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/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-rust.yml b/.github/workflows/test-rust.yml index c6901411167..17b6481a2bf 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,8 +4,22 @@ 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" @@ -20,8 +34,22 @@ 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" 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 cc606339a20..f55c87c2ae5 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -200,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/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 91835e19e3c..d360074ea4e 100644 --- a/Makefile +++ b/Makefile @@ -150,8 +150,8 @@ 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) @base_ref=$$($(RESOLVE_BASE)) && \ changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \ 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/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 f7c918a6827..cf7b3f8a38d 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -56,118 +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.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 }} - {{- 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 }} @@ -245,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 }} @@ -252,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 }} @@ -280,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/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/values.yaml b/helm/litellm-helm/values.yaml index 8dc7b967e11..fcee331a5aa 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -190,6 +190,48 @@ 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 @@ -222,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: @@ -243,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 @@ -319,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 c459512c7b9..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. @@ -443,3 +457,34 @@ ImplementationSpecific {{- 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 9cb6b07e77b..c06cc9583a0 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -61,6 +61,9 @@ 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 }} @@ -71,8 +74,11 @@ spec: - 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 .Values.gateway.metricsServer.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 @@ -83,6 +89,10 @@ spec: - 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 }} @@ -142,10 +152,53 @@ spec: 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 .Values.gateway.metricsServer.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 @@ -156,6 +209,11 @@ spec: - 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/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/values.yaml b/helm/litellm/values.yaml index b5d535c992d..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 @@ -284,6 +304,52 @@ gateway: 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 @@ -340,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/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 05c5aad9303..817df082d8c 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[] diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 91b4e4a7ba1..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.95" +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.95" +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 43b9ec1aac2..0e8e6e09a21 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,20 @@ # It is not intended for manual editing. version = 4 +[[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" @@ -418,7 +432,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", "axum-core", - "base64", + "base64 0.22.1", "bytes", "futures-util", "http 1.4.2", @@ -468,6 +482,12 @@ dependencies = [ "tracing", ] +[[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" @@ -542,6 +562,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" @@ -644,6 +673,21 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[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 = "const-oid" version = "0.10.2" @@ -696,7 +740,7 @@ dependencies = [ "ciborium", "clap", "criterion-plot", - "itertools", + "itertools 0.13.0", "num-traits", "oorandom", "page_size", @@ -716,7 +760,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" dependencies = [ "cast", - "itertools", + "itertools 0.13.0", ] [[package]] @@ -778,6 +822,56 @@ 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" @@ -790,6 +884,37 @@ version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +[[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" version = "0.10.7" @@ -841,6 +966,12 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + [[package]] name = "fastrand" version = "2.5.0" @@ -964,6 +1095,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 +1116,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 6.0.0", "rand_core 0.10.1", "wasm-bindgen", ] @@ -1220,7 +1363,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", @@ -1319,6 +1462,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 +1497,8 @@ checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", "hashbrown", + "serde", + "serde_core", ] [[package]] @@ -1365,6 +1516,15 @@ 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" @@ -1409,7 +1569,7 @@ name = "litellm-ai-gateway" version = "0.1.0" dependencies = [ "axum", - "base64", + "base64 0.22.1", "futures-channel", "futures-util", "litellm-config", @@ -1447,17 +1607,19 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "base64", + "base64 0.22.1", "rand 0.8.7", "reqwest", "rstest", "serde", "serde_json", + "serde_path_to_error", "sha2 0.10.9", "thiserror 2.0.19", "tokio", "tracing", "tracing-subscriber", + "url", ] [[package]] @@ -1469,6 +1631,7 @@ dependencies = [ "litellm-ai-gateway", "litellm-core", "litellm-python-interop", + "litellm-token-counter", "pyo3", "pyo3-async-runtimes", "serde", @@ -1489,6 +1652,22 @@ 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" @@ -1507,6 +1686,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" @@ -1535,6 +1730,12 @@ dependencies = [ "unicase", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "mio" version = "1.2.2" @@ -1546,6 +1747,38 @@ dependencies = [ "windows-sys 0.61.2", ] +[[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" @@ -1576,6 +1809,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" @@ -1604,6 +1859,18 @@ dependencies = [ "winapi", ] +[[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" @@ -1850,6 +2117,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" @@ -1863,10 +2136,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" @@ -1888,6 +2171,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" @@ -1897,6 +2190,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" @@ -1922,6 +2224,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" @@ -1979,7 +2292,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", @@ -2361,12 +2674,36 @@ 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 = "subtle" version = "2.6.1" @@ -2535,6 +2872,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" @@ -2773,6 +3143,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" @@ -2856,6 +3247,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" @@ -3081,6 +3481,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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 82de7f40069..f3e54e5b2aa 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" } @@ -39,6 +41,8 @@ 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" +url = "2.5.8" +criterion = "0.8.2" [profile.release] opt-level = 3 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..098ce071efc 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,7 @@ 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 { .. } => "AuthError", Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index ed41f1ff9e7..d1dd811f7ab 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -386,7 +386,7 @@ 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 { .. } => "AuthError", Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", 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..c22d05f5726 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -114,7 +114,8 @@ impl IntoResponse for MessagesRouteError { | Error::Connect(_) | Error::InvalidResponse(_) | Error::InvalidType { .. } - | Error::MissingField(_) => ( + | Error::MissingField(_) + | Error::MissingApiKey { .. } => ( StatusCode::BAD_GATEWAY, "messages provider request 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/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index c0de7ff3977..b4ca88cf16f 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,6 +4,11 @@ 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 @@ -11,10 +16,13 @@ rand.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true +serde_path_to_error = "0.1" +tokio.workspace = true thiserror.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, optional = true } sha2.workspace = true +url.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 } diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index fc81f4fa029..108d2a48e30 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -43,3 +43,6 @@ 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 OCR_HTTP_TIMEOUT_SECS: u64 = 600; +pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; +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..0382314057f 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -17,6 +17,10 @@ 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("upstream request failed with status {status}: {body}")] Http { status: u16, body: String }, #[error("upstream network error: {0}")] @@ -36,6 +40,59 @@ pub enum Error { Unsupported(&'static str), } +#[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), + } + } +} + pub fn json_type_name(value: &serde_json::Value) -> &'static str { match value { serde_json::Value::Null => "null", @@ -46,3 +103,53 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str { serde_json::Value::Object(_) => "object", } } + +#[cfg(test)] +mod transport_tests { + use super::*; + + #[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..3a0896a4d5a 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -14,5 +14,6 @@ pub mod realtime; pub mod responses; pub mod router; pub mod routing_utils; +mod url_utils; pub use error::Error; 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..7dfb08b4dc2 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -0,0 +1,69 @@ +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 mistral; + +pub(crate) use mistral::MistralAdapter; + +/// 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; + } + }; +} + +pub(crate) use for_each_ocr_adapter; 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..36bfc036bae --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -0,0 +1,75 @@ +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::constants::OCR_CONNECT_TIMEOUT_SECS; +use crate::error::TransportError; + +#[derive(Clone)] +pub struct OcrClient { + provider_http: reqwest::Client, +} + +impl OcrClient { + pub fn new(provider_http: reqwest::Client) -> Result { + Ok(Self { provider_http }) + } + + #[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 + } + + #[cfg(test)] + pub(crate) fn for_test(provider_http: reqwest::Client) -> Self { + Self { provider_http } + } +} + +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/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..cd0a1dc6b17 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs @@ -0,0 +1,102 @@ +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}; + + #[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 document: OcrDocument = serde_json::from_value( + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + ) + .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); + } + + #[test] + fn request_mapping_filters_unknown_fields() { + let params: MistralOcrParams = serde_json::from_value(json!({"unknown": true})).unwrap(); + let document: OcrDocument = serde_json::from_value( + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + ) + .unwrap(); + let result = + serde_json::to_value(transform_ocr_request("model", document, ¶ms).unwrap()) + .unwrap(); + assert!(result.get("unknown").is_none()); + } + + #[test] + fn response_preserves_provider_fields() { + let response: MistralOcrResponse = serde_json::from_value(json!({ + "pages":[{"index":0,"markdown":"hello","header":"head","confidence_scores":{"mean":0.99}}], + "model":"returned-model", + "usage_info":{"pages_processed":1,"future_counter":5}, + "future_response_field":"kept" + })) + .unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["header"], "head"); + assert_eq!(result["usage_info"]["future_counter"], 5); + assert_eq!(result["future_response_field"], "kept"); + assert_eq!(result["model"], "returned-model"); + } + + #[test] + fn response_rejects_null_pages() { + assert!(serde_json::from_value::(json!({"pages":null})).is_err()); + } +} 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..170ef5f68a7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/mod.rs @@ -0,0 +1 @@ +pub(crate) mod mistral; 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..50793aad566 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -0,0 +1,42 @@ +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), +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum OcrResponseError { + #[error("invalid OCR response field: {path}")] + ResponseField { path: String }, +} + +#[derive(Debug, Error)] +pub enum OcrError { + #[error("{0}")] + Request(#[from] OcrRequestError), + #[error("{0}")] + Response(#[from] OcrResponseError), + #[error("{0}")] + Transport(#[from] TransportError), + #[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::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..78c86aeaa0e 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,2 +1,21 @@ +mod adapters; +pub mod client; +mod codecs; +pub mod error; +mod handler; +pub mod hooks; +mod prepare; +mod registry; pub mod transformation; pub mod types; +pub mod wire; + +pub use client::{OcrClient, ocr}; +pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; + +#[cfg(test)] +#[path = "../../tests/ocr/support.rs"] +pub(crate) mod test_support; +#[cfg(test)] +#[path = "../../tests/ocr.rs"] +pub(crate) mod 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..cff40b7de50 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -0,0 +1,142 @@ +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; + +#[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) 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) +} + +#[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..9bae8153ee8 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -0,0 +1,53 @@ +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, +} + +impl OcrProvider { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Mistral => "mistral", + } + } +} + +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, + value => return Err(Error::InvalidProvider(value.to_string())), + }; + match typed_provider { + OcrProvider::Mistral => Ok((provider.model.to_string(), OcrAdapterKind::Mistral)), + } +} diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 62299faf9ed..ac4f10bf15b 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -1,7 +1,7 @@ use crate::Error; use serde_json::{Map, Value}; -use super::types::{OcrRequestData, OcrResponseData}; +use super::types::{LiteLLMOcrResponse, OcrRequestData}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrAuthStrategy { @@ -49,14 +49,14 @@ pub trait OcrProviderConfig: Sync { &self, model: &str, response_json: Value, - ) -> Result; + ) -> Result; fn transform_ocr_response_with_params( &self, model: &str, response_json: Value, _optional_params: &Map, - ) -> Result { + ) -> Result { self.transform_ocr_response(model, response_json) } diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 71cdb232a87..02b5330188c 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,6 +1,14 @@ +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::constants::OCR_HTTP_TIMEOUT_SECS; + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct OcrRequestData { pub data: Value, @@ -8,31 +16,145 @@ pub struct OcrRequestData { } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct OcrResponseData { +#[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, + }, +} + +#[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_base: Option, + pub extra_headers: Vec<(String, String)>, + pub timeout: Duration, +} + +impl Default for OcrConnection { + fn default() -> Self { + Self { + api_key: None, + api_base: None, + extra_headers: Vec::new(), + timeout: Duration::from_secs(OCR_HTTP_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(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, + 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 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..f37c06a01ac --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -0,0 +1,154 @@ +use crate::ocr::error::OcrRequestError; +use crate::ocr::error::OcrResponseError; +use std::time::Duration; + +use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest}; +use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument}; +use crate::Error; +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, + 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 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_base: nonblank(wire.api_base), + extra_headers: headers, + timeout: timeout.unwrap_or(defaults.timeout), + }; + Ok(LiteLLMOcrRequest { + connection, + ..request + }) +} + +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/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index d15c032f0bc..2af25a5639a 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -2,7 +2,7 @@ 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 crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -440,7 +440,7 @@ fn transform_document_intelligence_response( model: &str, response_json: Value, preserve_native_response: bool, -) -> Result { +) -> Result { let response = response_json .as_object() .ok_or_else(|| Error::InvalidType { @@ -488,7 +488,7 @@ fn transform_document_intelligence_response( }) .collect(); - Ok(OcrResponseData { + Ok(LiteLLMOcrResponse { usage_info: Some(json!({ "pages_processed": pages.len(), "doc_size_bytes": null, @@ -521,7 +521,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -599,7 +599,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { transform_document_intelligence_response(model, response_json, false) } @@ -608,7 +608,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { model: &str, response_json: Value, optional_params: &Map, - ) -> Result { + ) -> Result { transform_document_intelligence_response( model, response_json, @@ -718,7 +718,7 @@ mod tests { }) } - fn assert_native_fields_preserved(response: &OcrResponseData, operation: &Value) { + fn assert_native_fields_preserved(response: &LiteLLMOcrResponse, operation: &Value) { let analyze_result = &operation["analyzeResult"]; assert_eq!(response.extra_fields["content"], analyze_result["content"]); diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 11e8fe7db18..044fc587c22 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -1,6 +1,6 @@ use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData}; use serde_json::{Map, Value}; const SUPPORTED_OCR_PARAMS: &[&str] = &[ @@ -107,7 +107,7 @@ impl OcrProviderConfig for MistralOcrConfig { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { let response_object = response_json .as_object() .ok_or_else(|| Error::InvalidType { @@ -128,7 +128,7 @@ impl OcrProviderConfig for MistralOcrConfig { let document_annotation = response_object.get("document_annotation").cloned(); let usage_info = response_object.get("usage_info").cloned(); - Ok(OcrResponseData { + Ok(LiteLLMOcrResponse { pages, model, document_annotation, @@ -178,7 +178,10 @@ pub fn transform_ocr_request( } #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub fn transform_ocr_response(model: &str, response_json: Value) -> Result { +pub fn transform_ocr_response( + model: &str, + response_json: Value, +) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs index b8507541e18..f025887846f 100644 --- a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs @@ -6,7 +6,7 @@ use serde_json::{Map, Value, json}; use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData}; pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai"; pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY"; @@ -283,7 +283,7 @@ fn build_pages(result: &Map) -> Vec { pub fn transform_reducto_response( model: &str, response_json: Value, -) -> Result { +) -> Result { let response = response_json .as_object() .ok_or_else(|| Error::InvalidType { @@ -311,7 +311,7 @@ pub fn transform_reducto_response( "credits": usage.get("credits").cloned().unwrap_or(Value::Null), })); - Ok(OcrResponseData { + Ok(LiteLLMOcrResponse { pages: build_pages(result), model: model.to_string(), document_annotation: None, @@ -341,7 +341,7 @@ impl OcrProviderConfig for ReductoParseV3Config { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { transform_reducto_response(model, response_json) } @@ -383,7 +383,7 @@ impl OcrProviderConfig for ReductoParseLegacyConfig { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { transform_reducto_response(model, response_json) } 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 index c324de8cb45..b0e5a278d0b 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,6 +1,6 @@ use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -226,7 +226,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -301,7 +301,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { &self, model: &str, response_json: Value, - ) -> Result { + ) -> Result { let response = response_json .as_object() .ok_or_else(|| Error::InvalidType { @@ -339,7 +339,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { .get("usage_info") .cloned() .or_else(|| response.get("usage").cloned()); - Ok(OcrResponseData { + Ok(LiteLLMOcrResponse { pages, model: object .get("model") 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..982dca0dbe3 --- /dev/null +++ b/litellm-rust/crates/core/src/url_utils.rs @@ -0,0 +1,95 @@ +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 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"); + } +} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs new file mode 100644 index 00000000000..d1828dfb816 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -0,0 +1,227 @@ +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(), + 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(), + 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..d9720019419 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -0,0 +1,106 @@ +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 { + OcrClient::for_test(reqwest::Client::new()) +} + +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(), + 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/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..407475dedce 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -41,6 +41,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidRequest(_) | Error::InvalidType { .. } | Error::MissingField(_) + | Error::MissingApiKey { .. } | 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/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/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 5a461801b62..157aed12a83 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -45,9 +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, @@ -154,6 +156,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "smtp_email", "deepeval", "s3_v2", + "pointfive", "aws_sqs", "vector_store_pre_call_hook", "dotprompt", @@ -439,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 @@ -475,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", 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/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..baabfad6852 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 @@ -23,7 +23,7 @@ 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 .redis_cache import RedisCache, log_redis_failure if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -177,8 +177,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, @@ -217,8 +219,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 +254,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, @@ -339,8 +345,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,7 +365,9 @@ 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): @@ -372,7 +386,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 +426,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 +457,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/redis_cache.py b/litellm/caching/redis_cache.py index 106c1580110..6b93529e456 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -14,6 +14,7 @@ import functools import hashlib import inspect import json +import logging import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar @@ -391,10 +392,23 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +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) + + 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.""" if breaker.is_open(): - raise Exception(f"Redis circuit breaker is open — skipping {name}") + raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open — skipping {name}") return _swallowed_redis_failures.get() @@ -440,7 +454,7 @@ def _run_under_circuit_breaker_sync( 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) return result diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index c3c18e3d009..5a6debc4af5 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1201,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. diff --git a/litellm/constants.py b/litellm/constants.py index d545890dd59..028c08a691e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1654,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" @@ -1682,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"))) @@ -1998,6 +2000,9 @@ 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 + # 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/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 8c7f4557992..421a3c84d75 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -54,18 +54,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 @@ -777,8 +781,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) @@ -854,8 +869,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) @@ -890,8 +916,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/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index c40b90cee25..85bfcc6e7ed 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -378,6 +378,27 @@ }, "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_logger.py b/litellm/integrations/custom_logger.py index 8f03e08f02d..d445a3adf14 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 @@ -897,10 +898,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac 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 +909,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: 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 +951,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/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/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 996ff9c75af..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 @@ -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) - await run_aws_signing(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", 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 eacc3e4860a..aa7d6ca1699 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -2,6 +2,7 @@ ## Helper utilities import copy import logging +import re from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal @@ -21,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. 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/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 92b32d32dc0..edd2e88f95c 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -17,6 +17,7 @@ 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", } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e6b2bb164ef..6ee68ab21c5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -17,7 +17,7 @@ from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast from httpx import Response -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm import ( @@ -42,6 +42,7 @@ 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, @@ -64,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 ( @@ -89,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 @@ -120,7 +127,6 @@ from litellm.types.utils import ( CachingDetails, CallTypes, CostBreakdown, - CostResponseTypes, CustomPricingLiteLLMParams, DynamicPromptManagementParamLiteral, EmbeddingResponse, @@ -183,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 @@ -204,7 +211,7 @@ if TYPE_CHECKING: from litellm.integrations.otel.logger import OpenTelemetryV2 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 BasePassthroughConfig + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, LoggedRelayResponse try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -288,6 +295,9 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: # 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 @@ -473,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, @@ -1211,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( @@ -2381,7 +2400,7 @@ class Logging(LiteLLMLoggingBaseClass): self, raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: 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, @@ -4377,6 +4396,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): @@ -5065,6 +5092,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): @@ -5358,23 +5389,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( @@ -5632,7 +5663,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 @@ -6281,6 +6312,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( 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/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index d70530534da..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, Iterator, 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: 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 cf9604a0fd5..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 @@ -210,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] @@ -1004,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. @@ -1030,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 db23929e0c3..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 @@ -1641,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) @@ -1996,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 @@ -2248,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 @@ -2371,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 70562748c6a..1de2533d514 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -179,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) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index ca520116606..87c4ec8938e 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -73,8 +73,13 @@ _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: - return user_agent.startswith("claude-cli/") + """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: @@ -1656,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"), @@ -1671,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. @@ -1680,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/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/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 039c462b38a..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,7 +14,10 @@ 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 @@ -146,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/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/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index adbf2e126fb..20180c5cfa2 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, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm.types.utils import CallTypes from ..base_utils import BaseLLMModelInfo @@ -7,9 +16,68 @@ 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 BasePassthroughConfig(BaseLLMModelInfo): @@ -23,8 +91,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 +126,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 +156,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,21 +165,21 @@ 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 _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: 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/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 96804a1fa62..f52c1cec6a8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -6,7 +6,7 @@ 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 @@ -14,7 +14,8 @@ from threading import Lock 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 @@ -31,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 @@ -52,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 @@ -129,6 +172,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): "aws_sts_endpoint", "aws_bedrock_runtime_endpoint", "aws_external_id", + "aws_session_tags", ] def _get_ssl_verify(self, ssl_verify: bool | str | None = None): @@ -146,7 +190,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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. """ @@ -156,7 +200,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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: """ @@ -231,6 +275,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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, ): """ @@ -267,6 +312,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): (aws_external_id, "AWS_EXTERNAL_ID"), ) ) + session_tags: Final = _canonical_aws_session_tags(aws_session_tags) verbose_logger.debug( "in get credentials\n" @@ -279,7 +325,8 @@ class BaseAWSLLM(SignsRequestsWithAWS): "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, @@ -290,6 +337,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_web_identity_token is not None, aws_sts_endpoint, aws_external_id, + session_tags, ) args: Final = { @@ -303,6 +351,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): "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, } @@ -345,6 +394,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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, ), ) @@ -989,6 +1039,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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 @@ -1041,16 +1092,9 @@ class BaseAWSLLM(SignsRequestsWithAWS): # 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, @@ -1060,6 +1104,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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 @@ -1083,16 +1128,9 @@ class BaseAWSLLM(SignsRequestsWithAWS): # 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. @@ -1127,6 +1165,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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]: """ @@ -1153,6 +1192,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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, ) @@ -1168,6 +1208,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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 @@ -1198,6 +1239,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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( @@ -1207,6 +1249,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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) @@ -1243,14 +1286,9 @@ class BaseAWSLLM(SignsRequestsWithAWS): **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) @@ -1469,6 +1507,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): "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( @@ -1487,6 +1526,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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, @@ -1632,6 +1672,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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( @@ -1645,6 +1686,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): 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) 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/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 6d48ff3f07c..d397420cb17 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -357,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 @@ -375,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, ) ) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index be4f0f32689..cb2c70e74c8 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -93,6 +93,7 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( "aws_web_identity_token", "aws_sts_endpoint", "aws_external_id", + "aws_session_tags", ) @@ -1663,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/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index be766eaedd0..7efdfd3cebb 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -87,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: @@ -117,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 diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e1f0fc9e7d3..12e515e19a8 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -751,7 +751,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 +771,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/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/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/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/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/transformation.py b/litellm/llms/openai/responses/transformation.py index 926de3e8854..833ae206024 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,4 +1,4 @@ -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 @@ -15,6 +15,10 @@ 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 @@ -40,7 +44,7 @@ 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}) @@ -106,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. @@ -293,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 @@ -378,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: @@ -421,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: """ @@ -620,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]: """ 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/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/main.py b/litellm/main.py index 819626cc986..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, ), ) @@ -5550,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 @@ -8595,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"] @@ -8712,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: @@ -8789,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) @@ -8966,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 0d2eda93323..b7726290f0e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30295,7 +30295,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, @@ -30340,7 +30340,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, @@ -30432,7 +30432,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, @@ -30478,7 +30478,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, @@ -31474,7 +31474,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 @@ -31526,7 +31526,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 @@ -31578,7 +31578,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 }, @@ -31629,7 +31629,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 }, @@ -33709,13 +33709,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, @@ -39969,6 +39970,45 @@ "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-oss-120b": { "input_cost_per_token": 3.7e-08, "litellm_provider": "openrouter", @@ -48615,6 +48655,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, @@ -48625,6 +48666,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, @@ -48635,6 +48677,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, @@ -48663,6 +48706,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, @@ -48719,6 +48763,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, @@ -48729,6 +48774,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, @@ -57212,6 +57258,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 + } } ] }, @@ -58590,6 +58644,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, @@ -58602,6 +58657,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, @@ -58614,6 +58670,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, @@ -58626,6 +58683,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, @@ -58666,6 +58724,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, @@ -58678,6 +58737,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, @@ -58690,6 +58750,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, @@ -58702,6 +58763,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, @@ -58714,6 +58776,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, @@ -58736,6 +58799,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, @@ -58748,6 +58812,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, @@ -58758,6 +58823,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, @@ -58770,6 +58836,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, @@ -58789,7 +58856,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, 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/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 4bac65125b4..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, ) @@ -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. diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 7379126983a..789b2ffaef4 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -6,10 +6,18 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone 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, @@ -117,6 +125,7 @@ class _OAuthCredentialAccessToken(TypedDict): class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): + identity_binding_proof: ReadOnly[str] type: str refresh_token: str expires_at: str @@ -1393,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. @@ -1409,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 @@ -1628,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", "") @@ -1677,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( @@ -1709,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 ) @@ -1742,6 +1772,10 @@ async def resolve_valid_user_oauth_token( grant: Final = oauth_grant_state(cred) if cred is None or grant == "absent": return None + 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 prisma_client is None: @@ -1782,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 @@ -1816,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( diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cab4b6c161a..42fbe82531c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -57,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, @@ -139,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. @@ -149,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 @@ -169,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, @@ -210,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) @@ -225,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 @@ -234,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) @@ -547,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. @@ -588,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", @@ -616,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, ) @@ -854,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 @@ -884,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( @@ -908,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, @@ -930,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: @@ -1015,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 @@ -1051,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, @@ -1070,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( @@ -1097,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( @@ -1137,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( @@ -2134,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( @@ -2310,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 @@ -2401,11 +2478,6 @@ 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.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], @@ -2413,6 +2485,9 @@ async def _build_oauth_protected_resource_response( "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/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 3d94fa345d0..f4889008e94 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -411,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 @@ -443,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 @@ -729,11 +729,15 @@ async def _flow_target( 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 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 = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive" + state: Final = ( + "interactive" + if server.is_gateway_managed_oauth2 and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" + else "m2m" + ) return state, server diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 6aaa00ae415..ff30ef99ebd 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,16 @@ Usage with curl:: http://localhost:4000/mcp/atlassian_mcp """ -from typing import TYPE_CHECKING, Final +import json +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final +from starlette.requests import HTTPConnection from starlette.types import Message, Send 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 +118,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. @@ -144,37 +238,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 +307,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 +338,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 +358,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) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d7f238f142c..af25b0e919a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -80,6 +80,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 record_auth_resolution from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, @@ -108,12 +109,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, @@ -2441,6 +2444,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") @@ -3832,13 +3837,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, @@ -3853,11 +3866,14 @@ 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. @@ -3960,6 +3976,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) @@ -4032,6 +4049,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, @@ -4086,6 +4104,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, 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/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/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/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f8f1c563811..95129d9aeed 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -49,7 +49,11 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_server_name, _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode ) -from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug +from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + MCPAuthDiagnostics, + MCPDebug, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, get_passthrough_www_authenticate, @@ -2044,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 @@ -4472,13 +4502,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: @@ -4677,6 +4709,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": @@ -4813,6 +4846,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/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4d5cdcc8003..f0af17ab818 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -5384,7 +5384,7 @@ "additionalProperties": { "type": "string" }, - "description": "Git source reference", + "description": "Plugin source reference", "title": "Source", "type": "object" }, @@ -5411,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": [ @@ -5509,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" }, @@ -5653,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" }, @@ -5721,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": { @@ -5731,6 +5749,16 @@ } }, "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" } }, "summary": "Get Marketplace", @@ -5788,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": { @@ -5923,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": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c22bb76629d..2cbff128635 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -497,6 +497,9 @@ 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. @@ -700,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", @@ -929,10 +933,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", @@ -1124,6 +1128,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 @@ -2427,6 +2432,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)) @@ -3726,6 +3738,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): """ 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/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a71a1993064..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, @@ -3450,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 f78c4221f5a..be65c3b39ec 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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 @@ -2003,9 +2005,20 @@ def get_model_from_request( 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, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index bfccb703e76..20ab9904f46 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -91,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, @@ -183,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]: @@ -2650,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, @@ -2685,6 +2726,7 @@ 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: @@ -2710,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), ) @@ -2768,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 @@ -3131,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 03d01ff7f66..4045857a237 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. @@ -548,7 +548,7 @@ lite --base-url https://your-proxy.example.com configure claude --api-key sk-... claude ``` -With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way 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 (the ones whose id contains `claude` or `anthropic`) 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, which has to be on `/v1/models` for the key. 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 and sends no thinking parameters for it, so either name the group like a Claude model id or append `[1m]` to opt into the 1M window. 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 +With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way 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, which has to be on `/v1/models` for the key. 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 diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py index 83fbd323924..9d329f8d8f2 100644 --- a/litellm/proxy/client/cli/commands/configure.py +++ b/litellm/proxy/client/cli/commands/configure.py @@ -1,16 +1,23 @@ """`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" import os -import re 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, context_secret_vault, get_stored_api_key from .claude_settings import ( STARTING_MODEL_ROLE, @@ -30,14 +37,16 @@ from .claude_settings import ( settings_file_owners, unconfigure_claude_settings, ) -from .pi import ListingFailure, PiSyncError, fetch_model_ids +from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing from .up import ensure_fresh_login _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_PICKER_FILTER: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +_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 " @@ -65,7 +74,16 @@ def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeC return ApiKeyHelper(resolve_api_key_helper(base_url)), stored -def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, tuple[str, ...]]: +@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[ClaudeCredential, _Listing]: """Every configure path begins the same way: the local ownership check first, so a `lite up` session is refused before any login prompt or request, then the credential, then the listing.""" settings_path: Final = claude_settings_path(os.environ) @@ -88,21 +106,28 @@ def _listing_error(base_url: str, error: PiSyncError) -> str: 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) -> tuple[str, ...]: - listed: Final = fetch_model_ids(base_url, key) +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 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: ClaudeCredential, listed: Sequence[str], model: str | None) -> None: +def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Listing, model: str | None) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] - if model is not None and model not in listed: + 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( @@ -113,29 +138,32 @@ def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequ configure_claude_settings( base_url, credential, - _model_choice(model), + _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_FILTER.search(listed_model)) + 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." if isinstance(credential, StaticToken) else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it." ) click.echo( - f"Starting model: {model} ({STARTING_MODEL_ROLE}); switch any time with /model." - if model is not None + 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." ) click.echo( - f"/model will list {in_picker} of the proxy's {len(listed)} models (Claude Code shows only ids containing " - "'claude' or 'anthropic')." + 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 isinstance(credential, StaticToken) and settings_path.is_symlink(): @@ -173,8 +201,10 @@ def interactive_configure( targets: Final = pick_targets() if _CLAUDE_TARGET not in targets: return - credential, listed = _start(ctx, None) - _apply_claude(ctx, credential, listed, pick_model(listed)) + 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) @@ -218,8 +248,8 @@ def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. Assumes the proxy is already running. """ - credential, listed = _start(ctx, api_key) - _apply_claude(ctx, credential, listed, model) + credential, listing = _start(ctx, api_key) + _apply_claude(ctx, credential, listing, model) @unconfigure_group.command(name="claude") diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 70c89a853e3..9810e81ae36 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -13,10 +13,11 @@ 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" @@ -51,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): @@ -69,17 +83,18 @@ class _ModelGroupList(BaseModel): data: tuple[_ModelGroup, ...] -def fetch_model_ids( +def fetch_model_listing( base_url: str, api_key: str, *, get: Callable[..., requests.Response] = requests.get, -) -> tuple[str, ...] | PiSyncError: + 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}"}, # mutable-ok: requests headers require a dict + headers={"Authorization": f"Bearer {api_key}", **headers}, # mutable-ok: requests headers require a dict timeout=10, ) except requests.RequestException as e: @@ -94,10 +109,21 @@ def fetch_model_ids( 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) - ids: Final = tuple(dict.fromkeys(model.id for model in listing.data)) - if not ids: + models: Final = tuple(dict.fromkeys(listing.data)) + if not models: return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY) - return ids + 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: + 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({}) @@ -222,11 +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/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_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/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/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_url_settings.py b/litellm/proxy/db/db_url_settings.py index d1e4b3e92b8..e93bc96da69 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -52,6 +52,7 @@ 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, @@ -358,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 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/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/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/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 747cfa526d9..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( 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 a074f02f4e8..0339cf4dfea 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -59,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. @@ -666,7 +670,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if response_has_hidden_params(response): priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) additional_headers: Final = ensure_response_additional_headers(response) - additional_headers["x-litellm-priority"] = priority or "default" + 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" return response 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/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c6c3dde4b6e..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,6 +27,7 @@ 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 ( @@ -1228,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, @@ -1475,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) @@ -1505,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: @@ -1631,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: @@ -1814,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] @@ -1866,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( @@ -3856,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, @@ -3922,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: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..00406ad436e 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -27,6 +27,16 @@ from litellm.proxy.db.db_spend_update_writer import ( 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, @@ -34,6 +44,7 @@ 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 ( @@ -71,8 +82,43 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( 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, + ) -> None: + super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) + self.spend_event_producer = spend_event_producer + + 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, @@ -503,6 +549,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 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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index da033dc2276..ad4687e95db 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -31,6 +31,7 @@ from litellm.constants import ( 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, @@ -83,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,}$") @@ -810,16 +807,6 @@ def apply_missing_session_id_policy( ) -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 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/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 2aa7fdc7393..4c97bbaf5de 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2278,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( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0f19e9ce149..94d2b773e14 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -111,6 +111,7 @@ from litellm.router_utils.auto_router_model_naming import ( 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, @@ -897,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, @@ -1650,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 @@ -1704,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, @@ -2037,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, @@ -2221,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/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index bf6ee57b4fc..face515ef88 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -34,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 @@ -53,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, @@ -78,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 @@ -120,6 +123,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. @@ -412,7 +433,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, @@ -1499,6 +1520,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, @@ -1508,30 +1537,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"} 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/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 30ff47ac9f9..0f24acb8bb4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,16 @@ 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 @@ -51,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, @@ -151,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: @@ -233,7 +244,7 @@ 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 @@ -268,7 +279,7 @@ 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 ( @@ -366,8 +377,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, @@ -412,6 +426,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, @@ -454,7 +469,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, @@ -577,6 +592,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: @@ -928,6 +948,17 @@ def cleanup_router_config_variables(): 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: if prisma_client is None: return @@ -1363,6 +1394,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() @@ -2442,16 +2477,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 @@ -3744,13 +3790,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 @@ -3762,16 +3845,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: @@ -4078,6 +4161,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 @@ -4903,7 +4988,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: " @@ -6639,6 +6726,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"], @@ -8632,6 +8727,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( @@ -9222,7 +9318,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." @@ -10469,6 +10567,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": @@ -10555,6 +10671,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( @@ -10603,6 +10724,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( @@ -17655,6 +17781,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"], @@ -17687,10 +17873,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) @@ -17721,6 +17907,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 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 05c5aad9303..817df082d8c 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[] diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index bf5eadcd85d..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 @@ -210,6 +211,7 @@ 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 @@ -237,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]] = {} @@ -315,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), } @@ -1272,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)) @@ -1314,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") @@ -1355,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( 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..da79328fa59 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 @@ -3099,7 +3104,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 +3133,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 +4654,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 a21d761996f..f0f38358cf0 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,21 @@ 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, 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 +39,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 +337,15 @@ def _sl_attribution_fallback( return standard_logging_payload.get(field) or "" +def _looks_like_model_name(model: str) -> bool: + return len(model) <= MAX_SPEND_LOG_MODEL_NAME_LENGTH and not any(char.isspace() for char in model) + + 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 +443,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 +554,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 @@ -843,7 +866,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": @@ -860,7 +883,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: @@ -1093,7 +1116,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)] @@ -1165,7 +1188,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): @@ -1242,14 +1265,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") @@ -1272,7 +1299,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) @@ -1292,7 +1319,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 @@ -1316,7 +1343,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 "{}" @@ -1364,7 +1391,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 00ccad33b6d..253494b02f4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -68,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 @@ -124,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, @@ -3902,6 +3909,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. @@ -3996,6 +4008,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). @@ -4004,13 +4017,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, ) @@ -4079,6 +4092,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")) ) @@ -5796,15 +5811,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.", @@ -6124,6 +6144,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: @@ -6135,6 +6167,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. @@ -6404,48 +6469,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 ### @@ -7547,6 +7577,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 @@ -7554,6 +7590,10 @@ 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( @@ -7955,6 +7995,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, @@ -7979,31 +8101,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/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/main.py b/litellm/responses/main.py index 7a210cdd970..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 @@ -407,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): @@ -1078,6 +1110,7 @@ def responses( 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") @@ -1186,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: diff --git a/litellm/router.py b/litellm/router.py index b4f48fa9496..9e6db66db83 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -53,6 +53,7 @@ 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, @@ -99,6 +100,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( 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, @@ -150,6 +152,7 @@ from litellm.router_utils.common_utils import ( 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, @@ -228,6 +231,7 @@ from litellm.types.router import ( CredentialLiteLLMParams, CustomRoutingStrategyBase, Deployment, + DeploymentModelListingInfo, DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, @@ -941,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], ...] = () @@ -5199,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): @@ -5235,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, @@ -5755,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( **{ @@ -10057,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: """ @@ -10225,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: @@ -11018,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) @@ -11795,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 @@ -13345,8 +13396,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 diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index d605b43e42a..1a4764c291e 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -361,6 +361,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 @@ -442,11 +455,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: @@ -461,7 +489,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/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index faafcea404a..7519f4c5156 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -25,7 +25,7 @@ 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 ( @@ -36,9 +36,11 @@ from litellm.constants import ( 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 ( @@ -46,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 ( @@ -56,6 +59,7 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionImageObject, ChatCompletionTextObject, + ResponsesAPIResponse, ) from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, @@ -364,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 @@ -387,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 @@ -486,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, @@ -591,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, @@ -601,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( @@ -1630,6 +1691,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: @@ -1924,16 +1989,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 """ @@ -1944,6 +2005,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( @@ -1953,26 +2015,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, @@ -2004,34 +2070,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 @@ -2040,6 +2114,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, @@ -2427,7 +2528,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: @@ -3248,6 +3349,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 @@ -3349,6 +3462,7 @@ 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) + 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 @@ -3358,7 +3472,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: @@ -3369,7 +3483,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) @@ -3557,7 +3671,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: @@ -3591,7 +3706,7 @@ class ComplexityRouter(CustomLogger): ) 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 diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index d28924c69b2..1f1b5a5cc4b 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -971,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'." ), ) @@ -985,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'." ), @@ -1292,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." ), ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 14e6592e1fd..0b73f4e31a7 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,3 +1,4 @@ +import logging from collections.abc import Mapping, Sequence from typing import Final @@ -6,6 +7,7 @@ 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 @@ -87,16 +89,22 @@ def _least_busy( def _warn_unreadable(model_group: str, error: Exception) -> None: - verbose_router_logger.warning( - "least-busy routing could not read the shared in-flight counts for %s, " - "falling back to this worker's own counts: %s", - model_group, + 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: - verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error) + 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): diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index c1c6d25beca..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 @@ -256,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/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/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/bedrock.py b/litellm/types/llms/bedrock.py index 9f93886a9c6..76756ac35bb 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1107,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/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/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 84ffd50eea1..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 @@ -173,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 @@ -225,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).""" @@ -250,7 +280,23 @@ class MCPServer(BaseModel): @property def advertises_gateway_authorization_server(self) -> bool: """Whether named discovery should advertise the aggregate gateway authorization server.""" - return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + 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: 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/router.py b/litellm/types/router.py index 6b707a544a2..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 @@ -21,6 +21,7 @@ if TYPE_CHECKING: 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 ( @@ -288,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 @@ -612,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/utils.py b/litellm/types/utils.py index ab0cc5f959c..58b940227f8 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, @@ -3373,7 +3374,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 diff --git a/litellm/utils.py b/litellm/utils.py index 5110c42ee43..a765e1b1246 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, @@ -278,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 @@ -1195,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, @@ -1662,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 @@ -1676,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__ @@ -1843,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 @@ -1941,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, @@ -1998,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: @@ -2197,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 @@ -2967,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 @@ -3095,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: @@ -6982,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): @@ -7002,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): @@ -7029,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 ############################ @@ -8850,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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0d2eda93323..b7726290f0e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30295,7 +30295,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, @@ -30340,7 +30340,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, @@ -30432,7 +30432,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, @@ -30478,7 +30478,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, @@ -31474,7 +31474,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 @@ -31526,7 +31526,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 @@ -31578,7 +31578,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 }, @@ -31629,7 +31629,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 }, @@ -33709,13 +33709,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, @@ -39969,6 +39970,45 @@ "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-oss-120b": { "input_cost_per_token": 3.7e-08, "litellm_provider": "openrouter", @@ -48615,6 +48655,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, @@ -48625,6 +48666,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, @@ -48635,6 +48677,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, @@ -48663,6 +48706,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, @@ -48719,6 +48763,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, @@ -48729,6 +48774,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, @@ -57212,6 +57258,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 + } } ] }, @@ -58590,6 +58644,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, @@ -58602,6 +58657,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, @@ -58614,6 +58670,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, @@ -58626,6 +58683,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, @@ -58666,6 +58724,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, @@ -58678,6 +58737,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, @@ -58690,6 +58750,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, @@ -58702,6 +58763,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, @@ -58714,6 +58776,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, @@ -58736,6 +58799,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, @@ -58748,6 +58812,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, @@ -58758,6 +58823,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, @@ -58770,6 +58836,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, @@ -58789,7 +58856,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, diff --git a/pyproject.toml b/pyproject.toml index 04f2f3fd1dd..d33d693f794 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,7 +67,7 @@ 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.95", + "litellm-proxy-extras==0.4.96", "litellm-enterprise==0.1.66", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", diff --git a/schema.prisma b/schema.prisma index 05c5aad9303..817df082d8c 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[] 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 923a4ca2da8..6fcbdc2f500 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -258,6 +258,133 @@ 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 aa2c3d558e2..2b235c2bad5 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -213,6 +213,12 @@ locals { # 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" @@ -253,7 +259,7 @@ locals { 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 ? { @@ -272,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 ---------- @@ -298,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" @@ -316,16 +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, - local.metrics_env, - ) - secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) - mountPoints = local.metrics_mount_points + 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 @@ -342,7 +412,7 @@ resource "aws_ecs_task_definition" "gateway" { }, local.gateway_proxy_overrides, ) - ], local.gateway_metrics_container)) + ], local.gateway_metrics_container, local.collector_container)) dynamic "volume" { for_each = local.metrics_enabled ? [1] : [] 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/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 667f6db63c9..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 @@ -729,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/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index ffdc448f5fe..018d01f75a8 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -327,6 +327,7 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ key.Metadata = metadata if _, err := c.UpdateKey(key); err != nil { + d.Partial(true) return diag.FromErr(fmt.Errorf("error updating key: %s", err)) } 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_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_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/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/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index f0fc3d892e2..e1099fe0a62 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -510,6 +510,37 @@ async def _call(session: ClientSession, tool_id: str, a: int = 3, b: int = 4) -> 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" @@ -527,13 +558,31 @@ class TestProxyMcpAuthorizationScope: _assert_unauthorized(await _call(ungranted, restricted_id)) @pytest.mark.asyncio - async def test_no_mcp_servers_sentinel_hides_every_tool(self, proxy_server_url: str) -> None: + 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"] - async with _scoped_session(proxy_server_url, "sk-none") as session: - assert await _search(session, "add") == {} - _assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": tool_id})) - _assert_unauthorized(await _call(session, tool_id)) + 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: @@ -581,24 +630,7 @@ class TestProxyMcpAuthorizationScope: @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: - async with httpx.AsyncClient() as client: - response = await client.post( - f"{proxy_server_url}/mcp/proxy", - headers={ - "Accept": "application/json, text/event-stream", - **({"Authorization": f"Bearer {key}"} if key else {}), - }, - json={ - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2025-03-26", - "capabilities": {}, - "clientInfo": {"name": "auth-test", "version": "1"}, - }, - }, - ) + response = await _raw_initialize(proxy_server_url, key) assert response.status_code == 401, response.text @pytest.mark.asyncio @@ -646,25 +678,28 @@ class TestProxyMcpAuthorizationScope: @pytest.mark.asyncio async def test_proxy_scope_exception_returns_iserror_and_emits_failure_log(self, proxy_server_url: str) -> None: - async with _scoped_session( + 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"}, - ) as session: - result = await session.call_tool("call_tool", {"tool_id": "denied-scope", "arguments": {}}) - 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"] + ) + 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( 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_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/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/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..eae8bbfdaff 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,102 @@ 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 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 6b2df118611..bcaa58c9c40 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1013,3 +1013,20 @@ async def test_breaker_metrics_track_state_and_failure_class(): 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 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 b07e5876e8b..713330a8280 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -12,6 +12,7 @@ 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 @@ -1686,3 +1687,211 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N 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) 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/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_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/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 7c1445d79c9..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 @@ -2,11 +2,12 @@ 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, @@ -25,6 +26,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( 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" @@ -1442,7 +1445,7 @@ class TestFlattenTopLevelSchemaCombinators: assert schema == snapshot -class TestToolWithFlattenedParameters: +class TestToolWithSanitizedParameters: def _anyof_tool(self): return { "type": "function", @@ -1469,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"] @@ -1486,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 = { @@ -1497,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", @@ -1510,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: 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_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_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 6aa77745e3d..2f6339dcdbb 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3,7 +3,8 @@ import contextlib import datetime import os import sys -from typing import Literal +from collections.abc import Callable +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6801,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_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 3be0bae4120..76092e96307 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,69 @@ 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" 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 626b8a63b20..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 @@ -1554,3 +1554,52 @@ def test_stream_chunk_builder_reads_role_from_first_frame_with_choices() -> None 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_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 7da04d12569..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,6 +1,7 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function import asyncio +import base64 import importlib import threading import time @@ -25,6 +26,8 @@ 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 @@ -1558,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/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 91652c89092..e1b39c4ba13 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -2267,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_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..c8007acf70f --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -0,0 +1,617 @@ +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" + + await logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=[stream.encode()], provider_config=AzureAIPassthroughConfig() + ) + 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/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_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 cbf160c451f..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 @@ -57,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 @@ -105,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 ( 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 8831fd9061c..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 @@ -7,12 +7,15 @@ 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 @@ -562,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/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 4aad2846537..e5a460e2f1a 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1041,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 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 2c7e476e9a8..c5b8e7ecc9d 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -17,7 +17,7 @@ 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 ( @@ -2408,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 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 3f03305423a..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 @@ -616,6 +618,60 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): 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) # # --------------------------------------------------------------------------- # 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..3d4ba264c1c 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,68 @@ 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() 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/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_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index c5902b32a06..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): @@ -2022,6 +2025,86 @@ class TestFlattenToolSchemaCombinatorsWiring: 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 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/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/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/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/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 20f4719e4bf..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 @@ -7193,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( @@ -7209,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"), @@ -7223,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 @@ -7303,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, ) @@ -7333,7 +7335,11 @@ class TestAggregateGatewayDcrChallenge: (_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/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_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 355b3bfd30e..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 @@ -628,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) @@ -642,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 @@ -1184,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", @@ -1309,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" @@ -1348,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" @@ -1609,3 +1612,95 @@ 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 763200c3709..5a99139a67f 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 @@ -7651,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.""" @@ -7794,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() @@ -7914,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, @@ -8128,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 @@ -10921,3 +11047,97 @@ 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() 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 73a52a8d2e8..8aa4cfc5619 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 @@ -833,7 +833,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 @@ -844,7 +844,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, ) @@ -2044,3 +2044,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..7d46bb0237a 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,17 @@ 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 from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_DEBUG_REQUEST_HEADER, MCPDebug, + MCPAuthDiagnostics, ) @@ -166,77 +172,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 +204,90 @@ class TestWrapSendWithDebugHeaders: asyncio.run(wrapped(body_msg)) assert captured[0] == body_msg + + +@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" 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 9b6eaba3177..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 @@ -3,6 +3,7 @@ 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 @@ -1866,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 @@ -2013,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, @@ -2498,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, @@ -2610,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, @@ -5624,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: 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 46fef83092d..adf985e9a21 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 @@ -12567,3 +12567,112 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, 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) 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_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/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/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5bfef2b6445..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) diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index aaf630ad29b..cdf1f897707 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -565,6 +565,38 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): ) +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( @@ -3497,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): @@ -3516,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_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 0fd19f518d9..38a0e85c1ea 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): """ @@ -3826,3 +3902,17 @@ def test_team_disable_logging_stays_proxy_admin_only(): 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" 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 8e6c41761cc..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 @@ -7249,3 +7249,46 @@ async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_a 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/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py index a1808383982..ac7408339fe 100644 --- a/tests/test_litellm/proxy/client/cli/test_configure_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -89,7 +89,7 @@ class TestConfigureClaudeWithAVirtualKey: 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 len(responses.calls) == 1 + 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): @@ -358,3 +358,74 @@ class TestUnconfigureClaude: 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_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 1c2da514ee2..03e5d7dd197 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -13,6 +13,7 @@ from litellm.proxy.client.cli.commands.pi import ( PiSyncError, fetch_model_ids, fetch_model_limits, + fetch_model_listing, models_json_path, provider_block, sync_models_json, @@ -50,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") 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_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/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_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_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_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/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index b06d87ac67e..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(): """ @@ -3131,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/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 0cd6b4ede9c..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 @@ -1918,3 +1918,30 @@ async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): ) 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 6d382370f5f..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 @@ -6236,3 +6236,51 @@ async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): ) 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 eff892f2d80..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 @@ -5,14 +5,20 @@ 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 @@ -2096,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/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_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 c02f886fc31..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 @@ -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"), 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_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/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 a1cfd973b75..addc952af14 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 @@ -5022,7 +5022,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: @@ -5121,7 +5125,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" @@ -5181,7 +5189,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" @@ -5231,6 +5243,97 @@ class TestAzureRouterModelStreamingKeepalive: 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 @@ -5263,3 +5366,88 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): 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/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 dafe686c4f2..deb7289d2d1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -253,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 # --------------------------------------------------------------------------- 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_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_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 0fb9b1a6d88..300cac5cdb2 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 @@ -880,7 +880,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 +922,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 +957,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 +1003,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 +1060,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 = [ @@ -1315,7 +1315,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 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 de6c7c2a40a..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,13 +1,24 @@ +import json +import math from typing import Final import pytest -import litellm.proxy.proxy_server as proxy_server +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 estimate_request_max_cost, 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", @@ -56,6 +67,16 @@ async def test_non_exempt_llm_route_still_reserves_budget(): 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}), @@ -139,3 +160,172 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): ) 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_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..671a8ae63fc 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", [ @@ -6777,3 +6828,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_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5a79560b972..df113197ec6 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, { @@ -587,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 @@ -605,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 @@ -625,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 @@ -633,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 @@ -660,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 @@ -677,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 @@ -695,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 @@ -713,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 @@ -741,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 @@ -754,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, ): @@ -809,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 @@ -831,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. @@ -922,6 +990,88 @@ 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), + ], +) +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(): @@ -1412,7 +1562,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, ): @@ -1476,7 +1626,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 @@ -1486,7 +1636,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 @@ -1495,7 +1645,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 @@ -1506,22 +1656,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" @@ -1554,7 +1704,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, ): @@ -1590,7 +1740,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, ): @@ -1654,7 +1804,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, ): @@ -1686,7 +1836,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, ): @@ -1709,13 +1859,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 @@ -1747,7 +1897,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 @@ -1766,7 +1916,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 @@ -2182,7 +2332,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 @@ -2262,7 +2412,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, ): @@ -2290,7 +2440,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, ): @@ -2312,7 +2462,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, ): @@ -2343,7 +2493,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 @@ -2428,7 +2578,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, ): @@ -2460,7 +2610,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, ): @@ -2568,7 +2718,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, ): @@ -3164,7 +3314,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_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_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_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_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index b78ec7dcff6..646596c5b88 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 @@ -887,6 +888,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 +916,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 +931,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 +990,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 +1005,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 @@ -1960,6 +2014,122 @@ 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", [ @@ -2038,3 +2208,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/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index dc30798df55..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 @@ -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/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/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/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/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 51103297c58..dddaee71a63 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -5,6 +5,10 @@ Tests the rule-based complexity scoring and tier assignment logic. """ import asyncio +from collections.abc import AsyncIterator +import json +from copy import deepcopy +from functools import partial import logging import sys import time @@ -12,6 +16,7 @@ from typing import Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import pytest +import httpx from pydantic import ValidationError import litellm @@ -66,6 +71,8 @@ from litellm.types.router import ( LiteLLM_Params, TaggedPreRoutingStrategy, ) +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler requires_semantic_router = pytest.mark.skipif( @@ -194,7 +201,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.""" @@ -2482,6 +2496,366 @@ 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.""" @@ -2940,6 +3314,29 @@ 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( @@ -7601,11 +7998,382 @@ _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", [ @@ -13662,6 +14430,7 @@ class _OutputCeilingRecorder(CustomLogger): 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", diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index 9efa526fc02..c2fa41f4ca8 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -1,9 +1,11 @@ +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" @@ -185,3 +187,24 @@ def test_calls_without_a_deployment_are_ignored() -> 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_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/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/test_component_entrypoint.py b/tests/test_litellm/test_component_entrypoint.py index 09837d2b233..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 = { @@ -330,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() @@ -367,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", ) @@ -387,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 8f8a7640c08..f610821e06a 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -18,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, @@ -128,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 = { 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 5cb52295f94..f71225c6fc5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2394,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 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_router.py b/tests/test_litellm/test_router.py index f5b1f79ad73..def67ccf88b 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -19,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 @@ -42,7 +44,7 @@ from litellm.router import ( _is_retriable_anthropic_status, ) from litellm.router_strategy import simple_shuffle -from litellm.types.router import DeploymentTypedDict, RetryPolicy +from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, RetryPolicy def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -5049,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 @@ -8173,6 +8210,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=[ @@ -14137,6 +14372,49 @@ 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", @@ -15062,3 +15340,53 @@ def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_i {"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_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_utils.py b/tests/test_litellm/test_utils.py index 8b186be43e5..7753cb7d770 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,9 +1,12 @@ 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 @@ -15,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, @@ -6208,3 +6212,135 @@ def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_with "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/types/test_router.py b/tests/test_litellm/types/test_router.py index fd933a9d993..7fcb76b638f 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -1,6 +1,7 @@ import logging import pytest +from pydantic import ValidationError from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, @@ -122,3 +123,26 @@ 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_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 index 02d274bb405..b0c75d9d2f5 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,24 +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 pytest_collection_modifyitems(items): - rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - if not rust_enabled: - skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") +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: - item.add_marker(skip) + 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 + 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 index d5b1fce1139..1293de9ee0e 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -1,31 +1,34 @@ import json import threading +from collections.abc import Generator +from dataclasses import dataclass 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 +@dataclass(frozen=True, slots=True) +class RecordedOCRRequest: + body: object + + @pytest.fixture -def ocr_server(): - requests = [] +def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest]]]: + requests: Final[list[RecordedOCRRequest]] = [] class Handler(BaseHTTPRequestHandler): - def do_POST(self): + 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"]))), - } + RecordedOCRRequest( + body=json.loads(self.rfile.read(int(self.headers["Content-Length"]))), + ) ) - if self.headers.get("User-Agent", "").startswith("python-httpx"): - self.send_response(418) - self.end_headers() - return - response = json.dumps( + response: Final = json.dumps( { "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], "model": "mistral-ocr-latest", @@ -38,11 +41,11 @@ def ocr_server(): self.end_headers() self.wfile.write(response) - def log_message(self, format, *args): + def log_message(self, format: str, *args: object) -> None: pass - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + 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 @@ -52,21 +55,29 @@ def ocr_server(): thread.join() -def test_ocr_with_rust_extension(ocr_server): +def test_native_ocr_with_compiled_rust_extension( + ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]], +) -> None: server, requests = ocr_server - host, port = server.server_address + address: Final = server.server_address + host: Final = str(address[0]) + port: Final = int(address[1]) - response = litellm.ocr( - model="mistral/mistral-ocr-latest", + 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.pages[0].markdown == "native OCR response" + 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"}, + assert requests[0].body == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 76ac60a6453..e2a08a40bcb 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2325,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/package-lock.json b/ui/litellm-dashboard/package-lock.json index 44360e94392..1f459bc50ea 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -11402,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": { 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)/hooks/organizations/useOrganizations.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts index 960afe7392c..5448b7182d7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts @@ -81,7 +81,7 @@ describe("useOrganizations", () => { userRole: "Admin", token: "test-token", userEmail: "test@example.com", - premiumUser: false, + premiumUser: true, disabledPersonalKeyCreation: null, showSSOBanner: false, }); @@ -181,7 +181,7 @@ describe("useOrganizations", () => { userRole: "Admin", token: null, userEmail: "test@example.com", - premiumUser: false, + premiumUser: true, disabledPersonalKeyCreation: null, showSSOBanner: false, }); @@ -197,6 +197,26 @@ describe("useOrganizations", () => { expect(organizationListCall).not.toHaveBeenCalled(); }); + it("does not call the organization API when the session is not premium", async () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(organizationListCall).not.toHaveBeenCalled(); + }); + it("should not execute query when userId is missing", async () => { // Mock missing userId mockUseAuthorized.mockReturnValue({ @@ -205,7 +225,7 @@ describe("useOrganizations", () => { userRole: "Admin", token: "test-token", userEmail: "test@example.com", - premiumUser: false, + premiumUser: true, disabledPersonalKeyCreation: null, showSSOBanner: false, }); @@ -229,7 +249,7 @@ describe("useOrganizations", () => { userRole: null, token: "test-token", userEmail: "test@example.com", - premiumUser: false, + premiumUser: true, disabledPersonalKeyCreation: null, showSSOBanner: false, }); @@ -253,7 +273,7 @@ describe("useOrganizations", () => { userRole: null, token: null, userEmail: "test@example.com", - premiumUser: false, + premiumUser: true, disabledPersonalKeyCreation: null, showSSOBanner: false, }); @@ -335,7 +355,7 @@ describe("useOrganization", () => { userRole: "Admin", token: "test-token", userEmail: "test@example.com", - premiumUser: false, + premiumUser: true, disabledPersonalKeyCreation: null, showSSOBanner: false, }); @@ -356,6 +376,26 @@ describe("useOrganization", () => { expect(result.current.isLoading).toBe(false); }); + it("does not call the organization info API when the session is not premium", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useOrganization("org-1"), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(organizationInfoCall).not.toHaveBeenCalled(); + }); + it("falls through to the detail API call when no cached list contains the organization", async () => { (organizationInfoCall as any).mockResolvedValue(mockOrganizations[0]); queryClient.setQueryData(organizationKeys.list({ filters: { org_id: "org-2" } }), [mockOrganizations[1]]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts index 734c1986f8f..98053fdf038 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts @@ -11,9 +11,10 @@ export interface OrganizationListFilters { } export const useOrganizations = (filters?: OrganizationListFilters): UseQueryResult => { - const { accessToken, userId, userRole } = useAuthorized(); + const { accessToken, userId, userRole, premiumUser } = useAuthorized(); const orgId = filters?.org_id || null; const orgAlias = filters?.org_alias || null; + const hasSession = Boolean(accessToken && userId && userRole); return useQuery({ queryKey: organizationKeys.list( orgId || orgAlias @@ -21,16 +22,16 @@ export const useOrganizations = (filters?: OrganizationListFilters): UseQueryRes : {}, ), queryFn: async () => await organizationListCall(accessToken!, orgId, orgAlias), - enabled: Boolean(accessToken && userId && userRole), + enabled: hasSession && premiumUser === true, }); }; export const useOrganization = (organizationID?: string) => { const queryClient = useQueryClient(); - const { accessToken } = useAuthorized(); + const { accessToken, premiumUser } = useAuthorized(); return useQuery({ queryKey: organizationKeys.detail(organizationID!), - enabled: Boolean(accessToken && organizationID), + enabled: Boolean(accessToken && organizationID) && premiumUser === true, queryFn: async () => { if (!accessToken || !organizationID) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx index 1058d4d0466..6507855d0c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.test.tsx @@ -20,21 +20,25 @@ const DEFAULT_PROPS = { onSuccess: vi.fn(), }; -const URL_PLACEHOLDER = "https://github.com/org/repo or https://gitlab.com/org/repo"; +const URL_PLACEHOLDER = "https://github.com/org/repo or https://bucket.s3.amazonaws.com/my-skill.zip"; const SUBPATH_PLACEHOLDER = "plugins/my-skill"; +const SHA256_PLACEHOLDER = "64 hex characters"; +const S3_ZIP_URL = "https://skills-bucket.s3.us-east-1.amazonaws.com/plugins/s3-skill-1.0.0.zip"; +const DIGEST = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; describe("AddPluginForm", () => { beforeEach(() => { vi.clearAllMocks(); }); - it("renders the host-agnostic repository URL input and subfolder field", () => { + it("renders the host-agnostic source URL input and subfolder field, hiding the digest until a zip is entered", () => { renderWithProviders(); - expect(screen.getByText("Repository URL")).toBeInTheDocument(); + expect(screen.getByText("Source URL")).toBeInTheDocument(); expect(screen.getByPlaceholderText(URL_PLACEHOLDER)).toBeInTheDocument(); expect(screen.getByText("Subfolder path (Optional)")).toBeInTheDocument(); expect(screen.getByPlaceholderText(SUBPATH_PLACEHOLDER)).toBeInTheDocument(); + expect(screen.queryByPlaceholderText(SHA256_PLACEHOLDER)).not.toBeInTheDocument(); }); it("shows GitHub repo preview for a plain repo URL", async () => { @@ -255,6 +259,85 @@ describe("AddPluginForm", () => { }); }); + it("shows a zip archive preview, disables the subfolder field, and reveals the digest field", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + + expect(await screen.findByText(/Zip archive/)).toBeInTheDocument(); + expect(screen.getByPlaceholderText(SUBPATH_PLACEHOLDER)).toBeDisabled(); + expect(screen.getByText("A zip archive is installed as a whole, so this field is disabled")).toBeInTheDocument(); + expect(screen.getByPlaceholderText(SHA256_PLACEHOLDER)).toBeInTheDocument(); + expect((screen.getByPlaceholderText("my-skill") as HTMLInputElement).value).toBe("s3-skill-1-0-0"); + }); + + it("submits an archive source without a digest when the field is left empty", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await submit(); + + await waitFor(() => { + expect(mockRegister).toHaveBeenCalledWith( + "sk-test", + expect.objectContaining({ source: { source: "archive", url: S3_ZIP_URL } }), + ); + }); + }); + + it("submits an archive source pinned to the lowercased digest", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await act(async () => { + fireEvent.change(screen.getByPlaceholderText(SHA256_PLACEHOLDER), { + target: { value: ` ${DIGEST.toUpperCase()} ` }, + }); + }); + await submit(); + + await waitFor(() => { + expect(mockRegister).toHaveBeenCalledWith( + "sk-test", + expect.objectContaining({ source: { source: "archive", url: S3_ZIP_URL, sha256: DIGEST } }), + ); + }); + }); + + it("drops the digest once the archive URL changes so a stale checksum is never sent for a new file", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await act(async () => { + fireEvent.change(screen.getByPlaceholderText(SHA256_PLACEHOLDER), { target: { value: DIGEST } }); + }); + const otherZipUrl = S3_ZIP_URL.replace("1.0.0", "1.1.0"); + await typeUrl(otherZipUrl); + + expect(screen.getByPlaceholderText(SHA256_PLACEHOLDER)).toHaveValue(""); + await submit(); + + await waitFor(() => { + expect(mockRegister).toHaveBeenCalledWith( + "sk-test", + expect.objectContaining({ source: { source: "archive", url: otherZipUrl } }), + ); + }); + }); + + it("blocks submission and shows the digest error for a malformed sha256", async () => { + renderWithProviders(); + + await typeUrl(S3_ZIP_URL); + await act(async () => { + fireEvent.change(screen.getByPlaceholderText(SHA256_PLACEHOLDER), { target: { value: "not-a-digest" } }); + }); + await submit(); + + expect(await screen.findByText("SHA-256 must be a 64-character hex digest")).toBeInTheDocument(); + expect(mockRegister).not.toHaveBeenCalled(); + }); + it("surfaces the backend error message when registration fails", async () => { mockRegister.mockRejectedValueOnce(new Error("Plugin 'claude-code' already exists")); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx index 70669ffc0cc..e420038b124 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx @@ -26,6 +26,7 @@ import { parseKeywords, parseSkillSource, isValidSubPath, + isValidSha256, SkillSourcePreview, } from "@/components/claude_code_plugins/helpers"; import { PluginAuthor, PluginSource, SkillRegisterRequest } from "@/components/claude_code_plugins/types"; @@ -39,13 +40,14 @@ interface AddPluginFormProps { } const addPluginShape = { - skillUrl: z.string().min(1, "Please enter a repository URL"), + skillUrl: z.string().min(1, "Please enter a repository or zip archive URL"), subPath: z .string() .refine( (value) => !value || isValidSubPath(value), "Subfolder must be a relative path like plugins/my-skill (letters, numbers, dots, hyphens, underscores)", ), + sha256: z.string().refine(isValidSha256, "SHA-256 must be a 64-character hex digest"), name: z .string() .min(1, "Please enter skill name") @@ -69,6 +71,7 @@ type AddPluginFormValues = z.infer; const EMPTY_VALUES: AddPluginFormValues = { skillUrl: "", subPath: "", + sha256: "", name: "", domain: "", namespace: "", @@ -89,11 +92,19 @@ const buildAuthor = (values: AddPluginFormValues): PluginAuthor | undefined => { return email ? { name, email } : { name }; }; +const archiveUrlOf = (preview: SkillSourcePreview | null): string | undefined => + preview?.parsed.source === "archive" ? preview.parsed.url : undefined; + +const withArchiveDigest = (source: PluginSource, sha256: string): PluginSource => { + const digest = sha256.trim(); + return source.source === "archive" && digest ? { ...source, sha256: digest.toLowerCase() } : source; +}; + const buildRegisterRequest = (values: AddPluginFormValues, source: PluginSource): SkillRegisterRequest => { const author = buildAuthor(values); return { name: values.name.trim(), - source, + source: withArchiveDigest(source, values.sha256), ...(values.version ? { version: values.version.trim() } : {}), ...(values.description ? { description: values.description.trim() } : {}), ...(author ? { author } : {}), @@ -115,6 +126,16 @@ const PREDEFINED_CATEGORIES = [ "Documentation", ]; +const SUB_PATH_LOCK_REASON = { + "git-subdir": "The URL already points to a subfolder, so this field is disabled", + archive: "A zip archive is installed as a whole, so this field is disabled", +} as const; + +type SubPathLock = keyof typeof SUB_PATH_LOCK_REASON; + +const subPathLockFor = (source: PluginSource["source"] | undefined): SubPathLock | null => + source === "git-subdir" || source === "archive" ? source : null; + const labelWithHint = (label: string, hint: string): React.ReactNode => ( <> {label} @@ -129,15 +150,18 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT const form = useZodForm(addPluginSchema, { defaultValues: EMPTY_VALUES }); const [isSubmitting, setIsSubmitting] = useState(false); const [urlPreview, setUrlPreview] = useState(null); - const [urlEncodesSubdir, setUrlEncodesSubdir] = useState(false); + const [subPathLock, setSubPathLock] = useState(null); const recomputePreview = (skillUrl: string, subPath: string) => { - const encodesSubdir = parseSkillSource(skillUrl)?.parsed.source === "git-subdir"; - setUrlEncodesSubdir(encodesSubdir); - if (encodesSubdir && form.getValues("subPath")) { + const lock = subPathLockFor(parseSkillSource(skillUrl)?.parsed.source); + setSubPathLock(lock); + if (lock && form.getValues("subPath")) { form.setValue("subPath", ""); } - const preview = parseSkillSource(skillUrl, encodesSubdir ? undefined : subPath); + const preview = parseSkillSource(skillUrl, lock ? undefined : subPath); + if (archiveUrlOf(preview) !== archiveUrlOf(urlPreview) && form.getValues("sha256")) { + form.setValue("sha256", ""); + } setUrlPreview(preview); if (preview && !form.getValues("name")) { form.setValue("name", preview.suggestedName); @@ -151,7 +175,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT } if (!urlPreview) { - toast.error("Please enter a valid repository URL"); + toast.error("Please enter a valid repository or zip archive URL"); return; } @@ -176,7 +200,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT toast.success("Skill registered successfully"); form.reset(EMPTY_VALUES); setUrlPreview(null); - setUrlEncodesSubdir(false); + setSubPathLock(null); onSuccess(); onClose(); } catch (error) { @@ -190,7 +214,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT const handleCancel = () => { form.reset(EMPTY_VALUES); setUrlPreview(null); - setUrlEncodesSubdir(false); + setSubPathLock(null); onClose(); }; @@ -207,15 +231,15 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT control={form.control} name="skillUrl" label={labelWithHint( - "Repository URL", - "Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host. E.g. github.com/org/repo, gitlab.com/org/repo, or github.com/org/repo/tree/main/my-skill", + "Source URL", + "Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host (e.g. github.com/org/repo or github.com/org/repo/tree/main/my-skill), or an HTTPS link to a .zip archive of the skill hosted on S3 or any static file server.", )} > {({ ref, onChange, ...field }) => ( { onChange(event); @@ -232,9 +256,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT "Subfolder path (Optional)", "Path within the repository where the skill lives (e.g., plugins/my-skill). Leave empty if the skill is at the repo root.", )} - description={ - urlEncodesSubdir ? "The URL already points to a subfolder, so this field is disabled" : undefined - } + description={subPathLock ? SUB_PATH_LOCK_REASON[subPathLock] : undefined} > {({ ref, onChange, ...field }) => ( = ({ visible, onClose, accessT onChange(event); recomputePreview(form.getValues("skillUrl"), event.target.value); }} - disabled={urlEncodesSubdir} + disabled={subPathLock !== null} /> )} + {urlPreview?.parsed.source === "archive" && ( + + {({ ref, ...field }) => ( + + )} + + )} + {urlPreview && (
Detected: {urlPreview.label} diff --git a/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx b/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx index 2a1530cc352..1104e2c7485 100644 --- a/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/SkillHubTableColumns.tsx @@ -26,7 +26,7 @@ function getSkillSourceLink(skill: Plugin): { url: string; label: string } | nul const url = src.path ? `${src.url}/tree/main/${src.path}` : src.url; return { url, label: url.replace("https://github.com/", "") }; } - if (src?.source === "url" && src.url) { + if ((src?.source === "url" || src?.source === "archive") && src.url) { return { url: src.url, label: src.url.replace(/^https?:\/\//, "") }; } return null; diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx index 3ef7dee01b0..a4d76e2cc06 100644 --- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx @@ -97,6 +97,24 @@ describe("LoggingCallbacksTable", () => { expect(onDelete).toHaveBeenCalledWith(callback); }); + it("shows a read-only label instead of the actions menu for runtime-only callback rows", () => { + render( + , + ); + expect(screen.getByTestId("callback-actions-langfuse-success")).toBeInTheDocument(); + expect(screen.queryByTestId("callback-actions-datadog-success")).not.toBeInTheDocument(); + expect(screen.getAllByText("Read only")).toHaveLength(1); + }); + // Regression: `/get_callbacks` returns the same `name` twice when a // callback is registered for both success and failure (e.g. `generic_api` // → POST to spend-log on both 200 and 4xx/5xx). The UI used to ignore diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx index 2263fe03b3d..950ef5de6e4 100644 --- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx @@ -50,6 +50,16 @@ interface CallbackRowActionsProps { } function CallbackRowActions({ callback, onTest, onEdit, onDelete }: CallbackRowActionsProps) { + if (callback.read_only) { + return ( + + Read only + + ); + } return ( ({ + default: () => ({ + token: "test-token", + accessToken: "test-token", + userId: "test-user", + userEmail: "test-user@example.com", + userRole: "Admin", + premiumUser: true, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }), +})); + vi.mock("./common_components/budget_duration_dropdown", () => { const BudgetDurationDropdown = ({ value, onChange }: { value: string | null; onChange: (value: string) => void }) => ( ); +const SKILLS_HINT = + "Enabled skills are visible to every key. Grant disabled (private) Claude Code plugins to this key here."; + +export const KeyAgentAndSkillFields = ({ + control, + accessToken, +}: { + control: Control; + accessToken: string; +}) => ( + <> + + {({ value, onChange }) => ( + + )} + + + + {({ value, onChange }) => ( + + )} + + +); + export const KeyBudgetNumberField = ({ control, name, diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 8a42ecff50c..522af5a85ad 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -366,6 +366,45 @@ describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { }); }); +describe("KeyInfoView handleKeyUpdate skills", () => { + it("should forward the skills the edit form supplies into object_permission and drop the form key", async () => { + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + skills: ["private-skill"], + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.object_permission.skills).toEqual(["private-skill"]); + expect(sentPayload).not.toHaveProperty("skills"); + }); + + it("should send an explicit empty skills list when the form clears every skill", async () => { + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + skills: [], + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.object_permission.skills).toEqual([]); + }); +}); + describe("KeyInfoView handleKeyUpdate budget_duration", () => { it("should send a canonical budget_duration through unchanged", async () => { renderView(true); diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index f24fb6e5a86..fec8e749143 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -46,6 +46,7 @@ export interface KeyEditFormValues { mcp_servers_and_groups?: McpServersAndGroups; mcp_tool_permissions?: Record; agents_and_groups?: AgentsAndGroups; + skills?: string[]; organization_id?: string | null; team_id?: string | null; logging_settings?: unknown[]; @@ -102,6 +103,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => agents: keyData.object_permission?.agents || [], accessGroups: keyData.object_permission?.agent_access_groups || [], }, + skills: keyData.object_permission?.skills || [], organization_id: keyData.organization_id, team_id: keyData.team_id, logging_settings: extractLoggingSettings(keyData.metadata), @@ -148,6 +150,7 @@ export const keyEditFormSchema = z.object({ mcp_servers_and_groups: z.custom(), mcp_tool_permissions: z.custom | undefined>(), agents_and_groups: z.custom(), + skills: z.custom(), organization_id: z.custom(), team_id: z.custom(), logging_settings: z.custom(), @@ -196,6 +199,7 @@ export const toSubmittedValues = ( mcp_servers_and_groups: values.mcp_servers_and_groups, mcp_tool_permissions: values.mcp_tool_permissions, agents_and_groups: values.agents_and_groups, + skills: values.skills, organization_id: values.organization_id, team_id: values.team_id, logging_settings: values.logging_settings, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 35e9268e1c2..b2c2a381b42 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -59,6 +59,7 @@ vi.mock("../networking", async () => { agents: [], }), getAgentAccessGroups: vi.fn().mockResolvedValue([]), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), }; }); @@ -135,6 +136,14 @@ vi.mock("../agent_management/AgentSelector", () => ({ ), })); +vi.mock("../skills/SkillSelector", () => ({ + default: ({ onChange }: { onChange: (selected: string[]) => void }) => ( + + ), +})); + vi.mock("../common_components/AccessGroupSelector", () => ({ default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( { mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, mcp_tool_permissions: {}, agents_and_groups: { agents: [], accessGroups: [] }, + skills: [], organization_id: null, team_id: null, logging_settings: [], @@ -2241,6 +2251,36 @@ describe("KeyEditView", () => { expect(onSubmitMock.mock.calls[0][0].agents_and_groups.agents).toEqual(["agent-1"]); }); + it("carries a picked skill into the payload", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + await userEvent.click(screen.getByRole("button", { name: "pick skill" })); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].skills).toEqual(["private-skill"]); + }); + + it("preloads the stored skills into the payload when the selector is left untouched", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock, { + ...MOCK_KEY_DATA, + object_permission: { ...MOCK_KEY_DATA.object_permission, skills: ["stored-skill"] }, + } as KeyResponse); + await screen.findByRole("button", { name: /save changes/i }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].skills).toEqual(["stored-skill"]); + }); + it("carries an added logging integration into the payload", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); renderForPayload(onSubmitMock); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index b1bbbeed6f9..e464dfe5008 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -15,7 +15,6 @@ import { FormField } from "@/components/shared/form/FormField"; import React, { useEffect, useRef, useState } from "react"; import { hasCapability } from "../../utils/capabilities"; import { isProxyAdminRole, rolesWithWriteAccess } from "../../utils/roles"; -import AgentSelector from "../agent_management/AgentSelector"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; @@ -32,9 +31,8 @@ import { modelSentinelOptions, parseAllowedRoutes, } from "./keyEditFieldNormalizers"; -import { KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls"; +import { KeyAgentAndSkillFields, KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls"; import { - AgentsAndGroups, KeyEditFormValues, keyEditFormSchema, McpServersAndGroups, @@ -762,16 +760,7 @@ export function KeyEditView({ />
- - {({ value, onChange }) => ( - - )} - + ({ + JsonViewer: ({ data }: { data: unknown }) => 
{JSON.stringify(data)}
, +})); + +describe("ClassifierAuditView", () => { + it("separates and copies the provider input, source request, and returned verdict", async () => { + const user = userEvent.setup(); + const input = { system: "classification rubric", messages: [{ role: "user", content: "classify this" }] }; + render( + , + ); + const classifier = within(screen.getByRole("region", { name: "Classifier input" })); + expect(classifier.getByText(/classification rubric/)).toBeInTheDocument(); + expect(classifier.queryByText(/source-only/)).not.toBeInTheDocument(); + expect( + within(screen.getByRole("region", { name: "Originating request, credentials masked" })).getByText(/source-only/), + ).toBeInTheDocument(); + expect( + within(screen.getByRole("region", { name: "Classifier response" })).getByText(/a greeting/), + ).toBeInTheDocument(); + await user.click(classifier.getByRole("button", { name: "Copy Classifier input" })); + expect(await navigator.clipboard.readText()).toBe(JSON.stringify(input, null, 2)); + }); + + it("does not present legacy source messages as captured classifier input", () => { + render(); + expect(screen.getAllByText("Not captured or message logging disabled")).toHaveLength(3); + expect(screen.queryByRole("button", { name: "Copy Classifier input" })).not.toBeInTheDocument(); + }); + + it("labels truncated input without marking a complete source request as truncated", () => { + render( + , + ); + expect(within(screen.getByRole("region", { name: "Classifier input" })).getByRole("status")).toHaveTextContent( + "This stored copy is truncated", + ); + expect( + within(screen.getByRole("region", { name: "Originating request, credentials masked" })).queryByRole("status"), + ).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/ClassifierAuditView.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/ClassifierAuditView.tsx new file mode 100644 index 00000000000..913700efa9b --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/ClassifierAuditView.tsx @@ -0,0 +1,52 @@ +import CopyButton from "@/components/shared/CopyButton"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import type { ReactNode } from "react"; +import { JsonViewer } from "./JsonViewer"; + +interface ClassifierAuditViewProps { + request: Record; + response: unknown; +} + +export function ClassifierAuditView({ request, response }: ClassifierAuditViewProps) { + return ( +
+ + Provider request payload. A cached call or disabled message logging may have no capture. + + + Comparison only. This source request was not appended to the classifier input. + + + The returned verdict and any explanation supplied by the classifier. Later routing rules may change the tier. + +
+ ); +} + +function AuditField({ title, value, children }: { title: string; value: unknown; children: ReactNode }) { + const serialized = JSON.stringify(value); + const truncated = serialized?.includes("litellm_truncated") ?? false; + + return ( + + + {title} + {value != null && } + + +

{children}

+ {truncated && ( +

+ This stored copy is truncated. The complete payload is unavailable from the configured log storage. +

+ )} + {value == null ? ( +

Not captured or message logging disabled

+ ) : ( + + )} + + + ); +} diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx similarity index 89% rename from ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx rename to ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx index e0d6a1d41e0..721525268b3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx @@ -272,6 +272,59 @@ describe("LogDetailContent", () => { expect(screen.queryByRole("tab", { name: "Request" })).not.toBeInTheDocument(); }); + it.each(["object", "serialized", "messages only", "null captures"])( + "preserves request inspection and copying for classifier logs with %s data", + async (shape) => { + const user = userEvent.setup(); + const messages = [{ role: "user", content: "legacy classifier prompt" }]; + const request = { + messages, + temperature: 0.5, + ...(shape === "null captures" ? { classifier_input: null, originating_request_masked: null } : {}), + }; + const storedRequest = shape === "serialized" ? JSON.stringify(request) : request; + const logEntry: Partial = { + call_type: "acompletion", + messages, + proxy_server_request: shape === "messages only" ? undefined : storedRequest, + metadata: { status: "success", internal_call_origin: "autorouter_classifier" }, + }; + render(); + + expect(screen.getByText("legacy classifier prompt")).toBeInTheDocument(); + expect(screen.queryByRole("region", { name: "Classifier input" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: "JSON", exact: true })); + await user.click(screen.getByRole("button", { name: "Copy JSON", exact: true })); + expect(await navigator.clipboard.readText()).toBe( + JSON.stringify(shape === "messages only" ? messages : request, null, 2), + ); + }, + ); + + it.each([ + ["classifier_input", "acompletion"], + ["originating_request_masked", "acompletion"], + ["classifier_input", "aresponses"], + ["originating_request_masked", "responses"], + ])("shows partial classifier audits when only %s is captured for %s", (field, callType) => { + render( + , + ); + + expect(screen.getByRole("region", { name: "Classifier input" })).toBeInTheDocument(); + expect(screen.getByRole("region", { name: "Originating request, credentials masked" })).toBeInTheDocument(); + expect(screen.getByText("Not captured or message logging disabled")).toBeInTheDocument(); + expect(screen.queryByText("Request & Response")).not.toBeInTheDocument(); + }); + it("should display Request and Response tabs when JSON view is selected", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 186cb66592b..b5a3412a96f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -48,6 +48,8 @@ import { } from "./constants"; import { ToolsSection } from "../ToolsSection"; import { PrettyMessagesView } from "./PrettyMessagesView"; +import { ClassifierAuditView } from "./ClassifierAuditView"; +import { AUTOROUTER_CLASSIFIER_ORIGIN } from "./ClassifyTag"; export interface LogDetailContentProps { logEntry: LogEntry; @@ -68,6 +70,12 @@ export function LogDetailContent({ logEntry, isLoadingDetails = false, accessTok const metadata = logEntry.metadata || {}; const hasError = metadata.status === "failure"; const errorInfo = hasError ? metadata.error_information : null; + const isClassifier = + metadata.internal_call_origin === AUTOROUTER_CLASSIFIER_ORIGIN && + ["completion", "acompletion", "responses", "aresponses"].includes(logEntry.call_type); + const rawRequest = formatData(logEntry.proxy_server_request || logEntry.messages); + const hasClassifierAudit = + isClassifier && (rawRequest?.classifier_input != null || rawRequest?.originating_request_masked != null); const hasMessages = checkHasMessages(logEntry.messages); const hasResponse = checkHasResponse(logEntry.response); @@ -88,10 +96,6 @@ export function LogDetailContent({ logEntry, isLoadingDetails = false, accessTok // Vector store data const hasVectorStoreData = checkHasVectorStoreData(metadata); - const getRawRequest = () => { - return formatData(logEntry.proxy_server_request || logEntry.messages); - }; - const getFormattedResponse = () => { if (hasError && errorInfo) { return { @@ -196,11 +200,15 @@ export function LogDetailContent({ logEntry, isLoadingDetails = false, accessTok Loading request & response data... - ) : ( + ) : null} + {!isLoadingDetails && hasClassifierAudit && ( + + )} + {!isLoadingDetails && !hasClassifierAudit && ( rawRequest} getFormattedResponse={getFormattedResponse} logEntry={logEntry} /> @@ -541,22 +549,20 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: )} - {metadata?.attempted_retries !== undefined && metadata?.attempted_retries !== null ? ( - metadata.attempted_retries > 0 ? ( - <> - {metadata.attempted_retries} - {metadata.max_retries !== undefined && metadata.max_retries !== null - ? ` / ${metadata.max_retries}` - : ""} - - ) : ( - - None - - ) - ) : ( - "-" + {metadata?.attempted_retries != null && metadata.attempted_retries > 0 && ( + <> + {metadata.attempted_retries} + {metadata.max_retries !== undefined && metadata.max_retries !== null + ? ` / ${metadata.max_retries}` + : ""} + )} + {metadata?.attempted_retries != null && metadata.attempted_retries <= 0 && ( + + None + + )} + {metadata?.attempted_retries == null && "-"} @@ -602,16 +608,10 @@ function RequestResponseSection({ const totalTokens = promptTokens + completionTokens; const costBreakdown = logEntry.metadata?.cost_breakdown; const useCostBreakdown = costBreakdown?.input_cost !== undefined && costBreakdown?.output_cost !== undefined; - const inputCost = useCostBreakdown - ? costBreakdown!.input_cost ?? 0 - : totalTokens > 0 - ? (totalSpend * promptTokens) / totalTokens - : 0; - const outputCost = useCostBreakdown - ? costBreakdown!.output_cost ?? 0 - : totalTokens > 0 - ? (totalSpend * completionTokens) / totalTokens - : 0; + const estimatedInputCost = totalTokens > 0 ? (totalSpend * promptTokens) / totalTokens : 0; + const estimatedOutputCost = totalTokens > 0 ? (totalSpend * completionTokens) / totalTokens : 0; + const inputCost = useCostBreakdown ? costBreakdown!.input_cost ?? 0 : estimatedInputCost; + const outputCost = useCostBreakdown ? costBreakdown!.output_cost ?? 0 : estimatedOutputCost; return (
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d62ce2b675..6826cded6f5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2110,12 +2110,17 @@ export interface paths { * - 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 * ``` */ @@ -2152,8 +2157,8 @@ export interface paths { * @description 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 @@ -2163,7 +2168,7 @@ export interface paths { * * 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) @@ -2229,7 +2234,7 @@ export interface paths { * * 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) @@ -23713,6 +23718,15 @@ export interface components { /** @description The decision record this request would have written to its log row */ routing_decision: components["schemas"]["StandardLoggingRoutingDecision"]; }; + /** AwsSessionTag */ + AwsSessionTag: { + /** Key */ + Key: string; + /** Value */ + Value: string; + } & { + [key: string]: unknown; + }; /** BaseLitellmParams */ BaseLitellmParams: { /** @@ -26357,6 +26371,31 @@ export interface components { * independently of the response-cache backend in `litellm_settings.cache_params`. */ CoordinationRedisParams: { + /** + * Aws Iam Auth + * @description enable AWS ElastiCache IAM authentication + */ + aws_iam_auth?: boolean | string | null; + /** + * Aws Iam Cache Name + * @description AWS ElastiCache cache name + */ + aws_iam_cache_name?: string | null; + /** + * Aws Iam Region + * @description AWS region for ElastiCache IAM authentication + */ + aws_iam_region?: string | null; + /** + * Aws Iam Serverless + * @description the ElastiCache cache is serverless rather than a self-designed cluster + */ + aws_iam_serverless?: boolean | string | null; + /** + * Aws Iam User Name + * @description AWS ElastiCache IAM user name + */ + aws_iam_user_name?: string | null; /** * Host * @description Redis hostname @@ -29324,6 +29363,8 @@ export interface components { models?: string[] | null; /** Search Tools */ search_tools?: string[] | null; + /** Skills */ + skills?: string[] | null; /** Vector Stores */ vector_stores?: string[] | null; }; @@ -29377,6 +29418,8 @@ export interface components { * @default [] */ search_tools: string[] | null; + /** Skills */ + skills?: string[] | null; /** * Vector Stores * @default [] @@ -29538,6 +29581,8 @@ export interface components { aws_secret_access_key?: string | null; /** Aws Session Name */ aws_session_name?: string | null; + /** Aws Session Tags */ + aws_session_tags?: components["schemas"]["AwsSessionTag"][] | null; /** Aws Session Token */ aws_session_token?: string | null; /** Aws Sts Endpoint */ @@ -33627,7 +33672,7 @@ export interface components { name: string; /** * Source - * @description Git source reference + * @description Plugin source reference */ source: { [key: string]: string; @@ -34876,7 +34921,7 @@ export interface components { * @description 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. */ RegisterPluginRequest: { /** @description Plugin author */ @@ -34918,10 +34963,11 @@ export interface components { namespace?: string | null; /** * Source - * @description Git source reference. Supported formats: + * @description Plugin source reference. Supported formats: * - GitHub: {'source': 'github', 'repo': 'org/repo'} * - Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'} * - Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'} + * - Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''} */ source: { [key: string]: string; @@ -35210,7 +35256,7 @@ export interface components { reasoning_override_min_score?: number | null; /** * Reminder Markers - * @description 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. + * @description 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 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. */ reminder_markers?: components["schemas"]["ReminderMarkerPair"][] | null; /** @@ -38200,10 +38246,11 @@ export interface components { namespace?: string | null; /** * Source - * @description Git source reference. Supported formats: + * @description Plugin source reference. Supported formats: * - GitHub: {'source': 'github', 'repo': 'org/repo'} * - Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'} * - Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'} + * - Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': ''} */ source: { [key: string]: string; @@ -39731,6 +39778,8 @@ export interface components { aws_secret_access_key?: string | null; /** Aws Session Name */ aws_session_name?: string | null; + /** Aws Session Tags */ + aws_session_tags?: components["schemas"]["AwsSessionTag"][] | null; /** Aws Session Token */ aws_session_token?: string | null; /** Aws Sts Endpoint */ @@ -43191,7 +43240,9 @@ export interface operations { }; get_marketplace_claude_code_marketplace_json_get: { parameters: { - query?: never; + query?: { + key?: string | null; + }; header?: never; path?: never; cookie?: never; @@ -43207,6 +43258,15 @@ export interface operations { "application/json": unknown; }; }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; }; }; list_plugins_claude_code_plugins_get: { @@ -49399,7 +49459,7 @@ export interface operations { parameters: { query: { /** @description Specify the service being hit. */ - service: ("slack_budget_alerts" | "langfuse" | "langfuse_otel" | "slack" | "ms_teams" | "openmeter" | "webhook" | "email" | "braintrust" | "datadog" | "datadog_llm_observability" | "generic_api" | "arize" | "galileo" | "newrelic" | "sqs") | string; + service: ("slack_budget_alerts" | "langfuse" | "langfuse_otel" | "slack" | "ms_teams" | "openmeter" | "webhook" | "email" | "braintrust" | "datadog" | "datadog_llm_observability" | "generic_api" | "arize" | "galileo" | "newrelic" | "pointfive" | "sqs") | string; }; header?: never; path?: never; diff --git a/ui/litellm-dashboard/src/utils/entityLinks.test.ts b/ui/litellm-dashboard/src/utils/entityLinks.test.ts index 47161a903ed..231413b213a 100644 --- a/ui/litellm-dashboard/src/utils/entityLinks.test.ts +++ b/ui/litellm-dashboard/src/utils/entityLinks.test.ts @@ -2,7 +2,29 @@ import { describe, expect, it, vi } from "vitest"; vi.mock("@/components/networking", () => ({ serverRootPath: "" })); -import { modelGroupHref } from "./entityLinks"; +import { modelGroupHref, teamDetailHref, userDetailHref } from "./entityLinks"; + +describe("userDetailHref", () => { + it("targets the users page filtered to the encoded user id", () => { + expect(userDetailHref("user-1")).toMatch(/\/users\?user=user-1$/); + expect(userDetailHref("a b/c")).toMatch(/\?user=a%20b%2Fc$/); + }); + + it("returns no href for the proxy admin placeholder, which has no user page", () => { + expect(userDetailHref("default_user_id")).toBeUndefined(); + }); +}); + +describe("teamDetailHref", () => { + it("targets the teams page filtered to the encoded team id", () => { + expect(teamDetailHref("team-1")).toMatch(/\/teams\?team=team-1$/); + expect(teamDetailHref("a b/c")).toMatch(/\?team=a%20b%2Fc$/); + }); + + it("returns no href for the Admin UI session team, which has no team page", () => { + expect(teamDetailHref("litellm-dashboard")).toBeUndefined(); + }); +}); describe("modelGroupHref", () => { it("targets the models page filtered to the encoded model group", () => { diff --git a/ui/litellm-dashboard/src/utils/entityLinks.ts b/ui/litellm-dashboard/src/utils/entityLinks.ts index b8c70bdda48..33f2aa34976 100644 --- a/ui/litellm-dashboard/src/utils/entityLinks.ts +++ b/ui/litellm-dashboard/src/utils/entityLinks.ts @@ -1,3 +1,4 @@ +import { DEFAULT_PROXY_ADMIN_USER_ID, UI_TEAM_ID } from "@/utils/sentinels"; import { uiHref } from "@/utils/uiHref"; const MODEL_GRANT_SENTINELS: ReadonlySet = new Set([ @@ -6,7 +7,8 @@ const MODEL_GRANT_SENTINELS: ReadonlySet = new Set([ "no-default-models", ]); -export function teamDetailHref(teamId: string): string { +export function teamDetailHref(teamId: string): string | undefined { + if (teamId === UI_TEAM_ID) return undefined; return `${uiHref("teams")}?team=${encodeURIComponent(teamId)}`; } @@ -14,7 +16,8 @@ export function keyDetailHref(keyToken: string): string { return `${uiHref("api-keys")}?key=${encodeURIComponent(keyToken)}`; } -export function userDetailHref(userId: string): string { +export function userDetailHref(userId: string): string | undefined { + if (userId === DEFAULT_PROXY_ADMIN_USER_ID) return undefined; return `${uiHref("users")}?user=${encodeURIComponent(userId)}`; } diff --git a/ui/litellm-dashboard/src/utils/sentinels.ts b/ui/litellm-dashboard/src/utils/sentinels.ts new file mode 100644 index 00000000000..da6d09dfdd7 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/sentinels.ts @@ -0,0 +1,3 @@ +export const DEFAULT_PROXY_ADMIN_USER_ID = "default_user_id"; + +export const UI_TEAM_ID = "litellm-dashboard"; diff --git a/uv.lock b/uv.lock index 0fe787645a2..cedc505acdf 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-06T00:40:30.433549Z" +exclude-newer = "2026-09-07T23:09:03.362777Z" exclude-newer-span = "P3D" [manifest] @@ -4772,7 +4772,7 @@ source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.95" +version = "0.4.96" source = { editable = "litellm-proxy-extras" } [[package]]