diff --git a/.github/workflows/create_daily_oss_agent_shin_branch.yml b/.github/workflows/create_daily_oss_agent_shin_branch.yml new file mode 100644 index 00000000000..d6118f3b53c --- /dev/null +++ b/.github/workflows/create_daily_oss_agent_shin_branch.yml @@ -0,0 +1,47 @@ +name: Create Daily oss-agent-shin Branch + +on: + schedule: + - cron: "0 0 * * *" # Runs every day at midnight UTC + workflow_dispatch: # Allow manual trigger + +jobs: + create-oss-agent-shin-branch: + if: github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + permissions: + contents: write + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Create daily oss-agent-shin branch + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + # Configure Git user + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + + # Generate branch name with MM_DD_YYYY format + BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')" + echo "Creating branch: $BRANCH_NAME" + + # Fetch all branches + git fetch --all + + # Check if the branch already exists + if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then + echo "Branch $BRANCH_NAME already exists. Skipping creation." + else + echo "Creating new branch: $BRANCH_NAME" + # Create the new branch from main + git checkout -b $BRANCH_NAME origin/main + # Push the new branch + git push origin $BRANCH_NAME + echo "Successfully created and pushed branch: $BRANCH_NAME" + fi diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 1439b2c07f7..118408f7463 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -7,6 +7,7 @@ on: - litellm_internal_staging - litellm_oss_branch - "litellm_**" + workflow_dispatch: permissions: contents: read @@ -42,3 +43,16 @@ jobs: workers: 2 reruns: 2 artifact-name: proxy-endpoints + + # Behavior-pinning tests for litellm/proxy/proxy_server.py. Owns its + # own job (not a path on the proxy-endpoints job above) so its budget + # is independent and its coverage artifact is uploaded separately. + # See: https://www.notion.so/36c43b8acdab81ee845fd5365128a2fc + proxy-server: + uses: ./.github/workflows/_test-unit-base.yml + with: + test-path: tests/test_litellm/proxy/proxy_server + workers: 4 + reruns: 2 + timeout-minutes: 60 + artifact-name: proxy-server diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml index 155445acdf6..57ff746c9c8 100644 --- a/.github/workflows/test_server_root_path.yml +++ b/.github/workflows/test_server_root_path.yml @@ -101,6 +101,31 @@ jobs: docker logs litellm-test exit 1 + - name: Setup Node for Playwright + uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 + with: + node-version: "20" + + - name: Install UI deps and Chromium + working-directory: ui/litellm-dashboard + run: | + npm ci + npx playwright install --with-deps chromium + + - name: Run SERVER_ROOT_PATH redirect e2e + working-directory: ui/litellm-dashboard + env: + SERVER_ROOT_PATH: ${{ matrix.root_path }} + run: npx playwright test --config=e2e_tests/serverRootPath.config.ts + + - name: Upload Playwright artifacts on failure + if: failure() + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: playwright-trace-${{ strategy.job-index }} + path: ui/litellm-dashboard/test-results/ + retention-days: 7 + - name: Cleanup if: always() run: | diff --git a/AGENTS.md b/AGENTS.md index e99bf79d783..a41fc4268d9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,6 +2,19 @@ This document provides comprehensive instructions for AI agents working in the LiteLLM repository. +## Confidentiality: Customer and Company Names in Code + +The codebase is public. Before writing **any** third-party organization name into this repository — in source code, file or directory names, docstrings, comments, tests, fixtures, mock payloads, error messages, log lines, commit messages, or PR descriptions — pause and check: + +**Already in the codebase** (OpenAI, Anthropic, Google, Azure, Bedrock, Fireworks, and other established LLM providers / integrations) — fine to use. Quick check: `git grep -i ""` — if it returns hits in real code (not just your current diff), the name is established. + +**Anything else** — customers, prospects, partners, new vendor integrations, observability tools, infra vendors, or any organization name that does not already appear in the repo. STOP and surface it to the user. Ask for explicit consent before writing the name into any file, commit message, or PR description. Do not write it speculatively and clean up later. Do not substitute a placeholder and proceed. Do not assume it is safe because it "looks like" a public company. The user must approve first. + +**What to do instead of a customer-specific reference:** +- If you find yourself reaching for a customer name — real or fake — step back. The code shouldn't be customer-specific in the first place. Generalize the feature, or capture the customer motivation in internal docs (Notion / Linear / the internal staging PR description), never in the repo. +- Frame changes by the capability they add, not the customer who asked for it ("add per-team Bedrock guardrail routing", not "add routing for $CUSTOMER"). +- Standard "fake value" markers (`example.com`, `localhost`, `127.0.0.1`, `test@example.com`) and abstract identifiers (`team_a`, `user_1`, `tenant_x`) are fine — those are not customer stand-ins. + ## OVERVIEW LiteLLM is a unified interface for 100+ LLMs that: diff --git a/CLAUDE.md b/CLAUDE.md index baf23c90148..b9a336b8f40 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,6 +2,19 @@ This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. +## Confidentiality: Customer and Company Names in Code + +The codebase is public. Before writing **any** third-party organization name into this repository — in source code, file or directory names, docstrings, comments, tests, fixtures, mock payloads, error messages, log lines, commit messages, or PR descriptions — pause and check: + +**Already in the codebase** (OpenAI, Anthropic, Google, Azure, Bedrock, Fireworks, and other established LLM providers / integrations) — fine to use. Quick check: `git grep -i ""` — if it returns hits in real code (not just your current diff), the name is established. + +**Anything else** — customers, prospects, partners, new vendor integrations, observability tools, infra vendors, or any organization name that does not already appear in the repo. STOP and surface it to the user. Ask for explicit consent before writing the name into any file, commit message, or PR description. Do not write it speculatively and clean up later. Do not substitute a placeholder and proceed. Do not assume it is safe because it "looks like" a public company. The user must approve first. + +**What to do instead of a customer-specific reference:** +- If you find yourself reaching for a customer name — real or fake — step back. The code shouldn't be customer-specific in the first place. Generalize the feature, or capture the customer motivation in internal docs (Notion / Linear / the internal staging PR description), never in the repo. +- Frame changes by the capability they add, not the customer who asked for it ("add per-team Bedrock guardrail routing", not "add routing for $CUSTOMER"). +- Standard "fake value" markers (`example.com`, `localhost`, `127.0.0.1`, `test@example.com`) and abstract identifiers (`team_a`, `user_1`, `tenant_x`) are fine — those are not customer stand-ins. + ## Documentation Documentation lives in a separate repository: [BerriAI/litellm-docs](https://github.com/BerriAI/litellm-docs). It is served at [docs.litellm.ai](https://docs.litellm.ai). Do not create or edit documentation files in this repository — open doc PRs against `BerriAI/litellm-docs` instead. diff --git a/backend/Dockerfile b/backend/Dockerfile index c08014fc0ef..2cfdde8a517 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -12,17 +12,27 @@ USER root COPY --from=uvbin /uv /uvx /usr/local/bin/ -RUN apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile +# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE +# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi +# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. +RUN for i in 1 2 3; do \ + apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done # UV_COMPILE_BYTECODE=1 precompiles .pyc at install time → faster cold start. # UV_LINK_MODE=copy avoids hardlink warnings when uv installs from a # BuildKit cache mount (different filesystem). # UV_PYTHON_DOWNLOADS=0 force uv to use the apk-installed CPython instead of # silently pulling a managed interpreter. +# PRISMA_USE_GLOBAL_NODE explicit (matches default) so an env override can't +# silently re-enable nodeenv's Node download. ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ UV_COMPILE_BYTECODE=1 \ UV_PYTHON_DOWNLOADS=0 \ + PRISMA_USE_GLOBAL_NODE=true \ PATH="/app/.venv/bin:${PATH}" # Stage 1 — install dependencies only. @@ -58,7 +68,11 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root -RUN apk add --no-cache bash openssl tzdata python3 libsndfile libatomic +RUN for i in 1 2 3; do \ + apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done # wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532) with # /home/nonroot. We run the backend as that user diff --git a/cookbook/gollem_go_agent_framework/go.mod b/cookbook/gollem_go_agent_framework/go.mod index 89d9033aa22..a8dc9365d7f 100644 --- a/cookbook/gollem_go_agent_framework/go.mod +++ b/cookbook/gollem_go_agent_framework/go.mod @@ -1,5 +1,5 @@ module github.com/BerriAI/litellm/cookbook/gollem_go_agent_framework -go 1.25.1 +go 1.26.3 require github.com/fugue-labs/gollem v0.1.0 diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index aefe2a564bb..b9cd1be06ec 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -30,7 +30,7 @@ spec: checksum/config: {{ include (print $.Template.BasePath "/configmap-litellm.yaml") . | sha256sum }} {{- end }} {{- with .Values.podAnnotations }} - {{- toYaml . | nindent 8 }} + {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} labels: {{- include "litellm.labels" . | nindent 8 }} diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index df6d1345644..f3d62651d8f 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -377,3 +377,28 @@ tests: content: name: sidecar-tpl image: "ghcr.io/berriai/litellm-database:test" + - it: should support tpl in podAnnotations + template: deployment.yaml + set: + image: + repository: ghcr.io/berriai/litellm-database + tag: test + # Mirrors the real-world scenario this feature unblocks: + # user disables the built-in ConfigMap (and its built-in checksum/config + # annotation) and re-implements checksum/config themselves via tpl. + proxyConfigMap: + create: false + podAnnotations: + checksum/config: "{{ .Values.image.tag }}" + example.com/some-key: "{{ .Values.image.repository }}" + example.com/literal: "plain-string-value" + asserts: + - equal: + path: spec.template.metadata.annotations["checksum/config"] + value: "test" + - equal: + path: spec.template.metadata.annotations["example.com/some-key"] + value: "ghcr.io/berriai/litellm-database" + - equal: + path: spec.template.metadata.annotations["example.com/literal"] + value: "plain-string-value" diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2729babb6d6..8717e5b3fcd 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -55,22 +55,10 @@ COPY . . # Set non-root flag for build time consistency ENV LITELLM_NON_ROOT=true -# Stage the pre-built Admin UI from the checked-in Next.js static export. -# _experimental/out/ is regenerated as part of the release runbook. -# Restructure extensionless routes (foo.html -> foo/index.html) to match the layout -# proxy_server.py expects, and drop a readiness marker. RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \ cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \ cp /app/litellm/proxy/logo.jpg /var/lib/litellm/assets/logo.jpg && \ - ( cd /var/lib/litellm/ui && \ - for html_file in *.html; do \ - if [ "$html_file" != "index.html" ] && [ -f "$html_file" ]; then \ - folder_name="${html_file%.html}" && \ - mkdir -p "$folder_name" && \ - mv "$html_file" "$folder_name/index.html"; \ - fi; \ - done && \ - touch .litellm_ui_ready ) + touch /var/lib/litellm/ui/.litellm_ui_ready RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \ diff --git a/gateway/Dockerfile b/gateway/Dockerfile index a2ca3d3f83f..19c8a10fdfe 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -12,17 +12,27 @@ USER root COPY --from=uvbin /uv /uvx /usr/local/bin/ -RUN apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile +# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE +# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi +# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. +RUN for i in 1 2 3; do \ + apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done # UV_COMPILE_BYTECODE=1 precompiles .pyc at install time → faster cold start. # UV_LINK_MODE=copy avoids hardlink warnings when uv installs from a # BuildKit cache mount (different filesystem). # UV_PYTHON_DOWNLOADS=0 force uv to use the apk-installed CPython instead of # silently pulling a managed interpreter. +# PRISMA_USE_GLOBAL_NODE explicit (matches default) so an env override can't +# silently re-enable nodeenv's Node download. ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ UV_COMPILE_BYTECODE=1 \ UV_PYTHON_DOWNLOADS=0 \ + PRISMA_USE_GLOBAL_NODE=true \ PATH="/app/.venv/bin:${PATH}" # Stage 1 — install dependencies only. @@ -58,7 +68,11 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root -RUN apk add --no-cache bash openssl tzdata python3 libsndfile libatomic +RUN for i in 1 2 3; do \ + apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done # wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532) with # /home/nonroot. We run the proxy as that user. diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index e2faf42b766..4319907883e 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -56,16 +56,34 @@ app.kubernetes.io/component: ui {{- end -}} {{/* -Shared ServiceAccount name used by all three component Deployments. When -`serviceAccount.create` is true and `serviceAccount.name` is empty, default -to the chart fullname. When `create` is false, fall back to the provided -name or the namespace's `default` SA. +Per-component ServiceAccount name helpers. + +Each component (gateway, backend, ui) has its own SA config under +.Values.serviceAccounts.. When `create` is true and `name` is +empty the chart defaults to "-litellm-". When `create` +is false the chart uses the provided name, or the namespace `default` SA. */}} -{{- define "litellm.serviceAccountName" -}} -{{- if .Values.serviceAccount.create -}} -{{ default (include "litellm.fullname" .) .Values.serviceAccount.name }} +{{- define "litellm.gateway.serviceAccountName" -}} +{{- if .Values.serviceAccounts.gateway.create -}} +{{ default (include "litellm.gateway.fullname" .) .Values.serviceAccounts.gateway.name }} {{- else -}} -{{ default "default" .Values.serviceAccount.name }} +{{ default "default" .Values.serviceAccounts.gateway.name }} +{{- end -}} +{{- end -}} + +{{- define "litellm.backend.serviceAccountName" -}} +{{- if .Values.serviceAccounts.backend.create -}} +{{ default (include "litellm.backend.fullname" .) .Values.serviceAccounts.backend.name }} +{{- else -}} +{{ default "default" .Values.serviceAccounts.backend.name }} +{{- end -}} +{{- end -}} + +{{- define "litellm.ui.serviceAccountName" -}} +{{- if .Values.serviceAccounts.ui.create -}} +{{ default (include "litellm.ui.fullname" .) .Values.serviceAccounts.ui.name }} +{{- else -}} +{{ default "default" .Values.serviceAccounts.ui.name }} {{- end -}} {{- end -}} diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index e761409f8c4..3b59c58c8bf 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -19,7 +19,8 @@ spec: labels: {{- include "litellm.backend.selectorLabels" . | nindent 8 }} spec: - serviceAccountName: {{ include "litellm.serviceAccountName" . }} + serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }} + automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 935d432342e..05ea4052159 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -22,7 +22,8 @@ spec: labels: {{- include "litellm.gateway.selectorLabels" . | nindent 8 }} spec: - serviceAccountName: {{ include "litellm.serviceAccountName" . }} + serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }} + automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index f3dc2ae0236..92671388546 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -28,7 +28,7 @@ spec: app.kubernetes.io/component: migrations spec: restartPolicy: Never - serviceAccountName: {{ include "litellm.serviceAccountName" . }} + serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} diff --git a/helm/litellm/templates/serviceaccount.yaml b/helm/litellm/templates/serviceaccount.yaml index 3c998448ae5..a2fc52f47c0 100644 --- a/helm/litellm/templates/serviceaccount.yaml +++ b/helm/litellm/templates/serviceaccount.yaml @@ -1,13 +1,51 @@ -{{- if .Values.serviceAccount.create -}} +{{- $prev := false -}} +{{- if .Values.serviceAccounts.gateway.create -}} +{{- $prev = true }} apiVersion: v1 kind: ServiceAccount metadata: - name: {{ include "litellm.serviceAccountName" . }} + name: {{ include "litellm.gateway.serviceAccountName" . }} labels: {{- include "litellm.commonLabels" . | nindent 4 }} - {{- with .Values.serviceAccount.annotations }} + app.kubernetes.io/component: gateway + {{- with .Values.serviceAccounts.gateway.annotations }} annotations: {{- toYaml . | nindent 4 }} {{- end }} -automountServiceAccountToken: {{ .Values.serviceAccount.automount }} +automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }} +{{- end }} +{{- if .Values.serviceAccounts.backend.create }} +{{- if $prev }} +--- +{{- end }} +{{- $prev = true }} +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ include "litellm.backend.serviceAccountName" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: backend + {{- with .Values.serviceAccounts.backend.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }} +{{- end }} +{{- if .Values.serviceAccounts.ui.create }} +{{- if $prev }} +--- +{{- end }} +apiVersion: v1 +kind: ServiceAccount +metadata: + name: {{ include "litellm.ui.serviceAccountName" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: ui + {{- with .Values.serviceAccounts.ui.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }} {{- end }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index 549bf61a0dd..b40b44cca53 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -19,7 +19,8 @@ spec: labels: {{- include "litellm.ui.selectorLabels" . | nindent 8 }} spec: - serviceAccountName: {{ include "litellm.serviceAccountName" . }} + serviceAccountName: {{ include "litellm.ui.serviceAccountName" . }} + automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 92477616a9a..934661643bd 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -14,16 +14,33 @@ ingress: host: "" # optional; if set, becomes the rule's host tls: [] -# Shared ServiceAccount used by all three component Deployments. Set -# `create: true` to have the chart provision it (e.g. when wiring an EKS -# Pod Identity association by SA name). Set `name` to use an existing SA -# (chart-created or out-of-band). When both are empty / false, pods run -# with the namespace's `default` SA. -serviceAccount: - create: false - automount: true - annotations: {} - name: "" +# Per-component ServiceAccounts for gateway, backend, and ui. +# +# Each section mirrors the old shared serviceAccount shape. Set `create: +# true` to have the chart provision the SA (useful for EKS Pod Identity / +# GKE Workload Identity annotations). Set `name` to bind an existing SA. +# When both are unset the component pod runs with the namespace `default` SA. +# +# The UI SA deliberately defaults to `automount: false` — the static nginx +# container does not need the K8s API and should not carry a projected +# ServiceAccount token that a compromised container could use to call the +# cloud-provider metadata service or the K8s API. +serviceAccounts: + gateway: + create: false + automount: true + annotations: {} + name: "" + backend: + create: false + automount: true + annotations: {} + name: "" + ui: + create: false + automount: false + annotations: {} + name: "" # Pre-install / pre-upgrade Helm hook that runs `prisma migrate deploy` # against the writer database, creating the LiteLLM schema (tables that diff --git a/litellm/__init__.py b/litellm/__init__.py index 7c92623358d..56d516536e8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -225,6 +225,11 @@ use_chat_completions_url_for_anthropic_messages: bool = bool( route_all_chat_openai_to_responses: bool = ( os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true" ) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge +# When True, Gemini/Vertex Live setup is deferred until client `session.update`. +# Default False preserves historical behavior (auto-send setup on connect). +gemini_live_defer_setup: bool = ( + os.getenv("LITELLM_GEMINI_LIVE_DEFER_SETUP", "false").lower() == "true" +) use_legacy_interactions_schema: bool = ( os.getenv("LITELLM_USE_LEGACY_INTERACTIONS_SCHEMA", "false").lower() == "true" ) # When True, sends Api-Revision: 2026-05-07 to Google so responses use the legacy `outputs` diff --git a/litellm/constants.py b/litellm/constants.py index fb765c0226c..f72528eb170 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1147,6 +1147,7 @@ BEDROCK_CONVERSE_MODELS = [ "openai.gpt-oss-120b-1:0", "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6-v1:0", "anthropic.claude-opus-4-6-v1", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 98e00cf5788..9a4b158b622 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import from litellm.litellm_core_utils.llm_cost_calc.utils import ( CostCalculatorUtils, _generic_cost_per_character, + _get_regional_uplift_multiplier, _get_service_tier_cost_key, _parse_prompt_tokens_details, calculate_cost_component, @@ -132,6 +133,8 @@ _VIDEO_CALL_TYPES = frozenset( { CallTypes.create_video.value, CallTypes.acreate_video.value, + CallTypes.video_edit.value, + CallTypes.avideo_edit.value, CallTypes.video_remix.value, CallTypes.avideo_remix.value, } @@ -312,6 +315,10 @@ def cost_per_token( # noqa: PLR0915 audio_transcription_file_duration: float = 0.0, # for audio transcription calls - the file time in seconds ### SERVICE TIER ### service_tier: Optional[str] = None, # for OpenAI service tier pricing + ### DATA RESIDENCY ### + data_residency: Optional[ + str + ] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") response: Optional[Any] = None, ### REQUEST MODEL ### request_model: Optional[str] = None, # original request model for router detection @@ -412,9 +419,36 @@ def cost_per_token( # noqa: PLR0915 prompt_tokens_cost_usd_dollar: float = 0 completion_tokens_cost_usd_dollar: float = 0 model_cost_ref = litellm.model_cost + # Only callers that explicitly pass `custom_llm_provider` get the + # dedup/prefix-join treatment. When provider is omitted, preserve legacy + # behavior: `model_with_provider` stays equal to the raw `model` string + # (provider is detected below for downstream use only). + caller_supplied_provider = custom_llm_provider is not None + + # `model` is normally a string, but callers that mock the transport can pass + # non-string objects. Only run the string-based dedup/prefix-join when it is + # actually a string — e.g. a MagicMock's `.startswith()` is always truthy and + # its slices return new mocks, which would spin the dedup loop forever. + model_is_str = isinstance(model, str) + + # Router/proxy deployments may repeat the provider segment (e.g. model_name + # "openai/openai/gpt-5.5"). Strip duplicated `{provider}/` chains before joining. + if caller_supplied_provider and model_is_str: + _dup_prefix = f"{custom_llm_provider}/" + while model.startswith(_dup_prefix): + _remainder = model[len(_dup_prefix) :] + if _remainder.startswith(_dup_prefix): + model = _remainder + else: + break + model_with_provider = model - if custom_llm_provider is not None: - model_with_provider = custom_llm_provider + "/" + model + if caller_supplied_provider: + _prov_prefix = f"{custom_llm_provider}/" + if model_is_str and model.startswith(_prov_prefix): + model_with_provider = model + else: + model_with_provider = f"{custom_llm_provider}/{model}" if region_name is not None: model_with_provider_and_region = ( f"{custom_llm_provider}/{region_name}/{model}" @@ -425,6 +459,9 @@ def cost_per_token( # noqa: PLR0915 model_with_provider = model_with_provider_and_region else: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) + + assert custom_llm_provider is not None # caller-supplied or get_llm_provider + model_without_prefix = model model_parts = model.split("/", 1) if len(model_parts) > 1: @@ -493,6 +530,7 @@ def cost_per_token( # noqa: PLR0915 usage=usage_block, custom_llm_provider=custom_llm_provider, service_tier=service_tier, + data_residency=data_residency, ) return prompt_cost, completion_cost @@ -521,7 +559,10 @@ def cost_per_token( # noqa: PLR0915 or call_type == CallTypes.retrieve_batch ): return batch_cost_calculator( - usage=usage_block, model=model, custom_llm_provider=custom_llm_provider + usage=usage_block, + model=model, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, ) elif call_type == "atranscription" or call_type == "transcription": if _transcription_usage_has_token_details(usage_block): @@ -529,6 +570,7 @@ def cost_per_token( # noqa: PLR0915 model=model_without_prefix, usage=usage_block, service_tier=service_tier, + data_residency=data_residency, ) return openai_cost_per_second( @@ -579,7 +621,10 @@ def cost_per_token( # noqa: PLR0915 ) elif custom_llm_provider == "openai": return openai_cost_per_token( - model=model, usage=usage_block, service_tier=service_tier + model=model, + usage=usage_block, + service_tier=service_tier, + data_residency=data_residency, ) elif custom_llm_provider == "databricks": return databricks_cost_per_token(model=model, usage=usage_block) @@ -631,6 +676,7 @@ def cost_per_token( # noqa: PLR0915 usage=usage_block, custom_llm_provider=custom_llm_provider, service_tier=service_tier, + data_residency=data_residency, ) if ( @@ -1117,6 +1163,10 @@ def completion_cost( # noqa: PLR0915 litellm_logging_obj: Optional[LitellmLoggingObject] = None, ### SERVICE TIER ### service_tier: Optional[str] = None, # for OpenAI service tier pricing + ### DATA RESIDENCY ### + data_residency: Optional[ + str + ] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ) -> float: """ Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. @@ -1516,6 +1566,7 @@ def completion_cost( # noqa: PLR0915 combined_usage_object=cost_per_token_usage_object, custom_llm_provider=custom_llm_provider, litellm_model_name=model, + data_residency=data_residency, ) elif call_type == _MCP_CALL_TYPE: from litellm.proxy._experimental.mcp_server.cost_calculator import ( @@ -1600,6 +1651,7 @@ def completion_cost( # noqa: PLR0915 audio_transcription_file_duration=audio_transcription_file_duration, rerank_billed_units=rerank_billed_units, service_tier=service_tier, + data_residency=data_residency, response=completion_response, request_model=request_model_for_cost, ) @@ -1811,6 +1863,10 @@ def response_cost_calculator( litellm_logging_obj: Optional[LitellmLoggingObject] = None, ### SERVICE TIER ### service_tier: Optional[str] = None, # for OpenAI service tier pricing + ### DATA RESIDENCY ### + data_residency: Optional[ + str + ] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ) -> float: """ Returns @@ -1844,6 +1900,7 @@ def response_cost_calculator( router_model_id=router_model_id, litellm_logging_obj=litellm_logging_obj, service_tier=service_tier, + data_residency=data_residency, ) return response_cost except Exception as e: @@ -2202,6 +2259,7 @@ def batch_cost_calculator( model: str, custom_llm_provider: Optional[str] = None, model_info: Optional[ModelInfo] = None, + data_residency: Optional[str] = None, ) -> Tuple[float, float]: """ Calculate the cost of a batch job. @@ -2286,6 +2344,11 @@ def batch_cost_calculator( usage.completion_tokens * (output_cost_per_token) / 2 ) # batch cost is usually half of the regular token cost + uplift = _get_regional_uplift_multiplier(model_info, data_residency) + if uplift != 1.0: + total_prompt_cost *= uplift + total_completion_cost *= uplift + return total_prompt_cost, total_completion_cost @@ -2431,6 +2494,7 @@ def handle_realtime_stream_cost_calculation( combined_usage_object: Usage, custom_llm_provider: str, litellm_model_name: str, + data_residency: Optional[str] = None, ) -> float: """ Handles the cost calculation for realtime stream responses. @@ -2461,6 +2525,7 @@ def handle_realtime_stream_cost_calculation( model=model_name, usage=combined_usage_object, custom_llm_provider=custom_llm_provider, + data_residency=data_residency, ) except Exception: continue diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index b8cd04836c3..d48dba8e7bb 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -1,5 +1,7 @@ import os -from typing import TYPE_CHECKING, Any, Optional, Union +import threading +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Optional, Tuple, Union from litellm._logging import verbose_logger from litellm.integrations.arize import _utils @@ -8,8 +10,10 @@ from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig if TYPE_CHECKING: from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SpanProcessor from opentelemetry.trace import Span as _Span from opentelemetry.trace import SpanKind + from opentelemetry.trace import Tracer from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry from litellm.integrations.opentelemetry import ( @@ -21,20 +25,27 @@ if TYPE_CHECKING: OpenTelemetryConfig = _OpenTelemetryConfig Span = Union[_Span, Any] OpenTelemetry = _OpenTelemetry + LITELLM_TRACER_NAME: str else: Protocol = Any OpenTelemetryConfig = Any Span = Any + Tracer = Any TracerProvider = Any SpanKind = Any - # Import OpenTelemetry at runtime + SpanProcessor = Any try: - from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.integrations.opentelemetry import ( + LITELLM_TRACER_NAME, + OpenTelemetry, + ) except ImportError: + LITELLM_TRACER_NAME = "litellm" OpenTelemetry = None # type: ignore ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces" +_MAX_PROJECT_PROVIDERS = 64 class ArizePhoenixLogger(OpenTelemetry): # type: ignore @@ -48,37 +59,142 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore def _init_tracing(self, tracer_provider): """ - Override to always create a *private* TracerProvider for Arize Phoenix. + Override to create per-project TracerProviders (LRU-cached) for Arize Phoenix. The base ``OpenTelemetry._init_tracing`` falls back to the global TracerProvider when one already exists. That causes whichever integration initialises second to silently reuse the first one's exporter, so spans only reach one destination. - - By creating our own provider we guarantee Arize Phoenix always gets - its own exporter pipeline, regardless of initialisation order. """ - from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import SpanKind if tracer_provider is not None: - # Explicitly supplied (e.g. in tests) — honour it. - self.tracer = tracer_provider.get_tracer("litellm") + self._use_injected_tracer_provider = True + self._shared_span_processor = None + self.tracer = tracer_provider.get_tracer(LITELLM_TRACER_NAME) self.span_kind = SpanKind return - # Always create a dedicated provider — never touch the global one. - provider = TracerProvider(resource=self._get_litellm_resource(self.config)) - provider.add_span_processor(self._get_span_processor()) - self.tracer = provider.get_tracer("litellm") + self._use_injected_tracer_provider = False + self._project_providers: OrderedDict[str, TracerProvider] = OrderedDict() + self._project_providers_lock = threading.Lock() + self._shared_span_processor = self._get_span_processor() self.span_kind = SpanKind + + default_project = self._resolve_project_name({}) + self.tracer = self._get_tracer_for(default_project) verbose_logger.debug( - "ArizePhoenixLogger: Created dedicated TracerProvider " - "(endpoint=%s, exporter=%s)", + "ArizePhoenixLogger: Initialized per-project TracerProvider cache " + "(default_project=%s, endpoint=%s, exporter=%s)", + default_project, self.config.endpoint, self.config.exporter, ) + def flush_tracer_providers(self) -> None: + """ + Flush all cached per-project providers and the shared span processor. + + Call on graceful proxy shutdown. Do not call on LRU eviction — in-flight + spans may still reference evicted providers. + """ + if getattr(self, "_use_injected_tracer_provider", False): + return + + shared_processor = getattr(self, "_shared_span_processor", None) + if shared_processor is not None: + try: + shared_processor.force_flush() + except Exception as e: + verbose_logger.debug( + "ArizePhoenixLogger: shared span processor force_flush failed: %s", + e, + ) + + with getattr(self, "_project_providers_lock", threading.Lock()): + providers = list(getattr(self, "_project_providers", {}).values()) + + for provider in providers: + try: + provider.force_flush() + except Exception as e: + verbose_logger.debug( + "ArizePhoenixLogger: TracerProvider force_flush failed: %s", e + ) + + def _get_litellm_resource_for_project(self, project_name: str): + """ + Build an OTEL Resource with project routing attrs that win over env detector. + + Phoenix uses ``openinference.project.name``; Arize AX uses ``model_id`` and + ``service.name``. Project attrs are merged last so OTEL_RESOURCE_ATTRIBUTES + from init does not pin every provider to one project. + """ + from opentelemetry.sdk.resources import OTELResourceDetector, Resource + + project_attributes: dict[str, str] = { + "openinference.project.name": project_name, + "model_id": project_name, + "service.name": project_name, + } + deployment_environment = getattr(self.config, "deployment_environment", None) + if deployment_environment is not None: + project_attributes["deployment.environment"] = deployment_environment + + env_resource = OTELResourceDetector().detect() + project_resource = Resource.create(project_attributes) # type: ignore[arg-type] + return env_resource.merge(project_resource) + + def _build_tracer_provider_for_project(self, project_name: str) -> TracerProvider: + """Create a TracerProvider for *project_name* (caller holds no cache lock).""" + from opentelemetry.sdk.trace import TracerProvider + + provider = TracerProvider( + resource=self._get_litellm_resource_for_project(project_name) + ) + provider.add_span_processor(self._shared_span_processor) + return provider + + def _get_tracer_for(self, project_name: str) -> Tracer: + """Return a tracer for *project_name*, creating/caching a provider on miss.""" + if getattr(self, "_use_injected_tracer_provider", False): + return self.tracer + + with self._project_providers_lock: + if project_name in self._project_providers: + self._project_providers.move_to_end(project_name) + return self._project_providers[project_name].get_tracer( + LITELLM_TRACER_NAME + ) + + # OTELResourceDetector().detect() is synchronous; build outside the lock so + # concurrent requests for other projects are not blocked on cache misses. + new_provider = self._build_tracer_provider_for_project(project_name) + + with self._project_providers_lock: + if project_name in self._project_providers: + self._project_providers.move_to_end(project_name) + return self._project_providers[project_name].get_tracer( + LITELLM_TRACER_NAME + ) + + if len(self._project_providers) >= _MAX_PROJECT_PROVIDERS: + self._project_providers.popitem(last=False) + + self._project_providers[project_name] = new_provider + return new_provider.get_tracer(LITELLM_TRACER_NAME) + + def _resolve_tracer_for_kwargs(self, kwargs: dict) -> Tuple[str, Tracer]: + """Resolve project name once and return the matching tracer.""" + project_name = self._resolve_project_name(kwargs) + return project_name, self._get_tracer_for(project_name) + + def get_tracer_to_use_for_request(self, kwargs: dict) -> Tracer: + """Route guardrail/raw-request spans to the same per-project tracer as the request.""" + if getattr(self, "_use_injected_tracer_provider", False): + return self.tracer + return self._resolve_tracer_for_kwargs(kwargs)[1] + def _init_otel_logger_on_litellm_proxy(self): """ Override: Arize Phoenix should NOT overwrite the proxy's @@ -93,56 +209,109 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): - from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import ( - safe_set_attribute, - ) - _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) - - # Dynamic project name: check metadata first, then fall back to env var config - dynamic_project_name = ArizePhoenixLogger._get_dynamic_project_name(kwargs) - if dynamic_project_name: - safe_set_attribute(span, "openinference.project.name", dynamic_project_name) - else: - # Fall back to static config from env var - config = ArizePhoenixLogger.get_arize_phoenix_config() - if config.project_name: - safe_set_attribute( - span, "openinference.project.name", config.project_name - ) - return @staticmethod - def _get_dynamic_project_name(kwargs) -> Optional[str]: - """ - Retrieve dynamic Phoenix project name from request metadata. + def _normalize_project_name(name: Optional[str]) -> Optional[str]: + if name is None: + return None + normalized = str(name).strip() + return normalized if normalized else None - Users can set `metadata.phoenix_project_name` in their request to route - traces to different Phoenix projects dynamically. - """ - standard_logging_payload = kwargs.get("standard_logging_object") - if isinstance(standard_logging_payload, dict): - metadata = standard_logging_payload.get("metadata") + @staticmethod + def _iter_metadata_dicts_from_kwargs(kwargs: dict): + """Yield request metadata dicts; standard_logging_object before litellm_params.""" + for key in ("standard_logging_object", "litellm_params"): + found_key = kwargs.get(key) + if not isinstance(found_key, dict): + continue + metadata = found_key.get("metadata") if isinstance(metadata, dict): - project_name = metadata.get("phoenix_project_name") - if project_name: - return str(project_name) + yield metadata - # Also check litellm_params.metadata for SDK usage + @staticmethod + def _is_proxy_request(kwargs: dict) -> bool: + """True when the call is routed through the LiteLLM proxy. + + Proxy mode is determined solely by the server-set ``proxy_server_request`` + field in ``litellm_params``. Checking request metadata for + ``user_api_key_auth_metadata`` is intentionally avoided: that field is + user-supplied and would let an authenticated caller fake proxy-mode + detection to route their telemetry into arbitrary Arize/Phoenix projects. + """ litellm_params = kwargs.get("litellm_params") - if isinstance(litellm_params, dict): - metadata = litellm_params.get("metadata") or {} - else: - metadata = {} - if isinstance(metadata, dict): - project_name = metadata.get("phoenix_project_name") - if project_name: - return str(project_name) + return isinstance(litellm_params, dict) and bool( + litellm_params.get("proxy_server_request") + ) + @staticmethod + def _project_from_metadata_dict( + metadata: dict, metadata_key: str, *, proxy_mode: bool + ) -> Optional[str]: + """ + Read a Phoenix project field from proxy/SDK metadata. + + On the proxy, only ``user_api_key_auth_metadata`` (team/key config) may + select the project. SDK callers may still set project fields directly on + ``metadata``. + """ + auth_metadata = metadata.get("user_api_key_auth_metadata") + if isinstance(auth_metadata, dict): + project = ArizePhoenixLogger._normalize_project_name( + auth_metadata.get(metadata_key) + ) + if project: + return project + + if not proxy_mode: + return ArizePhoenixLogger._normalize_project_name( + metadata.get(metadata_key) + ) return None - def _get_phoenix_context(self, kwargs): + @staticmethod + def _metadata_project_from_kwargs(kwargs: dict, metadata_key: str) -> Optional[str]: + proxy_mode = ArizePhoenixLogger._is_proxy_request(kwargs) + for metadata in ArizePhoenixLogger._iter_metadata_dicts_from_kwargs(kwargs): + project = ArizePhoenixLogger._project_from_metadata_dict( + metadata, metadata_key, proxy_mode=proxy_mode + ) + if project: + return project + return None + + @staticmethod + def _resolve_project_name(kwargs: dict) -> str: + """ + Resolve the target Phoenix/Arize project for this request. + + Proxy priority: ``user_api_key_auth_metadata.phoenix_project_name_override``, + ``user_api_key_auth_metadata.phoenix_project_name``, env, then ``default``. + SDK priority: request metadata fields, then env, then ``default``. + """ + override = ArizePhoenixLogger._metadata_project_from_kwargs( + kwargs, "phoenix_project_name_override" + ) + if override: + return override + + phoenix_name = ArizePhoenixLogger._metadata_project_from_kwargs( + kwargs, "phoenix_project_name" + ) + if phoenix_name: + return phoenix_name + + env_name = ArizePhoenixLogger._normalize_project_name( + os.environ.get("PHOENIX_PROJECT_NAME") + or os.environ.get("ARIZE_PROJECT_NAME") + ) + if env_name: + return env_name + + return "default" + + def _get_phoenix_context(self, kwargs, tracer: Optional[Tracer] = None): """ Build a trace context for Phoenix's dedicated TracerProvider. @@ -159,11 +328,13 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore """ from opentelemetry import trace + if tracer is None: + tracer = self._resolve_tracer_for_kwargs(kwargs)[1] + litellm_params = kwargs.get("litellm_params", {}) or {} proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} headers = proxy_server_request.get("headers", {}) or {} - # Propagate distributed trace context if the caller sent a traceparent traceparent_ctx = ( self.get_traceparent_from_header(headers=headers) if headers.get("traceparent") @@ -173,10 +344,8 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore is_proxy_mode = bool(proxy_server_request) if is_proxy_mode: - # Create a parent span on Phoenix's own tracer so both parent - # and child are exported to Phoenix. start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time")) - parent_span = self.tracer.start_span( + parent_span = tracer.start_span( name="litellm_proxy_request", start_time=( self._to_ns(start_time_val) if start_time_val is not None else None @@ -187,100 +356,77 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore ctx = trace.set_span_in_context(parent_span) return ctx, parent_span - # SDK mode — no parent span needed return traceparent_ctx, None def _handle_success(self, kwargs, response_obj, start_time, end_time): - """ - Override to always create spans on ArizePhoenixLogger's dedicated TracerProvider. - - The base class's ``_get_span_context`` would find the parent span created by - the ``otel`` callback on the *global* TracerProvider. That span is invisible - in Phoenix (different exporter pipeline), so we ignore it and build our own - hierarchy via ``_get_phoenix_context``. - """ - from opentelemetry.trace import Status, StatusCode - - verbose_logger.debug( - "ArizePhoenixLogger: Logging kwargs: %s, OTEL config settings=%s", - kwargs, - self.config, + self._handle_phoenix_trace( + kwargs, response_obj, start_time, end_time, success=True ) - ctx, parent_span = self._get_phoenix_context(kwargs) - - # Create litellm_request span (child of our parent when in proxy mode) - span = self.tracer.start_span( - name=self._get_span_name(kwargs), - start_time=self._to_ns(start_time), - context=ctx, - ) - span.set_status(Status(StatusCode.OK)) - self.set_attributes(span, kwargs, response_obj) - - # Raw-request sub-span (if enabled) — must be created before - # ending the parent span so the hierarchy is valid. - self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) - span.end(end_time=self._to_ns(end_time)) - - # Guardrail span - self._create_guardrail_span(kwargs=kwargs, context=ctx) - - # Annotate and close our proxy parent span - if parent_span is not None: - parent_span.set_status(Status(StatusCode.OK)) - self.set_attributes(parent_span, kwargs, response_obj) - parent_span.end(end_time=self._to_ns(end_time)) - - # Metrics & cost recording - self._record_metrics(kwargs, response_obj, start_time, end_time) - - # Semantic logs - if self.config.enable_events: - self._emit_semantic_logs(kwargs, response_obj, span) - def _handle_failure(self, kwargs, response_obj, start_time, end_time): - """ - Override to always create failure spans on ArizePhoenixLogger's dedicated - TracerProvider. Mirrors ``_handle_success`` but sets ERROR status. - """ + self._handle_phoenix_trace( + kwargs, response_obj, start_time, end_time, success=False + ) + + def _handle_phoenix_trace( + self, + kwargs, + response_obj, + start_time, + end_time, + *, + success: bool, + ): from opentelemetry.trace import Status, StatusCode verbose_logger.debug( - "ArizePhoenixLogger: Failure - Logging kwargs: %s, OTEL config settings=%s", + "ArizePhoenixLogger: %s - kwargs: %s, OTEL config settings=%s", + "success" if success else "failure", kwargs, self.config, ) - ctx, parent_span = self._get_phoenix_context(kwargs) + _project_name, tracer = self._resolve_tracer_for_kwargs(kwargs) + ctx, parent_span = self._get_phoenix_context(kwargs, tracer=tracer) - # Create litellm_request span (child of our parent when in proxy mode) - span = self.tracer.start_span( + status = Status(StatusCode.OK if success else StatusCode.ERROR) + + span = tracer.start_span( name=self._get_span_name(kwargs), start_time=self._to_ns(start_time), context=ctx, ) - span.set_status(Status(StatusCode.ERROR)) + span.set_status(status) self.set_attributes(span, kwargs, response_obj) - self._record_exception_on_span(span=span, kwargs=kwargs) + if not success: + self._record_exception_on_span(span=span, kwargs=kwargs) + + if success: + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, span + ) span.end(end_time=self._to_ns(end_time)) - # Guardrail span self._create_guardrail_span(kwargs=kwargs, context=ctx) - # Annotate and close our proxy parent span if parent_span is not None: - parent_span.set_status(Status(StatusCode.ERROR)) + parent_span.set_status(status) self.set_attributes(parent_span, kwargs, response_obj) - self._record_exception_on_span(span=parent_span, kwargs=kwargs) + if not success: + self._record_exception_on_span(span=parent_span, kwargs=kwargs) parent_span.end(end_time=self._to_ns(end_time)) + if success: + self._record_metrics(kwargs, response_obj, start_time, end_time) + + if self.config.enable_events: + self._emit_semantic_logs(kwargs, response_obj, span) + @staticmethod def get_arize_phoenix_config() -> ArizePhoenixConfig: """ Retrieves the Arize Phoenix configuration based on environment variables. Returns: - ArizePhoenixConfig: A Pydantic model containing Arize Phoenix configuration. """ api_key = os.environ.get("PHOENIX_API_KEY", None) @@ -295,18 +441,15 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore protocol: Protocol = "otlp_http" if collector_endpoint: - # Parse the endpoint to determine protocol if collector_endpoint.startswith("grpc://") or ( ":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint ): endpoint = collector_endpoint protocol = "otlp_grpc" else: - # Phoenix Cloud endpoints (app.phoenix.arize.com) include the space in the URL if "app.phoenix.arize.com" in collector_endpoint: endpoint = collector_endpoint protocol = "otlp_http" - # For other HTTP endpoints, ensure they have the correct path elif "/v1/traces" not in collector_endpoint: if collector_endpoint.endswith("/v1"): endpoint = collector_endpoint + "/traces" @@ -318,7 +461,6 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore endpoint = collector_endpoint protocol = "otlp_http" else: - # If no endpoint specified, self hosted phoenix endpoint = "http://localhost:6006/v1/traces" protocol = "otlp_http" verbose_logger.debug( @@ -329,12 +471,11 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore if api_key is not None: otlp_auth_headers = f"Authorization=Bearer {api_key}" elif "app.phoenix.arize.com" in endpoint: - # Phoenix Cloud requires an API key raise ValueError( "PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com)." ) - project_name = os.environ.get("PHOENIX_PROJECT_NAME", "default") + project_name = os.environ.get("PHOENIX_PROJECT_NAME") or "default" return ArizePhoenixConfig( otlp_auth_headers=otlp_auth_headers, @@ -343,8 +484,6 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore project_name=project_name, ) - ## cannot suppress additional proxy server spans, removed previous methods. - async def async_health_check(self): config = self.get_arize_phoenix_config() diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index a961d4f9244..0f954eb1ce0 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -2,10 +2,17 @@ import asyncio import os import time from datetime import datetime -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, cast from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.datadog.datadog_handler import ( + get_datadog_env, + get_datadog_hostname, + get_datadog_pod_name, + get_datadog_service, +) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -15,9 +22,30 @@ from litellm.types.integrations.datadog_cost_management import ( ) from litellm.types.utils import StandardLoggingPayload +# Reserved tag keys whose values come from trusted sources (infra env, LiteLLM +# core payload fields, or proxy-controlled auth metadata). User-supplied +# request_tags / metadata cannot overwrite these, even when the key is +# allowlisted via cost_tag_keys, because that would let an authenticated caller +# spoof cost attribution (e.g. request_tags=["team:victim-team"]). +_RESERVED_TAG_KEYS: frozenset = frozenset( + { + "env", + "service", + "host", + "pod_name", + "provider", + "model", + "model_id", + "team", + "user", + "model_group", + } +) + class DatadogCostManagementLogger(CustomBatchLogger): - def __init__(self, **kwargs): + def __init__(self, cost_tag_keys: Optional[List[str]] = None, **kwargs): + self.cost_tag_keys: List[str] = list(cost_tag_keys) if cost_tag_keys else [] self.dd_api_key = os.getenv("DD_API_KEY") self.dd_app_key = os.getenv("DD_APP_KEY") self.dd_site = os.getenv("DD_SITE", "datadoghq.com") @@ -68,20 +96,21 @@ class DatadogCostManagementLogger(CustomBatchLogger): if not self.log_queue: return + batch_to_send = self.log_queue[:] + self.log_queue = [] + try: - # Aggregate costs from the batch - aggregated_entries = self._aggregate_costs(self.log_queue) - + aggregated_entries = self._aggregate_costs(batch_to_send) if not aggregated_entries: + verbose_logger.debug( + "Datadog Cost Management: batch produced no aggregable entries; " + "dropping %d log(s) from queue.", + len(batch_to_send), + ) return - - # Send to Datadog await self._upload_to_datadog(aggregated_entries) - - # Clear queue only on success (or if we decide to drop on failure) - # CustomBatchLogger clears queue in flush_queue, so we just process here - except Exception as e: + self.log_queue = batch_to_send + self.log_queue verbose_logger.exception( f"Datadog Cost Management: Error in async_send_batch: {str(e)}" ) @@ -151,45 +180,81 @@ class DatadogCostManagementLogger(CustomBatchLogger): return list(aggregator.values()) def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]: - from litellm.integrations.datadog.datadog_handler import ( - get_datadog_env, - get_datadog_hostname, - get_datadog_pod_name, - get_datadog_service, - ) - - tags = { + tags: Dict[str, str] = { "env": get_datadog_env(), "service": get_datadog_service(), "host": get_datadog_hostname(), "pod_name": get_datadog_pod_name(), } - # Add metadata as tags - metadata = log.get("metadata", {}) - if metadata: - # Add user info - # Add user info - if metadata.get("user_api_key_alias"): - tags["user"] = str(metadata["user_api_key_alias"]) + # Always-on canonical FOCUS dimensions from top-level payload fields. + # Non-sensitive and required for Datadog Custom Costs per-model attribution. + self._add_tag(tags, "provider", log.get("custom_llm_provider")) + self._add_tag(tags, "model", log.get("model")) + self._add_tag(tags, "model_id", log.get("model_id")) - # Add Team Tag - team_tag = ( - metadata.get("user_api_key_team_alias") - or metadata.get("team_alias") # type: ignore - or metadata.get("user_api_key_team_id") - or metadata.get("team_id") # type: ignore - ) + # cast because StandardLoggingMetadata is a TypedDict; we iterate it + # as a generic mapping below. + metadata: Dict[str, Any] = cast(Dict[str, Any], log.get("metadata") or {}) - if team_tag: - tags["team"] = str(team_tag) - # model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get() - model_group = metadata.get("model_group") # type: ignore[misc] - if model_group: - tags["model_group"] = str(model_group) + # Backwards-compat: team/user/model_group preserved regardless of allowlist. + if metadata.get("user_api_key_alias"): + tags["user"] = str(metadata["user_api_key_alias"]) + team_tag = ( + metadata.get("user_api_key_team_alias") + or metadata.get("team_alias") + or metadata.get("user_api_key_team_id") + or metadata.get("team_id") + ) + if team_tag: + tags["team"] = str(team_tag) + if metadata.get("model_group"): + tags["model_group"] = str(metadata["model_group"]) + + # Allowlist-gated: request_tags (split on `:`) and arbitrary metadata.*. + # Reserved keys are hard-blocked here regardless of allowlist membership — + # see _RESERVED_TAG_KEYS for the rationale. + if self.cost_tag_keys: + allow = set(self.cost_tag_keys) + for rt in log.get("request_tags") or []: + if not isinstance(rt, str) or ":" not in rt: + continue + k, _, v = rt.partition(":") + if k in allow and v: + self._set_custom_tag(tags, k, v) + for k, v in metadata.items(): + if k in allow and v is not None and not isinstance(v, (dict, list)): + self._set_custom_tag(tags, k, str(v)) + for nested_key in ("spend_logs_metadata", "requester_metadata"): + nested = metadata.get(nested_key) + if isinstance(nested, dict): + for k, v in nested.items(): + if ( + k in allow + and v is not None + and not isinstance(v, (dict, list)) + ): + self._set_custom_tag(tags, k, str(v)) return tags + @staticmethod + def _set_custom_tag(tags: Dict[str, str], key: str, value: str) -> None: + if key in _RESERVED_TAG_KEYS: + verbose_logger.debug( + "Datadog Cost Management: dropping user-supplied tag %r=%r — " + "key is reserved for trusted cost attribution.", + key, + value, + ) + return + tags[key] = value + + @staticmethod + def _add_tag(tags: Dict[str, str], key: str, value: Any) -> None: + if value: + tags[key] = str(value) + async def _upload_to_datadog(self, payload: List[Dict]): if not self.dd_api_key or not self.dd_app_key: return @@ -201,8 +266,6 @@ class DatadogCostManagementLogger(CustomBatchLogger): } # The API endpoint expects a list of objects directly in the body (file content behavior) - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - data_json = safe_dumps(payload) response = await self.async_client.put( diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index fcf40701e28..d7847027d7e 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -144,7 +144,26 @@ class DatadogMetricsLogger(CustomBatchLogger): } self.log_queue.append(series_llm_latency) - # 3. Request Count / Status Code + # 3. LiteLLM Overhead Latency Metric (total - llm_api time) + hidden_params = log.get("hidden_params", {}) or {} + litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + if litellm_overhead_time_ms is not None: + overhead_tags = self._extract_tags(log) # no status_code on latency metric + series_overhead: DatadogMetricSeries = { + "metric": "litellm.overhead.latency", + "type": 3, # gauge + "points": [ + { + "timestamp": timestamp, + "value": litellm_overhead_time_ms + / 1000, # convert ms → seconds + } + ], + "tags": overhead_tags, + } + self.log_queue.append(series_overhead) + + # 4. Request Count / Status Code series_count: DatadogMetricSeries = { "metric": "litellm.llm_api.request_count", "type": 1, # count diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index e99d5f23a4c..a598124f612 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -1,18 +1,29 @@ +import json import os -from typing import Any, Dict, List, Optional +import re +from typing import Any, Dict, List, Optional, Tuple, cast from pydantic import BaseModel, Field import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, + get_content_from_model_response, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.llms.openai import AllMessageValues + +GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai" +# Cap the in-memory buffer so persistent flush failures (e.g. Galileo +# unavailable, invalid credentials) cannot leak memory unboundedly. +GALILEO_MAX_IN_MEMORY_RECORDS = 1000 -# from here: https://docs.rungalileo.io/galileo/gen-ai-studio-products/galileo-observe/how-to/logging-data-via-restful-apis#structuring-your-records class LLMResponse(BaseModel): latency_ms: int status_code: int @@ -37,65 +48,190 @@ class GalileoObserve(CustomLogger): def __init__(self) -> None: self.in_memory_records: List[dict] = [] self.batch_size = 1 - self.base_url = os.getenv("GALILEO_BASE_URL", None) - self.project_id = os.getenv("GALILEO_PROJECT_ID", None) + self.api_key = os.getenv("GALILEO_API_KEY") + self.project_id = os.getenv("GALILEO_PROJECT_ID") + self.log_stream_id = os.getenv("GALILEO_LOG_STREAM_ID") + self.username = os.getenv("GALILEO_USERNAME") + self.password = os.getenv("GALILEO_PASSWORD") + self.base_url = self._normalize_base_url(os.getenv("GALILEO_BASE_URL")) + if self.api_key and not self.base_url: + self.base_url = GALILEO_CLOUD_API_BASE_URL + self.use_v2_api = bool(self.api_key) self.headers: Optional[Dict[str, str]] = None self.async_httpx_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) - pass - def set_galileo_headers(self): - # following https://docs.rungalileo.io/galileo/gen-ai-studio-products/galileo-observe/how-to/logging-data-via-restful-apis#logging-your-records + @staticmethod + def _normalize_base_url(base_url: Optional[str]) -> Optional[str]: + if base_url: + return base_url.rstrip("/") + return None - headers = { - "accept": "application/json", - "Content-Type": "application/x-www-form-urlencoded", - } - galileo_login_response = litellm.module_level_client.post( + def _is_configured(self) -> bool: + if not self.project_id or not self.base_url: + return False + if self.use_v2_api: + return bool(self.api_key) + return bool(self.username and self.password) + + async def async_set_galileo_headers(self) -> None: + galileo_login_response = await self.async_httpx_handler.post( url=f"{self.base_url}/login", - headers=headers, + headers={ + "accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + }, data={ - "username": os.getenv("GALILEO_USERNAME"), - "password": os.getenv("GALILEO_PASSWORD"), + "username": self.username, + "password": self.password, }, ) - + galileo_login_response.raise_for_status() access_token = galileo_login_response.json()["access_token"] - self.headers = { "accept": "application/json", "Content-Type": "application/json", "Authorization": f"Bearer {access_token}", } - def get_output_str_from_response(self, response_obj, kwargs): - output = None - if response_obj is not None and ( - kwargs.get("call_type", None) == "embedding" - or isinstance(response_obj, litellm.EmbeddingResponse) - ): - output = None - elif response_obj is not None and isinstance( - response_obj, litellm.ModelResponse - ): - output = response_obj["choices"][0]["message"].json() - elif response_obj is not None and isinstance( - response_obj, litellm.TextCompletionResponse - ): - output = response_obj.choices[0].text - elif response_obj is not None and isinstance( - response_obj, litellm.ImageResponse - ): - output = response_obj["data"] + async def _ensure_headers(self) -> bool: + if self.headers is not None: + return True - return output + if self.use_v2_api: + if not self.api_key: + return False + self.headers = { + "accept": "application/json", + "Content-Type": "application/json", + "Galileo-API-Key": self.api_key, + } + return True + + if not (self.username and self.password and self.base_url): + return False + + try: + await self.async_set_galileo_headers() + return True + except Exception as e: + verbose_logger.debug("Galileo Logger: failed to authenticate: %s", e) + return False + + @staticmethod + def _galileo_input_messages( + messages: Optional[List[Any]], input_text: str + ) -> List[Dict[str, str]]: + if not messages: + return [{"role": "user", "content": input_text}] + + galileo_messages: List[Dict[str, str]] = [] + for message in messages: + if not isinstance(message, dict): + continue + role = message.get("role") + if not role: + continue + galileo_messages.append( + { + "role": str(role), + "content": convert_content_list_to_str( + message=cast(AllMessageValues, message) + ), + } + ) + + if galileo_messages: + return galileo_messages + return [{"role": "user", "content": input_text}] + + @staticmethod + def _record_to_v2_span(record: Dict[str, Any]) -> Dict[str, Any]: + created_at = record.get("created_at", "") + if created_at and not re.search(r"(Z|[+-]\d{2}:?\d{2})$", created_at): + created_at = f"{created_at}Z" + + span: Dict[str, Any] = { + "type": "llm", + "name": record.get("node_type", "litellm"), + "created_at": created_at, + "input": GalileoObserve._galileo_input_messages( + record.get("messages"), record.get("input_text", "") + ), + "output": { + "role": "assistant", + "content": record.get("output_text", ""), + }, + "status_code": record.get("status_code", 200), + "model": record.get("model"), + "metrics": { + "duration_ns": int(record.get("latency_ms", 0)) * 1_000_000, + "num_input_tokens": record.get("num_input_tokens"), + "num_output_tokens": record.get("num_output_tokens"), + }, + } + if record.get("tags"): + span["tags"] = record["tags"] + return span + + def _get_ingest_request(self) -> Optional[Tuple[str, Dict[str, Any]]]: + if not self.base_url or not self.project_id: + return None + + # Snapshot the records to be sent into a new list so concurrent appends + # during the network round-trip (across the await points in + # flush_in_memory_records) aren't silently dropped when we later clear + # the in-memory buffer. + records = list(self.in_memory_records) + + if self.use_v2_api: + payload: Dict[str, Any] = { + "spans": [self._record_to_v2_span(record) for record in records], + "reliable": False, + } + if self.log_stream_id: + payload["log_stream_id"] = self.log_stream_id + return ( + f"{self.base_url}/v2/projects/{self.project_id}/spans", + payload, + ) + + return ( + f"{self.base_url}/projects/{self.project_id}/observe/ingest", + {"records": records}, + ) + + def get_output_str_from_response( + self, response_obj: Any, kwargs: Dict[str, Any] + ) -> Optional[str]: + if response_obj is None: + return None + if kwargs.get("call_type", None) == "embedding" or isinstance( + response_obj, litellm.EmbeddingResponse + ): + return None + if isinstance(response_obj, litellm.TextCompletionResponse): + return response_obj.choices[0].text + if isinstance(response_obj, litellm.ImageResponse): + return json.dumps(response_obj["data"], default=str) + if isinstance(response_obj, (litellm.ModelResponse, dict)): + return get_content_from_model_response(response_obj) + return None async def async_log_success_event( self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any ): verbose_logger.debug("On Async Success") + if not self._is_configured(): + verbose_logger.debug( + "Galileo Logger: skipping flush — set GALILEO_PROJECT_ID and " + "either GALILEO_API_KEY (hosted) or GALILEO_USERNAME/GALILEO_PASSWORD " + "(enterprise Observe)." + ) + return + _latency_ms = int((end_time - start_time).total_seconds() * 1000) _call_type = kwargs.get("call_type", "litellm") input_text = litellm.utils.get_formatted_prompt( @@ -125,26 +261,69 @@ class GalileoObserve(CustomLogger): ), # timestamp str constructed in "%Y-%m-%dT%H:%M:%S" format ) - # dump to dict request_dict = request_record.model_dump() + messages = kwargs.get("messages") + if messages: + request_dict["messages"] = messages self.in_memory_records.append(request_dict) + # Bound the buffer so persistent flush failures cannot grow it + # without limit. Drop the oldest records once we exceed the cap. + if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS: + dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS + self.in_memory_records = self.in_memory_records[ + -GALILEO_MAX_IN_MEMORY_RECORDS: + ] + verbose_logger.warning( + "Galileo Logger: in-memory buffer exceeded %s records; " + "dropped %s oldest record(s). Check Galileo connectivity/credentials.", + GALILEO_MAX_IN_MEMORY_RECORDS, + dropped, + ) + if len(self.in_memory_records) >= self.batch_size: await self.flush_in_memory_records() async def flush_in_memory_records(self): - verbose_logger.debug("flushing in memory records") - response = await self.async_httpx_handler.post( - url=f"{self.base_url}/projects/{self.project_id}/observe/ingest", - headers=self.headers, - json={"records": self.in_memory_records}, - ) + if not self.in_memory_records: + return - if response.status_code == 200: + # Capture the number of records that will be sent BEFORE any await so + # that concurrent appends made by other asyncio tasks during the + # network round-trip aren't silently dropped on the success-clear. + records_in_payload = len(self.in_memory_records) + + ingest_request = self._get_ingest_request() + if ingest_request is None: verbose_logger.debug( - "Galileo Logger:successfully flushed in memory records" + "Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID" ) - self.in_memory_records = [] + return + + if not await self._ensure_headers(): + verbose_logger.debug("Galileo Logger: could not set request headers") + return + + url, payload = ingest_request + verbose_logger.debug("flushing in memory records to %s", url) + + try: + response = await self.async_httpx_handler.post( + url=url, + headers=self.headers, + json=payload, + ) + except Exception as e: + verbose_logger.debug( + "Galileo Logger: failed to flush in memory records: %s", e + ) + return + + if response.is_success: + verbose_logger.debug( + "Galileo Logger: successfully flushed in memory records" + ) + del self.in_memory_records[:records_in_payload] else: verbose_logger.debug("Galileo Logger: failed to flush in memory records") verbose_logger.debug( @@ -152,6 +331,13 @@ class GalileoObserve(CustomLogger): response.text, response.status_code, ) + # Legacy enterprise auth caches a bearer token obtained from + # /login. If the request was rejected for auth reasons, drop the + # cached headers so the next flush re-authenticates instead of + # silently failing forever on a stale token. The v2 API key path + # uses a long-lived static key, so leave its headers in place. + if not self.use_v2_api and response.status_code in (401, 403): + self.headers = None async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("On Async Failure") diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ad9538ac171..b32803b5dfc 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -1,5 +1,7 @@ from typing import Optional +from litellm.llms.openai.data_residency import infer_openai_data_residency + # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls _OPTIONAL_KWARGS_KEYS = frozenset( @@ -103,6 +105,10 @@ def get_litellm_params( if litellm_trace_id is None: litellm_trace_id = _meta.get("trace_id") or _meta.get("session_id") + data_residency: Optional[str] = infer_openai_data_residency( + custom_llm_provider, api_base + ) + # Build base dict with explicit parameters (always included) litellm_params = { "acompletion": acompletion, @@ -112,6 +118,7 @@ def get_litellm_params( "verbose": verbose, "custom_llm_provider": custom_llm_provider, "api_base": api_base, + "data_residency": data_residency, "litellm_call_id": litellm_call_id, "model_alias_map": model_alias_map, "completion_call_id": completion_call_id, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 63fa0e64695..97266096ef9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1546,6 +1546,11 @@ class Logging(LiteLLMLoggingBaseClass): if self.optional_params else None ), + "data_residency": ( + self.litellm_params.get("data_residency") + if hasattr(self, "litellm_params") and self.litellm_params + else None + ), } except Exception as e: # error creating kwargs for cost calculation debug_info = StandardLoggingModelCostFailureDebugInformation( @@ -3905,31 +3910,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 endpoint=arize_phoenix_config.endpoint, headers=arize_phoenix_config.otlp_auth_headers, ) - if arize_phoenix_config.project_name: - existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") - # Add openinference.project.name attribute - if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" - ) - else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={arize_phoenix_config.project_name}" - ) - - # Set Phoenix project name from environment variable - phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) - if phoenix_project_name: - existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") - # Add openinference.project.name attribute - if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={phoenix_project_name}" - ) - else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={phoenix_project_name}" - ) # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 59d0465e6d4..882561ed2e8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -9,6 +9,7 @@ from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, CompletionTokensDetailsWrapper, + DataResidency, ImageResponse, ModelInfo, PassthroughCallTypes, @@ -29,6 +30,9 @@ _IMAGE_RESPONSE_CALL_TYPES = frozenset( } ) +# Pre-resolved DataResidency enum values for fast membership checks +_VALID_DATA_RESIDENCIES = frozenset(r.value for r in DataResidency) + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: @@ -617,11 +621,46 @@ def _calculate_input_cost( return prompt_cost +def _get_regional_uplift_multiplier( + model_info: ModelInfo, data_residency: Optional[str] +) -> float: + """ + Resolve the per-model regional-processing uplift multiplier for a given + data-residency region. + + OpenAI applies a flat percentage uplift (e.g. +10%) on all token costs for + requests served from a regionalized hostname (eu./us.api.openai.com). The + multiplier is stored on the model entry as + ``regional_processing_uplift_multiplier_`` (e.g. 1.10). + + Returns 1.0 (no uplift) when ``data_residency`` is ``None`` or when the + model has no multiplier configured for the given region. + """ + if data_residency is None: + return 1.0 + residency = data_residency.lower() + if residency not in _VALID_DATA_RESIDENCIES: + return 1.0 + multiplier = model_info.get(f"regional_processing_uplift_multiplier_{residency}") + if multiplier is None: + return 1.0 + try: + return float(cast(float, multiplier)) + except (TypeError, ValueError): + verbose_logger.exception( + "Invalid regional_processing_uplift_multiplier_%s for model; " + "defaulting to 1.0", + residency, + ) + return 1.0 + + def generic_cost_per_token( # noqa: PLR0915 model: str, usage: Usage, custom_llm_provider: str, service_tier: Optional[str] = None, + data_residency: Optional[str] = None, ) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -631,6 +670,8 @@ def generic_cost_per_token( # noqa: PLR0915 Input: - model: str, the model name without provider prefix - usage: LiteLLM Usage block, containing anthropic caching information + - data_residency: optional OpenAI data-residency region (e.g. "eu", "us"), + used to apply the per-model regional-processing uplift multiplier. Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -781,6 +822,14 @@ def generic_cost_per_token( # noqa: PLR0915 ) completion_cost += float(image_tokens) * _output_cost_per_image_token + ## REGIONAL DATA-RESIDENCY UPLIFT + # Applied as a flat multiplier across all token costs for the request + # when the upstream is a regionalized OpenAI host (eu./us.api.openai.com). + uplift = _get_regional_uplift_multiplier(model_info, data_residency) + if uplift != 1.0: + prompt_cost *= uplift + completion_cost *= uplift + return prompt_cost, completion_cost diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index f169f86079a..03341ce0c9b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3997,7 +3997,7 @@ def _convert_to_bedrock_tool_call_invoke( for tool in tool_calls: if "function" in tool: tool_id = tool["id"] - name = tool["function"].get("name", "") + name = make_valid_bedrock_tool_name(tool["function"].get("name", "")) arguments = tool["function"].get("arguments", "") if not arguments or not arguments.strip(): @@ -5323,16 +5323,10 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 def make_valid_bedrock_tool_name(input_tool_name: str) -> str: - """ - Replaces any invalid characters in the input tool name with underscores - and ensures the resulting string is a valid identifier for Bedrock tools - """ + """Normalize tool names to Bedrock pattern [a-zA-Z][a-zA-Z0-9_-]*.""" def replace_invalid(char): - """ - Bedrock tool names only supports alpha-numeric characters and underscores - """ - if char.isalnum() or char == "_": + if char.isalnum() or char in ("_", "-"): return char return "_" @@ -5492,7 +5486,7 @@ def _bedrock_tools_pt( raw_name = f"litellm_unnamed_tool_{tool_idx}" # related issue: https://github.com/BerriAI/litellm/issues/5007 - # Bedrock tool names must satisfy regular expression pattern: [a-zA-Z][a-zA-Z0-9_]* ensure this is true + # Bedrock tool names must satisfy pattern: [a-zA-Z][a-zA-Z0-9_-]* name = make_valid_bedrock_tool_name(input_tool_name=raw_name) if _tool_description: # bedrock doesn't accept empty "" or None descriptions description = _tool_description diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c4528ff74e3..33bb6d7d2ea 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -86,6 +86,12 @@ class RealTimeStreaming: # When a text message is blocked, hold the guardrail reason so the next # response.create can be rewritten to include the failure context. self._pending_guardrail_message: Optional[str] = None + # Track whether session.created has already been sent to the client + # (e.g. synthetic event in deferred setup mode). + self._session_created_sent_to_client: bool = False + # Track whether we have already sent the guardrail turn-detection update + # that disables provider auto-response for transcription guardrails. + self._guardrail_turn_detection_update_sent: bool = False _SESSION_EVENT_TYPES = frozenset(["session.created", "session.updated"]) _AUDIO_FORMAT_MAP: Dict[str, Dict[str, Any]] = { @@ -248,40 +254,82 @@ class RealTimeStreaming: ## SYNC LOGGING executor.submit(self.logging_obj.success_handler(self.messages)) - async def _send_to_backend(self, message: str) -> None: + async def _send_to_backend(self, message: str) -> bool: """Send a message to the backend WebSocket. If a provider_config is set the message is first passed through transform_realtime_request so that provider-specific translation (e.g. dropping session.update for Vertex AI) is applied even for guardrail-injected messages. + + Returns True if at least one message was actually delivered to the + backend, False if the provider transformation produced no output and + the message was effectively dropped. """ if self.provider_config: transformed = self.provider_config.transform_realtime_request( message, self.model, self.session_configuration_request ) + sent = False for msg in transformed: + # Send first; only cache the setup payload once the backend + # has actually accepted it. Caching before send would leave + # ``session_configuration_request`` populated after a failed + # send, causing subsequent client session.update messages to + # be treated as "subsequent" and dropped even though the + # backend never received the original setup. await self.backend_ws.send(msg) # type: ignore[union-attr, attr-defined] - else: - await self.backend_ws.send(message) # type: ignore[union-attr, attr-defined] + self._cache_session_configuration_request(msg) + sent = True + return sent + await self.backend_ws.send(message) # type: ignore[union-attr, attr-defined] + return True + + def _cache_session_configuration_request(self, transformed_message: str) -> None: + """Store setup payload once sent to backend. + + Updates the cached setup on every successful setup send so follow-up + ``session.update`` messages (which produce a merged setup with new + ``generationConfig`` / ``systemInstruction`` / etc.) are reflected in + the cache used by downstream readers (``transform_session_created_event``, + ``return_new_content_delta_events`` modality lookup, ...). + """ + try: + message_obj = json.loads(transformed_message) + if "setup" in message_obj: + self.session_configuration_request = transformed_message + except (json.JSONDecodeError, TypeError): + return def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" + turn_detection: Dict[str, Any] = { + "type": "server_vad", + "create_response": False, + } if self._backend_uses_beta_protocol: - session: Dict[str, Any] = { - "turn_detection": {"create_response": False}, - } + session: Dict[str, Any] = {"turn_detection": turn_detection} else: session = { "type": "realtime", - "audio": { - "input": { - "turn_detection": {"create_response": False}, - } - }, + "audio": {"input": {"turn_detection": turn_detection}}, } return json.dumps({"type": "session.update", "session": session}) + async def _maybe_send_guardrail_turn_detection_update(self) -> None: + """Disable provider auto-response once when transcription guardrails are enabled.""" + if self._guardrail_turn_detection_update_sent: + return + if not self._has_audio_transcription_guardrails(): + return + sent = await self._send_to_backend(self._make_disable_auto_response_message()) + # Only mark as sent when the provider transformation actually delivered + # the update to the backend. Otherwise (e.g. Gemini drops session.update + # after the initial setup), leave the flag unset so future opportunities + # — such as a duplicate session.created — can retry. + if sent: + self._guardrail_turn_detection_update_sent = True + def _has_realtime_guardrails(self) -> bool: """Return True if any callback is registered for realtime guardrail event types.""" from litellm.integrations.custom_guardrail import CustomGuardrail @@ -320,12 +368,20 @@ class RealTimeStreaming: self, transcript: str, item_id: Optional[str] = None, + pre_block_backend_message: Optional[str] = None, ) -> bool: """ Run registered guardrails on a completed speech transcription. Returns True if blocked (synthetic warning already sent to client). Returns False if clean (caller should send response.create to the backend). + + ``pre_block_backend_message`` (if provided) is sent to the backend + BEFORE any of the guardrail's own backend messages when a block is + triggered. This is needed for protocol contracts that require a + specific message to be sent first — e.g. Gemini Live requires a + matching ``toolResponse`` immediately after a ``toolCall`` before any + other client messages can be accepted. """ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks @@ -385,6 +441,13 @@ class RealTimeStreaming: getattr(callback, "realtime_violation_message", None) or safe_msg ) + # Deliver any caller-supplied backend message FIRST so that + # protocol contracts requiring a specific ordering (e.g. + # Gemini Live's mandatory ``toolResponse`` after a + # ``toolCall``) are honored before the guardrail's own + # clientContent / cancel messages are sent. + if pre_block_backend_message is not None: + await self._send_to_backend(pre_block_backend_message) # Cancel any in-progress LLM response (e.g. VAD auto-response). await self._send_to_backend(json.dumps({"type": "response.cancel"})) # Send the policy violation hint (shows as small gray status text in UI). @@ -480,16 +543,34 @@ class RealTimeStreaming: else [transformed_response] ) for event in events: + is_session_created_event = ( + isinstance(event, dict) and event.get("type") == "session.created" + ) + if is_session_created_event: + if self._session_created_sent_to_client: + # A synthetic session.created (with placeholder defaults) was + # already forwarded to the client when we connected. The + # provider's real session.created (e.g. emitted from Gemini + # `setupComplete`) carries the authoritative modalities/model + # from the client's session.update. Re-emit it as + # `session.updated` so the client learns the corrected + # configuration without seeing two `session.created` events. + event = {**event, "type": "session.updated"} + else: + self._session_created_sent_to_client = True event_str = json.dumps(event) - ## For audio/VAD guardrail path: forward session.created first, then inject. - if ( - isinstance(event, dict) - and event.get("type") == "session.created" - and self._has_audio_transcription_guardrails() - ): + ## For audio/VAD guardrail path: forward the (possibly retyped) + ## session.created first, then invoke the one-time guardrail + ## turn-detection update. ``_maybe_send_guardrail_turn_detection_update`` + ## is idempotent (gated by ``_guardrail_turn_detection_update_sent``), + ## so duplicate session.created events — including those emitted + ## after a synthetic session.created from ``llm_http_handler`` in + ## deferred-setup mode — still get a single chance to inject the + ## update if a prior attempt was dropped by the provider transform. + if is_session_created_event and self._has_audio_transcription_guardrails(): self.store_message(event_str) await self.websocket.send_text(event_str) - await self._send_to_backend(self._make_disable_auto_response_message()) + await self._maybe_send_guardrail_turn_detection_update() continue ## GUARDRAIL: run on transcription events in provider_config path too if ( @@ -564,10 +645,19 @@ class RealTimeStreaming: try: raw_response = await self.backend_ws.recv( # type: ignore[union-attr] decode=False - ) # improves performance + ) except TypeError: raw_response = await self.backend_ws.recv() # type: ignore[union-attr, assignment] + if isinstance(raw_response, bytes): + try: + raw_response = raw_response.decode("utf-8") + except UnicodeDecodeError: + verbose_logger.warning( + "Received non-UTF-8 binary frame from backend, skipping." + ) + continue + if self.provider_config: try: await self._handle_provider_config_message(raw_response) @@ -783,12 +873,13 @@ class RealTimeStreaming: item["content"] = new_content return item - async def client_ack_messages(self): + async def client_ack_messages(self): # noqa: PLR0915 try: while True: message = await self.websocket.receive_text() ## GUARDRAIL: intercept conversation.item.create for text-based injection. + guardrail_turn_detection_injected = False try: msg_obj = json.loads(message) msg_type = msg_obj.get("type") @@ -796,7 +887,68 @@ class RealTimeStreaming: if msg_type == "conversation.item.create": # Check user text messages for prompt injection item = msg_obj.get("item", {}) - if item.get("role") == "user": + # Check function_call_output first so a client cannot + # bypass the tool-result guardrail by also setting + # role="user" on a function_call_output item. + if item.get("type") == "function_call_output": + # Tool results are client-controlled and fed to the + # model; check them with the same guardrail used for + # user text so an attacker cannot smuggle blocked + # content into a function_call_output. + output = item.get("output", "") + output_text = ( + output + if isinstance(output, str) + else json.dumps(output) + ) + if output_text: + # Build the sanitized function_call_output up + # front so we can hand it to the guardrail + # runner as the pre-block message. Providers + # that pair every toolCall with a toolResponse + # (e.g. Gemini/Vertex Live) require the + # toolResponse to arrive BEFORE any other + # client message — otherwise the guardrail's + # own clientContent would violate the + # pending-tool-call protocol contract and the + # backend could close the connection before + # the sanitized response ever lands. Dropping + # the blocked item outright would similarly + # leave such providers waiting indefinitely. + # The sanitized payload carries no blocked + # content — only a generic policy marker. + sanitized_msg = json.dumps( + { + **msg_obj, + "item": { + **item, + "output": json.dumps( + { + "error": "Tool output blocked by content policy", + } + ), + }, + } + ) + blocked = await self.run_realtime_guardrails( + output_text, + pre_block_backend_message=sanitized_msg, + ) + if blocked: + # ``_pending_guardrail_message`` is + # intentionally NOT set here. That flag + # exists to swallow the reflexive + # ``response.create`` an OpenAI client + # sends immediately after a user text + # message. In a tool-calling flow the + # client may not send a ``response.create`` + # at all (e.g. Gemini SDKs auto-respond), + # so leaving the flag set would + # incorrectly drop an unrelated + # ``response.create`` from a later + # interaction turn. + continue + elif item.get("role") == "user": content_list = item.get("content", []) texts = [ c.get("text", "") @@ -824,6 +976,89 @@ class RealTimeStreaming: self._pending_guardrail_message = None continue + ## GUARDRAIL: Inject turn_detection into first session.update + # if needed. Done BEFORE the GA remap so the injected + # ``create_response`` rides along with any client-provided + # turn_detection fields (e.g. silence_duration_ms) into the + # nested ``audio.input.turn_detection`` path produced by the + # remap. Doing this after the remap would create a separate + # minimal root-level ``turn_detection`` and silently drop + # the client's nested settings. + if ( + msg_type == "session.update" + and self.session_configuration_request is None + and not self._guardrail_turn_detection_update_sent + and self._has_audio_transcription_guardrails() + ): + session = msg_obj.setdefault("session", {}) + if isinstance(session, dict): + existing_td = session.get("turn_detection") + if not isinstance(existing_td, dict): + existing_td = {} + existing_td["create_response"] = False + session["turn_detection"] = existing_td + message = json.dumps(msg_obj) + guardrail_turn_detection_injected = True + verbose_logger.debug( + "Injected turn_detection into first session.update for audio transcription guardrails" + ) + + ## GUARDRAIL: Force ``create_response`` to False in any + # client-provided ``turn_detection`` so a later + # ``session.update`` cannot re-enable VAD auto-response + # and bypass the transcription guardrail after the + # initial disable. Covers both the flat beta key and the + # nested GA ``audio.input.turn_detection`` shape, since + # the GA remap below also accepts either form. Skipped + # when the injection block above already ran for this + # message, to avoid redundant double-serialization. + if ( + msg_type == "session.update" + and not guardrail_turn_detection_injected + and self._has_audio_transcription_guardrails() + ): + session = msg_obj.get("session") + if isinstance(session, dict): + td_overridden = False + flat_td = session.get("turn_detection") + flat_td_present = flat_td is not None + if flat_td_present: + if not isinstance(flat_td, dict): + flat_td = {} + if flat_td.get("create_response") is not False: + flat_td["create_response"] = False + session["turn_detection"] = flat_td + td_overridden = True + nested_td_present = False + audio = session.get("audio") + if isinstance(audio, dict): + audio_input = audio.get("input") + if isinstance(audio_input, dict): + nested_td = audio_input.get("turn_detection") + if nested_td is not None: + nested_td_present = True + if not isinstance(nested_td, dict): + nested_td = {} + if ( + nested_td.get("create_response") + is not False + ): + nested_td["create_response"] = False + audio_input["turn_detection"] = nested_td + td_overridden = True + # Symmetric with the first-update injection block: + # if the client omitted turn_detection entirely on + # a subsequent session.update, still inject the + # ``create_response: False`` override so the + # transcription guardrail cannot be re-enabled by + # any downstream merge that drops the original + # disable. + if not flat_td_present and not nested_td_present: + session["turn_detection"] = {"create_response": False} + td_overridden = True + if td_overridden: + message = json.dumps(msg_obj) + # GA compatibility: remap beta-style session fields only when # the upstream is in GA mode. Beta upstreams expect the flat # session shape unchanged. @@ -841,17 +1076,20 @@ class RealTimeStreaming: pass ## LOGGING + # Log after any in-place modifications (GA remap, guardrail + # turn_detection injection) so audit logs reflect what we + # actually forward to the backend. self.store_input(message=message) - ## FORWARD TO BACKEND - if self.provider_config: - message = self.provider_config.transform_realtime_request( - message, self.model - ) - for msg in message: - await self.backend_ws.send(msg) # type: ignore[union-attr] - else: - await self.backend_ws.send(message) # type: ignore[union-attr] + ## FORWARD TO BACKEND + # Only mark the guardrail turn_detection update as sent after the + # backend actually accepted the message. Setting the flag earlier + # would permanently disable the injection if ``_send_to_backend`` + # raised — neither this loop nor + # ``_maybe_send_guardrail_turn_detection_update`` would retry. + sent = await self._send_to_backend(message) + if guardrail_turn_detection_injected and sent: + self._guardrail_turn_detection_update_sent = True except Exception as e: verbose_logger.debug(f"Error in client ack messages: {e}") diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index d7803455b4a..4928dd08386 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -146,6 +146,37 @@ class SensitiveDataMasker: return masked_data +_default_masker = SensitiveDataMasker() + + +def mask_sensitive_keys( + data: Dict[str, Any], sensitive_fields: Set[str] +) -> Dict[str, Any]: + """Return a new dict with values masked for keys listed in ``sensitive_fields``. + + Unlike :meth:`SensitiveDataMasker.mask_dict`, this does exact key-name + matching (not segment matching), so callers explicitly enumerate which + fields to mask. Non-string and None values are passed through unchanged. + + Values shorter than ``visible_prefix + visible_suffix`` (8 by default) + fall outside :meth:`SensitiveDataMasker._mask_value`'s partial-reveal + range and are replaced with a fixed-length all-mask string, so a short + credential is never returned verbatim. + """ + masked: Dict[str, Any] = {} + mask_char = _default_masker.mask_char + min_visible = _default_masker.visible_prefix + _default_masker.visible_suffix + for key, value in data.items(): + if value is not None and key in sensitive_fields and isinstance(value, str): + if len(value) < min_visible: + masked[key] = mask_char * len(value) if value else value + else: + masked[key] = _default_masker._mask_value(value) + else: + masked[key] = value + return masked + + # Usage example: """ masker = SensitiveDataMasker() diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fa7faf3035d..4642201ca67 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -59,6 +59,8 @@ FUNCTION_CALL_ATTRIBUTE = "function_call" _SYNC_ITER_EXHAUSTED = object() +_GCHUNK_FIELDS: frozenset = frozenset(GChunk.__annotations__) + def _next_sync_or_exhausted(it: Any) -> Any: """ @@ -181,6 +183,30 @@ class CustomStreamWrapper: self.created: Optional[int] = None self._last_returned_hidden_params: Optional[dict] = None + _cached_logging_provider = self.logging_obj.model_call_details.get( + "custom_llm_provider", None + ) + self._cached_logging_llm_provider: Optional[str] = _cached_logging_provider + _effective_model = model or "" + if ( + custom_llm_provider == "openai" + and custom_llm_provider != _cached_logging_provider + ): + _effective_model = "{}/{}".format( + _cached_logging_provider, _effective_model + ) + self._cached_model_name: str = _effective_model + + # Snapshot assumes self._hidden_params is populated from litellm_params + # at init and never mutated during the stream. If that ever changes, + # this cache must be removed. + self._base_hidden_params: Dict[str, Any] = { + **self._hidden_params, + "response_cost": None, + } + + self._post_streaming_hooks: Optional[List] = None + def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS @@ -681,29 +707,16 @@ class CustomStreamWrapper: def model_response_creator( self, chunk: Optional[dict] = None, hidden_params: Optional[dict] = None ): - _model = self.model - _received_llm_provider = self.custom_llm_provider - _logging_obj_llm_provider = self.logging_obj.model_call_details.get("custom_llm_provider", None) # type: ignore - if ( - _received_llm_provider == "openai" - and _received_llm_provider != _logging_obj_llm_provider - ): - _model = "{}/{}".format(_logging_obj_llm_provider, _model) + _model = self._cached_model_name + _logging_obj_llm_provider = self._cached_logging_llm_provider + if chunk is None: - chunk = {} + args: Dict[str, Any] = {"model": _model} else: - # pop model keyword chunk.pop("model", None) - - chunk_dict = {} - for key, value in chunk.items(): - if key != "stream": - chunk_dict[key] = value - - args = { - "model": _model, - **chunk_dict, - } + args = {"model": _model} + if chunk: + args.update({k: v for k, v in chunk.items() if k != "stream"}) model_response = ModelResponseStream(**args) if self.response_id is not None: @@ -717,15 +730,23 @@ class CustomStreamWrapper: model_response.created = self.created else: self.created = model_response.created + + # Spread order is load-bearing: _base_hidden_params (model_id, api_base, ...) + # must win over both caller-supplied hidden_params and the computed + # custom_llm_provider/created_at values, so it comes last. if hidden_params is not None: - model_response._hidden_params = hidden_params - model_response._hidden_params["custom_llm_provider"] = _logging_obj_llm_provider - model_response._hidden_params["created_at"] = time.time() - model_response._hidden_params = { - **model_response._hidden_params, - **self._hidden_params, - "response_cost": None, - } + model_response._hidden_params = { + **hidden_params, + "custom_llm_provider": _logging_obj_llm_provider, + "created_at": time.time(), + **self._base_hidden_params, + } + else: + model_response._hidden_params = { + "custom_llm_provider": _logging_obj_llm_provider, + "created_at": time.time(), + **self._base_hidden_params, + } if ( len(model_response.choices) > 0 @@ -1627,7 +1648,17 @@ class CustomStreamWrapper: from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import CallTypes - # Get request kwargs from logging object + if self._post_streaming_hooks is None: + self._post_streaming_hooks = [ + cb + for cb in litellm.callbacks + if isinstance(cb, CustomLogger) + and hasattr(cb, "async_post_call_streaming_deployment_hook") + ] + + if not self._post_streaming_hooks: + return chunk + request_data = self.logging_obj.model_call_details call_type_str = self.logging_obj.call_type @@ -1636,18 +1667,14 @@ class CustomStreamWrapper: except ValueError: typed_call_type = None - # Call hooks for all callbacks - for callback in litellm.callbacks: - if isinstance(callback, CustomLogger) and hasattr( - callback, "async_post_call_streaming_deployment_hook" - ): - result = await callback.async_post_call_streaming_deployment_hook( - request_data=request_data, - response_chunk=chunk, - call_type=typed_call_type, - ) - if result is not None: - chunk = result + for callback in self._post_streaming_hooks: + result = await callback.async_post_call_streaming_deployment_hook( + request_data=request_data, + response_chunk=chunk, + call_type=typed_call_type, + ) + if result is not None: + chunk = result return chunk except Exception as e: @@ -1888,17 +1915,15 @@ class CustomStreamWrapper: response = self._add_mcp_list_tools_to_first_chunk(response) self.sent_first_chunk = True - if hasattr( - response, "usage" - ): # remove usage from chunk, only send on final chunk - # Convert the object to a dictionary + # ModelResponseStream declares `usage` as a field, so + # hasattr(response, "usage") is always True — must check + # `is not None` to avoid running this path on every chunk. + if getattr(response, "usage", None) is not None: obj_dict = response.model_dump() - # Remove an attribute (e.g., 'attr2') if "usage" in obj_dict: del obj_dict["usage"] - # Create a new object without the removed attribute response = self.model_response_creator( chunk=obj_dict, hidden_params=response._hidden_params ) @@ -2398,10 +2423,7 @@ def generic_chunk_has_all_required_fields(chunk: dict) -> bool: :param chunk: The dictionary to check. :return: True if all required fields are present, False otherwise. """ - _all_fields = GChunk.__annotations__ - - decision = all(key in _all_fields for key in chunk) - return decision + return all(key in _GCHUNK_FIELDS for key in chunk) def convert_generic_chunk_to_model_response_stream( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0b56eb86d9c..57609cfcd26 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -337,13 +337,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _supports_effort_level(model: str, level: str) -> bool: - """Check ``supports_{level}_reasoning_effort`` in the model map. + def _supports_model_capability(model: str, key: str) -> bool: + """Check a boolean capability ``key`` in the model map. Strips bedrock/vertex prefixes so a provider-routed Claude still resolves to the Anthropic model-map entry. """ - key = f"supports_{level}_reasoning_effort" try: if _supports_factory( model=model, @@ -372,8 +371,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): except Exception: pass try: - import litellm - for cand in candidates: if cand in litellm.model_cost and ( litellm.model_cost[cand].get(key) is True @@ -383,6 +380,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): pass return False + @staticmethod + def _supports_effort_level(model: str, level: str) -> bool: + """Check ``supports_{level}_reasoning_effort`` in the model map.""" + return AnthropicConfig._supports_model_capability( + model, f"supports_{level}_reasoning_effort" + ) + @staticmethod def _validate_effort_for_model(model: str, effort: Optional[str]) -> Optional[str]: """Return ``None`` if ``effort`` is allowed on ``model``, else an error message.""" @@ -400,7 +404,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): @staticmethod def _model_supports_effort_param(model: str) -> bool: - """Whether the model accepts ``output_config.effort`` at all.""" + """Whether the model accepts ``output_config.effort`` at all. + + A model qualifies if its map entry advertises ``supports_output_config`` + or any ``supports_*_reasoning_effort`` flag. The two are independent + signals: e.g. Claude Opus 4.5 supports ``output_config`` without + advertising a non-default (max/xhigh) effort level. + """ + if AnthropicConfig._supports_model_capability(model, "supports_output_config"): + return True return any( AnthropicConfig._supports_effort_level(model, level) for level in ("low", "minimal", "medium", "high", "xhigh", "max") @@ -1793,7 +1805,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self._ensure_context_management_beta_header( headers, optional_params["context_management"] ) - if optional_params.get("output_format") is not None: + output_config = optional_params.get("output_config") + if optional_params.get("output_format") is not None or ( + isinstance(output_config, dict) and output_config.get("format") is not None + ): self._ensure_beta_header( headers, ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value ) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 15f404d3f53..f94232fa451 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -427,8 +427,13 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value ) - # Check for structured outputs - if optional_params.get("output_format") is not None: + # Check for structured outputs. Anthropic's newer request shape nests + # the schema under output_config.format; the older top-level + # output_format remains supported for backwards compatibility. + output_config = optional_params.get("output_config") + if optional_params.get("output_format") is not None or ( + isinstance(output_config, dict) and output_config.get("format") is not None + ): beta_values.add( ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value ) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 4fc1ae960b8..e1ac1858912 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,3 +1,5 @@ +import asyncio +import hashlib import json import os from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast @@ -449,6 +451,25 @@ class BaseAzureLLM(BaseOpenAILLM): ] = None client_initialization_params: dict = locals() client_initialization_params["is_async"] = _is_async + _lp = litellm_params or {} + _ad_provider = _lp.get("azure_ad_token_provider") + _ad_token = _lp.get("azure_ad_token") + _client_secret = _lp.get("client_secret") + _azure_password = _lp.get("azure_password") + client_initialization_params["azure_ad_token"] = ( + hashlib.sha256(_ad_token.encode()).hexdigest() + if isinstance(_ad_token, str) + else None + ) + client_initialization_params["azure_ad_token_provider"] = ( + f"provider_id={id(_ad_provider) if callable(_ad_provider) else None}" + f"|tenant_id={_lp.get('tenant_id')}" + f"|client_id={_lp.get('client_id')}" + f"|client_secret={hashlib.sha256(_client_secret.encode()).hexdigest() if isinstance(_client_secret, str) else None}" + f"|azure_username={_lp.get('azure_username')}" + f"|azure_password={hashlib.sha256(_azure_password.encode()).hexdigest() if isinstance(_azure_password, str) else None}" + f"|azure_scope={_lp.get('azure_scope')}" + ) if client is None: cached_client = self.get_cached_openai_client( client_initialization_params=client_initialization_params, @@ -474,8 +495,29 @@ class BaseAzureLLM(BaseOpenAILLM): if self._is_azure_v1_api_version(api_version): # Extract only params that OpenAI client accepts # Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview" - v1_params = { - "api_key": azure_client_params.get("api_key"), + # The OpenAI client accepts a callable for `api_key` and re-invokes it + # on every request (via `_refresh_api_key`), so passing + # `azure_ad_token_provider` directly preserves Azure AD token refresh + # behavior that the regular AzureOpenAI client provides. + v1_api_key: Optional[Union[str, Callable[[], Any]]] = ( + azure_client_params.get("api_key") + or azure_client_params.get("azure_ad_token_provider") + or azure_client_params.get("azure_ad_token") + ) + if _is_async is True and callable(v1_api_key): + # AsyncOpenAI expects an async provider; wrap the sync provider + # returned by azure-identity. Offload to a thread so a token + # refresh (blocking HTTP call to AAD on cache miss) does not + # stall the event loop. + _sync_provider = v1_api_key + + async def _async_v1_api_key() -> str: + return await asyncio.to_thread(_sync_provider) + + v1_api_key = _async_v1_api_key + + v1_params: Dict[str, Any] = { + "api_key": v1_api_key, "base_url": f"{api_base}/openai/v1/", } if "timeout" in azure_client_params: diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index bec25916c4b..5f35a58ce1f 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -108,10 +108,9 @@ class BaseConfig(ABC): return type_to_response_format_param(response_format=response_format) def is_thinking_enabled(self, non_default_params: dict) -> bool: - return ( - non_default_params.get("thinking", {}).get("type") == "enabled" - or non_default_params.get("reasoning_effort") is not None - ) + return (non_default_params.get("thinking") or {}).get( + "type" + ) == "enabled" or non_default_params.get("reasoning_effort") is not None def is_max_tokens_in_request(self, non_default_params: dict) -> bool: """ diff --git a/litellm/llms/base_llm/managed_resources/utils.py b/litellm/llms/base_llm/managed_resources/utils.py index 59f5ff0d845..6e30b6cb252 100644 --- a/litellm/llms/base_llm/managed_resources/utils.py +++ b/litellm/llms/base_llm/managed_resources/utils.py @@ -177,8 +177,14 @@ def extract_model_id_from_unified_id( if decoded_id: unified_id = decoded_id - # Extract model ID - match = re.search(r"model_id,([^;]+)", unified_id) + # Extract model ID. Anchor to a field boundary (start of string or + # after `;`) so this regex doesn't substring-match the `model_id,` + # inside file_id encodings' `llm_output_file_model_id,` + # field — that would feed the deployment UUID as a model candidate + # into the team-access check and 403 every team-BYOK file attach + # with `Tried to access ` (LIT-3244 patch/1.86.0 second-order + # finding). + match = re.search(r"(?:^|;)model_id,([^;]+)", unified_id) if match: return match.group(1).strip() diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index d5531a532b9..0f239b4ad45 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Union import httpx +from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents from litellm.types.realtime import ( RealtimeResponseTransformInput, RealtimeResponseTypedDict, @@ -69,6 +70,20 @@ class BaseRealtimeConfig(ABC): ) -> Optional[str]: # message sent to setup the realtime session return None + def transform_session_created_event( + self, + model: str, + logging_session_id: str, + session_configuration_request: Optional[str] = None, + ) -> Optional[Union[dict, OpenAIRealtimeStreamSessionEvents]]: + """ + Optional hook for providers that defer session setup until client `session.update`. + + Return an OpenAI-compatible `session.created` payload when the proxy should + emit a synthetic event immediately after backend websocket connection. + """ + return None + @abstractmethod def transform_realtime_response( self, diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index 87289ad6a0c..9b4cf777280 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -321,6 +321,23 @@ class BaseVideoConfig(ABC): "video get character is not supported for this provider" ) + def get_video_edit_prefetch_params( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Optional[Tuple[str, Dict]]: + """ + Return (url, body) for a pre-fetch HTTP call that must be made before + transform_video_edit_request, or None if no pre-fetch is required. + + Providers that need to retrieve the source video before constructing the + edit request (e.g. Vertex AI) should override this method. The handler + uses the existing shared httpx client so the call is properly async. + """ + return None + def transform_video_edit_request( self, prompt: str, @@ -329,6 +346,7 @@ class BaseVideoConfig(ABC): litellm_params: GenericLiteLLMParams, headers: dict, extra_body: Optional[Dict[str, Any]] = None, + prefetched_source_data: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Transform the video edit request into a URL and JSON data. @@ -343,6 +361,7 @@ class BaseVideoConfig(ABC): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict] = None, ) -> VideoObject: raise NotImplementedError("video edit is not supported for this provider") diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 44ba1ce3c86..9b9b96aae04 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -157,8 +157,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): def _get_agent_runtime_arn(self, model: str) -> str: """ Extract ARN from model string - model = "agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" - returns: "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" + model = "agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp" + returns: "arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp" """ parts = model.split("/", 1) if len(parts) != 2 or parts[0] != "agentcore": @@ -170,7 +170,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): def _extract_region_from_arn(self, arn: str) -> str: """ Extract region from ARN - arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC + arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp returns: us-west-2 """ parts = arn.split(":") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index efc890d9ee2..d58d2e27595 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -30,6 +30,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, _bedrock_tools_pt, + make_valid_bedrock_tool_name, ) from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, @@ -77,6 +78,7 @@ from ..common_utils import ( get_anthropic_beta_from_headers, get_bedrock_tool_name, is_claude_4_5_on_bedrock, + normalize_bedrock_opus_output_config_effort, ) # Computer use tool prefixes supported by Bedrock @@ -447,10 +449,20 @@ class AmazonConverseConfig(BaseConfig): value=reasoning_effort, llm_provider="bedrock_converse", ) + existing_output_config = optional_params.get("output_config") + if not isinstance(existing_output_config, dict): + existing_output_config = {} + existing_output_config.setdefault("effort", mapped_effort) + normalize_bedrock_opus_output_config_effort( + model=model, + output_config=existing_output_config, + ) + mapped_effort = existing_output_config["effort"] self._validate_anthropic_adaptive_effort( model=model, effort=mapped_effort ) - optional_params["output_config"] = {"effort": mapped_effort} + optional_params["output_config"] = existing_output_config + optional_params["_output_config_normalized"] = True @staticmethod def _validate_anthropic_adaptive_effort(model: str, effort: str) -> None: @@ -595,7 +607,9 @@ class AmazonConverseConfig(BaseConfig): elif isinstance(tool_choice, dict): # only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html specific_tool = SpecificToolChoiceBlock( - name=tool_choice.get("function", {}).get("name", "") + name=make_valid_bedrock_tool_name( + tool_choice.get("function", {}).get("name", "") + ) ) return ToolChoiceValuesBlock(tool=specific_tool) else: @@ -1198,6 +1212,12 @@ class AmazonConverseConfig(BaseConfig): self, optional_params: dict, model: str ) -> Tuple[dict, dict, dict, Optional[OutputConfigBlock]]: """Prepare and separate request parameters.""" + # Consume the internal ``_output_config_normalized`` marker set by + # ``_handle_reasoning_effort_parameter`` so it does not linger on the + # caller's ``optional_params`` after the transformation returns. + anthropic_output_config_already_normalized = bool( + optional_params.pop("_output_config_normalized", False) + ) # Filter out exception objects before deepcopy to prevent deepcopy failures # Exceptions should not be stored in optional_params (this is a defensive fix) cleaned_params = filter_exceptions_from_params(optional_params) @@ -1216,8 +1236,17 @@ class AmazonConverseConfig(BaseConfig): # Anthropic-only ``output_config`` (snake_case) — re-attached to # ``additionalModelRequestFields`` for Anthropic models below. The - # Bedrock-native ``outputConfig`` (camelCase) is handled separately. + # structured-output ``format`` subfield is consumed into Bedrock's + # native ``outputConfig`` (camelCase), which is handled separately. anthropic_output_config = inference_params.pop("output_config", None) + output_config_format = None + if isinstance(anthropic_output_config, dict): + anthropic_output_config = dict(anthropic_output_config) + candidate_output_config_format = anthropic_output_config.pop("format", None) + if isinstance(candidate_output_config_format, dict): + output_config_format = candidate_output_config_format + if not anthropic_output_config: + anthropic_output_config = None # Extract requestMetadata before processing other parameters request_metadata = inference_params.pop("requestMetadata", None) @@ -1227,6 +1256,30 @@ class AmazonConverseConfig(BaseConfig): output_config: Optional[OutputConfigBlock] = inference_params.pop( "outputConfig", None ) + base_model = BedrockModelInfo.get_base_model(model) + if ( + output_config is None + and output_config_format is not None + and output_config_format.get("type") == "json_schema" + and base_model.startswith("anthropic") + and self._supports_native_structured_outputs( + model, self.custom_llm_provider + ) + ): + output_config = self._create_output_config_for_response_format( + json_schema=output_config_format.get("schema"), + name=output_config_format.get("name"), + description=output_config_format.get("description"), + ) + elif output_config is None and output_config_format is not None: + litellm.verbose_logger.warning( + "Bedrock Converse: dropping `output_config.format` for model=%s — " + "model does not advertise `supports_native_structured_output` in " + "model_prices_and_context_window.json. The schema will not be " + "enforced; pass `response_format` to use the synthetic tool-call " + "fallback.", + model, + ) # keep supported params in 'inference_params', and set all model-specific params in 'additional_request_params' additional_request_params = { @@ -1272,7 +1325,6 @@ class AmazonConverseConfig(BaseConfig): if anthropic_output_config is not None and isinstance( anthropic_output_config, dict ): - base_model = BedrockModelInfo.get_base_model(model) if base_model.startswith("anthropic"): if ( litellm.drop_params is True @@ -1283,6 +1335,11 @@ class AmazonConverseConfig(BaseConfig): model, ) else: + if not anthropic_output_config_already_normalized: + normalize_bedrock_opus_output_config_effort( + model=model, + output_config=anthropic_output_config, + ) effort = anthropic_output_config.get("effort") if effort is not None: self._validate_anthropic_adaptive_effort( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index d9599b8b9c4..a13336b6c88 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -16,8 +16,11 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + convert_bedrock_invoke_output_format_to_inline_schema, get_anthropic_beta_from_headers, + normalize_bedrock_opus_output_config_effort, normalize_tool_input_schema_types_for_bedrock_invoke, + pop_bedrock_invoke_output_config_format, remove_custom_field_from_tools, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER @@ -75,6 +78,17 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Use a model name that forces tool-based approach model = "claude-3-sonnet-20240229" + # Clamp ``reasoning_effort`` to the Bedrock effort ceiling before the + # parent mapping converts it to ``output_config.effort`` and the + # downstream effort gate runs. Mirrors the converse path's + # ``_handle_reasoning_effort_parameter`` and the messages path's + # ``_clamp_adaptive_reasoning_effort_for_bedrock`` so adaptive Claude + # requests degrade ``xhigh`` -> ``max`` rather than 400-ing on + # models like Opus 4.6 that don't natively advertise xhigh. + self._clamp_adaptive_reasoning_effort_for_bedrock( + model=original_model, params=non_default_params + ) + optional_params = AnthropicConfig.map_openai_params( self, non_default_params, @@ -88,6 +102,27 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): return optional_params + @staticmethod + def _clamp_adaptive_reasoning_effort_for_bedrock(model: str, params: dict) -> None: + """Lower ``reasoning_effort`` to the Bedrock effort ceiling before mapping. + + Bedrock's adaptive Claude models accept the OpenAI-style + ``reasoning_effort`` tier, but the request validator can reject tiers + the model does not natively advertise (e.g. ``xhigh`` on Opus 4.6). + Clamp the raw tier to the model's + ``bedrock_output_config_effort_ceiling`` so Claude Code "goal mode" + keeps working. Non-adaptive models and models without a ceiling are + left untouched. + """ + if not AnthropicConfig._is_adaptive_thinking_model(model): + return + effort = params.get("reasoning_effort") + if not isinstance(effort, str): + return + clamped = {"effort": effort} + normalize_bedrock_opus_output_config_effort(model=model, output_config=clamped) + params["reasoning_effort"] = clamped["effort"] + def transform_request( self, model: str, @@ -157,6 +192,13 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): for k, v in optional_params.items() if k not in self.aws_authentication_params } + output_config = filtered_params.get("output_config") + if isinstance(output_config, dict): + filtered_params["output_config"] = dict(output_config) + normalize_bedrock_opus_output_config_effort( + model=model, + output_config=filtered_params["output_config"], + ) filtered_params = self._normalize_bedrock_tool_search_tools(filtered_params) anthropic_request = AnthropicConfig.transform_request( @@ -170,7 +212,20 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) - anthropic_request.pop("output_format", None) + output_format = anthropic_request.pop("output_format", None) + output_config_format = pop_bedrock_invoke_output_config_format( + anthropic_request + ) + if output_format: + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=output_format, + request_body=anthropic_request, + ) + elif output_config_format: + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=output_config_format, + request_body=anthropic_request, + ) if not ( _supports_factory( model=model, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4f4729e4019..bdc5da321c6 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -34,6 +34,15 @@ class BedrockError(BaseLLMException): # Lazy import cache to avoid circular imports and performance impact _get_model_info = None +BedrockOutputConfigEffort = Literal["low", "medium", "high", "max", "xhigh"] +_BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER: Dict[BedrockOutputConfigEffort, int] = { + "low": 0, + "medium": 1, + "high": 2, + "max": 3, + "xhigh": 4, +} + def get_cached_model_info(): """ @@ -51,6 +60,79 @@ def get_cached_model_info(): return _get_model_info +@functools.lru_cache(maxsize=1) +def _get_local_model_cost_map() -> Dict: + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + + return GetModelCostMap.load_local_model_cost_map() + + +def pop_bedrock_invoke_output_config_format(request_body: Dict) -> Optional[Dict]: + """ + Remove and return Anthropic's nested ``output_config.format`` field. + + Bedrock Invoke paths convert the schema to inline message text. Any remaining + ``output_config`` keys, such as ``effort``, are left in place. + """ + output_config = request_body.get("output_config") + if not isinstance(output_config, dict): + return None + + output_format = output_config.pop("format", None) + if not output_config: + request_body.pop("output_config", None) + + if isinstance(output_format, dict): + return output_format + return None + + +def convert_bedrock_invoke_output_format_to_inline_schema( + output_format: Dict, + request_body: Dict, +) -> None: + """ + Embed an Anthropic structured-output schema into the last user message. + + Bedrock Invoke does not support ``output_format`` directly, so the schema is + appended to the final user message for prompt-engineered structured output. + The caller's ``messages`` list, message dict, and content list are not + mutated; a fresh ``messages`` list with a copied final user message is + written back to ``request_body``. + """ + schema = output_format.get("schema") + if not schema: + return + + messages = request_body.get("messages") + if not isinstance(messages, list) or not messages: + return + + last_user_idx = None + for i in range(len(messages) - 1, -1, -1): + message = messages[i] + if isinstance(message, dict) and message.get("role") == "user": + last_user_idx = i + break + + if last_user_idx is None: + return + + original = messages[last_user_idx] + content = original.get("content", []) + schema_block = {"type": "text", "text": json.dumps(schema)} + if isinstance(content, str): + new_content = [{"type": "text", "text": content}, schema_block] + elif isinstance(content, list): + new_content = [*content, schema_block] + else: + return + + new_messages = list(messages) + new_messages[last_user_idx] = {**original, "content": new_content} + request_body["messages"] = new_messages + + def remove_custom_field_from_tools(request_body: dict) -> None: """ Remove ``custom`` field from each tool in the request body. @@ -603,6 +685,62 @@ def is_claude_4_5_on_bedrock(model: str) -> bool: return any(pattern in model_lower for pattern in claude_4_5_patterns) +def normalize_bedrock_opus_output_config_effort(model: str, output_config: Any) -> None: + """ + Normalize Anthropic ``output_config.effort`` values for Bedrock Opus ids. + + Bedrock's Claude Opus request validator can accept a narrower effort + vocabulary than Anthropic's compatibility surface. The Bedrock ceiling is + read from ``model_prices_and_context_window.json`` via + ``bedrock_output_config_effort_ceiling``. + + Mutates ``output_config`` in place so callers can accept Claude Code's + ``xhigh`` input without forwarding a provider-invalid value. + """ + if not isinstance(output_config, dict): + return + + effort = output_config.get("effort") + if effort not in _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER: + return + + ceiling = _get_bedrock_output_config_effort_ceiling(model) + if ceiling is None: + return + + if ( + _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER[effort] + > _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER[ceiling] + ): + output_config["effort"] = ceiling + + +def _get_bedrock_output_config_effort_ceiling( + model: str, +) -> Optional[BedrockOutputConfigEffort]: + try: + model_info = get_cached_model_info()( + model=model, + custom_llm_provider="bedrock", + ) + except Exception: + return None + + ceiling = model_info.get("bedrock_output_config_effort_ceiling") + if isinstance(ceiling, str) and ceiling in _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER: + return ceiling # type: ignore[return-value] + + model_cost_key = model_info.get("key") + if not isinstance(model_cost_key, str): + return None + + local_model_info = _get_local_model_cost_map().get(model_cost_key, {}) + ceiling = local_model_info.get("bedrock_output_config_effort_ceiling") + if isinstance(ceiling, str) and ceiling in _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER: + return ceiling # type: ignore[return-value] + return None + + # Import after standalone functions to avoid circular imports from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 69b61298d33..b223f4534fa 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -32,10 +32,13 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + convert_bedrock_invoke_output_format_to_inline_schema, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, is_claude_4_5_on_bedrock, + normalize_bedrock_opus_output_config_effort, normalize_tool_input_schema_types_for_bedrock_invoke, + pop_bedrock_invoke_output_config_format, remove_custom_field_from_tools, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER @@ -450,145 +453,15 @@ class AmazonAnthropicClaudeMessagesConfig( else: anthropic_messages_request.pop("context_management", None) - def _convert_output_format_to_inline_schema( - self, - output_format: Dict, - anthropic_messages_request: Dict, - ) -> None: - """ - Convert Anthropic output_format to inline schema in message content. - - Bedrock Invoke doesn't support the output_format parameter, so we embed - the schema directly into the user message content as text instructions. - - This approach adds the schema to the last user message, instructing the model - to respond in the specified JSON format. - - Args: - output_format: The output_format dict with 'type' and 'schema' - anthropic_messages_request: The request dict to modify in-place - - Ref: https://aws.amazon.com/blogs/machine-learning/structured-data-response-with-amazon-bedrock-prompt-engineering-and-tool-use/ - """ - import json - - # Extract schema from output_format - schema = output_format.get("schema") - if not schema: - return - - # Get messages from the request - messages = anthropic_messages_request.get("messages", []) - if not messages: - return - - # Find the last user message - last_user_message_idx = None - for idx in range(len(messages) - 1, -1, -1): - if messages[idx].get("role") == "user": - last_user_message_idx = idx - break - - if last_user_message_idx is None: - return - - last_user_message = messages[last_user_message_idx] - content = last_user_message.get("content", []) - - # Ensure content is a list - if isinstance(content, str): - content = [{"type": "text", "text": content}] - last_user_message["content"] = content - - # Add schema as text content to the message - schema_text = {"type": "text", "text": json.dumps(schema)} - content.append(schema_text) - - def transform_anthropic_messages_request( + def _get_bedrock_invoke_anthropic_beta_headers( self, model: str, messages: List[Dict], anthropic_messages_optional_request_params: Dict, - litellm_params: GenericLiteLLMParams, headers: dict, - ) -> Dict: - anthropic_messages_request = AnthropicMessagesConfig.transform_anthropic_messages_request( - self=self, - model=model, - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params=litellm_params, - headers=headers, - ) - ######################################################### - ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### - ######################################################### - - # 1. anthropic_version is required for all claude models - if "anthropic_version" not in anthropic_messages_request: - anthropic_messages_request["anthropic_version"] = ( - self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION - ) - - # 2. `stream` is not allowed in request body for bedrock invoke - if "stream" in anthropic_messages_request: - anthropic_messages_request.pop("stream", None) - - # 3. `model` is not allowed in request body for bedrock invoke - if "model" in anthropic_messages_request: - anthropic_messages_request.pop("model", None) - - injected_thinking_for_clear_thinking = ( - self._ensure_thinking_for_clear_thinking_context_management( - anthropic_messages_request=anthropic_messages_request, - model=model, - ) - ) - - # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models) - self._remove_ttl_from_cache_control( - anthropic_messages_request=anthropic_messages_request, model=model - ) - - # 5. Convert `output_format` to inline schema (Bedrock invoke doesn't support output_format) - output_format = anthropic_messages_request.pop("output_format", None) - if output_format: - self._convert_output_format_to_inline_schema( - output_format=output_format, - anthropic_messages_request=anthropic_messages_request, - ) - - # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, - # but older models do not — strip it to avoid request rejection. - # Ref: https://github.com/BerriAI/litellm/issues/22797 - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model) - ): - if anthropic_messages_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) - - # 5b. Remove `custom` field from tools (Bedrock doesn't support it) - # Claude Code sends `custom: {defer_loading: true}` on tool definitions, - # which causes Bedrock to reject the request with "Extra inputs are not permitted" - # Ref: https://github.com/BerriAI/litellm/issues/22847 - remove_custom_field_from_tools(anthropic_messages_request) - normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_messages_request) - ensure_bedrock_anthropic_messages_tool_names(anthropic_messages_request) - - # 6. AUTO-INJECT beta headers based on features used + anthropic_messages_request: Dict, + injected_thinking_for_clear_thinking: bool, + ) -> List[str]: anthropic_model_info = AnthropicModelInfo() tools = anthropic_messages_optional_request_params.get("tools") messages_typed = cast(List[AllMessageValues], messages) @@ -651,6 +524,160 @@ class AmazonAnthropicClaudeMessagesConfig( dropped_user_betas, ) + return filtered_betas + + def _strip_unsupported_bedrock_invoke_fields( + self, + anthropic_messages_request: Dict, + ) -> Dict: + allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS + stripped = sorted(k for k in anthropic_messages_request if k not in allowed) + if stripped: + verbose_logger.debug( + "Bedrock Invoke: stripping unsupported top-level request fields: %s", + stripped, + ) + return {k: v for k, v in anthropic_messages_request.items() if k in allowed} + + @staticmethod + def _clamp_adaptive_reasoning_effort_for_bedrock( + model: str, optional_params: Dict + ) -> None: + """Lower ``reasoning_effort`` to the Bedrock effort ceiling before validation. + + The shared ``/v1/messages`` effort gate rejects tiers a model does not + natively support (e.g. ``xhigh`` on Opus 4.6). Bedrock's chat paths instead + clamp the tier to the model's ``bedrock_output_config_effort_ceiling`` so + Claude Code "goal mode" keeps working; mirror that here so the messages + path degrades ``xhigh`` -> ``max`` rather than 400-ing. Non-adaptive models + and models without a ceiling are left untouched. + """ + if not AnthropicModelInfo._is_adaptive_thinking_model(model): + return + effort = optional_params.get("reasoning_effort") + if not isinstance(effort, str): + return + clamped = {"effort": effort} + normalize_bedrock_opus_output_config_effort(model=model, output_config=clamped) + optional_params["reasoning_effort"] = clamped["effort"] + + def transform_anthropic_messages_request( + self, + model: str, + messages: List[Dict], + anthropic_messages_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Dict: + self._clamp_adaptive_reasoning_effort_for_bedrock( + model=model, + optional_params=anthropic_messages_optional_request_params, + ) + anthropic_messages_request = AnthropicMessagesConfig.transform_anthropic_messages_request( + self=self, + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + ######################################################### + ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### + ######################################################### + + # 1. anthropic_version is required for all claude models + if "anthropic_version" not in anthropic_messages_request: + anthropic_messages_request["anthropic_version"] = ( + self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION + ) + + # 2. `stream` is not allowed in request body for bedrock invoke + if "stream" in anthropic_messages_request: + anthropic_messages_request.pop("stream", None) + + # 3. `model` is not allowed in request body for bedrock invoke + if "model" in anthropic_messages_request: + anthropic_messages_request.pop("model", None) + + injected_thinking_for_clear_thinking = ( + self._ensure_thinking_for_clear_thinking_context_management( + anthropic_messages_request=anthropic_messages_request, + model=model, + ) + ) + + # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models) + self._remove_ttl_from_cache_control( + anthropic_messages_request=anthropic_messages_request, model=model + ) + + # 5. Convert structured-output params to inline schema. + # Bedrock Invoke doesn't support top-level `output_format`; its + # accepted `output_config` subset is also narrower than Anthropic's, so + # consume the newer `output_config.format` shape here instead of + # forwarding it as an unknown nested key. + existing_output_config = anthropic_messages_request.get("output_config") + if isinstance(existing_output_config, dict): + anthropic_messages_request["output_config"] = dict(existing_output_config) + output_format = anthropic_messages_request.pop("output_format", None) + output_config_format = pop_bedrock_invoke_output_config_format( + anthropic_messages_request + ) + if output_format: + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=output_format, + request_body=anthropic_messages_request, + ) + elif output_config_format: + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=output_config_format, + request_body=anthropic_messages_request, + ) + normalize_bedrock_opus_output_config_effort( + model=model, + output_config=anthropic_messages_request.get("output_config"), + ) + + # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, + # but older models do not — strip it to avoid request rejection. + # Ref: https://github.com/BerriAI/litellm/issues/22797 + if not ( + _supports_factory( + model=model, + custom_llm_provider="bedrock", + key="supports_output_config", + ) + or AnthropicConfig._model_supports_effort_param(model) + ): + if anthropic_messages_request.pop("output_config", None) is not None: + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` for " + "model=%s — neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) + + # 5b. Remove `custom` field from tools (Bedrock doesn't support it) + # Claude Code sends `custom: {defer_loading: true}` on tool definitions, + # which causes Bedrock to reject the request with "Extra inputs are not permitted" + # Ref: https://github.com/BerriAI/litellm/issues/22847 + remove_custom_field_from_tools(anthropic_messages_request) + normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_messages_request) + ensure_bedrock_anthropic_messages_tool_names(anthropic_messages_request) + + # 6. AUTO-INJECT beta headers based on features used + filtered_betas = self._get_bedrock_invoke_anthropic_beta_headers( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + headers=headers, + anthropic_messages_request=anthropic_messages_request, + injected_thinking_for_clear_thinking=injected_thinking_for_clear_thinking, + ) + if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas @@ -669,16 +696,9 @@ class AmazonAnthropicClaudeMessagesConfig( # Catches Anthropic-only extensions (output_config, speed, mcp_servers, ...) # and any future additions Claude Code may start sending. ``context_management`` # has already been pre-filtered to its Bedrock-supported subset above. - allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS - stripped = sorted(k for k in anthropic_messages_request if k not in allowed) - if stripped: - verbose_logger.debug( - "Bedrock Invoke: stripping unsupported top-level request fields: %s", - stripped, - ) - anthropic_messages_request = { - k: v for k, v in anthropic_messages_request.items() if k in allowed - } + anthropic_messages_request = self._strip_unsupported_bedrock_invoke_fields( + anthropic_messages_request + ) return anthropic_messages_request diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c9ab3c648ac..941fe59e825 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5316,6 +5316,28 @@ class BaseLLMHTTPHandler: ) if _session_config: realtime_streaming.session_configuration_request = _session_config + + # For providers that defer setup until client session.update, optionally + # send synthetic session.created to unblock clients waiting on connect. + if not provider_config.requires_session_configuration(): + synthetic_session = provider_config.transform_session_created_event( + model=model, + logging_session_id=logging_obj.litellm_trace_id, + session_configuration_request=None, + ) + if synthetic_session is not None: + synthetic_session_str = json.dumps(synthetic_session) + # Record before sending so the synthetic session.created is + # captured in the session log alongside provider-driven + # events; without this it would be silently absent from + # success_handler / async_success_handler payloads. + realtime_streaming.store_message(synthetic_session_str) + await websocket.send_text(synthetic_session_str) + realtime_streaming._session_created_sent_to_client = True + verbose_logger.debug( + "Sent synthetic session.created to client to unblock connection" + ) + await realtime_streaming.bidirectional_forward() except websockets.exceptions.InvalidStatusCode as e: # type: ignore @@ -6538,6 +6560,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: @@ -6620,6 +6643,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: @@ -6712,6 +6736,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -6783,6 +6808,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -6866,6 +6892,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -6923,6 +6950,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -6999,6 +7027,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -7009,27 +7038,49 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, data = video_provider_config.transform_video_edit_request( - prompt=prompt, + prefetched_source_data = None + prefetch_params = video_provider_config.get_video_edit_prefetch_params( video_id=video_id, api_base=api_base, litellm_params=litellm_params, headers=headers, - extra_body=extra_body, - ) - - logging_obj.pre_call( - input=prompt, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": url, - "headers": headers, - "video_id": video_id, - }, ) + if prefetch_params is not None: + prefetch_url, prefetch_body = prefetch_params + try: + prefetch_resp = sync_httpx_client.post( + url=prefetch_url, + headers=headers, + json=prefetch_body, + timeout=timeout, + ) + prefetch_resp.raise_for_status() + except Exception as e: + raise self._handle_error(e=e, provider_config=video_provider_config) + prefetched_source_data = prefetch_resp.json() try: + url, data = video_provider_config.transform_video_edit_request( + prompt=prompt, + video_id=video_id, + api_base=api_base, + litellm_params=litellm_params, + headers=headers, + extra_body=extra_body, + prefetched_source_data=prefetched_source_data, + ) + + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": url, + "headers": headers, + "video_id": video_id, + }, + ) + response = sync_httpx_client.post( url=url, headers=headers, @@ -7041,6 +7092,7 @@ class BaseLLMHTTPHandler: raw_response=response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + request_data=data, ) except Exception as e: raise self._handle_error(e=e, provider_config=video_provider_config) @@ -7071,6 +7123,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -7081,27 +7134,49 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, data = video_provider_config.transform_video_edit_request( - prompt=prompt, + prefetched_source_data = None + prefetch_params = video_provider_config.get_video_edit_prefetch_params( video_id=video_id, api_base=api_base, litellm_params=litellm_params, headers=headers, - extra_body=extra_body, - ) - - logging_obj.pre_call( - input=prompt, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": url, - "headers": headers, - "video_id": video_id, - }, ) + if prefetch_params is not None: + prefetch_url, prefetch_body = prefetch_params + try: + prefetch_resp = await async_httpx_client.post( + url=prefetch_url, + headers=headers, + json=prefetch_body, + timeout=timeout, + ) + prefetch_resp.raise_for_status() + except Exception as e: + raise self._handle_error(e=e, provider_config=video_provider_config) + prefetched_source_data = prefetch_resp.json() try: + url, data = video_provider_config.transform_video_edit_request( + prompt=prompt, + video_id=video_id, + api_base=api_base, + litellm_params=litellm_params, + headers=headers, + extra_body=extra_body, + prefetched_source_data=prefetched_source_data, + ) + + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": url, + "headers": headers, + "video_id": video_id, + }, + ) + response = await async_httpx_client.post( url=url, headers=headers, @@ -7113,6 +7188,7 @@ class BaseLLMHTTPHandler: raw_response=response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + request_data=data, ) except Exception as e: raise self._handle_error(e=e, provider_config=video_provider_config) @@ -7160,6 +7236,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -7234,6 +7311,7 @@ class BaseLLMHTTPHandler: api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -7445,6 +7523,7 @@ class BaseLLMHTTPHandler: api_key=api_key, headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 4378db06358..cf1fc75ef10 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -3,8 +3,10 @@ This file contains the transformation logic for the Gemini realtime API. """ import json +from collections import OrderedDict from typing import Any, Dict, List, Optional, Union, cast +import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -29,6 +31,7 @@ from litellm.types.llms.openai import ( OpenAIRealtimeDoneEvent, OpenAIRealtimeEvents, OpenAIRealtimeEventTypes, + OpenAIRealtimeFunctionCallArgumentsDone, OpenAIRealtimeOutputItemDone, OpenAIRealtimeResponseAudioDone, OpenAIRealtimeResponseContentPartAdded, @@ -36,10 +39,12 @@ from litellm.types.llms.openai import ( OpenAIRealtimeResponseDoneObject, OpenAIRealtimeResponseTextDone, OpenAIRealtimeStreamResponseBaseObject, + OpenAIRealtimeStreamResponseOutputItem, OpenAIRealtimeStreamResponseOutputItemAdded, OpenAIRealtimeStreamSession, OpenAIRealtimeStreamSessionEvents, OpenAIRealtimeTurnDetection, + ResponsesAPIStreamEvents, ) from litellm.types.llms.vertex_ai import ( GeminiResponseModalities, @@ -56,15 +61,43 @@ from litellm.utils import get_empty_usage from ..common_utils import encode_unserializable_types, get_api_key_from_env -MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Dict[str, OpenAIRealtimeEventTypes] = { +MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Dict[ + str, Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents] +] = { "setupComplete": OpenAIRealtimeEventTypes.SESSION_CREATED, "serverContent.generationComplete": OpenAIRealtimeEventTypes.RESPONSE_TEXT_DONE, "serverContent.turnComplete": OpenAIRealtimeEventTypes.RESPONSE_DONE, "serverContent.interrupted": OpenAIRealtimeEventTypes.RESPONSE_DONE, + "toolCall": ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, +} + +# Top-level keys in a Gemini realtime message that map_openai_event knows how +# to handle. Other keys (e.g. ``usageMetadata``) can appear alongside these as +# siblings and must be skipped by the main transform loop — otherwise +# map_openai_event raises ``ValueError`` and the WebSocket session terminates. +_KNOWN_GEMINI_TOP_LEVEL_KEYS: set = { + map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT } class GeminiRealtimeConfig(BaseRealtimeConfig): + # Cap the LRU of in-flight tool calls so long sessions with many tool + # calls don't grow the dict without bound. Sized large enough to cover + # bursts of pending tool responses; the oldest entry is evicted when a + # new call beyond the cap arrives. + _TOOL_CALL_ID_TO_NAME_MAX = 256 + + def __init__(self): + super().__init__() + # Store call_id → function_name mapping for tool call round-trip + self._tool_call_id_to_name: "OrderedDict[str, str]" = OrderedDict() + # Buffer ``usageMetadata`` that Gemini Live emits as a standalone + # frame (between turns) so the next ``response.done`` attributes the + # tokens consumed. Without this an authenticated client can drive + # tool-call or normal turns whose token usage is recorded as zero, + # bypassing spend and budget accounting. + self._pending_usage_metadata: Optional[dict] = None + def validate_environment( self, headers: dict, model: str, api_key: Optional[str] = None ) -> dict: @@ -190,10 +223,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) vertex_gemini_config = VertexGeminiConfig() - optional_params["generationConfig"]["tools"] = ( - vertex_gemini_config._map_function( - value=value, optional_params=optional_params - ) + # Tools should be at the top level of setup, not inside generationConfig + optional_params["tools"] = vertex_gemini_config._map_function( + value=value, optional_params=optional_params ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} @@ -214,6 +246,272 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): optional_params.pop("generationConfig") return optional_params + @staticmethod + def _extract_turn_detection(session: dict) -> Optional[dict]: + """Extract turn_detection from a session.update payload. + + Handles both the flat beta shape (``session.turn_detection``) and the + GA shape (``session.audio.input.turn_detection``). + """ + if not isinstance(session, dict): + return None + td = session.get("turn_detection") + if isinstance(td, dict): + return td + audio = session.get("audio") + if isinstance(audio, dict): + input_cfg = audio.get("input") + if isinstance(input_cfg, dict): + td = input_cfg.get("turn_detection") + if isinstance(td, dict): + return td + return None + + @staticmethod + def _normalize_session_payload_for_mapping(session: dict) -> dict: + """Normalize GA-remapped session fields back to their beta keys. + + ``map_openai_params`` only recognises the flat OpenAI-beta key names + (``modalities``, ``input_audio_transcription``, ``turn_detection``). + For GA clients the upstream shim renames these into the nested GA + schema (``output_modalities``, ``audio.input.transcription``, + ``audio.input.turn_detection``), which would otherwise be silently + dropped here. Surface them back at the top level so the existing + mapping logic picks them up without duplicating provider-specific + knowledge of the GA schema in ``map_openai_params``. + """ + if not isinstance(session, dict): + return session + + normalized = dict(session) + + if "modalities" not in normalized and "output_modalities" in normalized: + normalized["modalities"] = normalized["output_modalities"] + + audio = normalized.get("audio") + if isinstance(audio, dict): + input_cfg = audio.get("input") + if isinstance(input_cfg, dict): + if ( + "input_audio_transcription" not in normalized + and "transcription" in input_cfg + ): + normalized["input_audio_transcription"] = input_cfg["transcription"] + + extracted_turn_detection = GeminiRealtimeConfig._extract_turn_detection( + normalized + ) + if extracted_turn_detection is not None and not isinstance( + normalized.get("turn_detection"), dict + ): + normalized["turn_detection"] = extracted_turn_detection + + return normalized + + def _handle_session_update( + self, + json_message: dict, + model: str, + session_configuration_request: Optional[str], + ) -> List[str]: + """ + Handle session.update by sending setup to Gemini. + + On the FIRST session.update (when session_configuration_request is None), + the full setup with all configuration is sent. + + Subsequent session.update messages are forwarded as a follow-up setup + with the new fields merged into the original setup. Gemini Live treats + a follow-up BidiGenerateContentSetup as a full session replacement + rather than a partial merge, so we carry forward the previous setup + (tools, generationConfig, inputAudioTranscription, systemInstruction, + ...) and overlay the new fields on top. This preserves the old + behavior where clients could refine the session via session.update + (e.g. add tools after the auto-setup on connect), and also keeps the + guardrail-driven turn_detection update working. + """ + session_payload = json_message.get("session") or {} + # Normalize GA-remapped fields (``output_modalities``, + # nested ``audio.input.transcription``, + # ``audio.input.turn_detection``) back to their flat beta keys so + # ``map_openai_params`` picks them up. Without this, GA clients' + # explicit modality / transcription / turn-detection settings + # would be silently dropped because ``map_openai_params`` only + # recognises the flat OpenAI-beta key names. + session_payload = self._normalize_session_payload_for_mapping(session_payload) + new_overrides = self.map_openai_params( + optional_params={}, non_default_params=session_payload + ) + + if session_configuration_request is None: + generation_config = new_overrides.setdefault("generationConfig", {}) + generation_config.setdefault("responseModalities", ["AUDIO"]) + new_overrides.setdefault("inputAudioTranscription", {}) + new_overrides["model"] = f"models/{model}" + verbose_logger.debug( + "Gemini Realtime: Sending initial setup with tools to backend" + ) + return [json.dumps({"setup": new_overrides})] + + if not new_overrides: + verbose_logger.debug( + "Gemini Realtime: Ignoring session.update (no mappable fields)" + ) + return [] + + try: + original_setup = cast( + BidiGenerateContentSetup, + json.loads(session_configuration_request).get("setup", {}), + ) + except (json.JSONDecodeError, AttributeError): + original_setup = {} + + # Deep-merge ``generationConfig`` and ``realtimeInputConfig`` so a + # partial session.update (e.g. only ``temperature`` or only + # ``modalities``) does not silently drop unrelated sub-keys + # (``responseModalities``, ``maxOutputTokens``, ...) from the original + # setup. + follow_up_setup: BidiGenerateContentSetup = { + **original_setup, + **new_overrides, + "model": f"models/{model}", + } + original_generation_config = original_setup.get("generationConfig") + new_generation_config = new_overrides.get("generationConfig") + if isinstance(original_generation_config, dict) and isinstance( + new_generation_config, dict + ): + follow_up_setup["generationConfig"] = { + **original_generation_config, + **new_generation_config, + } + original_realtime_input_config = original_setup.get("realtimeInputConfig") + new_realtime_input_config = new_overrides.get("realtimeInputConfig") + if isinstance(original_realtime_input_config, dict) and isinstance( + new_realtime_input_config, dict + ): + merged_realtime_input_config = { + **original_realtime_input_config, + **new_realtime_input_config, + } + # Deep-merge ``automaticActivityDetection`` so a partial VAD + # update (e.g. the guardrail-injected ``disabled: True`` from + # ``create_response: False``) does not silently drop unrelated + # knobs like ``silenceDurationMs`` / ``prefixPaddingMs`` from + # the original setup. + original_automatic_activity_detection = original_realtime_input_config.get( + "automaticActivityDetection" + ) + new_automatic_activity_detection = new_realtime_input_config.get( + "automaticActivityDetection" + ) + if isinstance(original_automatic_activity_detection, dict) and isinstance( + new_automatic_activity_detection, dict + ): + merged_realtime_input_config["automaticActivityDetection"] = { + **original_automatic_activity_detection, + **new_automatic_activity_detection, + } + follow_up_setup["realtimeInputConfig"] = cast( + BidiGenerateContentRealtimeInputConfig, + merged_realtime_input_config, + ) + verbose_logger.debug( + "Gemini Realtime: Forwarding session.update as follow-up setup" + ) + return [json.dumps({"setup": follow_up_setup})] + + def _handle_conversation_item(self, json_message: dict) -> List[str]: + """ + Handle conversation.item.create for user text or function call output. + + Converts OpenAI format to Gemini's clientContent (for user text) or + toolResponse (for function outputs). + """ + item = json_message.get("item", {}) + item_type = item.get("type") + + # Handle function call output (tool response) + if item_type == "function_call_output": + return self._handle_function_call_output(item) + + # Handle regular text content + return self._handle_user_text_content(item) + + def _handle_function_call_output(self, item: dict) -> List[str]: + """Transform function_call_output to Gemini toolResponse format.""" + call_id = item.get("call_id", "") + output = item.get("output", "{}") + + verbose_logger.debug( + f"Gemini Realtime: Transforming function_call_output for call_id={call_id}" + ) + + # Parse the output to get the result. Gemini's + # functionResponses[].response field is a Struct, so it must be a + # dict; wrap any non-dict (primitives, lists, invalid JSON) under a + # `result` key. + try: + parsed_output = json.loads(output) if isinstance(output, str) else output + except json.JSONDecodeError: + parsed_output = output + output_dict = ( + parsed_output + if isinstance(parsed_output, dict) + else {"result": parsed_output} + ) + + # Look up the function name from stored mapping. Keep the entry so a + # client SDK that retries function_call_output (or sends it twice for + # the same tool call) still produces a Gemini toolResponse with the + # required ``name`` field; refresh the LRU position so an active + # call_id stays warm across long sessions. + function_name = self._tool_call_id_to_name.get(call_id) + if function_name: + self._tool_call_id_to_name.move_to_end(call_id) + else: + verbose_logger.warning( + f"Gemini Realtime: Function name not found for call_id={call_id}. " + "This may cause Gemini to reject the response." + ) + + # Build Gemini toolResponse format + function_response = { + "id": call_id, + "response": output_dict, + } + if function_name: + function_response["name"] = function_name + + tool_response_message = { + "toolResponse": {"functionResponses": [function_response]} + } + + return [json.dumps(tool_response_message)] + + def _handle_user_text_content(self, item: dict) -> List[str]: + """Transform user text content to Gemini clientContent format.""" + content_list = item.get("content", []) + text_parts = [ + c.get("text", "") + for c in content_list + if isinstance(c, dict) and c.get("type") == "input_text" + ] + text = " ".join(filter(None, text_parts)) + if not text: + return [] + + # Build clientContent message with turns (proper Gemini Live API format) + client_content_message = { + "clientContent": { + "turns": [{"role": "user", "parts": [{"text": text}]}], + "turnComplete": True, + } + } + + return [json.dumps(client_content_message)] + def transform_realtime_request( self, message: str, @@ -233,55 +531,42 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): messages: List[str] = [] msg_type = json_message.get("type") - ## HANDLE SESSION UPDATE — translate to Gemini setup; no realtime_input needed ## + ## HANDLE SESSION UPDATE — translate to Gemini setup ## if msg_type == "session.update": - client_session_configuration_request = self.map_openai_params( - optional_params={}, non_default_params=json_message["session"] + return self._handle_session_update( + json_message, model, session_configuration_request ) - client_session_configuration_request["model"] = f"models/{model}" - messages.append(json.dumps({"setup": client_session_configuration_request})) - return messages ## HANDLE response.create — Gemini responds automatically; nothing to forward ## if msg_type == "response.create": return [] - ## HANDLE INPUT AUDIO BUFFER ## + ## HANDLE conversation.item.create — extract user text or function call output ## + if msg_type == "conversation.item.create": + return self._handle_conversation_item(json_message) + + ## HANDLE INPUT AUDIO BUFFER - use realtimeInput for audio streaming ## if msg_type == "input_audio_buffer.append": realtime_input_dict["audio"] = HttpxBlobType( mimeType=self.get_audio_mime_type(), data=json_message["audio"] ) - ## HANDLE conversation.item.create — extract actual user text ## - elif msg_type == "conversation.item.create": - item = json_message.get("item", {}) - content_list = item.get("content", []) - text_parts = [ - c.get("text", "") - for c in content_list - if isinstance(c, dict) and c.get("type") == "input_text" - ] - text = " ".join(filter(None, text_parts)) - if not text: - return [] - realtime_input_dict["text"] = text - else: - # Unknown/unsupported OpenAI event type — drop silently rather than - # forwarding raw JSON as text input to the model. - return [] - if len(realtime_input_dict) != 1: - raise ValueError( - f"Only one argument can be set, got {len(realtime_input_dict)}:" - f" {list(realtime_input_dict.keys())}" + realtime_input_dict = cast( + BidiGenerateContentRealtimeInput, + encode_unserializable_types( + cast(Dict[str, object], realtime_input_dict) + ), ) - realtime_input_dict = cast( - BidiGenerateContentRealtimeInput, - encode_unserializable_types(cast(Dict[str, object], realtime_input_dict)), - ) - - messages.append(json.dumps({"realtime_input": realtime_input_dict})) - return messages + gemini_msg = json.dumps({"realtimeInput": realtime_input_dict}) + verbose_logger.debug( + "Gemini Realtime: Sending audio realtimeInput to backend" + ) + messages.append(gemini_msg) + return messages + # Unknown/unsupported OpenAI event type — drop silently rather than + # forwarding raw JSON as text input to the model. + return [] def transform_session_created_event( self, @@ -300,7 +585,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): generation_config = ( session_configuration_request_dict.get("generationConfig", {}) or {} ) - gemini_modalities = generation_config.get("responseModalities", ["TEXT"]) + gemini_modalities = generation_config.get("responseModalities", ["AUDIO"]) _modalities = [ modality.lower() for modality in cast(List[str], gemini_modalities) ] @@ -352,18 +637,18 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): delta_type: ALL_DELTA_TYPES, session_configuration_request: Optional[str] = None, ) -> List[OpenAIRealtimeEvents]: - if session_configuration_request is None: - raise ValueError( - "session_configuration_request is required for Gemini API calls" - ) - - session_configuration_request_dict: BidiGenerateContentSetup = json.loads( - session_configuration_request - ).get("setup", {}) + session_configuration_request_dict: BidiGenerateContentSetup = {} + if session_configuration_request is not None: + try: + session_configuration_request_dict = json.loads( + session_configuration_request + ).get("setup", {}) + except json.JSONDecodeError: + session_configuration_request_dict = {} generation_config = session_configuration_request_dict.get( "generationConfig", {} ) - gemini_modalities = generation_config.get("responseModalities", ["TEXT"]) + gemini_modalities = generation_config.get("responseModalities", ["AUDIO"]) _modalities = [ modality.lower() for modality in cast(List[str], gemini_modalities) ] @@ -576,6 +861,86 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): returned_items.append(response_output_item_done) return returned_items + def _consume_usage_metadata_for_response_done(self, frame: dict) -> Optional[dict]: + """Return the ``usageMetadata`` to attribute to a ``response.done``. + + Gemini Live emits ``usageMetadata`` either alongside the closing + frame (``serverContent.turnComplete`` / ``toolCall``) or as a + standalone frame between turns. The standalone form would otherwise + be discarded by the no-op branch in ``transform_realtime_response`` + and the consumed tokens silently dropped from spend/budget + accounting. ``_pending_usage_metadata`` buffers any such standalone + frames so the next emitted ``response.done`` carries the deferred + token counts. + + Returns the in-frame ``usageMetadata`` if present (and clears the + buffer since the in-frame counts are the authoritative attribution + for this turn), otherwise returns the buffered counts. ``None`` is + returned when neither is available so the caller can fall back to + ``get_empty_usage()``. + """ + # ``pop`` (rather than ``get``) so a single Gemini frame containing + # multiple closing keys (e.g. both ``toolCall`` and + # ``serverContent.turnComplete``) cannot attribute the same + # ``usageMetadata`` to two ``response.done`` events and double-count + # tokens in spend/budget accounting. + in_frame = frame.pop("usageMetadata", None) if isinstance(frame, dict) else None + if isinstance(in_frame, dict): + self._pending_usage_metadata = None + return in_frame + buffered = self._pending_usage_metadata + self._pending_usage_metadata = None + return buffered + + def transform_tool_call_events( + self, + tool_call_message: dict, + response_id: Optional[str] = None, + output_item_id: Optional[str] = None, + ) -> List[OpenAIRealtimeFunctionCallArgumentsDone]: + """ + Transform Gemini toolCall message to OpenAI function call events. + + Converts Gemini's functionCalls format to OpenAI's response.function_call_arguments.done events. + Also stores call_id → name mapping for later use in function_call_output responses. + """ + function_calls = tool_call_message.get("functionCalls", []) + resolved_response_id = response_id or f"resp_{uuid.uuid4()}" + resolved_output_item_id = output_item_id or f"item_{uuid.uuid4()}" + + verbose_logger.debug( + f"Gemini Realtime: Transforming {len(function_calls)} tool call(s) to OpenAI format" + ) + + events: List[OpenAIRealtimeFunctionCallArgumentsDone] = [] + for idx, fc in enumerate(function_calls): + call_id = fc.get("id", "") + name = fc.get("name", "") + + # Store call_id → name mapping for round-trip. Use an LRU so + # repeated function_call_output lookups (retries) still hit, while + # sessions with many tool calls don't grow the dict unboundedly. + if call_id and name: + self._tool_call_id_to_name[call_id] = name + self._tool_call_id_to_name.move_to_end(call_id) + while len(self._tool_call_id_to_name) > self._TOOL_CALL_ID_TO_NAME_MAX: + self._tool_call_id_to_name.popitem(last=False) + + events.append( + OpenAIRealtimeFunctionCallArgumentsDone( + type="response.function_call_arguments.done", + event_id=f"event_{uuid.uuid4()}", + response_id=resolved_response_id, + item_id=f"{resolved_output_item_id}_tool_{idx}", + output_index=idx, + call_id=call_id, + name=name, + arguments=json.dumps(fc.get("args", {})), + ) + ) + + return events + @staticmethod def get_nested_value(obj: dict, path: str) -> Any: keys = path.split(".") @@ -681,14 +1046,20 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "generationConfig", {} ) temperature = generation_config.get("temperature") - max_output_tokens = generation_config.get("max_output_tokens") - gemini_modalities = generation_config.get("responseModalities", ["TEXT"]) + max_output_tokens = generation_config.get("maxOutputTokens") + gemini_modalities = generation_config.get("responseModalities", ["AUDIO"]) _modalities = [ modality.lower() for modality in cast(List[str], gemini_modalities) ] - if "usageMetadata" in message: + resolved_usage_metadata = self._consume_usage_metadata_for_response_done( + cast(dict, message) + ) + if resolved_usage_metadata is not None: _chat_completion_usage = VertexGeminiConfig._calculate_usage( - completion_response=message, + completion_response=cast( + BidiGenerateContentServerMessage, + {**cast(dict, message), "usageMetadata": resolved_usage_metadata}, + ), ) else: _chat_completion_usage = get_empty_usage() @@ -716,7 +1087,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if temperature is not None: response_done_event["response"]["temperature"] = temperature if max_output_tokens is not None: - response_done_event["response"]["max_output_tokens"] = max_output_tokens + response_done_event["response"]["max_output_tokens"] = cast( + int, max_output_tokens + ) return response_done_event @@ -808,13 +1181,18 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): def map_openai_event( self, key: str, - value: dict, + value: Any, current_delta_type: Optional[ALL_DELTA_TYPES], - json_message: dict, - ) -> OpenAIRealtimeEventTypes: - model_turn_event = value.get("modelTurn") - generation_complete_event = value.get("generationComplete") - openai_event: Optional[OpenAIRealtimeEventTypes] = None + ) -> Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents]: + if isinstance(value, dict): + model_turn_event = value.get("modelTurn") + generation_complete_event = value.get("generationComplete") + else: + model_turn_event = None + generation_complete_event = None + openai_event: Optional[ + Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents] + ] = None if model_turn_event: # check if model turn event openai_event = self.map_model_turn_event(model_turn_event) elif generation_complete_event: @@ -822,15 +1200,27 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): delta_type=current_delta_type ) else: - # Check if this key or any nested key matches our mapping - for map_key, openai_event in MAP_GEMINI_FIELD_TO_OPENAI_EVENT.items(): - if map_key == key or ( - "." in map_key - and GeminiRealtimeConfig.get_nested_value(json_message, map_key) - is not None - ): - openai_event = openai_event + # Check if this key or any nested key matches our mapping. Use a + # distinct loop variable so we don't shadow ``openai_event`` and + # leak the last dict value when no entry matches. Scope dotted-key + # lookups to the current ``key``/``value`` pair — checking the + # whole ``json_message`` would let a sibling key (e.g. + # ``serverContent.turnComplete``) misclassify the event currently + # being processed (e.g. ``toolCall``). + for map_key, candidate_event in MAP_GEMINI_FIELD_TO_OPENAI_EVENT.items(): + if map_key == key: + openai_event = candidate_event break + if "." in map_key: + prefix, _, nested_path = map_key.partition(".") + if ( + prefix == key + and isinstance(value, dict) + and GeminiRealtimeConfig.get_nested_value(value, nested_path) + is not None + ): + openai_event = candidate_event + break if openai_event is None: raise ValueError(f"Unknown openai event: {key}, value: {value}") return openai_event @@ -854,6 +1244,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): message_str = str(message) raise ValueError(f"Invalid JSON message: {message_str}") + verbose_logger.debug( + "Realtime Response Transform: Gemini frame keys=%s", + ( + sorted(json_message.keys()) + if isinstance(json_message, dict) + else type(json_message).__name__ + ), + ) + logging_session_id = logging_obj.litellm_trace_id current_output_item_id = realtime_response_transform_input[ @@ -913,32 +1312,44 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) # If serverContent only contained transcription(s) and no model - # content, return early — the main loop would fail on unknown keys. + # content, mark it as already handled so the main loop skips it + # (map_openai_event would raise on an unknown serverContent + # subkey). Fall through so sibling top-level keys such as + # ``toolCall`` are still processed in the main loop. _model_content_keys = { "modelTurn", "turnComplete", "interrupted", "generationComplete", } - if not any(k in server_content for k in _model_content_keys): - return { - "response": returned_message, - "current_output_item_id": current_output_item_id, - "current_response_id": current_response_id, - "current_delta_chunks": current_delta_chunks, - "current_conversation_id": current_conversation_id, - "current_item_chunks": current_item_chunks, - "current_delta_type": current_delta_type, - "session_configuration_request": session_configuration_request, - } + server_content_handled = not any( + k in server_content for k in _model_content_keys + ) + else: + server_content_handled = False - for key, value in json_message.items(): + tool_call_handled = False + # Snapshot the items so handlers below can safely mutate + # ``json_message`` (e.g. ``_consume_usage_metadata_for_response_done`` + # pops ``usageMetadata`` to prevent a single frame from attributing + # the same token counts to two ``response.done`` events). + for key, value in list(json_message.items()): + # Skip sibling metadata keys (e.g. ``usageMetadata``) that can + # accompany a primary payload like ``toolCall`` or ``serverContent``. + # ``map_openai_event`` raises ValueError on unknown keys, which + # would otherwise terminate the WebSocket session. + if key not in _KNOWN_GEMINI_TOP_LEVEL_KEYS: + continue + # serverContent was a transcription-only payload already emitted + # above; skip it here so map_openai_event doesn't raise on the + # missing model-content subkeys. + if key == "serverContent" and server_content_handled: + continue # Check if this key or any nested key matches our mapping openai_event = self.map_openai_event( key=key, value=value, current_delta_type=current_delta_type, - json_message=json_message, ) if openai_event == OpenAIRealtimeEventTypes.SESSION_CREATED: @@ -947,8 +1358,226 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): logging_session_id, realtime_response_transform_input["session_configuration_request"], ) - session_configuration_request = json.dumps(transformed_message) returned_message.append(transformed_message) + elif openai_event == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE: + # Handle toolCall from Gemini. If the payload has no function + # calls, emit nothing — an orphaned response.created/done pair + # with no output items would confuse OpenAI-compatible clients. + # Mark the key as intentionally consumed (mirroring + # ``server_content_handled``) so any sibling keys in the same + # frame are still processed by the rest of the loop and the + # post-loop guard doesn't treat the no-op as fatal. + if not value.get("functionCalls"): + tool_call_handled = True + continue + + if current_conversation_id is None: + current_conversation_id = f"conv_{uuid.uuid4()}" + + # Extract session-level response metadata once so both + # response.created and response.done can include matching + # modalities/temperature/max_output_tokens fields. + session_setup: BidiGenerateContentSetup = {} + if session_configuration_request is not None: + try: + session_setup = json.loads(session_configuration_request).get( + "setup", {} + ) + except (json.JSONDecodeError, TypeError): + session_setup = {} + tool_call_generation_config = ( + session_setup.get("generationConfig", {}) or {} + ) + tool_call_modalities = [ + modality.lower() + for modality in cast( + List[str], + tool_call_generation_config.get( + "responseModalities", ["AUDIO"] + ), + ) + ] + + # Emit response.created preamble if this is the first event in the response + if current_response_id is None: + current_response_id = f"resp_{uuid.uuid4()}" + current_output_item_id = f"item_{uuid.uuid4()}" + + # Mirror the audio/text path: include modalities, + # temperature, and max_output_tokens on response.created so + # spec-compliant clients see consistent response metadata + # regardless of whether the response starts with content or + # a tool call. + returned_message.append( + { + "type": "response.created", + "event_id": f"event_{uuid.uuid4()}", + "response": { + "object": "realtime.response", + "id": current_response_id, + "status": "in_progress", + "output": [], + "conversation_id": current_conversation_id, + "modalities": tool_call_modalities, + "temperature": tool_call_generation_config.get( + "temperature" + ), + "max_output_tokens": tool_call_generation_config.get( + "maxOutputTokens" + ), + }, + } + ) + + tool_call_events = self.transform_tool_call_events( + value, + response_id=current_response_id, + output_item_id=current_output_item_id, + ) + # Emit output_item.added and conversation.item.created for each function call + for idx, tool_call in enumerate(tool_call_events): + item_id = tool_call["item_id"] + function_call_item: OpenAIRealtimeStreamResponseOutputItem = { + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "call_id": tool_call["call_id"], + "name": tool_call["name"], + "arguments": tool_call["arguments"], + } + # response.output_item.added + returned_message.append( + OpenAIRealtimeStreamResponseOutputItemAdded( + type="response.output_item.added", + event_id=f"event_{uuid.uuid4()}", + response_id=current_response_id, + output_index=idx, + item={ + **function_call_item, + "status": "in_progress", + "arguments": "", + }, + ) + ) + # response.function_call_arguments.delta — Gemini delivers + # the full arguments string in a single toolCall frame + # rather than streaming partial chunks, so emit one delta + # carrying the complete payload before the matching + # ``.done`` event. Spec-compliant OpenAI Realtime SDK + # clients accumulate ``delta.delta`` and rely on at least + # one delta before ``.done``. + returned_message.append( + cast( + OpenAIRealtimeEvents, + { + "type": "response.function_call_arguments.delta", + "event_id": f"event_{uuid.uuid4()}", + "response_id": current_response_id, + "item_id": item_id, + "output_index": idx, + "call_id": tool_call["call_id"], + "delta": tool_call["arguments"], + }, + ) + ) + # response.function_call_arguments.done + returned_message.append(tool_call) + # response.output_item.done — pass a fresh copy so + # downstream handlers that mutate the item dict (e.g. the + # beta-protocol translator) don't corrupt the references + # used by sibling events sharing the same function_call_item. + returned_message.append( + OpenAIRealtimeOutputItemDone( + type="response.output_item.done", + event_id=f"event_{uuid.uuid4()}", + response_id=current_response_id, + output_index=idx, + item={**function_call_item}, + ) + ) + # conversation.item.created + returned_message.append( + OpenAIRealtimeConversationItemCreated( + type="conversation.item.created", + event_id=f"event_{uuid.uuid4()}", + item={**function_call_item}, + ) + ) + + # response.done - close the response so clients can submit tool + # results. Mirror the non-tool-call RESPONSE_DONE path: if Gemini + # delivered ``usageMetadata`` alongside this ``toolCall`` frame, + # propagate the real token counts so spend/budget accounting + # records the tokens consumed by the tool-call turn. Standalone + # ``usageMetadata`` frames emitted in a separate WebSocket frame + # are buffered on the instance so the next ``response.done`` + # picks them up (otherwise an authenticated client could drive + # tool-call turns whose token usage is recorded as zero, + # bypassing budgets). Falls back to an empty usage block when + # neither is available (OpenAI-compatible clients expect + # ``usage`` to always be present on response.done). + resolved_tool_call_usage_metadata = ( + self._consume_usage_metadata_for_response_done(json_message) + ) + if resolved_tool_call_usage_metadata is not None: + _tool_call_chat_completion_usage = ( + VertexGeminiConfig._calculate_usage( + completion_response=cast( + BidiGenerateContentServerMessage, + { + **json_message, + "usageMetadata": resolved_tool_call_usage_metadata, + }, + ), + ) + ) + else: + _tool_call_chat_completion_usage = get_empty_usage() + tool_call_responses_api_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + _tool_call_chat_completion_usage, + ) + tool_call_done_event = OpenAIRealtimeDoneEvent( + type="response.done", + event_id=f"event_{uuid.uuid4()}", + response=OpenAIRealtimeResponseDoneObject( + id=current_response_id, + object="realtime.response", + status="completed", + output=[ + { + "id": te["item_id"], + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "call_id": te["call_id"], + "name": te["name"], + "arguments": te["arguments"], + } + for te in tool_call_events + ], + conversation_id=current_conversation_id, + modalities=tool_call_modalities, + usage=tool_call_responses_api_usage.model_dump(), + ), + ) + tool_call_temperature = tool_call_generation_config.get("temperature") + if tool_call_temperature is not None: + tool_call_done_event["response"][ + "temperature" + ] = tool_call_temperature + tool_call_max_output_tokens = tool_call_generation_config.get( + "maxOutputTokens" + ) + if tool_call_max_output_tokens is not None: + tool_call_done_event["response"]["max_output_tokens"] = cast( + int, tool_call_max_output_tokens + ) + returned_message.append(tool_call_done_event) + # Reset IDs so the next model turn (after tool results) starts a + # fresh response with its own response.created preamble. + current_output_item_id = None + current_response_id = None elif openai_event == OpenAIRealtimeEventTypes.RESPONSE_DONE: transformed_response_done_event = self.transform_response_done_event( message=BidiGenerateContentServerMessage(**json_message), # type: ignore @@ -958,16 +1587,37 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): output_items=None, ) returned_message.append(transformed_response_done_event) + # Reset IDs so a subsequent turn (e.g. a `toolCall` arriving in + # a later WebSocket frame after `turnComplete`) starts a fresh + # response with its own `response.created` preamble instead of + # reusing the just-completed response ID. + current_output_item_id = None + current_response_id = None elif ( openai_event == OpenAIRealtimeEventTypes.RESPONSE_TEXT_DELTA or openai_event == OpenAIRealtimeEventTypes.RESPONSE_TEXT_DONE or openai_event == OpenAIRealtimeEventTypes.RESPONSE_AUDIO_DELTA or openai_event == OpenAIRealtimeEventTypes.RESPONSE_AUDIO_DONE ): + # Pass the locally-updated state (rather than the original + # input snapshot) so that prior iterations of this loop — + # e.g. a tool-call or response.done that just reset + # current_response_id/current_output_item_id to None — are + # honoured by the modality handler. + _modality_input: RealtimeResponseTransformInput = { + **realtime_response_transform_input, + "current_output_item_id": current_output_item_id, + "current_response_id": current_response_id, + "current_conversation_id": current_conversation_id, + "current_delta_chunks": current_delta_chunks, + "current_item_chunks": current_item_chunks, + "current_delta_type": current_delta_type, + "session_configuration_request": session_configuration_request, + } _returned_message = self.handle_openai_modality_event( openai_event, json_message, - realtime_response_transform_input, + _modality_input, delta_type="text" if "text" in openai_event.value else "audio", ) returned_message.extend(_returned_message["returned_message"]) @@ -979,6 +1629,41 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): else: raise ValueError(f"Unknown openai event: {openai_event}") if len(returned_message) == 0: + # A frame whose only top-level keys are sibling metadata (e.g. + # a standalone ``{"usageMetadata": {...}}`` emitted by Gemini + # Live between turns) is not an error — there is just nothing + # to forward to the OpenAI-shaped client. Returning the + # unchanged state keeps the WebSocket alive; raising would + # terminate the session for a benign no-op frame. + # serverContent already consumed by the transcription handler is + # a benign no-op for downstream — treat it like a metadata-only + # key when deciding whether to raise. + unhandled_known_keys = [ + key + for key in json_message + if key in _KNOWN_GEMINI_TOP_LEVEL_KEYS + and not (key == "serverContent" and server_content_handled) + and not (key == "toolCall" and tool_call_handled) + ] + # Buffer standalone usage metadata so the next response.done can + # attribute the token counts. Without this, an authenticated + # client driving turns whose usageMetadata is emitted in a + # separate frame would have those tokens recorded as zero spend, + # bypassing budget enforcement. + standalone_usage_metadata = json_message.get("usageMetadata") + if isinstance(standalone_usage_metadata, dict): + self._pending_usage_metadata = standalone_usage_metadata + if not unhandled_known_keys: + return { + "response": returned_message, + "current_output_item_id": current_output_item_id, + "current_response_id": current_response_id, + "current_delta_chunks": current_delta_chunks, + "current_conversation_id": current_conversation_id, + "current_item_chunks": current_item_chunks, + "current_delta_type": current_delta_type, + "session_configuration_request": session_configuration_request, + } if isinstance(message, bytes): message_str = message.decode("utf-8", errors="replace") else: @@ -993,6 +1678,13 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): transformed_message=returned_message, current_item_chunks=current_item_chunks, ) + + for msg in returned_message: + event_type = msg.get("type") if isinstance(msg, dict) else "unknown" + verbose_logger.debug( + "Realtime Response Transform: OpenAI event=%s", event_type + ) + return { "response": returned_message, "current_output_item_id": current_output_item_id, @@ -1005,7 +1697,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): } def requires_session_configuration(self) -> bool: - return True + # Default behavior is backwards-compatible: send setup on connect. + # Opt-in to deferred setup for tool-injection flow via: + # litellm.gemini_live_defer_setup = True + return not litellm.gemini_live_defer_setup def session_configuration_request(self, model: str) -> str: """ diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 9714c8a3923..77a95bfa5ab 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -581,12 +581,23 @@ class GeminiVideoConfig(BaseVideoConfig): raise NotImplementedError("video get character is not supported for Gemini") def transform_video_edit_request( - self, prompt, video_id, api_base, litellm_params, headers, extra_body=None + self, + prompt, + video_id, + api_base, + litellm_params, + headers, + extra_body=None, + prefetched_source_data=None, ): raise NotImplementedError("video edit is not supported for Gemini") def transform_video_edit_response( - self, raw_response, logging_obj, custom_llm_provider=None + self, + raw_response, + logging_obj, + custom_llm_provider=None, + request_data=None, ): raise NotImplementedError("video edit is not supported for Gemini") diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index 32b71a43afa..6935cafd0d9 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -19,7 +19,10 @@ def cost_router(call_type: CallTypes) -> Literal["cost_per_token", "cost_per_sec def cost_per_token( - model: str, usage: Usage, service_tier: Optional[str] = None + model: str, + usage: Usage, + service_tier: Optional[str] = None, + data_residency: Optional[str] = None, ) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -27,6 +30,9 @@ def cost_per_token( Input: - model: str, the model name without provider prefix - usage: LiteLLM Usage block, containing anthropic caching information + - data_residency: optional OpenAI data-residency region (e.g. "eu", "us"), + inferred from api_base. Applies the model's regional-processing + uplift multiplier when set. Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -37,6 +43,7 @@ def cost_per_token( usage=usage, custom_llm_provider="openai", service_tier=service_tier, + data_residency=data_residency, ) # ### Non-cached text tokens # non_cached_text_tokens = usage.prompt_tokens diff --git a/litellm/llms/openai/data_residency.py b/litellm/llms/openai/data_residency.py new file mode 100644 index 00000000000..7162f70ca5f --- /dev/null +++ b/litellm/llms/openai/data_residency.py @@ -0,0 +1,41 @@ +""" +Helpers for resolving OpenAI data-residency (regional processing) from an +api_base URL. + +OpenAI enforces hostname-per-region for projects with geography restrictions +enabled and rejects requests sent to the wrong host, so the api_base hostname +is the authoritative signal of which region a request was processed in. +""" + +from typing import Dict, Optional +from urllib.parse import urlparse + +# Mapping of OpenAI regional hostnames to the corresponding data-residency +# value used by the cost calculator. See +# https://developers.openai.com/api/docs/pricing for the regional-processing +# uplift these hostnames trigger. +_OPENAI_REGIONAL_HOSTS: Dict[str, str] = { + "eu.api.openai.com": "eu", + "us.api.openai.com": "us", +} + + +def infer_openai_data_residency( + custom_llm_provider: Optional[str], api_base: Optional[str] +) -> Optional[str]: + """ + Derive the OpenAI data-residency region from an api_base URL. + + Returns ``"eu"`` for the EU regional host, ``"us"`` for the US regional + host, and ``None`` for the default global host, any non-OpenAI provider, + or any non-OpenAI URL. + """ + if custom_llm_provider != "openai" or not api_base: + return None + try: + host = urlparse(api_base).hostname + except (TypeError, ValueError): + return None + if not host: + return None + return _OPENAI_REGIONAL_HOSTS.get(host.lower()) diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 2d165a7d7df..520a42e9dd1 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -534,6 +534,7 @@ class OpenAIVideoConfig(BaseVideoConfig): litellm_params: GenericLiteLLMParams, headers: dict, extra_body: Optional[Dict[str, Any]] = None, + prefetched_source_data: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: original_video_id = extract_original_video_id(video_id) url = f"{api_base.rstrip('/')}/edits" @@ -547,6 +548,7 @@ class OpenAIVideoConfig(BaseVideoConfig): raw_response: httpx.Response, logging_obj: Any, custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict] = None, ) -> VideoObject: video_obj = VideoObject(**raw_response.json()) if custom_llm_provider and video_obj.id: diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 4f84816a2bc..b1723f494ec 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -623,12 +623,23 @@ class RunwayMLVideoConfig(BaseVideoConfig): raise NotImplementedError("video get character is not supported for RunwayML") def transform_video_edit_request( - self, prompt, video_id, api_base, litellm_params, headers, extra_body=None + self, + prompt, + video_id, + api_base, + litellm_params, + headers, + extra_body=None, + prefetched_source_data=None, ): raise NotImplementedError("video edit is not supported for RunwayML") def transform_video_edit_response( - self, raw_response, logging_obj, custom_llm_provider=None + self, + raw_response, + logging_obj, + custom_llm_provider=None, + request_data=None, ): raise NotImplementedError("video edit is not supported for RunwayML") diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 2b4746b174e..ea4dbccc8c8 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -14,6 +14,7 @@ Auth: OAuth2 Bearer token (not an API key). import json from typing import List, Optional +from litellm import verbose_logger from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig @@ -26,6 +27,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): """ def __init__(self, access_token: str, project: str, location: str) -> None: + super().__init__() self._access_token = access_token self._project = project self._location = location @@ -138,6 +140,62 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # Request translation # ------------------------------------------------------------------ + def _vertex_model_path(self, model: str) -> str: + """Return the fully-qualified Vertex AI model resource path.""" + return ( + f"projects/{self._project}" + f"/locations/{self._location}" + f"/publishers/google/models/{model}" + ) + + def _build_vertex_ai_setup_config(self, model: str, session_params: dict) -> dict: + """Build Vertex AI setup configuration with proper model path and defaults.""" + # Normalize GA-remapped fields (``output_modalities``, nested + # ``audio.input.transcription``, ``audio.input.turn_detection``) back to + # their flat beta keys so ``map_openai_params`` picks them up. Without + # this, GA clients' explicit modality / transcription / turn-detection + # settings would be silently dropped because ``map_openai_params`` only + # recognises the flat OpenAI-beta key names. + session_params = self._normalize_session_payload_for_mapping(session_params) + setup_config = self.map_openai_params( + optional_params={}, non_default_params=session_params + ) + + # Use full Vertex AI model path + setup_config["model"] = self._vertex_model_path(model) + + # Add Vertex AI specific defaults if not provided + generation_config = setup_config.setdefault("generationConfig", {}) + generation_config.setdefault("responseModalities", ["AUDIO"]) + + # Ensure Vertex defaults for realtimeInputConfig apply even when + # the client provided a partial ``turn_detection`` (e.g. only + # ``silence_duration_ms``). ``map_automatic_turn_detection`` sets + # ``disabled=True`` whenever ``create_response`` is absent or + # ``False``. Force ``disabled=False`` only when the client did + # not explicitly request ``create_response: False`` — that path + # is how transcription guardrails suppress automatic responses, + # and overriding it here would silently bypass the guardrail. + # Vertex Live has no "VAD on, no auto-response" mode, so callers + # that need that behaviour must accept that VAD is off. + client_turn_detection = session_params.get("turn_detection") + client_disabled_auto_response = ( + isinstance(client_turn_detection, dict) + and client_turn_detection.get("create_response") is False + ) + realtime_input_config = setup_config.setdefault("realtimeInputConfig", {}) + automatic_detection = realtime_input_config.setdefault( + "automaticActivityDetection", {} + ) + if not client_disabled_auto_response: + automatic_detection["disabled"] = False + automatic_detection.setdefault("silenceDurationMs", 800) + + setup_config.setdefault("inputAudioTranscription", {}) + setup_config.setdefault("outputAudioTranscription", {}) + + return setup_config + def transform_realtime_request( self, message: str, @@ -147,16 +205,50 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): """ Translate OpenAI realtime client messages to Vertex AI format. - ``session.update`` is intentionally ignored (returns []) because - Vertex AI only accepts a single ``setup`` message at the start of - the connection — sending a second one causes a 1007 close error. - The initial setup (sent automatically before bidirectional_forward) - already includes AUDIO modality and server VAD, so there is nothing - more to configure. + On the first ``session.update`` (when no setup has been sent yet) the + full ``BidiGenerateContentSetup`` is built with Vertex AI's model path + and forwarded. Any later ``session.update`` is dropped: Vertex AI + documents ``setup`` as the first-and-only client message, and a second + ``setup`` closes the connection with a 1007 policy error. """ json_message = json.loads(message) - if json_message.get("type") == "session.update": - # Do not forward as a second setup — Vertex AI rejects it. + msg_type = json_message.get("type") + + if msg_type == "session.update": + if session_configuration_request is None: + setup_config = self._build_vertex_ai_setup_config( + model, json_message.get("session") or {} + ) + gemini_setup_msg = json.dumps({"setup": setup_config}) + + verbose_logger.debug( + "Vertex AI Realtime: Sending initial setup with tools to backend" + ) + return [gemini_setup_msg] + + # A follow-up session.update can't be forwarded as a second setup + # (Vertex Live closes the WebSocket with 1007). If this drop is + # silencing the audio-transcription guardrail's create_response + # disable, surface a warning so operators know the model will + # auto-respond before the guardrail can gate it on Vertex AI. + client_turn_detection = GeminiRealtimeConfig._extract_turn_detection( + json_message.get("session") or {} + ) + if ( + isinstance(client_turn_detection, dict) + and client_turn_detection.get("create_response") is False + ): + verbose_logger.warning( + "Vertex AI Realtime: Dropping subsequent session.update " + "(turn_detection.create_response=False) — Vertex Live " + "rejects a second setup message. Audio-transcription " + "guardrails cannot suppress the model's auto-response on " + "Vertex AI in non-deferred mode." + ) + else: + verbose_logger.debug( + "Vertex AI Realtime: Ignoring session.update (setup already sent)" + ) return [] return super().transform_realtime_request( diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index ed6176cef05..b84966354b8 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -40,6 +40,29 @@ else: BaseLLMException = Any +def _build_vertex_video_usage_from_request_data( + request_data: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + """Build usage metadata (duration, resolution) for video cost calculation.""" + usage_data: Dict[str, Any] = {} + if not request_data: + return usage_data + + parameters = request_data.get("parameters", {}) + duration = ( + parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS + ) + if duration is not None: + try: + usage_data["duration_seconds"] = float(duration) + except (ValueError, TypeError): + pass + res = parameters.get("resolution") + if res is not None and str(res).strip() != "": + usage_data["video_resolution"] = str(res).strip().lower() + return usage_data + + def _convert_image_to_vertex_format(image_file) -> Dict[str, str]: """ Convert image file to Vertex AI format with base64 encoding and MIME type. @@ -363,23 +386,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): id=video_id, object="video", status="processing", model=model ) - usage_data: Dict[str, Any] = {} - if request_data: - parameters = request_data.get("parameters", {}) - duration = ( - parameters.get("durationSeconds") - or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS - ) - if duration is not None: - try: - usage_data["duration_seconds"] = float(duration) - except (ValueError, TypeError): - pass - res = parameters.get("resolution") - if res is not None and str(res).strip() != "": - usage_data["video_resolution"] = str(res).strip().lower() - - video_obj.usage = usage_data + video_obj.usage = _build_vertex_video_usage_from_request_data(request_data) return video_obj def transform_video_status_retrieve_request( @@ -647,15 +654,123 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): def transform_video_get_character_response(self, raw_response, logging_obj): raise NotImplementedError("video get character is not supported for Vertex AI") + def get_video_edit_prefetch_params( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Tuple[str, Dict]: + """Return the fetchPredictOperation URL and body needed to retrieve the source video.""" + return self.transform_video_status_retrieve_request( + video_id=video_id, + api_base=api_base, + litellm_params=litellm_params, + headers=headers, + ) + def transform_video_edit_request( - self, prompt, video_id, api_base, litellm_params, headers, extra_body=None - ): - raise NotImplementedError("video edit is not supported for Vertex AI") + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, + extra_body: Optional[Dict[str, Any]] = None, + prefetched_source_data: Optional[Dict[str, Any]] = None, + ) -> Tuple[str, Dict]: + """ + Build a predictLongRunning edit request from the pre-fetched source video. + + The actual fetchPredictOperation HTTP call is hoisted into the handler so + it can use the shared async/sync httpx client instead of blocking the loop. + """ + if prefetched_source_data is None: + raise ValueError( + "prefetched_source_data is required for Vertex AI video edit. " + "Ensure get_video_edit_prefetch_params is called by the handler." + ) + + if not prefetched_source_data.get("done", False): + raise ValueError( + "Source video generation is not complete yet. " + "Check the video status before editing." + ) + + videos = prefetched_source_data.get("response", {}).get("videos", []) + if not videos: + raise ValueError("No videos found in the completed operation. Cannot edit.") + + source_video = videos[0] + video_input: Dict[str, Any] = {} + if "gcsUri" in source_video: + video_input["gcsUri"] = source_video["gcsUri"] + elif "bytesBase64Encoded" in source_video: + video_input["bytesBase64Encoded"] = source_video["bytesBase64Encoded"] + video_input["mimeType"] = source_video.get("mimeType", "video/mp4") + else: + raise ValueError( + "Source video has neither gcsUri nor bytesBase64Encoded. Cannot edit." + ) + + operation_name = extract_original_video_id(video_id) + model = self.extract_model_from_operation_name(operation_name) or "" + + instance_dict: Dict[str, Any] = {"prompt": prompt, "video": video_input} + request_data: Dict[str, Any] = {"instances": [instance_dict]} + + if extra_body: + extra_body_copy = dict(extra_body) + nested_params = extra_body_copy.pop("parameters", None) + vertex_params: Dict[str, Any] = {} + if isinstance(nested_params, dict): + vertex_params.update(nested_params) + vertex_params.update(extra_body_copy) + if vertex_params: + request_data["parameters"] = vertex_params + + edit_url = f"{api_base.rstrip('/')}/{model}:predictLongRunning" + return edit_url, request_data def transform_video_edit_response( - self, raw_response, logging_obj, custom_llm_provider=None - ): - raise NotImplementedError("video edit is not supported for Vertex AI") + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: Optional[str] = None, + request_data: Optional[Dict] = None, + ) -> VideoObject: + """ + Transform the Veo video edit response. + + Veo returns the same operation response as video generation: + {"name": "projects/.../operations/OPERATION_ID"} + + usage includes duration_seconds and optional video_resolution from the + edit request parameters for cost calculation. + """ + response_data = raw_response.json() + + operation_name = response_data.get("name") + if not operation_name: + raise ValueError(f"No operation name in Veo edit response: {response_data}") + + model = self.extract_model_from_operation_name(operation_name) or "" + + if custom_llm_provider: + video_id = encode_video_id_with_provider( + operation_name, custom_llm_provider, model + ) + else: + video_id = operation_name + + video_obj = VideoObject( + id=video_id, + object="video", + status="processing", + model=model, + ) + video_obj.usage = _build_vertex_video_usage_from_request_data(request_data) + return video_obj def transform_video_extension_request( self, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 62e576ea0f7..ce6d4ac824c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -731,7 +731,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "anthropic.claude-haiku-4-5@20251001": { @@ -755,7 +754,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_streaming": true, "supports_native_structured_output": true }, @@ -926,8 +924,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -952,8 +949,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, @@ -977,12 +973,12 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, @@ -1009,10 +1005,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "global.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, @@ -1039,10 +1035,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "us.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, @@ -1069,10 +1065,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "eu.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, @@ -1098,10 +1094,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "au.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, @@ -1127,10 +1123,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -1158,10 +1154,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "anthropic.claude-mythos-preview": { "input_cost_per_token": 0, @@ -1175,8 +1171,8 @@ "supports_vision": true, "supports_prompt_caching": false, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_output_config": true }, "global.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -1204,10 +1200,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "us.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, @@ -1235,10 +1231,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "eu.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, @@ -1265,10 +1261,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "au.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, @@ -1295,10 +1291,165 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "au.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, @@ -1326,9 +1477,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "global.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, @@ -1356,9 +1506,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "us.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1386,9 +1535,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "eu.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1415,9 +1563,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "au.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1444,9 +1591,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "jp.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1473,9 +1619,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -1504,8 +1649,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -1537,7 +1681,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, "supports_native_structured_output": true }, "anthropic.claude-v1": { @@ -1788,7 +1931,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "apac.anthropic.claude-3-sonnet-20240229-v1:0": { @@ -1834,8 +1976,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "assemblyai/best": { "input_cost_per_second": 3.333e-05, @@ -1877,7 +2018,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "azure/ada": { @@ -1965,10 +2105,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_output_config": true }, "azure_ai/claude-opus-4-6": { "input_cost_per_token": 5e-06, @@ -1995,9 +2135,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-7": { "input_cost_per_token": 5e-06, @@ -2025,9 +2164,35 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 159, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-opus-4-8": { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -2092,8 +2257,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "azure/computer-use-preview": { "input_cost_per_token": 3e-06, @@ -9467,8 +9631,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 + "supports_web_search": true }, "claude-3-haiku-20240307": { "cache_creation_input_token_cost": 3e-07, @@ -9486,8 +9649,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 264 + "supports_vision": true }, "claude-3-opus-20240229": { "cache_creation_input_token_cost": 1.875e-05, @@ -9506,8 +9668,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 395 + "supports_vision": true }, "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, @@ -9532,8 +9693,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-4-sonnet-20250514": { "cache_creation_input_token_cost": 3.75e-06, @@ -9563,8 +9723,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 + "supports_web_search": true }, "claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.75e-06, @@ -9593,8 +9752,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 346 + "supports_vision": true }, "claude-sonnet-4-5-20250929": { "cache_creation_input_token_cost": 3.75e-06, @@ -9624,8 +9782,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 346 + "supports_web_search": true }, "claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, @@ -9653,8 +9810,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -9678,8 +9834,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -9705,8 +9860,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-1-20250805": { "cache_creation_input_token_cost": 1.875e-05, @@ -9733,8 +9887,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-20250514": { "cache_creation_input_token_cost": 1.875e-05, @@ -9761,8 +9914,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-5-20251101": { "cache_creation_input_token_cost": 6.25e-06, @@ -9786,11 +9938,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "claude-opus-4-5": { "cache_creation_input_token_cost": 6.25e-06, @@ -9814,11 +9965,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "claude-opus-4-6": { "cache_creation_input_token_cost": 6.25e-06, @@ -9846,13 +9996,12 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "supports_max_reasoning_effort": true }, "claude-opus-4-6-20260205": { "cache_creation_input_token_cost": 6.25e-06, @@ -9880,13 +10029,12 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -9916,12 +10064,11 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "claude-opus-4-7-20260416": { "cache_creation_input_token_cost": 6.25e-06, @@ -9951,12 +10098,45 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true + "supports_output_config": true + }, + "claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1, + "fast": 2.0 + }, + "supports_output_config": true }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-05-14", @@ -9987,8 +10167,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "cloudflare/@cf/meta/llama-2-7b-chat-fp16": { "input_cost_per_token": 1.923e-06, @@ -11233,8 +11412,8 @@ "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_output_config": true }, "databricks/databricks-claude-sonnet-4": { "input_cost_per_token": 2.9999900000000002e-06, @@ -13400,7 +13579,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": { @@ -13528,8 +13706,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "eu.anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -13554,8 +13731,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "eu.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -13584,8 +13760,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "eu.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, @@ -13615,7 +13790,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "eu.meta.llama3-2-1b-instruct-v1:0": { @@ -14958,7 +15132,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 1.5e-06, "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://ai.google.dev/gemini-api/docs/models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -17901,22 +18075,9 @@ }, "github_copilot/claude-haiku-4.5": { "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 16000, - "max_tokens": 16000, - "mode": "chat", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/claude-opus-4.5": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 16000, - "max_tokens": 16000, + "max_input_tokens": 200000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "supported_endpoints": [ "/v1/chat/completions" @@ -17924,7 +18085,22 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_vision": true, - "supports_minimal_reasoning_effort": true + "supports_reasoning": true + }, + "github_copilot/claude-opus-4.5": { + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_output_config": true }, "github_copilot/claude-opus-4.6-fast": { "litellm_provider": "github_copilot", @@ -17939,6 +18115,22 @@ "supports_parallel_function_calling": true, "supports_vision": true }, + "github_copilot/claude-opus-4.7": { + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_reasoning": true + }, "github_copilot/claude-opus-41": { "litellm_provider": "github_copilot", "max_input_tokens": 80000, @@ -17965,16 +18157,33 @@ }, "github_copilot/claude-sonnet-4.5": { "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 16000, - "max_tokens": 16000, + "max_input_tokens": 200000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "chat", "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true + }, + "github_copilot/claude-sonnet-4.6": { + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_reasoning": true }, "github_copilot/gemini-2.5-pro": { "litellm_provider": "github_copilot", @@ -17984,7 +18193,25 @@ "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_reasoning": true + }, + "github_copilot/gemini-3-flash-preview": { + "litellm_provider": "github_copilot", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true }, "github_copilot/gemini-3-pro-preview": { "litellm_provider": "github_copilot", @@ -17996,13 +18223,30 @@ "supports_parallel_function_calling": true, "supports_vision": true }, + "github_copilot/gemini-3.1-pro-preview": { + "litellm_provider": "github_copilot", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true + }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "supports_function_calling": true + "supports_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-3.5-turbo-0613": { "litellm_provider": "github_copilot", @@ -18010,7 +18254,10 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "supports_function_calling": true + "supports_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4": { "litellm_provider": "github_copilot", @@ -18018,7 +18265,22 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "supports_function_calling": true + "supports_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] + }, + "github_copilot/gpt-4-0125-preview": { + "litellm_provider": "github_copilot", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true }, "github_copilot/gpt-4-0613": { "litellm_provider": "github_copilot", @@ -18026,16 +18288,22 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "supports_function_calling": true + "supports_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4-o-preview": { "litellm_provider": "github_copilot", - "max_input_tokens": 64000, + "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "supports_function_calling": true, - "supports_parallel_function_calling": true + "supports_parallel_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4.1": { "litellm_provider": "github_copilot", @@ -18046,7 +18314,10 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4.1-2025-04-14": { "litellm_provider": "github_copilot", @@ -18057,68 +18328,89 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-41-copilot": { "litellm_provider": "github_copilot", - "mode": "completion" + "mode": "chat" }, "github_copilot/gpt-4o": { "litellm_provider": "github_copilot", - "max_input_tokens": 64000, + "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4o-2024-05-13": { "litellm_provider": "github_copilot", - "max_input_tokens": 64000, + "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4o-2024-08-06": { "litellm_provider": "github_copilot", - "max_input_tokens": 64000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true - }, - "github_copilot/gpt-4o-2024-11-20": { - "litellm_provider": "github_copilot", - "max_input_tokens": 64000, + "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supported_endpoints": [ + "/v1/chat/completions" + ] + }, + "github_copilot/gpt-4o-2024-11-20": { + "litellm_provider": "github_copilot", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4o-mini": { "litellm_provider": "github_copilot", - "max_input_tokens": 64000, + "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "supports_function_calling": true, - "supports_parallel_function_calling": true + "supports_parallel_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-4o-mini-2024-07-18": { "litellm_provider": "github_copilot", - "max_input_tokens": 64000, + "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "supports_function_calling": true, - "supports_parallel_function_calling": true + "supports_parallel_function_calling": true, + "supported_endpoints": [ + "/v1/chat/completions" + ] }, "github_copilot/gpt-5": { "litellm_provider": "github_copilot", @@ -18137,14 +18429,19 @@ }, "github_copilot/gpt-5-mini": { "litellm_provider": "github_copilot", - "max_input_tokens": 128000, + "max_input_tokens": 264000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_reasoning": true }, "github_copilot/gpt-5.1": { "litellm_provider": "github_copilot", @@ -18177,7 +18474,7 @@ }, "github_copilot/gpt-5.2": { "litellm_provider": "github_copilot", - "max_input_tokens": 128000, + "max_input_tokens": 264000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -18188,11 +18485,27 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true + }, + "github_copilot/gpt-5.2-codex": { + "litellm_provider": "github_copilot", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_reasoning": true }, "github_copilot/gpt-5.3-codex": { "litellm_provider": "github_copilot", - "max_input_tokens": 128000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -18202,25 +18515,96 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_reasoning": true + }, + "github_copilot/gpt-5.4": { + "litellm_provider": "github_copilot", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "github_copilot/gpt-5.4-mini": { + "litellm_provider": "github_copilot", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "github_copilot/gpt-5.5": { + "litellm_provider": "github_copilot", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "github_copilot/oswe-vscode-prime": { + "litellm_provider": "github_copilot", + "max_input_tokens": 264000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_vision": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true }, "github_copilot/text-embedding-3-small": { "litellm_provider": "github_copilot", "max_input_tokens": 8191, "max_tokens": 8191, - "mode": "embedding" + "mode": "embedding", + "supported_endpoints": [ + "/v1/embeddings" + ] }, "github_copilot/text-embedding-3-small-inference": { "litellm_provider": "github_copilot", "max_input_tokens": 8191, "max_tokens": 8191, - "mode": "embedding" + "mode": "embedding", + "supported_endpoints": [ + "/v1/embeddings" + ] }, "github_copilot/text-embedding-ada-002": { "litellm_provider": "github_copilot", "max_input_tokens": 8191, "max_tokens": 8191, - "mode": "embedding" + "mode": "embedding", + "supported_endpoints": [ + "/v1/embeddings" + ] }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", @@ -18438,7 +18822,7 @@ "output_cost_per_token": 2.5e-05, "supports_function_calling": true, "supports_vision": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "gmi/anthropic/claude-sonnet-4.5": { "input_cost_per_token": 3e-06, @@ -18738,7 +19122,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "global.anthropic.claude-sonnet-4-20250514-v1:0": { @@ -18768,8 +19151,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "global.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.25e-06, @@ -18792,7 +19174,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "global.amazon.nova-2-lite-v1:0": { @@ -19014,6 +19395,8 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19087,6 +19470,8 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19160,6 +19545,8 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19231,6 +19618,8 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -19272,6 +19661,8 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -19293,6 +19684,8 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -19581,6 +19974,8 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_priority": 1e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20284,6 +20679,8 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21206,6 +21603,8 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -21612,6 +22011,8 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21693,6 +22094,8 @@ "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, @@ -22878,7 +23281,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "jp.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -22901,7 +23303,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "crusoe/deepseek-ai/DeepSeek-R1-0528": { @@ -26843,8 +27244,7 @@ "supports_computer_use": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-3.7-sonnet": { "input_cost_per_image": 0.0048, @@ -26860,8 +27260,7 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-opus-4": { "input_cost_per_image": 0.0048, @@ -26880,8 +27279,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-opus-4.1": { "input_cost_per_image": 0.0048, @@ -26901,8 +27299,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-sonnet-4": { "input_cost_per_image": 0.0048, @@ -26925,8 +27322,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-sonnet-4.6": { "cache_creation_input_token_cost": 3.75e-06, @@ -26950,9 +27346,7 @@ "supports_reasoning": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_minimal_reasoning_effort": true + "supports_vision": true }, "openrouter/anthropic/claude-opus-4.5": { "cache_creation_input_token_cost": 6.25e-06, @@ -26967,12 +27361,11 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "openrouter/anthropic/claude-opus-4.6": { "cache_creation_input_token_cost": 6.25e-06, @@ -26991,9 +27384,7 @@ "supports_reasoning": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_minimal_reasoning_effort": true + "supports_vision": true }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -27016,8 +27407,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-haiku-4.5": { "cache_creation_input_token_cost": 1.25e-06, @@ -27035,8 +27425,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 346 + "supports_vision": true }, "openrouter/anthropic/claude-opus-4.7": { "cache_creation_input_token_cost": 6.25e-06, @@ -27058,8 +27447,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346 + "supports_xhigh_reasoning_effort": true }, "openrouter/bytedance/ui-tars-1.5-7b": { "input_cost_per_token": 1e-07, @@ -28243,10 +28631,10 @@ "supports_tool_choice": true }, "openrouter/xiaomi/mimo-v2-flash": { - "input_cost_per_token": 9e-08, - "output_cost_per_token": 2.9e-07, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 0.0, + "cache_read_input_token_cost": 1e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -28256,7 +28644,43 @@ "supports_tool_choice": true, "supports_reasoning": true, "supports_vision": false, - "supports_prompt_caching": false + "supports_prompt_caching": true + }, + "openrouter/xiaomi/mimo-v2.5-pro": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": false, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "openrouter/xiaomi/mimo-v2.5": { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2e-06, + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 8e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true, + "supports_response_schema": true, + "supports_prompt_caching": true }, "openrouter/z-ai/glm-4.7": { "input_cost_per_token": 4e-07, @@ -28987,14 +29411,16 @@ "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "supports_output_config": true }, "perplexity/anthropic/claude-opus-4-7": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "supports_output_config": true }, "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", @@ -29002,7 +29428,7 @@ "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "perplexity/anthropic/claude-sonnet-4-5": { "litellm_provider": "perplexity", @@ -31255,7 +31681,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "us.anthropic.claude-3-5-sonnet-20240620-v1:0": { @@ -31383,8 +31808,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "us.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, @@ -31416,7 +31840,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0": { @@ -31442,7 +31865,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "au.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -31464,7 +31886,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "us.anthropic.claude-opus-4-20250514-v1:0": { @@ -31490,8 +31911,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "us.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.875e-06, @@ -31512,15 +31932,15 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "global.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, @@ -31541,15 +31961,15 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "eu.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, @@ -31569,15 +31989,15 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "us.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -31606,8 +32026,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "us.deepseek.r1-v1:0": { "input_cost_per_token": 1.35e-06, @@ -32150,13 +32569,13 @@ "output_cost_per_token": 2.5e-05, "supports_assistant_prefill": true, "supports_computer_use": true, - "supports_minimal_reasoning_effort": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_output_config": true }, "vercel_ai_gateway/anthropic/claude-opus-4.6": { "cache_creation_input_token_cost": 6.25e-06, @@ -32176,7 +32595,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "vercel_ai_gateway/anthropic/claude-sonnet-4": { "cache_creation_input_token_cost": 3.75e-06, @@ -33184,8 +33603,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-3-haiku": { "input_cost_per_token": 2.5e-07, @@ -33288,8 +33706,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -33343,14 +33760,13 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "vertex_ai/claude-opus-4-5@20251101": { "cache_creation_input_token_cost": 6.25e-06, @@ -33370,15 +33786,14 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_output_config": true }, "vertex_ai/claude-opus-4-6": { "cache_creation_input_token_cost": 6.25e-06, @@ -33404,9 +33819,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-4-6@default": { "cache_creation_input_token_cost": 6.25e-06, @@ -33432,9 +33846,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -33461,9 +33874,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-4-7@default": { "cache_creation_input_token_cost": 6.25e-06, @@ -33490,9 +33901,63 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-opus-4-8@default": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true }, "vertex_ai/claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.75e-06, @@ -33540,13 +34005,12 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "vertex_ai/claude-sonnet-4-5@20250929": { "cache_creation_input_token_cost": 3.75e-06, @@ -33598,8 +34062,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-sonnet-4": { "cache_creation_input_token_cost": 3.75e-06, @@ -33628,8 +34091,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-sonnet-4@20250514": { "cache_creation_input_token_cost": 3.75e-06, @@ -33658,8 +34120,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/mistralai/codestral-2@001": { "input_cost_per_token": 3e-07, @@ -40652,13 +41113,12 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "duckduckgo/search": { "litellm_provider": "duckduckgo", @@ -40974,7 +41434,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_pdf_input": true }, @@ -40997,7 +41456,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_pdf_input": true } diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index d39a0dda152..9484922833a 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -71,6 +71,11 @@ class BasePassthroughUtils: request_headers.pop("content-length", None) request_headers.pop("host", None) + custom_header_names = {header_name.lower() for header_name in headers} + for header_name in list(request_headers.keys()): + if header_name.lower() in custom_header_names: + request_headers.pop(header_name, None) + # Combine request headers with custom headers headers = {**request_headers, **headers} 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 708ec7f1176..97d3a8cf5cc 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 @@ -118,15 +118,19 @@ class MCPRequestHandler: return b"{}" request.body = mock_body # type: ignore + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415 + get_request_route, + ) + + request_route = get_request_route(request) # Only OAuth metadata routes registered under /.well-known/ are public. - # Match on request.url.path (path-only, exact prefix) so the substring - # cannot be smuggled via query string, hostname, or a deeper URL segment. - if request.url.path.startswith("/.well-known/"): + if request_route.startswith("/.well-known/"): validated_user_api_key_auth = UserAPIKeyAuth() elif ( not litellm_api_key and MCPRequestHandler._target_servers_delegate_auth_to_upstream( # noqa: E501 - path=request.url.path, mcp_servers=mcp_servers + path=request_route, mcp_servers=mcp_servers ) ): # Operator opted this oauth2 server into upstream-delegated auth @@ -174,7 +178,7 @@ class MCPRequestHandler: "401", "403", ) and MCPRequestHandler._target_servers_use_oauth2( - path=request.url.path, mcp_servers=mcp_servers + path=request_route, mcp_servers=mcp_servers ): verbose_logger.debug( "MCP OAuth2: target server is OAuth2-mode, treating " @@ -562,25 +566,32 @@ class MCPRequestHandler: ) ) + key_access_group_extras = ( + await MCPRequestHandler._get_key_access_group_mcp_server_extras( + user_api_key_auth + ) + ) + ######################################################### # Calculate key/team allowed servers using inheritance and intersection logic ######################################################### - allowed_mcp_servers: List[str] = [] - has_lower_level_mcp_restrictions = ( - len(allowed_mcp_servers_for_key) > 0 - or len(allowed_mcp_servers_for_team) > 0 - ) - if len(allowed_mcp_servers_for_team) > 0: - if len(allowed_mcp_servers_for_key) > 0: - # Key has its own MCP permissions - use intersection with team permissions - for _mcp_server in allowed_mcp_servers_for_key: - if _mcp_server in allowed_mcp_servers_for_team: - allowed_mcp_servers.append(_mcp_server) - else: - # Key has no MCP permissions - inherit from team - allowed_mcp_servers = allowed_mcp_servers_for_team + key_set = set(allowed_mcp_servers_for_key) + team_set = set(allowed_mcp_servers_for_team) + extras_set = set(key_access_group_extras) + + has_lower_level_mcp_restrictions = bool(key_set or team_set or extras_set) + + # 1. Team-gated base scope. + if not team_set: + base = key_set # no team restriction + elif not key_set: + base = team_set # key has no own perms → inherits team else: - allowed_mcp_servers = allowed_mcp_servers_for_key + base = key_set & team_set # both restrict → intersect + + # 2. Extend with access-group extras (LIT-3189 — bypasses team + # ceiling, gated by group's assigned_team_ids / assigned_key_ids). + allowed_mcp_servers: List[str] = list(base | extras_set) ######################################################### # Check end_user permissions if end_user_id is set @@ -873,6 +884,43 @@ class MCPRequestHandler: return True return False + @staticmethod + async def _get_key_access_group_mcp_server_extras( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[str]: + """ + Resolve the key's unified `access_group_ids` (LiteLLM_AccessGroupTable) to + MCP server IDs, gated by the access group's `assigned_team_ids` / + `assigned_key_ids`. These servers extend the team's MCP scope rather + than being capped by it. Tag-style `mcp_access_groups` (per-server tags) + are intentionally not handled here — they have no assignment fields and + remain subject to the team ceiling. + """ + if user_api_key_auth is None: + return [] + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.auth.auth_checks import ( + get_authorized_resources_from_key_access_groups, + ) + + raw_server_ids = await get_authorized_resources_from_key_access_groups( + valid_token=user_api_key_auth, + team_object=None, + resource_field="access_mcp_server_ids", + ) + if not raw_server_ids: + return [] + # Permission entries may be server_ids OR names/aliases — expand to ids. + return global_mcp_server_manager.expand_permission_list(raw_server_ids) + except Exception as e: + verbose_logger.warning( + f"Failed to get key access group MCP server extras: {str(e)}" + ) + return [] + @staticmethod async def _get_allowed_mcp_servers_for_key( user_api_key_auth: Optional[UserAPIKeyAuth] = None, @@ -944,42 +992,78 @@ class MCPRequestHandler: """ Get allowed MCP servers for a team. - Note: object_permission is automatically loaded by get_team_object() in main auth flow. + Unions two sources: + - Legacy team.object_permission (mcp_servers, mcp_access_groups, + mcp_tool_permissions). + - Unified team.access_group_ids → access_group.access_mcp_server_ids. + Mirrors the model-side pattern in can_team_access_model — the group + is already attached to the team, so the team relationship is itself + the gate (no assigned_team_ids check needed here). """ try: - # Get team object permission (already loaded in main auth flow) - object_permissions = await MCPRequestHandler._get_team_object_permission( - user_api_key_auth - ) - - if object_permissions is None: - return [] - - # Permission entries may be server_ids OR names/aliases — expand to ids. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.auth.auth_checks import ( + _get_mcp_server_ids_from_access_groups, + get_team_object, + ) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if ( + user_api_key_auth is None + or not user_api_key_auth.team_id + or prisma_client is None + ): + return [] + + team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( + team_id=user_api_key_auth.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if team_obj is None: + return [] + + team_access_group_servers = await _get_mcp_server_ids_from_access_groups( + access_group_ids=team_obj.access_group_ids or [], + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + object_permissions = team_obj.object_permission + if object_permissions is None: + return list(set(team_access_group_servers)) direct_mcp_servers = global_mcp_server_manager.expand_permission_list( object_permissions.mcp_servers or [] ) - # Get MCP servers from access groups - access_group_servers = ( + legacy_access_group_servers = ( await MCPRequestHandler._get_mcp_servers_from_access_groups( object_permissions.mcp_access_groups or [] ) ) - # servers referenced in tool permissions should also be accessible tool_perm_servers = list( global_mcp_server_manager.expand_tool_permissions( object_permissions.mcp_tool_permissions ).keys() ) - # Combine all lists - all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers + all_servers = ( + direct_mcp_servers + + legacy_access_group_servers + + tool_perm_servers + + team_access_group_servers + ) return list(set(all_servers)) except Exception as e: verbose_logger.warning( diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 652e284ed49..8324ba641a4 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,3 +1,4 @@ +import html as _html import json from typing import Any, Dict, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -618,8 +619,105 @@ async def token_endpoint( ) +# Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request +# redirects back to the configured redirect URI with ``error`` / +# ``error_description`` / ``error_uri`` query params and no ``code``. The MCP +# loopback flow funnels that response through this /callback endpoint, so +# the endpoint must accept either a successful (``code``+``state``) or an +# error response. Declaring ``code``/``state`` as required would cause +# FastAPI to reject the error response with a 422 before the handler runs, +# which strands the MCP client waiting on the loopback (see LIT-2750). + + +def _render_oauth_error_html(error: str, description: Optional[str]) -> HTMLResponse: + """Render an actionable HTML page for an IdP-reported OAuth error. + + Used when we cannot propagate the error back to the registered + ``redirect_uri`` (state missing or undecryptable). Returned with a 400 + status so the failure is observable to operators while still being a + human-readable page for the end user. + """ + safe_error = _html.escape(error or "unknown_error") + safe_description = _html.escape(description) if description else "" + description_html = f"

{safe_description}

" if safe_description else "" + body = ( + "" + "

Authentication failed

" + f"

Error: {safe_error}

" + f"{description_html}" + "

You can close this window and try again.

" + "" + ) + return HTMLResponse(body, status_code=400) + + @router.get("/callback") -async def callback(request: Request, code: str, state: str): +async def callback( + request: Request, + code: Optional[str] = None, + state: Optional[str] = None, + error: Optional[str] = None, + error_description: Optional[str] = None, + error_uri: Optional[str] = None, +): + """OAuth 2.0 authorization response handler for MCP loopback clients. + + Accepts either: + + - A successful authorization response (``code`` + ``state``), which is + forwarded back to the validated client ``redirect_uri`` with the + original (un-wrapped) ``state``. + - An error response (``error``[+``error_description``/``error_uri``]), per + RFC 6749 §4.1.2.1. When ``state`` is present and decodes to a trusted + ``redirect_uri``, the error params are propagated back to the client so + its OAuth library can surface them. Otherwise we render an HTML error + page so the user is not left on an opaque 422 / blank screen. + """ + # 1. IdP-reported error path (e.g. ``?error=access_denied``). + if error: + verbose_logger.info( + "MCP /callback received IdP error: error=%s, error_description=%s", + error, + error_description, + ) + if state: + try: + state_data = decode_state_hash(state) + original_state = state_data.get("original_state") + redirect_uri = _get_validated_client_redirect_uri(request, state_data) + except HTTPException: + # Untrusted/invalid client redirect_uri — surface inline rather + # than blindly forwarding the error to an attacker-controlled URL. + return _render_oauth_error_html(error, error_description) + except Exception: + # State could not be decrypted (expired key, tampered, etc.). + return _render_oauth_error_html(error, error_description) + + params: Dict[str, str] = {"error": error} + if error_description: + params["error_description"] = error_description + if error_uri: + params["error_uri"] = error_uri + if original_state is not None: + params["state"] = original_state + complete_returned_url = _append_query_params(redirect_uri, params) + return RedirectResponse(url=complete_returned_url, status_code=302) + + # No state — nothing to round-trip to. Show the user the error. + return _render_oauth_error_html(error, error_description) + + # 2. Neither success nor error parameters present — most likely a stray + # GET / dropped SSO redirect chain. Surface a 400 instead of 422. + if not code or not state: + missing = [ + name for name, value in (("code", code), ("state", state)) if not value + ] + return _render_oauth_error_html( + "invalid_request", + f"Missing authorization {' and '.join(repr(m) for m in missing)} parameter(s).", + ) + + # 3. Successful authorization response. try: state_data = decode_state_hash(state) original_state = state_data["original_state"] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 004f33e630a..9046d522280 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -255,6 +255,7 @@ class KeyManagementRoutes(str, enum.Enum): # team spend-log viewing SPEND_LOGS = "/spend/logs" + SPEND_LOGS_V2 = "/spend/logs/v2" class LiteLLMRoutes(enum.Enum): @@ -548,6 +549,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, KeyManagementRoutes.SPEND_LOGS.value, + KeyManagementRoutes.SPEND_LOGS_V2.value, KeyManagementRoutes.KEY_RESET_SPEND.value, KeyManagementRoutes.KEY_ALIASES.value, ] @@ -599,6 +601,7 @@ class LiteLLMRoutes(enum.Enum): "/spend/tags", "/spend/calculate", "/spend/logs", + "/spend/logs/v2", "/spend/logs/ui", "/spend/logs/session/ui", "/cost/estimate", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 14f198e0f12..38976f79aa3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -619,6 +619,9 @@ async def common_checks( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, ) + # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. + _reject_clientside_metadata_tags_check(general_settings, request_body, route) + # If this is a free model, skip all budget checks if not skip_budget_checks: # 3. If team is in budget @@ -660,6 +663,14 @@ async def common_checks( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, ) + if valid_token is not None: + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=request_body, + user_api_key_dict=valid_token, + ) + with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"): await _tag_max_budget_check( request_body=request_body, @@ -709,7 +720,6 @@ async def common_checks( # noqa: PLR0915 await _check_end_user_budget(end_user_obj=end_user_object, route=route) _enforce_user_param_check(general_settings, request, request_body, route) - _reject_clientside_metadata_tags_check(general_settings, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route) _guardrail_modification_check(request_body, team_object) @@ -1765,19 +1775,39 @@ async def _cache_team_object( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): - key = "team_id:{}".format(team_id) - ## CACHE REFRESH TIME! team_table.last_refreshed_at = time.time() + # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( - key=key, + key="team_id:{}".format(team_id), value=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, model_type=LiteLLM_TeamTableCachedObj, ) + # Invalidate the alias-keyed cache so the JWT auth path with + # `team_alias_jwt_field` (which reads via `get_team_object_by_alias`) + # doesn't keep serving the pre-mutation team after every team-write + # endpoint (team_model_add, team_model_delete, update_team, etc.). + # + # Why DELETE and not WRITE: `team_alias` has no UNIQUE constraint in + # schema.prisma. Writing this cache from the generic refresh path + # would let a team admin who renamed their team to collide with + # another team's alias silently overwrite the cached team for + # JWT-by-alias auth (veria-ai review on #28739). Deleting forces the + # next reader through `get_team_object_by_alias`, which DOES enforce + # uniqueness (len(teams) > 1 raises HTTPException) before populating + # the cache from a verified single row. + if team_table.team_alias: + alias_key = "team_alias:{}".format(team_table.team_alias) + user_api_key_cache.delete_cache(key=alias_key) + if proxy_logging_obj is not None: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( + key=alias_key + ) + async def _cache_key_object( hashed_token: str, @@ -3143,44 +3173,40 @@ async def can_team_access_model( raise -async def _key_access_group_grants_model( - model: Union[str, List[str]], +async def get_authorized_resources_from_key_access_groups( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], - llm_router: Optional[Router], -) -> bool: + resource_field: Literal[ + "access_model_names", "access_mcp_server_ids", "access_agent_ids" + ], +) -> List[str]: """ - Returns True if the key's `access_group_ids` expand to models that grant - access to `model`. Used to let a key's access group override a team's - model restriction in `common_checks`. - - A key's access group only counts if the access group itself authorizes the - caller as an owner — that is, the group's `assigned_team_ids` includes the - key's `team_id`, or the group's `assigned_key_ids` includes the key's - token. This preserves the team-as-owner boundary (a team member cannot - escalate by naming a group assigned to a different team) while still - letting a group reach the key without first being added to the team's - `access_group_ids` list. + For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row + and contribute its `resource_field` only if the group authorizes the caller + as an owner — that is, the group's `assigned_team_ids` includes the key's + `team_id`, or the group's `assigned_key_ids` includes the key's token. This + preserves the team-as-owner boundary while still letting a group reach the + key without first being added to the team's `access_group_ids` list. """ if valid_token is None: - return False + return [] key_access_group_ids = list(valid_token.access_group_ids or []) if not key_access_group_ids: - return False + return [] from litellm.proxy.proxy_server import prisma_client as _prisma_client from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache if _prisma_client is None or _user_api_key_cache is None: - return False + return [] key_team_id = valid_token.team_id or ( team_object.team_id if team_object is not None else None ) key_token = valid_token.token - authorized_models: List[str] = [] + authorized_resources: List[str] = [] for ag_id in key_access_group_ids: try: ag = await get_access_object( @@ -3196,17 +3222,36 @@ async def _key_access_group_grants_model( ) key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) if team_authorized or key_authorized: - authorized_models.extend(ag.access_model_names or []) + authorized_resources.extend(getattr(ag, resource_field, []) or []) + return list(set(authorized_resources)) + + +async def _key_access_group_grants_model( + model: Union[str, List[str]], + valid_token: Optional[UserAPIKeyAuth], + team_object: Optional[LiteLLM_TeamTable], + llm_router: Optional[Router], +) -> bool: + """ + Returns True if the key's `access_group_ids` expand to models that grant + access to `model`. Used to let a key's access group override a team's + model restriction in `common_checks`. + """ + authorized_models = await get_authorized_resources_from_key_access_groups( + valid_token=valid_token, + team_object=team_object, + resource_field="access_model_names", + ) if not authorized_models: return False try: _can_object_call_model( model=model, llm_router=llm_router, - models=list(set(authorized_models)), - team_model_aliases=valid_token.team_model_aliases, - team_id=valid_token.team_id, + models=authorized_models, + team_model_aliases=valid_token.team_model_aliases if valid_token else None, + team_id=valid_token.team_id if valid_token else None, object_type="key", ) return True diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index c4dcca764b2..86265270357 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -213,6 +213,12 @@ _EXTRA_BANNED_OBSERVABILITY_PARAMS: FrozenSet[str] = frozenset( { "posthog_api_url", "phoenix_project_name", + "phoenix_project_name_override", + # Server-reserved: written exclusively by add_user_api_key_auth_to_request_metadata + # from the authenticated key's database record. A caller-supplied value + # would survive the server merge and let an authenticated user redirect + # their Arize/Phoenix telemetry into arbitrary projects. + "user_api_key_auth_metadata", "wandb_api_key", "weave_project_id", } @@ -498,9 +504,18 @@ def route_in_additonal_public_routes(current_route: str): def get_request_route(request: Request) -> str: """ - Helper to get the route from the request + Resolve the request route from the ASGI scope, with ``root_path`` stripped. - remove base url from path if set e.g. `/genai/chat/completions` -> `/chat/completions + Prefer this over ``request.url.path`` for any auth, ACL, routing, or + audit-log decision: Starlette reconstructs ``url.path`` by interpolating + the Host header into a URL string and re-parsing with ``urlsplit``, so a + malformed Host (e.g. ``localhost/?x=1``) collapses ``url.path`` to ``"/"`` + while FastAPI continues to dispatch on ``scope["path"]``. ``scope["path"]`` + is uvicorn's parse of the HTTP request line and matches the actual + handler, so it's the authoritative route. + + Also normalizes sub-path deployments by stripping ``scope["root_path"]`` + e.g. ``/genai/chat/completions`` -> ``/chat/completions``. """ try: scope = request.scope diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index d364b52c676..0f4aa37ba91 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -14,6 +14,9 @@ from litellm.utils import get_valid_models _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) +_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) + + def _check_wildcard_routing(model: str) -> bool: """ Returns True if a model is a provider wildcard. diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index b2878ba0ae6..a9519aa6cc5 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -62,7 +62,11 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES = ("/regenerate", "/reset_spend") class RouteChecks: @staticmethod - def should_call_route(route: str, valid_token: UserAPIKeyAuth): + def should_call_route( + route: str, + valid_token: UserAPIKeyAuth, + request: Optional[Request] = None, + ): """ Check if management route is disabled and raise exception """ @@ -77,13 +81,15 @@ class RouteChecks: # Check if Virtual Key is allowed to call the route - Applies to all Roles RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token + route=route, valid_token=valid_token, request=request ) return True @staticmethod def is_virtual_key_allowed_to_call_route( - route: str, valid_token: UserAPIKeyAuth + route: str, + valid_token: UserAPIKeyAuth, + request: Optional[Request] = None, ) -> bool: """ Raises Exception if Virtual Key is not allowed to call the route @@ -130,6 +136,21 @@ class RouteChecks: ): return True + # Method-aware carve-out: allow GET on the two + # read-only MCP-server discovery endpoints + # (`/v1/mcp/server` and `/v1/mcp/server/{server_id}`) + # so virtual keys with allowed_routes=["llm_api_routes"] + # can list/inspect MCP servers. The GET handlers in + # mcp_management_endpoints.py sanitize the response + # for restricted virtual keys (stripping url, + # headers, env, credentials). POST/PUT/DELETE on + # these paths are admin-only management writes and + # are intentionally not covered. + if RouteChecks._is_get_mcp_server_discovery_route( + route=route, request=request + ): + return True + # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: if RouteChecks._route_matches_wildcard_pattern( @@ -401,6 +422,31 @@ class RouteChecks: return True return False + @staticmethod + def _is_get_mcp_server_discovery_route( + route: str, request: Optional[Request] + ) -> bool: + """ + Returns True if `request` is a GET against one of the two read-only + MCP-server discovery paths: + + - GET `/v1/mcp/server` (list) + - GET `/v1/mcp/server/{server_id}` (single server, single segment) + + Multi-segment paths (`/v1/mcp/server/{id}/approve`, etc.) and any + non-GET method return False, so admin-only management writes on the + same path prefix are not reachable through this carve-out. + """ + if request is None or request.method.upper() != "GET": + return False + if route == "/v1/mcp/server": + return True + prefix = "/v1/mcp/server/" + if not route.startswith(prefix): + return False + remainder = route[len(prefix) :] + return bool(remainder) and "/" not in remainder + @staticmethod def is_management_route(route: str) -> bool: """ @@ -627,7 +673,11 @@ class RouteChecks: Returns: bool: True if `thread` or `assistant` is in the request path, False otherwise """ - if "thread" in request.url.path or "assistant" in request.url.path: + # Inline import — auth_utils participates in a proxy import cycle. + from .auth_utils import get_request_route # noqa: PLC0415 + + route = get_request_route(request) + if "thread" in route or "assistant" in route: return True return False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 03278633928..813b9826b37 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2200,7 +2200,9 @@ async def user_api_key_auth( user_api_key_auth_obj.budget_reservation = None ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## - RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj) + RouteChecks.should_call_route( + route=route, valid_token=user_api_key_auth_obj, request=request + ) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 166ef7a66d0..85165709957 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -15,6 +15,9 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.callback_utils import ( + sanitize_openai_provider_metadata, +) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, @@ -120,6 +123,9 @@ async def create_batch( # noqa: PLR0915 or get_custom_llm_provider_from_request_headers(request=request) or "openai" ) + if isinstance(data.get("metadata"), dict): + data["metadata"] = sanitize_openai_provider_metadata(data["metadata"]) + _create_batch_data = LiteLLMBatchCreateRequest(**data) # Apply team-level batch output expiry enforcement diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ef1d64335b4..6782208458e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -839,6 +839,7 @@ class ProxyBaseLLMRequestProcessing: "aget_run", "acancel_run", "adelete_run", + "apply_guardrail", ], version: Optional[str] = None, user_model: Optional[str] = None, @@ -1368,6 +1369,21 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, request_data=self.data, ) + if route_type == "aresponses": + # Streaming /v1/responses returns here without + # reaching the non-streaming ownership tail below. + # Wrap the SSE generator so container ownership is + # written once the upstream iterator finishes + # assembling ``completed_response`` — otherwise + # code-interpreter containers created during the + # stream stay unregistered and follow-up file API + # calls 403. Covers the background-polling path + # too, which loops ``body_iterator`` end-to-end. + selected_data_generator = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=response, + wrapped_generator=selected_data_generator, + user_api_key_dict=user_api_key_dict, + ) return await create_response( generator=selected_data_generator, media_type="text/event-stream", @@ -1483,8 +1499,93 @@ class ProxyBaseLLMRequestProcessing: await check_response_size_is_safe(response=response) + if route_type in {"aresponses", "aget_responses"}: + await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( + response=response, + user_api_key_dict=user_api_key_dict, + ) + return response + @staticmethod + async def _record_container_owners_from_responses_if_needed( + response: Any, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """Register code-interpreter containers so follow-up file APIs pass ownership checks.""" + from litellm.proxy.container_endpoints.ownership import ( + record_container_owners_from_responses_response, + ) + + if response is None: + return + + try: + await record_container_owners_from_responses_response( + response=response, + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_proxy_logger.exception( + "Container ownership recording failed after responses call: %s", + e, + ) + + @staticmethod + def _extract_completed_responses_response(stream_response: Any) -> Any: + """Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator. + + ``ResponsesAPIStreamingIterator`` stores the terminal stream event + (``response.completed`` / ``response.incomplete`` / ``response.failed``) + in ``completed_response``; the actual response body hangs off + that event's ``.response`` attribute. Some iterators store the + ``ResponsesAPIResponse`` directly. Handle both shapes so the + container-ownership recording path can walk ``.output`` either way. + """ + completed = getattr(stream_response, "completed_response", None) + if completed is None: + return None + response_obj = getattr(completed, "response", None) + if response_obj is not None: + return response_obj + return completed + + @staticmethod + async def _wrap_responses_stream_for_container_ownership( + original_stream_response: Any, + wrapped_generator: Any, + user_api_key_dict: UserAPIKeyAuth, + ): + """Forward SSE chunks, then record container ownership at stream end. + + Streaming ``/v1/responses`` short-circuits out of + ``base_process_llm_request`` before the non-streaming ownership + tail runs, so without this wrap the + ``LiteLLM_ManagedObjectTable`` row for any container created + during the stream is never written and follow-up file API calls + return 403. + """ + try: + async for chunk in wrapped_generator: + yield chunk + finally: + try: + completed_obj = ( + ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + original_stream_response + ) + ) + if completed_obj is not None: + await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( + response=completed_obj, + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_proxy_logger.exception( + "Container ownership recording failed after streaming responses call: %s", + e, + ) + async def base_passthrough_process_llm_request( self, request: Request, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 4995752d441..a65e737f248 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -317,7 +317,15 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 DatadogCostManagementLogger, ) - datadog_cost_management_obj = DatadogCostManagementLogger() + init_params = {} + if ( + "datadog_cost_management" in callback_specific_params + and isinstance( + callback_specific_params["datadog_cost_management"], dict + ) + ): + init_params = callback_specific_params["datadog_cost_management"] + datadog_cost_management_obj = DatadogCostManagementLogger(**init_params) imported_list.append(datadog_cost_management_obj) elif isinstance(callback, CustomLogger): imported_list.append(callback) @@ -409,11 +417,15 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str, def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]: - _metadata = request_data.get("metadata", None) - if not _metadata: - _metadata = request_data.get("litellm_metadata", None) - if not isinstance(_metadata, dict): - _metadata = {} + _metadata: Dict = {} + metadata_bucket = request_data.get("metadata") + litellm_metadata_bucket = request_data.get("litellm_metadata") + if isinstance(metadata_bucket, dict): + _metadata.update(metadata_bucket) + if isinstance(litellm_metadata_bucket, dict): + # Batch/file routes store proxy tracking in litellm_metadata while + # user-facing metadata stays in metadata; merge both for headers. + _metadata.update(litellm_metadata_bucket) headers = {} if "applied_guardrails" in _metadata: headers["x-litellm-applied-guardrails"] = ",".join( @@ -452,19 +464,103 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]: return headers +def get_metadata_variable_name_from_kwargs( + kwargs: dict, +) -> Literal["metadata", "litellm_metadata"]: + """ + Helper to return what the "metadata" field should be called in the request data + + - New endpoints return `litellm_metadata` + - Old endpoints return `metadata` + + Context: + - LiteLLM used `metadata` as an internal field for storing metadata + - OpenAI then started using this field for their metadata + - LiteLLM is now moving to using `litellm_metadata` for our metadata + """ + return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" + + +LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset( + { + "applied_policies", + "applied_guardrails", + "policy_sources", + "guardrails", + "guardrail_config", + "_guardrail_pipelines", + "_pipeline_managed_guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", + "pillar_response_headers", + "_pillar_response_headers_trusted", + "pillar_flagged", + "pillar_scanners", + "pillar_evidence", + "pillar_evidence_truncated", + "pillar_session_id_response", + "standard_logging_object", + "proxy_server_request", + "secret_fields", + } +) + + +def _get_or_create_proxy_metadata_bucket( + request_data: Dict, +) -> tuple[Literal["metadata", "litellm_metadata"], dict]: + """ + Return the proxy-internal metadata bucket for this request. + + Batch/file routes store proxy state in ``litellm_metadata`` so the OpenAI + ``metadata`` field can remain provider-safe (string values only). + """ + metadata_key = get_metadata_variable_name_from_kwargs(request_data) + metadata_bucket = request_data.get(metadata_key) + if not isinstance(metadata_bucket, dict): + metadata_bucket = {} + request_data[metadata_key] = metadata_bucket + return metadata_key, metadata_bucket + + +def sanitize_openai_provider_metadata( + metadata: Optional[Dict[str, Any]], +) -> Optional[Dict[str, str]]: + """ + Keep only provider-safe OpenAI metadata entries (string keys -> string values). + + Strips LiteLLM proxy-internal tracking fields that must not be forwarded to + OpenAI batch/file APIs. + """ + if not metadata: + return metadata + sanitized: Dict[str, str] = {} + for key, value in metadata.items(): + if key in LITELLM_PROXY_INTERNAL_METADATA_KEYS: + continue + if isinstance(value, str): + sanitized[key] = value + else: + verbose_proxy_logger.debug( + "sanitize_openai_provider_metadata: dropping key %r with non-string value of type %s", + key, + type(value).__name__, + ) + return sanitized or None + + def add_guardrail_to_applied_guardrails_header( request_data: Dict, guardrail_name: Optional[str] ): if guardrail_name is None: return - _metadata = request_data.get("metadata", None) or {} + _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) if "applied_guardrails" in _metadata: if guardrail_name not in _metadata["applied_guardrails"]: _metadata["applied_guardrails"].append(guardrail_name) else: _metadata["applied_guardrails"] = [guardrail_name] - # Ensure metadata is set back to request_data (important when metadata didn't exist) - request_data["metadata"] = _metadata def add_policy_to_applied_policies_header( @@ -478,14 +574,12 @@ def add_policy_to_applied_policies_header( """ if policy_name is None: return - _metadata = request_data.get("metadata", None) or {} + _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) if "applied_policies" in _metadata: if policy_name not in _metadata["applied_policies"]: _metadata["applied_policies"].append(policy_name) else: _metadata["applied_policies"] = [policy_name] - # Ensure metadata is set back to request_data (important when metadata didn't exist) - request_data["metadata"] = _metadata def add_policy_sources_to_metadata(request_data: Dict, policy_sources: Dict[str, str]): @@ -498,13 +592,12 @@ def add_policy_sources_to_metadata(request_data: Dict, policy_sources: Dict[str, """ if not policy_sources: return - _metadata = request_data.get("metadata", None) or {} + _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) existing = _metadata.get("policy_sources", {}) if not isinstance(existing, dict): existing = {} existing.update(policy_sources) _metadata["policy_sources"] = existing - request_data["metadata"] = _metadata def add_guardrail_response_to_standard_logging_object( @@ -527,23 +620,6 @@ def add_guardrail_response_to_standard_logging_object( return standard_logging_object -def get_metadata_variable_name_from_kwargs( - kwargs: dict, -) -> Literal["metadata", "litellm_metadata"]: - """ - Helper to return what the "metadata" field should be called in the request data - - - New endpoints return `litellm_metadata` - - Old endpoints return `metadata` - - Context: - - LiteLLM used `metadata` as an internal field for storing metadata - - OpenAI then started using this field for their metadata - - LiteLLM is now moving to using `litellm_metadata` for our metadata - """ - return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - - def process_callback( _callback: str, callback_type: str, environment_variables: dict ) -> dict: diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 2ce3fda6297..678ff289649 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -546,7 +546,10 @@ def _add_vector_store_id_from_path(request_data: dict, request: Request) -> None request_data: The request data dictionary to populate request: The FastAPI Request object """ - path = request.url.path + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + path = get_request_route(request) vector_store_match = re.search(r"/vector_stores/([^/]+)/", path) if vector_store_match: vector_store_id = vector_store_match.group(1) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 57de6c4a63d..e0015e112e1 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -117,6 +117,58 @@ async def _get_prisma_client(): return prisma_client +def _custom_llm_provider_from_responses_response( + response: Any, + default: str = "openai", +) -> str: + hidden_params: Dict[str, Any] = {} + if isinstance(response, dict): + hidden_params = response.get("_hidden_params") or {} + else: + hidden_params = getattr(response, "_hidden_params", None) or {} + + provider = hidden_params.get("custom_llm_provider") + if isinstance(provider, str) and provider: + return provider + return default + + +async def record_container_owners_from_responses_response( + response: Any, + user_api_key_dict: UserAPIKeyAuth, + custom_llm_provider: Optional[str] = None, +) -> None: + """Track containers created implicitly by code interpreter in /v1/responses.""" + container_ids = ( + ResponsesAPIRequestUtils.collect_container_ids_from_responses_response(response) + ) + if not container_ids: + return + + resolved_provider = ( + custom_llm_provider or _custom_llm_provider_from_responses_response(response) + ) + + for container_id in container_ids: + try: + await record_container_owner( + response={"id": container_id, "object": "container"}, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=resolved_provider, + ) + except Exception as e: + # Per-container errors (including ``HTTPException`` from + # conflicting/forbidden ownership rows) must not abort the + # batch — other containers in the same response should still + # get recorded so their follow-up file API calls don't 403. + verbose_proxy_logger.exception( + "Failed to record container ownership from responses output " + "for container_id=%s: %s", + container_id, + e, + ) + + async def record_container_owner( response: Any, user_api_key_dict: UserAPIKeyAuth, @@ -151,6 +203,8 @@ async def record_container_owner( file_object = _dump_response(response) file_object["custom_llm_provider"] = resolved_provider file_object["provider_container_id"] = original_container_id + # Prisma Python requires Json fields to be serialized as a JSON string. + file_object_json: str = json.dumps(file_object) prisma_client = await _get_prisma_client() if prisma_client is None: @@ -172,7 +226,7 @@ async def record_container_owner( where={"model_object_id": model_object_id}, data={ "unified_object_id": container_id, - "file_object": file_object, + "file_object": file_object_json, "updated_by": owner, }, ) @@ -181,7 +235,7 @@ async def record_container_owner( data={ "unified_object_id": container_id, "model_object_id": model_object_id, - "file_object": file_object, + "file_object": file_object_json, "file_purpose": CONTAINER_OBJECT_PURPOSE, "created_by": owner, "updated_by": owner, diff --git a/litellm/proxy/example_config_yaml/oai_misc_config.yaml b/litellm/proxy/example_config_yaml/oai_misc_config.yaml index 16cc69c19a5..551043ec76b 100644 --- a/litellm/proxy/example_config_yaml/oai_misc_config.yaml +++ b/litellm/proxy/example_config_yaml/oai_misc_config.yaml @@ -23,11 +23,11 @@ model_list: model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0 ######################################################### ########## batch specific params ######################## - s3_bucket_name: litellm-proxy + s3_bucket_name: litellm-proxy-123456789012 s3_region_name: us-west-2 s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY - aws_batch_role_arn: arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV + aws_batch_role_arn: arn:aws:iam::123456789012:role/service-role/AmazonBedrockExecutionRoleForAgents_EXAMPLE model_info: mode: batch diff --git a/litellm/proxy/example_config_yaml/otel_test_config.yaml b/litellm/proxy/example_config_yaml/otel_test_config.yaml index c05e2b1b5df..9c7937efba9 100644 --- a/litellm/proxy/example_config_yaml/otel_test_config.yaml +++ b/litellm/proxy/example_config_yaml/otel_test_config.yaml @@ -55,7 +55,7 @@ guardrails: litellm_params: guardrail: bedrock # supported values: "bedrock", "lakera" mode: "during_call" - guardrailIdentifier: ff6ujrregl1q + guardrailIdentifier: 4w3d1di3snt5 guardrailVersion: "DRAFT" - guardrail_name: "custom-pre-guard" litellm_params: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index e55f3b6e16b..e0e4bdcf4a4 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -10,7 +10,7 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Literal, Optional, Type, TypeVar, Union, cast from urllib.parse import urlparse -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Request from pydantic import BaseModel from litellm.proxy.common_utils.path_utils import safe_join @@ -2187,9 +2187,97 @@ async def test_custom_code_guardrail( ) +def _resolve_guardrail_input_type( + active_guardrail: CustomGuardrail, input_type: str +) -> Literal["request", "response"]: + """Return the effective input_type, auto-upgrading to 'response' for post_call guardrails.""" + if input_type == "request": + hook = getattr(active_guardrail, "event_hook", None) + if hook == GuardrailEventHooks.post_call or hook == "post_call": + return "response" + return "response" if input_type == "response" else "request" + + +def _patch_logging_obj_for_guardrail( + litellm_logging_obj: Any, request: ApplyGuardrailRequest +) -> None: + """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" + litellm_logging_obj.call_type = "pass_through_endpoint" + litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" + litellm_logging_obj.update_messages( + request.messages + if request.messages + else [{"role": "user", "content": request.text}] + ) + + +async def _emit_guardrail_success_logs( + proxy_logging_obj: Any, + litellm_logging_obj: Any, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: ApplyGuardrailResponse, + start_time: datetime, +) -> ApplyGuardrailResponse: + """Fire proxy and LiteLLM success hooks after a successful guardrail run. + + Each hook is wrapped defensively so a callback failure never prevents the + caller from receiving the guardrail response. Returns the (possibly + hook-modified) response. + """ + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as thread_pool_executor, + ) + + try: + modified = await proxy_logging_obj.post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + if isinstance(modified, ApplyGuardrailResponse): + response = modified + except Exception: + verbose_proxy_logger.exception("apply_guardrail: post_call_success_hook failed") + + # Build the logging payload after post_call_success_hook so that logged + # data matches what the caller actually receives if the hook modified + # the response. + response_for_logging = {"response": response.model_dump(exclude_none=True)} + + if litellm_logging_obj is not None: + end_time = datetime.now(timezone.utc) + try: + await litellm_logging_obj.async_success_handler( + result=response_for_logging, + start_time=start_time, + end_time=end_time, + cache_hit=False, + ) + except Exception: + verbose_proxy_logger.exception( + "apply_guardrail: async_success_handler failed" + ) + try: + thread_pool_executor.submit( + litellm_logging_obj.success_handler, + response_for_logging, + start_time, + end_time, + False, + ) + except Exception: + verbose_proxy_logger.exception( + "apply_guardrail: success_handler submit failed" + ) + + return response + + @router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse) @router.post("/apply_guardrail", response_model=ApplyGuardrailResponse) async def apply_guardrail( + fastapi_request: Request, request: ApplyGuardrailRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -2198,8 +2286,29 @@ async def apply_guardrail( This endpoint allows testing guardrails by applying them to custom text inputs. """ + import traceback + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as thread_pool_executor, + ) + from litellm.proxy.proxy_server import ( + general_settings, + proxy_config, + proxy_logging_obj, + version, + ) from litellm.proxy.utils import handle_exception_on_proxy + data: dict = { + "guardrail_name": request.guardrail_name, + "input": [request.text], + "messages": request.messages or [], + "metadata": {"route": "/apply_guardrail"}, + } + litellm_logging_obj = None + start_time = datetime.now(timezone.utc) + try: active_guardrail: Optional[CustomGuardrail] = ( GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( @@ -2212,23 +2321,25 @@ async def apply_guardrail( detail=f"Guardrail '{request.guardrail_name}' not found. Please ensure the guardrail is configured in your LiteLLM proxy.", ) - request_data: dict = {} - if request.messages: - request_data["messages"] = request.messages + request_processor = ProxyBaseLLMRequestProcessing(data=data) + data, litellm_logging_obj = ( + await request_processor.common_processing_pre_call_logic( + request=fastapi_request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type="apply_guardrail", + ) + ) - # Auto-detect input_type: if the caller didn't specify "response" but the - # guardrail only runs post_call (e.g. LLM-as-a-judge), use "response" so - # the test actually exercises the guardrail logic. - from litellm.types.guardrails import GuardrailEventHooks + if litellm_logging_obj is not None: + _patch_logging_obj_for_guardrail(litellm_logging_obj, request) - resolved_input_type = request.input_type - if resolved_input_type == "request": - hook = getattr(active_guardrail, "event_hook", None) - if hook == GuardrailEventHooks.post_call or hook == "post_call": - resolved_input_type = "response" - - _input_type: Literal["request", "response"] = ( - "response" if resolved_input_type == "response" else "request" + request_data: dict = {"messages": request.messages} if request.messages else {} + _input_type = _resolve_guardrail_input_type( + active_guardrail, request.input_type ) guardrailed_inputs = await active_guardrail.apply_guardrail( inputs={"texts": [request.text]}, @@ -2236,13 +2347,55 @@ async def apply_guardrail( input_type=_input_type, ) response_text = guardrailed_inputs.get("texts", []) - - return ApplyGuardrailResponse( + response = ApplyGuardrailResponse( response_text=response_text[0] if response_text else request.text ) except Exception as e: + if litellm_logging_obj is not None and not isinstance(e, HTTPException): + try: + await litellm_logging_obj.async_failure_handler( + exception=e, + traceback_exception=traceback.format_exc(), + ) + except Exception: + verbose_proxy_logger.exception( + "apply_guardrail: async_failure_handler failed" + ) + try: + thread_pool_executor.submit( + litellm_logging_obj.failure_handler, + e, + traceback.format_exc(), + ) + except Exception: + verbose_proxy_logger.exception( + "apply_guardrail: failure_handler submit failed" + ) + try: + transformed_exception = await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + ) + if isinstance(transformed_exception, Exception): + e = transformed_exception + except Exception: + verbose_proxy_logger.exception( + "apply_guardrail: post_call_failure_hook failed" + ) raise handle_exception_on_proxy(e) + # Success logging outside except so a hook error never triggers failure handlers. + response = await _emit_guardrail_success_logs( + proxy_logging_obj=proxy_logging_obj, + litellm_logging_obj=litellm_logging_obj, + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + start_time=start_time, + ) + return response + # Usage (dashboard) endpoints: overview, detail, logs router.include_router(guardrails_usage_router) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index ff3df11c448..ba3aee75047 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -151,7 +151,10 @@ async def test_endpoint(request: Request): dict: A dictionary containing the route of the request URL. """ # ping the proxy server to check if its healthy - return {"route": request.url.path} + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + return {"route": get_request_route(request)} @router.get( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2b840b5495e..7666b23f2af 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -333,8 +333,10 @@ def _get_metadata_variable_name(request: Request) -> str: For ALL other endpoints we call this "metadata" """ - path = request.url.path + # Inline imports — auth_utils/route_checks participate in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + path = get_request_route(request) if "thread" in path or "assistant" in path: return "litellm_metadata" @@ -1191,6 +1193,36 @@ class LiteLLMProxyRequestSetup: return tags + @staticmethod + def apply_key_tags_pre_auth( + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """Merge key metadata tags into request_data before _tag_max_budget_check.""" + key_metadata = user_api_key_dict.metadata + if not key_metadata: + return + + key_tags = key_metadata.get("tags") + if not key_tags or not isinstance(key_tags, list): + return + + _metadata_variable_name = get_metadata_variable_name_from_kwargs(request_data) + metadata = request_data.get(_metadata_variable_name) + if isinstance(metadata, str): + parsed = safe_json_loads(metadata) + metadata = parsed if isinstance(parsed, dict) else {} + request_data[_metadata_variable_name] = metadata + elif not isinstance(metadata, dict): + metadata = {} + request_data[_metadata_variable_name] = metadata + + existing_tags = metadata.get("tags") + metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=existing_tags if isinstance(existing_tags, list) else None, + tags_to_add=key_tags, + ) + @staticmethod def apply_client_tag_policy_pre_auth( request: Request, @@ -1511,10 +1543,16 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # spend_tracking_utils, streaming_iterator) read `body` to audit the # request; taking the snapshot here ensures they see cleaned metadata. # - # Exclude secret_fields (which contains raw_headers with Authorization - # tokens) from the snapshot — they must never be persisted in spend logs - # or any other audit trail. - _body_snapshot = {k: v for k, v in data.items() if k != "secret_fields"} + # Exclude: + # - secret_fields: contains raw_headers with Authorization tokens; must + # never be persisted in spend logs or any other audit trail. + # - proxy_server_request: already a key on `data` at this point (set + # earlier in this function); including it would make the snapshot + # self-reference — body.proxy_server_request.body would be the same + # dict as body, producing an infinite traversal loop for any consumer + # that walks the structure. + _body_snapshot_exclude = {"secret_fields", "proxy_server_request"} + _body_snapshot = {k: v for k, v in data.items() if k not in _body_snapshot_exclude} data["proxy_server_request"]["body"] = _body_snapshot # Snapshot the requester-supplied metadata for downstream consumers. diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 55eb321185c..0a26b23beff 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -19,6 +19,7 @@ from pydantic import BaseModel, Field import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import ( AUDIT_ACTIONS, LiteLLM_AuditLogs, @@ -34,6 +35,10 @@ from litellm.types.management_endpoints import ( router = APIRouter() +# Cache fields holding credentials. Masked on read so plaintext Redis / +# Sentinel passwords never leave the server in a GET response. +_CACHE_SENSITIVE_FIELDS: set = {"password", "sentinel_password"} + _REDACTED_VALUE = "***REDACTED***" @@ -295,7 +300,11 @@ async def get_cache_settings( else: decrypted_settings["redis_type"] = "node" - current_values = decrypted_settings + # Mask credential fields so the GET response never carries + # plaintext Redis / Sentinel passwords off the server. + current_values = mask_sensitive_keys( + decrypted_settings, _CACHE_SENSITIVE_FIELDS + ) # Update field values with current values for field in cache_fields: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e9d9c243e7c..431ff49c7ce 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1568,6 +1568,9 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 global_mcp_server_manager, ) + from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415 + get_request_route, + ) server_id = request.path_params.get("server_id", "") if server_id: @@ -1584,7 +1587,7 @@ if MCP_AVAILABLE: ): # For /token, require PKCE authorization_code; refresh_token # grants must NOT bypass auth (see comment above). - path_lower = (request.url.path or "").rstrip("/").lower() + path_lower = get_request_route(request).rstrip("/").lower() if path_lower.endswith("/token"): body_data = await _read_request_body(request=request) grant_type = (body_data or {}).get("grant_type", "") diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index f2d8ec8fb55..722fcd30033 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -51,6 +51,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import UpdateUsefulLinksRequest, ) from litellm.types.router import ( + SPECIAL_MODEL_INFO_PARAMS, Deployment, DeploymentTypedDict, LiteLLMParamsTypedDict, @@ -130,6 +131,32 @@ def update_db_model( updated_patch.model_info.model_dump(exclude_none=True) ) + # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI + # passes through (which today re-sends the OLD pricing on every save) cannot + # silently undo a litellm_params clear via .update(). + # + # Restricted to SPECIAL_MODEL_INFO_PARAMS (input/output cost per token/character + # and cache read/write costs) so this path cannot be used to null out privileged + # model_info fields like team_id or access groups. SPECIAL_MODEL_INFO_PARAMS are + # mirrored between litellm_params and model_info by Deployment.__init__, so the + # clear propagates to both blobs. + if updated_patch.litellm_params: + for field in updated_patch.litellm_params.model_fields_set: + if ( + field in SPECIAL_MODEL_INFO_PARAMS + and getattr(updated_patch.litellm_params, field) is None + ): + merged_deployment_dict["litellm_params"].pop(field, None) # type: ignore + merged_deployment_dict.get("model_info", {}).pop(field, None) + if updated_patch.model_info: + for field in updated_patch.model_info.model_fields_set: + if ( + field in SPECIAL_MODEL_INFO_PARAMS + and getattr(updated_patch.model_info, field) is None + ): + merged_deployment_dict["model_info"].pop(field, None) # type: ignore + merged_deployment_dict.get("litellm_params", {}).pop(field, None) # type: ignore + # convert to prisma compatible format prisma_compatible_model_dict = PrismaCompatibleUpdateDBModel() diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 43fdc9ae1cf..0d34974fbef 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -64,6 +64,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + _cache_team_object, allowed_route_check_inside_route, can_org_access_model, get_org_object, @@ -130,6 +131,33 @@ def _sanitize_for_log(value: Any) -> str: return text.replace("\r", "").replace("\n", "") +async def _refresh_cached_team( + team_row: Any, + user_api_key_cache: Any, + proxy_logging_obj: Any, +) -> None: + """ + Refresh the in-memory cached team object after a DB write. + + Every endpoint that mutates `litellm_teamtable` must call this so the + cached `LiteLLM_TeamTableCachedObj` used by `common_checks` stays in + sync. Without this, subsequent auth checks read a stale team and can + 403 on permissions the DB has already granted (or, symmetrically, + keep granting permissions the DB has already revoked). + + `team_row` is the Prisma row returned by `update`/`find_unique` on + `litellm_teamtable`. It is converted to `LiteLLM_TeamTableCachedObj` + via `model_dump()` to match the cache shape `_cache_team_object` + expects. + """ + await _cache_team_object( + team_id=team_row.team_id, + team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, @@ -1591,7 +1619,6 @@ async def update_team( # noqa: PLR0915 ``` """ try: - from litellm.proxy.auth.auth_checks import _cache_team_object from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, @@ -1861,7 +1888,13 @@ async def update_team( # noqa: PLR0915 await prisma_client.db.litellm_teamtable.update( where={"team_id": data.team_id}, data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, # type: ignore ) ) @@ -1874,9 +1907,8 @@ async def update_team( # noqa: PLR0915 verbose_proxy_logger.info( "Successfully updated team - %s, info", team_row.team_id ) - await _cache_team_object( - team_id=team_row.team_id, - team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + await _refresh_cached_team( + team_row=team_row, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -4569,7 +4601,11 @@ async def team_model_add( }' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -4603,9 +4639,21 @@ async def team_model_add( ) updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models) - # Update team + # Update team. `include` mirrors the relations the auth path consumes + # off the cached team object so that `_refresh_cached_team` doesn't + # null them out — see object_permission_utils.validate_key_search_tools_against_team + # and the MCP/agent authz paths, which treat a missing object_permission + # as "no team-level restriction". updated_team = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, data={"models": updated_models} + where={"team_id": data.team_id}, + data={"models": updated_models}, + include={"object_permission": True}, # type: ignore + ) + + await _refresh_cached_team( + team_row=updated_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return updated_team @@ -4640,7 +4688,11 @@ async def team_model_delete( }' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -4679,9 +4731,17 @@ async def team_model_delete( # Remove specified models updated_models = [m for m in current_models if m not in data.models] - # Update team + # Update team. See team_model_add for the rationale on `include`. updated_team = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, data={"models": updated_models} + where={"team_id": data.team_id}, + data={"models": updated_models}, + include={"object_permission": True}, # type: ignore + ) + + await _refresh_cached_team( + team_row=updated_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return updated_team diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index b7d5cc30c49..888ccc78188 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -2,7 +2,7 @@ ## Helper utils for the management endpoints (keys/users/teams) from datetime import datetime from functools import wraps -from typing import List, Optional, Tuple +from typing import Any, Callable, List, Optional, Tuple from fastapi import HTTPException, Request @@ -435,6 +435,85 @@ async def send_management_endpoint_alert( ) +async def _emit_management_endpoint_otel_span( + func: Callable, + kwargs: dict, + parent_otel_span: Any, + start_time: datetime, + end_time: datetime, + result: Any = None, + exception: Optional[Exception] = None, +) -> None: + """Stamp + end the parent OTEL SERVER span for a management endpoint. + + Routes the request/response (or exception) through the OTEL success/failure + hook. Falls back to ``func.__name__`` for the route when the handler has no + ``http_request`` param — endpoints like ``/key/generate`` never receive one, + and gating the hook on it leaked their SERVER span (created in auth, never + ended → never exported). Always emitting keeps both success and failure + paths consistent. + """ + from litellm.proxy.proxy_server import open_telemetry_logger + + if open_telemetry_logger is None: + return + + http_request: Optional[Request] = kwargs.get("http_request") + if http_request is not None: + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415 + get_request_route, + ) + + route = get_request_route(http_request) + request_body: dict = await _read_request_body(request=http_request) + else: + route = func.__name__ + request_body = {} + + _CREDENTIAL_FIELDS = frozenset( + { + "key", + "token", + "api_key", + "secret", + "password", + "access_token", + "refresh_token", + "private_key", + "service_account_key", + } + ) + + _response: Optional[dict] = None + if exception is None and result is not None: + try: + raw = dict(result) + _response = {k: v for k, v in raw.items() if k not in _CREDENTIAL_FIELDS} + except Exception: + _response = None + + logging_payload = ManagementEndpointLoggingPayload( + route=route, + request_data=request_body, + response=_response, + start_time=start_time, + end_time=end_time, + exception=exception, + ) + + if exception is None: + await open_telemetry_logger.async_management_endpoint_success_hook( + logging_payload=logging_payload, + parent_otel_span=parent_otel_span, + ) + else: + await open_telemetry_logger.async_management_endpoint_failure_hook( + logging_payload=logging_payload, + parent_otel_span=parent_otel_span, + ) + + def management_endpoint_wrapper(func): """ This wrapper does the following: @@ -446,13 +525,10 @@ def management_endpoint_wrapper(func): @wraps(func) async def wrapper(*args, **kwargs): start_time = datetime.now() - _http_request: Optional[Request] = None try: result = await func(*args, **kwargs) end_time = datetime.now() try: - if kwargs is None: - kwargs = {} user_api_key_dict: UserAPIKeyAuth = ( kwargs.get("user_api_key_dict") or UserAPIKeyAuth() ) @@ -462,31 +538,16 @@ def management_endpoint_wrapper(func): user_api_key_dict=user_api_key_dict, function_name=func.__name__, ) - _http_request = kwargs.get("http_request", None) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: - from litellm.proxy.proxy_server import open_telemetry_logger - - if open_telemetry_logger is not None: - if _http_request: - _route = _http_request.url.path - _request_body: dict = await _read_request_body( - request=_http_request - ) - _response = dict(result) if result is not None else None - - logging_payload = ManagementEndpointLoggingPayload( - route=_route, - request_data=_request_body, - response=_response, - start_time=start_time, - end_time=end_time, - ) - - await open_telemetry_logger.async_management_endpoint_success_hook( # type: ignore - logging_payload=logging_payload, - parent_otel_span=parent_otel_span, - ) + await _emit_management_endpoint_otel_span( + func=func, + kwargs=kwargs, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + result=result, + ) # Delete updated/deleted info from cache _delete_api_key_from_cache(kwargs=kwargs) @@ -502,38 +563,26 @@ def management_endpoint_wrapper(func): except Exception as e: end_time = datetime.now() - if kwargs is None: - kwargs = {} user_api_key_dict: UserAPIKeyAuth = ( kwargs.get("user_api_key_dict") or UserAPIKeyAuth() ) parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: - from litellm.proxy.proxy_server import open_telemetry_logger - - if open_telemetry_logger is not None: - _http_request = kwargs.get("http_request") - if _http_request: - _route = _http_request.url.path - _request_body: dict = await _read_request_body( - request=_http_request - ) - else: - _route = func.__name__ - _request_body = {} - - logging_payload = ManagementEndpointLoggingPayload( - route=_route, - request_data=_request_body, - response=None, + try: + await _emit_management_endpoint_otel_span( + func=func, + kwargs=kwargs, + parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, exception=e, ) - - await open_telemetry_logger.async_management_endpoint_failure_hook( # type: ignore - logging_payload=logging_payload, - parent_otel_span=parent_otel_span, + except Exception as otel_exc: + # Non-Blocking Exception - never let OTEL failures swallow + # the original management-endpoint exception. + verbose_logger.debug( + "Error emitting OTEL span in management endpoint wrapper failure path: %s", + str(otel_exc), ) raise e diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ce103f806e1..7ca28a5d4ac 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -44,6 +44,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( ) from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) from litellm.proxy.utils import is_known_model from litellm.proxy.vector_store_endpoints.utils import ( @@ -1123,6 +1124,9 @@ async def bedrock_proxy_route( _forward_headers=True, ) # dynamically construct pass-through endpoint based on incoming path setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data) + # SigV4 signs an exact payload; pass-through must send prepped.body, not json.dumps + # of a dict that hooks may mutate (logging_obj, metadata, etc.). + setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, prepped.body) received_value = await endpoint_func( request, fastapi_response, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index df52c0fe204..36a389d233d 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,7 +6,7 @@ import posixpath import traceback from base64 import b64encode from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple, Union, cast +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast from urllib.parse import urlencode, urlparse import httpx @@ -62,6 +62,7 @@ from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, PassthroughStandardLoggingPayload, ) @@ -735,6 +736,22 @@ async def pass_through_request( # noqa: PLR0915 str(url) ) + # SigV4-signed callers (e.g. Bedrock) attach the exact bytes that were + # signed via request.state; we must send those instead of re-encoding the + # parsed dict (hooks mutate it, breaking the signature / Content-Length). + # Tolerate request objects without `state` (test fixtures) and only honor + # values httpx accepts for `content=`. + _request_state = getattr(request, "state", None) + state_raw_body: Optional[Union[str, bytes]] = ( + getattr(_request_state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, None) + if _request_state is not None + else None + ) + if state_raw_body is not None and not isinstance( + state_raw_body, (str, bytes, bytearray) + ): + state_raw_body = None + # Skip body parsing for multipart requests - make_multipart_http_request will handle it # But if custom_body is provided (e.g., JSON parsed despite multipart content-type), use it is_multipart = ( @@ -883,12 +900,19 @@ async def pass_through_request( # noqa: PLR0915 ) ) else: + # SigV4-signed callers (Bedrock) supply the exact pre-signed bytes; + # otherwise httpx encodes the parsed JSON dict as before. + body_kwargs: Dict[str, Any] = ( + {"content": state_raw_body} + if state_raw_body is not None + else {"json": _parsed_body} + ) req = async_client.build_request( - "POST", + request.method, url, - json=_parsed_body, params=requested_query_params, headers=headers, + **body_kwargs, ) response = await async_client.send(req, stream=stream) @@ -917,17 +941,28 @@ async def pass_through_request( # noqa: PLR0915 status_code=response.status_code, ) - response = ( - await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( - request=request, - async_client=async_client, + if state_raw_body is not None: + # SigV4-signed callers (Bedrock) require the exact pre-signed bytes + # to be forwarded so the signature/Content-Length stay valid. + response = await async_client.request( + method=request.method, url=url, headers=headers, - requested_query_params=requested_query_params, - _parsed_body=_parsed_body, - forward_multipart=is_multipart, + params=requested_query_params, + content=state_raw_body, + ) + else: + response = ( + await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( + request=request, + async_client=async_client, + url=url, + headers=headers, + requested_query_params=requested_query_params, + _parsed_body=_parsed_body, + forward_multipart=is_multipart, + ) ) - ) verbose_proxy_logger.debug("response.headers= %s", response.headers) if _is_streaming_response(response) is True: @@ -1225,7 +1260,7 @@ async def _parse_request_data_by_content_type( def create_pass_through_route( endpoint, target: str, - custom_headers: Optional[dict] = None, + custom_headers: Optional[Mapping[str, Any]] = None, _forward_headers: Optional[bool] = False, _merge_query_params: Optional[bool] = False, dependencies: Optional[List] = None, @@ -1272,11 +1307,14 @@ def create_pass_through_route( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), subpath: str = "", # captures sub-paths when include_subpath=True ): + from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415 + get_request_route, + ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, ) - path = request.url.path + path = get_request_route(request) # Parse request data based on content type ( @@ -1335,9 +1373,12 @@ def create_pass_through_route( ) ) - # Ensure custom_headers is a dict + # Ensure custom_headers is a dict. Botocore returns a HeadersDict + # for SigV4-prepared requests, which is a Mapping but not a dict. headers_dict = ( - param_custom_headers if isinstance(param_custom_headers, dict) else {} + dict(param_custom_headers) + if isinstance(param_custom_headers, Mapping) + else {} ) # Ensure query_params and custom_body are dicts or None @@ -1380,6 +1421,8 @@ def create_pass_through_route( finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY) + if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY): + delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY) return endpoint_func diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4af7caead0e..8fbe6d97dbc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -241,7 +241,10 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.litellm_core_utils.sensitive_data_masker import ( + SensitiveDataMasker, + mask_sensitive_keys, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * @@ -990,6 +993,15 @@ _OPENAPI_HTTP_METHODS = { } +# Credentials surfaced by `/get/config/callbacks` in the alerting block: the +# full Slack incoming-webhook URL is itself a credential, and the SMTP +# password is a service password. Masked on read so plaintext never reaches +# the UI. Kept here at module scope to match the analogous +# `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO +# and cache endpoint files. +_ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} + + def _strip_operation_id_method_suffix(operation_id: str) -> str: base, separator, suffix = operation_id.rpartition("_") if separator and suffix in _OPENAPI_HTTP_METHODS: @@ -1059,6 +1071,7 @@ app = FastAPI( root_path=server_root_path, lifespan=proxy_startup_event, # type: ignore[reportGeneralTypeIssues] generate_unique_id_function=_generate_stable_operation_id, + strict_content_type=False, ) vertex_live_passthrough_vertex_base = VertexBase() @@ -14708,6 +14721,9 @@ async def get_config(): # noqa: PLR0915 value=env_variable, key=_var ) _slack_env_vars[_var] = _decrypted_value + _slack_env_vars = mask_sensitive_keys( + _slack_env_vars, _ALERTING_SENSITIVE_VARS + ) _alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types _all_alert_types = ( @@ -14744,6 +14760,7 @@ async def get_config(): # noqa: PLR0915 # decode + decrypt the value _decrypted_value = decrypt_value_helper(value=env_variable, key=_var) _email_env_vars[_var] = _decrypted_value + _email_env_vars = mask_sensitive_keys(_email_env_vars, _ALERTING_SENSITIVE_VARS) alerting_data.append( { diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index e3019801aae..36beb5e9aba 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1817,7 +1817,10 @@ async def ui_view_spend_logs( # noqa: PLR0915 ) try: - is_v2 = "/spend/logs/v2" in request.url.path + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + is_v2 = "/spend/logs/v2" in get_request_route(request) formats = ["%Y-%m-%d %H:%M:%S", "%Y-%m-%d"] if is_v2 else ["%Y-%m-%d %H:%M:%S"] def parse_date(date_str: str) -> datetime: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index db3ae9ad942..07e2ca71950 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -9,6 +9,7 @@ from pydantic.fields import FieldInfo import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -19,6 +20,16 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router = APIRouter() +# SSO secret fields returned by /get/sso_settings. These are masked on read so +# the UI can show "(set)" without ever transporting the plaintext OAuth secret +# off the server, matching the write-once + masked-on-read contract used for +# the HashiCorp Vault config override. +_SSO_SENSITIVE_FIELDS: Set[str] = { + "google_client_secret", + "microsoft_client_secret", + "generic_client_secret", +} + class IPAddress(BaseModel): ip: str @@ -728,8 +739,9 @@ async def get_sso_settings(): schema = TypeAdapter(SSOConfig).json_schema(by_alias=True) - # Convert to dict for response - sso_dict = sso_config.model_dump() + # Convert to dict for response, masking OAuth client secrets so plaintext + # is never sent to the UI. + sso_dict = mask_sensitive_keys(sso_config.model_dump(), _SSO_SENSITIVE_FIELDS) # Add descriptions to the response result = { diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 657b520b271..1221ccf119f 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -330,11 +330,16 @@ def is_allowed_to_call_vector_store_endpoint( provider_config.get_vector_store_endpoints_by_type() ) + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + request_route = get_request_route(request) + # Determine the permission type based on the request permission_type = None for endpoint in provider_vector_store_endpoints["read"]: if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request.url.path + endpoint[1], request_route ): permission_type = "read" break @@ -342,7 +347,7 @@ def is_allowed_to_call_vector_store_endpoint( if permission_type is None: for endpoint in provider_vector_store_endpoints["write"]: if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request.url.path + endpoint[1], request_route ): permission_type = "write" break @@ -392,10 +397,15 @@ def is_allowed_to_call_vector_store_files_endpoint( provider_config.get_vector_store_file_endpoints_by_type() ) + # Inline import — auth_utils participates in a proxy import cycle. + from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 + + request_route = get_request_route(request) + permission_type: Optional[str] = None for endpoint in provider_vector_store_endpoints.get("read", ()): if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request.url.path + endpoint[1], request_route ): permission_type = "read" break @@ -403,7 +413,7 @@ def is_allowed_to_call_vector_store_files_endpoint( if permission_type is None: for endpoint in provider_vector_store_endpoints.get("write", ()): if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request.url.path + endpoint[1], request_route ): permission_type = "write" break diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 35680889d86..e4c713f67c0 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -54,6 +54,7 @@ if TYPE_CHECKING: else: ResponseText = str # Fallback for ResponseText import from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams @@ -1139,6 +1140,9 @@ def responses( "aresponses": _is_async, "litellm_call_id": litellm_call_id, "model_info": kwargs.get("model_info"), + "data_residency": infer_openai_data_residency( + custom_llm_provider, litellm_params.api_base + ), "metadata": ( kwargs["litellm_metadata"] if "litellm_metadata" in kwargs @@ -2032,6 +2036,9 @@ def compact_responses( litellm_params={ **responses_api_request_params, "litellm_call_id": litellm_call_id, + "data_residency": infer_openai_data_residency( + custom_llm_provider, litellm_params.api_base + ), }, custom_llm_provider=custom_llm_provider, ) @@ -2129,6 +2136,11 @@ async def _aresponses_websocket( api_key=api_key, ) + litellm_params_dict["data_residency"] = infer_openai_data_residency( + _custom_llm_provider, + dynamic_api_base or litellm_params.api_base or litellm.api_base, + ) + litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 74e4d7a533a..46a2894bd10 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -738,6 +738,98 @@ class ResponsesAPIRequestUtils: model_id, ) + @staticmethod + def _collect_container_ids_from_annotations( + annotations: Any, + collected: set[str], + ) -> None: + if not annotations or not isinstance(annotations, list): + return + for ann in annotations: + ResponsesAPIRequestUtils._collect_container_ids_from_output_item( + ann, collected + ) + + @staticmethod + def _collect_container_ids_from_message_content( + content: Any, + collected: set[str], + ) -> None: + if not content: + return + if isinstance(content, list): + for part in content: + if isinstance(part, dict): + ResponsesAPIRequestUtils._collect_container_ids_from_annotations( + part.get("annotations"), + collected, + ) + else: + ResponsesAPIRequestUtils._collect_container_ids_from_annotations( + getattr(part, "annotations", None), + collected, + ) + + @staticmethod + def _collect_container_ids_from_output_item( + item: Any, + collected: set[str], + ) -> None: + """Collect managed or raw ``container_id`` values from one output item.""" + if item is None: + return + + if isinstance(item, dict): + cid = item.get("container_id") + if isinstance(cid, str) and cid: + collected.add(cid) + nested = item.get("code_interpreter_call") + if isinstance(nested, dict): + nc = nested.get("container_id") + if isinstance(nc, str) and nc: + collected.add(nc) + if item.get("type") == "message": + ResponsesAPIRequestUtils._collect_container_ids_from_message_content( + item.get("content"), + collected, + ) + return + + cid_attr = getattr(item, "container_id", None) + if isinstance(cid_attr, str) and cid_attr: + collected.add(cid_attr) + + nested_obj = getattr(item, "code_interpreter_call", None) + if nested_obj is not None: + ResponsesAPIRequestUtils._collect_container_ids_from_output_item( + nested_obj, collected + ) + + if getattr(item, "type", None) == "message": + ResponsesAPIRequestUtils._collect_container_ids_from_message_content( + getattr(item, "content", None), + collected, + ) + + @staticmethod + def collect_container_ids_from_responses_response(response: Any) -> list[str]: + """Return unique container IDs referenced in a Responses API payload.""" + if response is None: + return [] + + if isinstance(response, dict): + output = response.get("output", []) + else: + output = getattr(response, "output", []) or [] + + collected: set[str] = set() + if output: + for item in output: + ResponsesAPIRequestUtils._collect_container_ids_from_output_item( + item, collected + ) + return list(collected) + @staticmethod def _update_container_ids_in_response( responses_api_response: Union[ResponsesAPIResponse, Dict[str, Any]], diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index f70cfad7fb5..862ca13e7ba 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -52,11 +52,12 @@ PROVIDERS: List[Dict] = [ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", + "description": "Claude Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ + "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6", "claude-sonnet-4-6", diff --git a/litellm/types/integrations/datadog_cost_management.py b/litellm/types/integrations/datadog_cost_management.py index fe04f43ea03..08744d2f52e 100644 --- a/litellm/types/integrations/datadog_cost_management.py +++ b/litellm/types/integrations/datadog_cost_management.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, TypedDict +from typing import Dict, List, Optional, TypedDict from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams @@ -9,7 +9,7 @@ class DatadogCostManagementInitParams(StandardCustomLoggerInitParams): Init params for Datadog Cost Management """ - datadog_cost_management_params: Optional[Dict] = None + cost_tag_keys: Optional[List[str]] = None class DatadogFOCUSCostEntry(TypedDict): diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 1c4d31d21ad..bbb892a0276 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -39,7 +39,8 @@ class AnthropicOutputSchema(TypedDict, total=False): class AnthropicOutputConfig(TypedDict, total=False): """Configuration for controlling Claude's output behavior.""" - effort: Literal["high", "medium", "low"] + effort: Literal["high", "medium", "low", "xhigh", "max"] + format: AnthropicOutputSchema class AnthropicMessagesTool(TypedDict, total=False): diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index 9e3fea1bbbb..8763544facc 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -133,7 +133,7 @@ class BidiGenerateContentSetup(TypedDict, total=False): tools: List[Tools] """The tools to be used for the realtime session.""" - realtimeInputConfig: dict + realtimeInputConfig: BidiGenerateContentRealtimeInputConfig """The realtime config to be used for the realtime session.""" sessionResumption: dict diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index abe58199dfd..14114b22f39 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -79,7 +79,14 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import Annotated, Dict, Required, TypedDict, override +from typing_extensions import ( + Annotated, + Dict, + NotRequired, + Required, + TypedDict, + override, +) from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.responses.main import ( @@ -1935,6 +1942,7 @@ class OpenAIRealtimeStreamResponseOutputItemAdded(TypedDict): response_id: str output_index: int item: OpenAIRealtimeStreamResponseOutputItem + event_id: NotRequired[str] class OpenAIRealtimeStreamResponseBaseObject(TypedDict): @@ -2061,6 +2069,17 @@ class OpenAIRealtimeContentPartDone(TypedDict): type: Literal["response.content_part.done"] +class OpenAIRealtimeFunctionCallArgumentsDone(TypedDict): + type: Literal["response.function_call_arguments.done"] + event_id: str + response_id: str + item_id: str + output_index: int + call_id: str + name: str + arguments: str + + class OpenAIRealtimeOutputItemDone(TypedDict): event_id: str item: OpenAIRealtimeStreamResponseOutputItem @@ -2126,6 +2145,7 @@ OpenAIRealtimeEvents = Union[ OpenAIRealtimeResponseAudioDone, OpenAIRealtimeContentPartDone, OpenAIRealtimeOutputItemDone, + OpenAIRealtimeFunctionCallArgumentsDone, OpenAIRealtimeDoneEvent, ] diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 4a07fa5e849..3524a7eb7f7 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -7,6 +7,10 @@ from typing_extensions import TypedDict # JSON without a FastAPI `custom_body` parameter (which would consume the HTTP body). LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY = "litellm_pass_through_custom_body" +# Request.state key for programmatic pass-through callers that must preserve an +# exact byte/string body, such as AWS SigV4-signed requests. +LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY = "litellm_pass_through_raw_body" + class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" diff --git a/litellm/types/router.py b/litellm/types/router.py index 6601f552b52..ef7eb05d087 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -398,6 +398,8 @@ SPECIAL_MODEL_INFO_PARAMS = [ "output_cost_per_token", "input_cost_per_character", "output_cost_per_character", + "cache_read_input_token_cost", + "cache_creation_input_token_cost", ] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 282baff07fe..8f471b62b5e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -148,6 +148,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_xhigh_reasoning_effort: Optional[bool] supports_max_reasoning_effort: Optional[bool] supports_output_config: Optional[bool] + bedrock_output_config_effort_ceiling: Optional[ + Literal["low", "medium", "high", "max", "xhigh"] + ] class SearchContextCostPerQuery(TypedDict, total=False): @@ -219,6 +222,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_priority: Optional[ float ] # OpenAI priority service tier pricing + regional_processing_uplift_multiplier_eu: Optional[ + float + ] # OpenAI EU data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%) + regional_processing_uplift_multiplier_us: Optional[ + float + ] # OpenAI US data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%) output_cost_per_character: Optional[float] # only for vertex ai models output_cost_per_audio_token: Optional[float] output_cost_per_token_above_128k_tokens: Optional[ @@ -3601,6 +3610,20 @@ class ServiceTier(Enum): PRIORITY = "priority" +class DataResidency(Enum): + """ + OpenAI data-residency / regional-processing regions. + + Inferred from the OpenAI api_base host (eu.api.openai.com -> EU, + us.api.openai.com -> US). Used to apply the regional-processing + cost uplift (see ``regional_processing_uplift_multiplier_`` + on ModelInfo). + """ + + US = "us" + EU = "eu" + + LLMResponseTypes = Union[ ModelResponse, EmbeddingResponse, diff --git a/litellm/utils.py b/litellm/utils.py index 2ba6ef9cae8..5a9dccc089e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5942,6 +5942,12 @@ def _get_model_info_helper( # noqa: PLR0915 output_cost_per_token_priority=_model_info.get( "output_cost_per_token_priority", None ), + regional_processing_uplift_multiplier_eu=_model_info.get( + "regional_processing_uplift_multiplier_eu", None + ), + regional_processing_uplift_multiplier_us=_model_info.get( + "regional_processing_uplift_multiplier_us", None + ), output_cost_per_audio_token=_model_info.get( "output_cost_per_audio_token", None ), @@ -6030,6 +6036,9 @@ def _get_model_info_helper( # noqa: PLR0915 supports_max_reasoning_effort=_model_info.get( "supports_max_reasoning_effort", None ), + bedrock_output_config_effort_ceiling=_model_info.get( + "bedrock_output_config_effort_ceiling", None + ), supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get( "search_context_cost_per_query", None diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 2160514251a..a78a4e2225a 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -31,12 +31,20 @@ USER root COPY --from=uvbin /uv /uvx /usr/local/bin/ -RUN apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile +# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE +# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi +# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. +RUN for i in 1 2 3; do \ + apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ UV_COMPILE_BYTECODE=1 \ UV_PYTHON_DOWNLOADS=0 \ + PRISMA_USE_GLOBAL_NODE=true \ PATH="/app/.venv/bin:${PATH}" # Stage 1 — install third-party deps only (cached by pyproject.toml/uv.lock). @@ -78,7 +86,11 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root -RUN apk add --no-cache bash openssl tzdata python3 libsndfile libatomic +RUN for i in 1 2 3; do \ + apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \ + [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ + sleep 5; \ + done # wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532). The # Prisma engine binaries are dynamically linked against libssl/libcrypto, so diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 62553e46ac1..80c2f32dc70 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -731,7 +731,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "anthropic.claude-haiku-4-5@20251001": { @@ -755,7 +754,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_streaming": true, "supports_native_structured_output": true }, @@ -926,8 +924,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -952,8 +949,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, @@ -977,12 +973,12 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, @@ -1009,11 +1005,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "global.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, @@ -1040,11 +1035,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "us.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, @@ -1071,11 +1065,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "eu.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, @@ -1101,11 +1094,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "au.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, @@ -1131,11 +1123,10 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_output_config": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "bedrock_output_config_effort_ceiling": "max" }, "anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -1163,10 +1154,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "anthropic.claude-mythos-preview": { "input_cost_per_token": 0, @@ -1180,8 +1171,8 @@ "supports_vision": true, "supports_prompt_caching": false, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_output_config": true }, "global.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -1209,10 +1200,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "us.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, @@ -1240,10 +1231,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "eu.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, @@ -1270,10 +1261,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "au.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, @@ -1300,10 +1291,165 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "au.anthropic.claude-opus-4-8": { + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" }, "anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, @@ -1331,10 +1477,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "global.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, @@ -1362,10 +1506,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "us.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1393,10 +1535,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "eu.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1423,10 +1563,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "au.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1453,10 +1591,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "jp.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, @@ -1483,10 +1619,8 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -1515,8 +1649,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -1548,7 +1681,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, "supports_native_structured_output": true }, "anthropic.claude-v1": { @@ -1799,7 +1931,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "apac.anthropic.claude-3-sonnet-20240229-v1:0": { @@ -1845,8 +1976,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "assemblyai/best": { "input_cost_per_second": 3.333e-05, @@ -1888,7 +2018,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "azure/ada": { @@ -1976,10 +2105,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_output_config": true }, "azure_ai/claude-opus-4-6": { "input_cost_per_token": 5e-06, @@ -2006,10 +2135,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, "supports_output_config": true, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-7": { "input_cost_per_token": 5e-06, @@ -2037,9 +2164,35 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 159, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-opus-4-8": { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -2104,9 +2257,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "azure/computer-use-preview": { "input_cost_per_token": 3e-06, @@ -9480,8 +9631,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 + "supports_web_search": true }, "claude-3-haiku-20240307": { "cache_creation_input_token_cost": 3e-07, @@ -9499,8 +9649,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 264 + "supports_vision": true }, "claude-3-opus-20240229": { "cache_creation_input_token_cost": 1.875e-05, @@ -9519,8 +9668,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 395 + "supports_vision": true }, "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, @@ -9545,8 +9693,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-4-sonnet-20250514": { "cache_creation_input_token_cost": 3.75e-06, @@ -9576,8 +9723,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 + "supports_web_search": true }, "claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.75e-06, @@ -9606,8 +9752,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 346 + "supports_vision": true }, "claude-sonnet-4-5-20250929": { "cache_creation_input_token_cost": 3.75e-06, @@ -9637,8 +9782,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 346 + "supports_web_search": true }, "claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, @@ -9666,9 +9810,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -9692,8 +9834,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -9719,8 +9860,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-1-20250805": { "cache_creation_input_token_cost": 1.875e-05, @@ -9747,8 +9887,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-20250514": { "cache_creation_input_token_cost": 1.875e-05, @@ -9775,8 +9914,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "claude-opus-4-5-20251101": { "cache_creation_input_token_cost": 6.25e-06, @@ -9800,11 +9938,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "claude-opus-4-5": { "cache_creation_input_token_cost": 6.25e-06, @@ -9828,11 +9965,10 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "claude-opus-4-6": { "cache_creation_input_token_cost": 6.25e-06, @@ -9860,14 +9996,12 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, "supports_output_config": true, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true }, "claude-opus-4-6-20260205": { "cache_creation_input_token_cost": 6.25e-06, @@ -9895,13 +10029,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true, "supports_output_config": true }, "claude-opus-4-7": { @@ -9932,12 +10064,10 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true, "supports_output_config": true }, "claude-opus-4-7-20260416": { @@ -9968,12 +10098,44 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, "provider_specific_entry": { "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true, + "supports_output_config": true + }, + "claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1, + "fast": 2.0 + }, "supports_output_config": true }, "claude-sonnet-4-20250514": { @@ -10005,8 +10167,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "cloudflare/@cf/meta/llama-2-7b-chat-fp16": { "input_cost_per_token": 1.923e-06, @@ -11251,8 +11412,8 @@ "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, - "supports_minimal_reasoning_effort": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_output_config": true }, "databricks/databricks-claude-sonnet-4": { "input_cost_per_token": 2.9999900000000002e-06, @@ -13418,7 +13579,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": { @@ -13546,8 +13706,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "eu.anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -13572,8 +13731,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "eu.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -13602,8 +13760,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "eu.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, @@ -13633,7 +13790,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "eu.meta.llama3-2-1b-instruct-v1:0": { @@ -17960,7 +18116,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_vision": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "github_copilot/claude-opus-4.6-fast": { "litellm_provider": "github_copilot", @@ -18474,7 +18630,7 @@ "output_cost_per_token": 2.5e-05, "supports_function_calling": true, "supports_vision": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "gmi/anthropic/claude-sonnet-4.5": { "input_cost_per_token": 3e-06, @@ -18774,7 +18930,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "global.anthropic.claude-sonnet-4-20250514-v1:0": { @@ -18804,8 +18959,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "global.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.25e-06, @@ -18828,7 +18982,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "global.amazon.nova-2-lite-v1:0": { @@ -19050,6 +19203,8 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19123,6 +19278,8 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19196,6 +19353,8 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -19267,6 +19426,8 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -19308,6 +19469,8 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -19329,6 +19492,8 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -19617,6 +19782,8 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_priority": 1e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20320,6 +20487,8 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21242,6 +21411,8 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -21648,6 +21819,8 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21729,6 +21902,8 @@ "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, @@ -22914,7 +23089,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "jp.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -22937,7 +23111,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "crusoe/deepseek-ai/DeepSeek-R1-0528": { @@ -26946,8 +27119,7 @@ "supports_computer_use": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-3.7-sonnet": { "input_cost_per_image": 0.0048, @@ -26963,8 +27135,7 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-opus-4": { "input_cost_per_image": 0.0048, @@ -26983,8 +27154,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-opus-4.1": { "input_cost_per_image": 0.0048, @@ -27004,8 +27174,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-sonnet-4": { "input_cost_per_image": 0.0048, @@ -27028,8 +27197,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-sonnet-4.6": { "cache_creation_input_token_cost": 3.75e-06, @@ -27053,9 +27221,7 @@ "supports_reasoning": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_minimal_reasoning_effort": true + "supports_vision": true }, "openrouter/anthropic/claude-opus-4.5": { "cache_creation_input_token_cost": 6.25e-06, @@ -27070,12 +27236,11 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "openrouter/anthropic/claude-opus-4.6": { "cache_creation_input_token_cost": 6.25e-06, @@ -27094,9 +27259,7 @@ "supports_reasoning": true, "supports_max_reasoning_effort": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 346, - "supports_minimal_reasoning_effort": true + "supports_vision": true }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -27119,8 +27282,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "openrouter/anthropic/claude-haiku-4.5": { "cache_creation_input_token_cost": 1.25e-06, @@ -27138,8 +27300,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 346 + "supports_vision": true }, "openrouter/anthropic/claude-opus-4.7": { "cache_creation_input_token_cost": 6.25e-06, @@ -27161,8 +27322,7 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346 + "supports_xhigh_reasoning_effort": true }, "openrouter/bytedance/ui-tars-1.5-7b": { "input_cost_per_token": 1e-07, @@ -29143,7 +29303,7 @@ "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "perplexity/anthropic/claude-sonnet-4-5": { "litellm_provider": "perplexity", @@ -31396,7 +31556,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "us.anthropic.claude-3-5-sonnet-20240620-v1:0": { @@ -31524,8 +31683,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "us.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, @@ -31557,7 +31715,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0": { @@ -31583,7 +31740,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "au.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -31605,7 +31761,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, "us.anthropic.claude-opus-4-20250514-v1:0": { @@ -31631,8 +31786,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "us.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.875e-06, @@ -31653,15 +31807,15 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "global.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, @@ -31682,15 +31836,15 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "eu.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, @@ -31710,15 +31864,15 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "high" }, "us.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -31747,8 +31901,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "us.deepseek.r1-v1:0": { "input_cost_per_token": 1.35e-06, @@ -32291,13 +32444,13 @@ "output_cost_per_token": 2.5e-05, "supports_assistant_prefill": true, "supports_computer_use": true, - "supports_minimal_reasoning_effort": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_output_config": true }, "vercel_ai_gateway/anthropic/claude-opus-4.6": { "cache_creation_input_token_cost": 6.25e-06, @@ -32317,7 +32470,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "vercel_ai_gateway/anthropic/claude-sonnet-4": { "cache_creation_input_token_cost": 3.75e-06, @@ -33325,8 +33478,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-3-haiku": { "input_cost_per_token": 2.5e-07, @@ -33429,8 +33581,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -33484,14 +33635,13 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_output_config": true }, "vertex_ai/claude-opus-4-5@20251101": { "cache_creation_input_token_cost": 6.25e-06, @@ -33511,15 +33661,14 @@ "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, - "supports_minimal_reasoning_effort": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 159, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_output_config": true }, "vertex_ai/claude-opus-4-6": { "cache_creation_input_token_cost": 6.25e-06, @@ -33545,10 +33694,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_output_config": true, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-4-6@default": { "cache_creation_input_token_cost": 6.25e-06, @@ -33574,10 +33721,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_output_config": true, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -33604,9 +33749,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-4-7@default": { "cache_creation_input_token_cost": 6.25e-06, @@ -33633,9 +33776,63 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "tool_use_system_prompt_tokens": 346, - "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-opus-4-8@default": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "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, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true }, "vertex_ai/claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.75e-06, @@ -33683,14 +33880,12 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "vertex_ai/claude-sonnet-4-5@20250929": { "cache_creation_input_token_cost": 3.75e-06, @@ -33742,8 +33937,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-sonnet-4": { "cache_creation_input_token_cost": 3.75e-06, @@ -33772,8 +33966,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/claude-sonnet-4@20250514": { "cache_creation_input_token_cost": 3.75e-06, @@ -33802,8 +33995,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 + "supports_vision": true }, "vertex_ai/mistralai/codestral-2@001": { "input_cost_per_token": 3e-07, @@ -40805,14 +40997,12 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_output_config": true, - "supports_minimal_reasoning_effort": true + "supports_output_config": true }, "duckduckgo/search": { "litellm_provider": "duckduckgo", @@ -41128,7 +41318,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_pdf_input": true }, @@ -41151,7 +41340,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, "supports_pdf_input": true } diff --git a/pyproject.toml b/pyproject.toml index 8dedca241ad..6e84afad17e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.87.0" +version = "1.88.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -33,62 +33,66 @@ Homepage = "https://litellm.ai" Repository = "https://github.com/BerriAI/litellm" Documentation = "https://docs.litellm.ai" -# Optional extras retain exact pins because they are consumed by Docker images -# where exact reproducibility matters. The core SDK uses ranges so downstream -# consumers can coexist with other packages without forced downgrades. +# Optional extras use compatible ranges (like the core SDK above) so downstream +# consumers can coexist with other packages and pick up security patches without +# forking. Reproducibility for our Docker/CI comes from `uv.lock` (images install +# via `uv sync --frozen`). A few deps stay exact-pinned: litellm's own +# sub-packages and the opentelemetry trio move in lockstep, and grpcio is +# supply-chain-pinned to a vetted, aged release. [project.optional-dependencies] proxy = [ - "gunicorn==23.0.0", - "uvicorn==0.33.0", - "granian==2.5.7", - "uvloop==0.21.0; sys_platform != 'win32'", - "fastapi==0.124.4", - "backoff==2.2.1", - "pyyaml==6.0.3", - "rq==2.7.0", - "orjson==3.11.6", - "apscheduler==3.11.2", - "fastapi-sso==0.19.0", - "PyJWT==2.12.0", - "python-multipart==0.0.27", - "cryptography==46.0.7", - "pynacl==1.6.2", - "websockets==15.0.1", - "boto3==1.43.1", - "azure-identity==1.25.2", - "azure-storage-blob==12.28.0", - "mcp==1.26.0", + "gunicorn>=23.0.0,<24.0", + "uvicorn>=0.33.0,<1.0", + "granian>=2.7.4,<3.0", + "uvloop>=0.21.0,<1.0; sys_platform != 'win32'", + "fastapi>=0.136.3,<1.0", + "starlette>=1.0.1,<2.0", + "backoff>=2.2.1,<3.0", + "pyyaml>=6.0.3,<7.0", + "rq>=2.7.0,<3.0", + "orjson>=3.11.6,<4.0", + "apscheduler>=3.11.2,<4.0", + "fastapi-sso>=0.19.0,<1.0", + "PyJWT>=2.12.0,<3.0", + "python-multipart>=0.0.27,<1.0", + "cryptography>=46.0.7,<47.0", + "pynacl>=1.6.2,<2.0", + "websockets>=15.0.1,<16.0", + "boto3>=1.43.1,<2.0", + "azure-identity>=1.25.2,<2.0", + "azure-storage-blob>=12.28.0,<13.0", + "mcp>=1.26.0,<2.0", "litellm-proxy-extras==0.4.73", "litellm-enterprise==0.1.41", - "RestrictedPython==8.1", - "rich==13.9.4", - "polars==1.38.1", - "soundfile==0.12.1", - "pyroscope-io==0.8.16; sys_platform != 'win32'", - "pydantic-settings>=2.14.1", + "RestrictedPython>=8.1,<9.0", + "rich>=13.9.4,<14.0", + "polars>=1.38.1,<2.0", + "soundfile>=0.12.1,<1.0", + "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", + "pydantic-settings>=2.14.1,<3.0", ] extra_proxy = [ - "prisma==0.11.0", - "azure-identity==1.25.2", - "azure-keyvault-secrets==4.10.0", + "prisma>=0.11.0,<1.0", + "azure-identity>=1.25.2,<2.0", + "azure-keyvault-secrets>=4.10.0,<5.0", # Not in PyPI proxy extra. - "google-cloud-kms==2.24.2", - "google-cloud-iam==2.19.1", + "google-cloud-kms>=2.24.2,<3.0", + "google-cloud-iam>=2.19.1,<3.0", # Not in PyPI proxy extra. - "resend==2.23.0", - "redisvl==0.4.1; python_version < '3.14'", - "a2a-sdk==0.3.24", + "resend>=2.23.0,<3.0", + "redisvl>=0.4.1,<1.0; python_version < '3.14'", + "a2a-sdk>=0.3.24,<1.0", ] utils = [ # Not in Docker or PyPI proxy extra. - "numpydoc==1.8.0", + "numpydoc>=1.8.0,<2.0", ] -caching = ["diskcache==5.6.3"] +caching = ["diskcache>=5.6.3,<6.0"] semantic-router = [ - "semantic-router==0.1.12; python_version < '3.14'", - "aurelio-sdk==0.0.19; python_version < '3.14'", + "semantic-router>=0.1.15,<1.0; python_version < '3.14'", + "aurelio-sdk>=0.0.19,<1.0; python_version < '3.14'", ] -mlflow = ["mlflow==3.11.1"] +mlflow = ["mlflow>=3.11.1,<4.0"] grpc = [ # Newest non-yanked release older than the 30-day cutoff. "grpcio==1.78.0", @@ -101,28 +105,28 @@ stt-nvidia-riva = [ "audioread>=3.0.1", "numpy>=1.26.0", ] -google = ["google-cloud-aiplatform==1.133.0"] +google = ["google-cloud-aiplatform>=1.133.0,<2.0"] proxy-runtime = [ # Historically bundled in the proxy Docker images via requirements.txt. # Keep these in a dedicated extra so uv-based images preserve the same # feature surface without forcing the base SDK install to grow. - "google-cloud-aiplatform==1.133.0", - "google-genai==1.37.0", - "anthropic[vertex]==0.84.0", + "google-cloud-aiplatform>=1.133.0,<2.0", + "google-genai>=1.37.0,<2.0", + "anthropic[vertex]>=0.84.0,<1.0", "grpcio==1.78.0", - "prometheus-client==0.20.0", - "langfuse==2.59.7", + "prometheus-client>=0.20.0,<1.0", + "langfuse>=2.59.7,<3.0", "opentelemetry-api==1.28.0", "opentelemetry-sdk==1.28.0", "opentelemetry-exporter-otlp==1.28.0", - "ddtrace==2.19.0", - "sentry-sdk==2.21.0", - "mangum==0.17.0", - "azure-ai-contentsafety==1.0.0", - "azure-storage-file-datalake==12.20.0", - "pypdf==6.10.2; python_version < '3.14'", - "llm-sandbox==0.3.39", - "detect-secrets==1.5.0", + "ddtrace>=2.19.0,<3.0", + "sentry-sdk>=2.21.0,<3.0", + "mangum>=0.17.0,<1.0", + "azure-ai-contentsafety>=1.0.0,<2.0", + "azure-storage-file-datalake>=12.20.0,<13.0", + "pypdf>=6.10.2,<7.0; python_version < '3.14'", + "llm-sandbox>=0.3.39,<1.0", + "detect-secrets>=1.5.0,<2.0", ] [project.scripts] @@ -188,7 +192,7 @@ ci = [ "psycopg2-binary==2.9.11", "pytest-codspeed==4.3.0", "pytest-retry==1.7.0", - "pyarrow==22.0.0", + "pyarrow==23.0.1", "langchain==1.2.10", "lunary==1.4.36; python_version == '3.10'", "lunary==1.4.37; python_version >= '3.11'", @@ -253,7 +257,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.87.0" +version = "1.88.0" version_files = [ "pyproject.toml:^version", ] diff --git a/scripts/benchmark_model_response_creator.py b/scripts/benchmark_model_response_creator.py new file mode 100644 index 00000000000..881870d3854 --- /dev/null +++ b/scripts/benchmark_model_response_creator.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python3 +"""Tight microbenchmark for CustomStreamWrapper.model_response_creator. + +Calls model_response_creator() in a tight loop on a pre-built wrapper to +isolate per-call cost. Driving the full wrapper adds threadpool logging, +gc, and other noise that swamps microsecond-scale changes here. + +Example: + uv run python scripts/benchmark_model_response_creator.py --label baseline + uv run python scripts/benchmark_model_response_creator.py --label optimized +""" + +from __future__ import annotations + +import argparse +import gc +import json +import logging +import os +import statistics +import time +from dataclasses import asdict, dataclass +from typing import List +from unittest.mock import MagicMock + +os.environ.setdefault("LITELLM_LOG", "ERROR") +logging.getLogger("LiteLLM").setLevel(logging.ERROR) + +import litellm # noqa: E402 + +litellm.suppress_debug_info = True + +from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper, +) # noqa: E402 + + +def _make_logging_obj(provider: str) -> MagicMock: + logging_obj = MagicMock() + logging_obj.model_call_details = { + "custom_llm_provider": provider, + "litellm_params": {}, + } + logging_obj.call_type = "completion" + logging_obj.stream_options = None + logging_obj.messages = [{"role": "user", "content": "hi"}] + logging_obj.completion_start_time = None + logging_obj._llm_caching_handler = None + return logging_obj + + +def _make_wrapper(provider: str, model: str) -> CustomStreamWrapper: + return CustomStreamWrapper( + completion_stream=iter([]), + model=model, + logging_obj=_make_logging_obj(provider), + custom_llm_provider=provider, + ) + + +@dataclass +class Result: + label: str + scenario: str + iterations: int + elapsed_min_s: float + elapsed_median_s: float + per_call_us: float + calls_per_sec: float + + +SCENARIOS = { + "no_chunk": { + "description": "model_response_creator() — no chunk arg (most common path)", + "chunk_factory": lambda i: None, + }, + "text_chunk": { + "description": "model_response_creator(chunk={'text': '...'}) — text delta path", + "chunk_factory": lambda i: {"text": f"token{i}"}, + }, + "rich_chunk": { + "description": "model_response_creator(chunk={...}) — full chunk dict path", + "chunk_factory": lambda i: { + "id": f"id-{i}", + "object": "chat.completion.chunk", + "created": 1234567890, + }, + }, +} + + +def bench_no_chunk(wrapper: CustomStreamWrapper, iterations: int) -> float: + gc.collect() + gc.disable() + try: + start = time.perf_counter() + for _ in range(iterations): + wrapper.model_response_creator() + elapsed = time.perf_counter() - start + finally: + gc.enable() + return elapsed + + +def bench_with_chunk(wrapper: CustomStreamWrapper, factory, iterations: int) -> float: + # Pre-build chunks so we don't measure their construction cost. + chunks = [factory(i) for i in range(iterations)] + gc.collect() + gc.disable() + try: + start = time.perf_counter() + for chunk in chunks: + wrapper.model_response_creator(chunk=dict(chunk)) # copy because mutated + elapsed = time.perf_counter() - start + finally: + gc.enable() + return elapsed + + +def run_scenario( + label: str, + scenario_key: str, + iterations: int, + repeats: int, + warmup: int, +) -> Result: + spec = SCENARIOS[scenario_key] + wrapper = _make_wrapper(provider="anthropic", model="claude-3-5-sonnet") + + if scenario_key == "no_chunk": + runner = lambda: bench_no_chunk(wrapper, iterations) # noqa: E731 + else: + runner = lambda: bench_with_chunk( + wrapper, spec["chunk_factory"], iterations + ) # noqa: E731 + + for _ in range(warmup): + runner() + samples = [runner() for _ in range(repeats)] + + elapsed_min = min(samples) + elapsed_median = statistics.median(samples) + per_call_us = (elapsed_min * 1_000_000) / iterations + calls_per_sec = iterations / elapsed_min if elapsed_min > 0 else 0.0 + + return Result( + label=label, + scenario=scenario_key, + iterations=iterations, + elapsed_min_s=elapsed_min, + elapsed_median_s=elapsed_median, + per_call_us=per_call_us, + calls_per_sec=calls_per_sec, + ) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--label", required=True) + ap.add_argument("--iterations", type=int, default=200_000) + ap.add_argument("--warmup", type=int, default=2) + ap.add_argument("--repeats", type=int, default=8) + ap.add_argument("--json", dest="json_out") + args = ap.parse_args() + + print( + f"\n=== label={args.label} iterations={args.iterations:,} " + f"warmup={args.warmup} repeats={args.repeats} (min reported) ===" + ) + results: List[Result] = [] + for scenario in SCENARIOS: + r = run_scenario( + args.label, scenario, args.iterations, args.repeats, args.warmup + ) + results.append(r) + print( + f" {r.scenario:12s}: " + f"min={r.elapsed_min_s*1000:8.2f} ms " + f"median={r.elapsed_median_s*1000:8.2f} ms " + f"per-call={r.per_call_us:7.3f} μs " + f"calls/s={r.calls_per_sec:>12,.0f}" + ) + + if args.json_out: + with open(args.json_out, "w", encoding="utf-8") as f: + json.dump([asdict(r) for r in results], f, indent=2) + print(f"\nWrote {len(results)} results to {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_streaming_chunk_overhead.py b/scripts/benchmark_streaming_chunk_overhead.py new file mode 100644 index 00000000000..948be096bec --- /dev/null +++ b/scripts/benchmark_streaming_chunk_overhead.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +"""Benchmark CustomStreamWrapper per-chunk overhead. + +Drives CustomStreamWrapper directly with synthetic in-memory chunks for +Anthropic (GenericStreamingChunk), Bedrock Invoke (GenericStreamingChunk), +and Bedrock Converse (ModelResponseStream). A full proxy benchmark adds +FastAPI, HTTP, and TCP latency, which dilutes the per-chunk CPU signal. + +Example: + uv run python scripts/benchmark_streaming_chunk_overhead.py \\ + --streams 500 --chunks 200 --warmup 50 --repeats 5 +""" + +from __future__ import annotations + +import argparse +import asyncio +import gc +import json +import logging +import os +import statistics +import time +from dataclasses import asdict, dataclass +from typing import Callable, List, Optional +from unittest.mock import MagicMock + +# Silence litellm's "Provider List" warnings emitted by get_llm_provider +# when it sees synthetic model names — we're not exercising provider +# routing, only the per-chunk wrapper hot path. +os.environ.setdefault("LITELLM_LOG", "ERROR") +logging.getLogger("LiteLLM").setLevel(logging.ERROR) + +import litellm # noqa: E402 + +litellm.suppress_debug_info = True + +from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper, +) # noqa: E402 +from litellm.types.utils import ( # noqa: E402 + Delta, + GenericStreamingChunk as GChunk, + ModelResponseStream, + StreamingChoices, + Usage, +) + +# --------------------------------------------------------------------------- +# Synthetic chunk fixtures +# --------------------------------------------------------------------------- + + +def _make_logging_obj(provider: str) -> MagicMock: + logging_obj = MagicMock() + logging_obj.model_call_details = { + "custom_llm_provider": provider, + "litellm_params": {}, + } + logging_obj.call_type = "completion" + logging_obj.stream_options = None + logging_obj.messages = [{"role": "user", "content": "hi"}] + logging_obj.completion_start_time = None + logging_obj._llm_caching_handler = None + return logging_obj + + +def _make_generic_chunk( + text: str, + is_finished: bool = False, + finish_reason: str = "", + usage: Optional[dict] = None, +) -> GChunk: + return GChunk( + text=text, + is_finished=is_finished, + finish_reason=finish_reason, + usage=usage, + index=0, + tool_use=None, + ) + + +def _make_converse_chunk( + text: str = "", + finish_reason: str = "", + usage: Optional[Usage] = None, +) -> ModelResponseStream: + return ModelResponseStream( + choices=[ + StreamingChoices( + finish_reason=finish_reason or None, + index=0, + delta=Delta(content=text, role="assistant"), + ) + ], + id="msg-bench", + model="anthropic.claude-3-5-sonnet", + usage=usage, + ) + + +# --------------------------------------------------------------------------- +# Provider stream factories +# --------------------------------------------------------------------------- + + +def anthropic_chunks(n: int) -> List[GChunk]: + out: List[GChunk] = [_make_generic_chunk(f"tok{i} ") for i in range(n)] + out.append( + _make_generic_chunk( + "", + is_finished=True, + finish_reason="stop", + usage={"prompt_tokens": 10, "completion_tokens": n, "total_tokens": 10 + n}, + ) + ) + return out + + +def bedrock_invoke_chunks(n: int) -> List[GChunk]: + # Bedrock Invoke surfaces GChunk-shaped dicts, same shape as Anthropic. + return anthropic_chunks(n) + + +def bedrock_converse_chunks(n: int) -> List[ModelResponseStream]: + out: List[ModelResponseStream] = [ + _make_converse_chunk(f"tok{i} ") for i in range(n) + ] + out.append( + _make_converse_chunk( + text="", + finish_reason="stop", + usage=Usage(prompt_tokens=10, completion_tokens=n, total_tokens=10 + n), + ) + ) + return out + + +PROVIDERS: dict[str, tuple[str, Callable[[int], list]]] = { + "anthropic": ("anthropic", anthropic_chunks), + "bedrock_invoke": ("bedrock", bedrock_invoke_chunks), + "bedrock_converse": ("bedrock", bedrock_converse_chunks), +} + + +# --------------------------------------------------------------------------- +# Drive a single stream end-to-end +# --------------------------------------------------------------------------- + + +def _make_wrapper( + chunks: list, provider: str, async_stream: bool +) -> CustomStreamWrapper: + logging_obj = _make_logging_obj(provider) + if async_stream: + + async def _agen(): + for c in chunks: + yield c + + stream = _agen() + else: + stream = iter(chunks) + return CustomStreamWrapper( + completion_stream=stream, + model="claude-3-5-sonnet", + logging_obj=logging_obj, + custom_llm_provider=provider, + ) + + +def drive_sync(provider_key: str, chunks_per_stream: int, n_streams: int) -> float: + provider, factory = PROVIDERS[provider_key] + # Pre-build the chunk lists; we only measure wrapper iteration cost. + chunk_lists = [factory(chunks_per_stream) for _ in range(n_streams)] + gc.collect() + gc.disable() + try: + start = time.perf_counter() + for chunks in chunk_lists: + wrapper = _make_wrapper(chunks, provider, async_stream=False) + for _ in wrapper: + pass + elapsed = time.perf_counter() - start + finally: + gc.enable() + return elapsed + + +async def drive_async( + provider_key: str, chunks_per_stream: int, n_streams: int +) -> float: + provider, factory = PROVIDERS[provider_key] + chunk_lists = [factory(chunks_per_stream) for _ in range(n_streams)] + gc.collect() + gc.disable() + try: + start = time.perf_counter() + for chunks in chunk_lists: + wrapper = _make_wrapper(chunks, provider, async_stream=True) + async for _ in wrapper: + pass + elapsed = time.perf_counter() - start + finally: + gc.enable() + return elapsed + + +# --------------------------------------------------------------------------- +# Repeat × take-min runner +# --------------------------------------------------------------------------- + + +@dataclass +class Result: + label: str + provider: str + mode: str + streams: int + chunks_per_stream: int + total_chunks: int + elapsed_min_s: float + elapsed_median_s: float + per_chunk_us: float + chunks_per_sec: float + streams_per_sec: float + + +def run_case( + label: str, + provider_key: str, + mode: str, + chunks_per_stream: int, + n_streams: int, + repeats: int, + warmup: int, +) -> Result: + if mode == "sync": + # Warmup runs amortize import-time and JIT-y caches. + for _ in range(warmup): + drive_sync(provider_key, chunks_per_stream, max(1, n_streams // 10)) + samples = [ + drive_sync(provider_key, chunks_per_stream, n_streams) + for _ in range(repeats) + ] + elif mode == "async": + + async def _warm(): + for _ in range(warmup): + await drive_async( + provider_key, chunks_per_stream, max(1, n_streams // 10) + ) + + asyncio.run(_warm()) + samples = [ + asyncio.run(drive_async(provider_key, chunks_per_stream, n_streams)) + for _ in range(repeats) + ] + else: + raise ValueError(f"unknown mode {mode!r}") + + elapsed_min = min(samples) + elapsed_median = statistics.median(samples) + # Each stream emits chunks_per_stream text chunks + 1 finish/usage chunk. + total_chunks = n_streams * (chunks_per_stream + 1) + per_chunk_us = (elapsed_min * 1_000_000) / total_chunks + chunks_per_sec = total_chunks / elapsed_min if elapsed_min > 0 else 0.0 + streams_per_sec = n_streams / elapsed_min if elapsed_min > 0 else 0.0 + + return Result( + label=label, + provider=provider_key, + mode=mode, + streams=n_streams, + chunks_per_stream=chunks_per_stream, + total_chunks=total_chunks, + elapsed_min_s=elapsed_min, + elapsed_median_s=elapsed_median, + per_chunk_us=per_chunk_us, + chunks_per_sec=chunks_per_sec, + streams_per_sec=streams_per_sec, + ) + + +def format_result(r: Result) -> str: + return ( + f" {r.provider:18s} {r.mode:5s}: " + f"min={r.elapsed_min_s*1000:8.2f} ms " + f"median={r.elapsed_median_s*1000:8.2f} ms " + f"per-chunk={r.per_chunk_us:7.2f} μs " + f"chunks/s={r.chunks_per_sec:>10,.0f} " + f"streams/s={r.streams_per_sec:>8,.1f}" + ) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument( + "--label", required=True, help="Run label (e.g. baseline / optimized)" + ) + ap.add_argument("--streams", type=int, default=500, help="Streams per run") + ap.add_argument( + "--chunks", + type=int, + default=200, + help="Text chunks per stream (excl. finish chunk)", + ) + ap.add_argument("--warmup", type=int, default=2, help="Warmup runs") + ap.add_argument( + "--repeats", type=int, default=5, help="Measured runs (we report min)" + ) + ap.add_argument( + "--providers", + default="anthropic,bedrock_invoke,bedrock_converse", + help="Comma-separated provider list", + ) + ap.add_argument( + "--modes", + default="sync,async", + help="Comma-separated iteration modes (sync/async)", + ) + ap.add_argument( + "--json", dest="json_out", help="Write results as JSON to this path" + ) + args = ap.parse_args() + + providers = [p.strip() for p in args.providers.split(",") if p.strip()] + modes = [m.strip() for m in args.modes.split(",") if m.strip()] + + for p in providers: + if p not in PROVIDERS: + raise SystemExit(f"unknown provider {p!r}; choose from {list(PROVIDERS)}") + for m in modes: + if m not in {"sync", "async"}: + raise SystemExit(f"unknown mode {m!r}; choose from sync/async") + + print( + f"\n=== label={args.label} streams={args.streams} chunks/stream={args.chunks} " + f"warmup={args.warmup} repeats={args.repeats} (min reported) ===" + ) + results: List[Result] = [] + for provider_key in providers: + for mode in modes: + r = run_case( + label=args.label, + provider_key=provider_key, + mode=mode, + chunks_per_stream=args.chunks, + n_streams=args.streams, + repeats=args.repeats, + warmup=args.warmup, + ) + results.append(r) + print(format_result(r)) + + if args.json_out: + with open(args.json_out, "w", encoding="utf-8") as f: + json.dump([asdict(r) for r in results], f, indent=2) + print(f"\nWrote {len(results)} results to {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/tests/_vcr_conftest_common.py b/tests/_vcr_conftest_common.py index cb43f1abbdd..d08b87bd580 100644 --- a/tests/_vcr_conftest_common.py +++ b/tests/_vcr_conftest_common.py @@ -13,10 +13,12 @@ import os import re import socket import sys +import threading from collections import defaultdict from typing import Iterable import pytest +import vcr.matchers as _vcr_matchers from tests._vcr_redis_persister import ( MAX_EPISODES_PER_CASSETTE, @@ -29,11 +31,28 @@ from tests._vcr_redis_persister import ( patch_vcrpy_aiohttp_record_path, ) +# Force litellm to use its bundled model-cost-map backup instead of fetching it +# from raw.githubusercontent.com on import. Several VCR conftests reload litellm +# in an autouse fixture (``importlib.reload(litellm)``); ``litellm.__init__`` +# calls ``get_model_cost_map()`` which issues a live ``httpx.get`` unless this is +# set. While a cassette is active that fetch gets *recorded* as an extra episode +# (it was present in ~710 of ~1900 cached cassettes). For tests that then skip, +# it is the only recorded episode, so the persister refuses to save it (skipped +# tests don't persist) and the test re-records it live and is classified +# MISS:NOT_PERSISTED on every run. Pinning to the local backup removes the +# network call entirely, so skip tests record nothing (NOOP) and passing tests +# stop carrying a volatile github episode. This matches the established idiom in +# the unit-test suite, which sets the same flag (see e.g. +# tests/test_litellm/test_cost_calculator.py). ``setdefault`` so an explicit +# override still wins. +os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") + CASSETTE_CACHE_HIGH_WATER_FRACTION = 0.85 SAFE_BODY_MATCHER_NAME = "safe_body" KEY_FINGERPRINT_MATCHER_NAME = "key_fingerprint" +TOLERANT_QUERY_MATCHER_NAME = "tolerant_query" KEY_FINGERPRINT_HEADER = "x-litellm-key-fp" VCR_DIAG_DIR_ENV = "LITELLM_VCR_DIAG_DIR" @@ -73,6 +92,17 @@ def reset_vcr_diag_dir() -> None: pass +# CircleCI truncates a step's retrievable output to the last ~400 KB. The +# diagnostic log is emitted right *before* the final pytest summary line but +# *after* the VCR CLASSIFICATION SUMMARY, so an unbounded dump (the body/key +# matchers log one block per *episode comparison*, even on an eventual HIT) +# pushes the classification summary out of the retrievable window and makes +# misses impossible to read in CI. Dedupe identical blocks (the same mismatch +# is logged against every non-matching episode) and cap the total emitted size +# so the summary always survives. +VCR_DIAG_EMIT_MAX_LINES = 400 + + def emit_vcr_diagnostic_log(terminalreporter) -> None: directory = _vcr_diag_dir() if not os.path.isdir(directory): @@ -83,25 +113,56 @@ def emit_vcr_diagnostic_log(terminalreporter) -> None: return if not files: return - terminalreporter.write_sep("=", "VCR DIAGNOSTIC LOG", bold=True) - terminalreporter.write_line( - f" source dir: {directory} (also archived as a CI artifact)" - ) + + # Collect every line, tagged by source file, deduplicating identical lines + # (with an occurrence count) so the repeated per-episode mismatch blocks + # collapse to one representative each. + seen_counts: dict[str, int] = defaultdict(int) + ordered: list[tuple[str, str]] = [] # (source_file, line) + read_errors: list[str] = [] for name in files: path = os.path.join(directory, name) try: with open(path, "r", encoding="utf-8") as fh: content = fh.read() except OSError as exc: - terminalreporter.write_line( + read_errors.append( f" [failed to read {name}: {type(exc).__name__}: {exc}]" ) continue - if not content.strip(): - continue - terminalreporter.write_sep("-", name, bold=False) for line in content.splitlines(): - terminalreporter.write_line(line) + if not line.strip(): + continue + seen_counts[line] += 1 + if seen_counts[line] == 1: + ordered.append((name, line)) + + if not ordered and not read_errors: + return + + terminalreporter.write_sep("=", "VCR DIAGNOSTIC LOG", bold=True) + terminalreporter.write_line( + f" source dir: {directory} (deduplicated; full log archived as a CI artifact)" + ) + for line in read_errors: + terminalreporter.write_line(line) + + emitted = 0 + last_source = None + for name, line in ordered: + if emitted >= VCR_DIAG_EMIT_MAX_LINES: + terminalreporter.write_line( + f" ... {len(ordered) - emitted} more unique diagnostic line(s) " + "suppressed to keep the classification summary retrievable in CI." + ) + break + if name != last_source: + terminalreporter.write_sep("-", name, bold=False) + last_source = name + count = seen_counts.get(line, 1) + suffix = f" (x{count})" if count > 1 else "" + terminalreporter.write_line(line + suffix) + emitted += 1 terminalreporter.write_sep("=", bold=True) @@ -326,6 +387,302 @@ def _canonical_body(request) -> tuple[bytes, str]: return b"", pre_type +# --------------------------------------------------------------------------- +# Volatile-token body normalization (compare-time only). +# +# Many tests append a cache-buster to the request body so the *live* call +# isn't served from an upstream prompt/response cache during recording: +# ``f"...{time.time()}"``, ``f"...{uuid.uuid4()}"``. LiteLLM's own +# observability payloads (langfuse/otel) likewise carry per-call UUIDs and +# ISO-8601 timestamps. None of that affects what the test asserts (response +# shape, cost, caching behaviour), but it makes the request body differ on +# every run, so vcrpy never matches and the cassette keeps appending episodes +# until it overflows ``MAX_EPISODES_PER_CASSETTE`` and re-records live forever. +# +# We canonicalize these volatile substrings to fixed placeholders *only for +# matching* (in ``_safe_body_matcher``), never in what we store — so the +# cassette on disk keeps the real bytes for debuggability, and the +# normalization is applied symmetrically to both the incoming and the stored +# request. Because it's symmetric and compare-time, it can never mask a +# response-level discrepancy; it only changes which recorded episode is +# selected. This mirrors the existing SigV4 / multipart-boundary / b64-image +# normalizations already in this module, and means the already-bloated +# cassettes start replaying immediately without a flush + re-record. +_VCR_UUID_RE = re.compile( + rb"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}" +) +# ISO-8601 timestamps, e.g. ``2026-05-25T03:40:37.262045Z`` / +# ``2026-05-25T03:40:37+00:00``. +_VCR_ISO_TS_RE = re.compile( + rb"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:?\d{2})?" +) +# Unix epoch as 13-digit milliseconds, then 10-digit ``time.time()`` float, +# then 10-digit integer seconds. Anchored to ``1`` + 9/12 digits, which keeps +# them inside the 2001-2033 / 2001-2033 epoch windows and avoids matching +# ordinary identifiers. Order matters: the longer/float forms are substituted +# before the bare-integer form so the integer rule can't bite off a prefix. +_VCR_UNIX_MS_RE = re.compile(rb"(? bytes: + """Replace per-run cache-busters (UUIDs / timestamps) with placeholders. + + Compare-time only — see the module note above. Returns ``body`` unchanged + when it contains none of these patterns, so deterministic requests are + unaffected. + """ + if not body: + return body + body = _VCR_UUID_RE.sub(b"", body) + body = _VCR_ISO_TS_RE.sub(b"", body) + body = _VCR_UNIX_MS_RE.sub(b"", body) + body = _VCR_UNIX_FLOAT_RE.sub(b"", body) + body = _VCR_UNIX_INT_RE.sub(b"", body) + return body + + +# Hosts whose request body is a rotating credential exchange (a freshly signed +# JWT ``assertion=...`` or refresh-token grant). The body changes on every run +# and carries no information the test asserts on, so matching on +# method+scheme+host+port+path+query is sufficient — skip the body comparison. +_CREDENTIAL_EXCHANGE_HOSTS = ( + "oauth2.googleapis.com", + "sts.googleapis.com", + "accounts.google.com", + "metadata.google.internal", + "169.254.169.254", +) + + +def _request_host(request) -> str: + uri = getattr(request, "uri", None) or getattr(request, "url", "") or "" + uri = str(uri) + if "//" not in uri: + return "" + rest = uri.split("//", 1)[1] + return rest.split("/", 1)[0].split("@")[-1].split(":")[0].lower() + + +def _is_credential_exchange_request(request) -> bool: + return _request_host(request) in _CREDENTIAL_EXCHANGE_HOSTS + + +# Observability / telemetry backends LiteLLM logs to. A telemetry export is a +# snapshot of the *whole* call — fresh span/trace UUIDs, ISO-8601 timestamps, +# durations, token costs, the LiteLLM build SHA (``release``), and the recorded +# LLM response content — and tests often round-trip a fresh ``trace_id`` back +# through the backend's query API to verify logging happened. None of that is +# reproducible under deterministic replay, and none of it is what the test +# asserts on (it checks redaction / presence, or a locally-computed trace id). +# So for these hosts we match on method+scheme+host+port+path only: the +# expensive LLM call still matches normally and stays cached, while the cheap +# telemetry POST/GET replays from the recorded response. This is why the body +# and query matchers below both short-circuit for telemetry hosts. +_TELEMETRY_HOST_SUFFIXES = ( + "langfuse.com", + "arize.com", + "phoenix.arize.com", + "traceloop.com", + "braintrust.dev", + "comet.com", + "wandb.ai", + "honeycomb.io", + "signoz.io", +) + + +def _is_telemetry_request(request) -> bool: + host = _request_host(request) + if not host: + return False + return any(host == s or host.endswith("." + s) for s in _TELEMETRY_HOST_SUFFIXES) + + +# Nodeid of the test currently executing, set per-test by +# ``install_live_call_probe`` (runs in the autouse gate at setup). Used to +# decide whether an incidental telemetry POST should be recorded — see +# ``_should_drop_telemetry_record``. xdist workers are separate processes and +# tests run sequentially within a worker, so a plain module global is safe. +_current_test_nodeid: str = "" + +# Test files/dirs that legitimately record & replay telemetry HTTP (they assert +# on the outgoing observability payload or query the backend back). Identified +# by a substring of the test path. Everything else is treated as a non-telemetry +# test for which a telemetry call is incidental leakage (see below). +_TELEMETRY_TEST_PATH_MARKERS = ( + "langfuse", + "arize", + "phoenix", + "traceloop", + "braintrust", + "comet", + "wandb", + "honeycomb", + "signoz", + "otel", + "opentelemetry", + "telemetry", + "observability", + "logging", # tests/logging_callback_tests, logging_testing dirs +) + + +def _current_test_records_telemetry() -> bool: + nodeid = _current_test_nodeid.lower() + return any(marker in nodeid for marker in _TELEMETRY_TEST_PATH_MARKERS) + + +# Test paths that legitimately RECORD AND REPLAY a telemetry *export* POST and +# assert on its response. Only the pass-through proxy test does this: it +# forwards a client POST to Langfuse's ``/api/public/ingestion`` and asserts the +# upstream multi-status (207) it replays from the cassette. Every other +# telemetry test either mocks the export client and asserts on the mock (the +# langfuse e2e suite) or asserts on a read-back GET / an in-memory span exporter +# — for those the export POST is fire-and-forget and must not be recorded (see +# ``_should_drop_telemetry_record``). +_TELEMETRY_EXPORT_REPLAY_TEST_MARKERS = ("pass_through",) + + +def _current_test_replays_telemetry_export() -> bool: + nodeid = _current_test_nodeid.lower() + return any(m in nodeid for m in _TELEMETRY_EXPORT_REPLAY_TEST_MARKERS) + + +def _is_telemetry_export_request(request) -> bool: + """A telemetry *export* — a span/trace/event ingestion call, always a POST + to an observability host. Read-backs (verifying a trace landed) are GETs.""" + if not _is_telemetry_request(request): + return False + return str(getattr(request, "method", "") or "").upper() == "POST" + + +# Thread-local "we are inside Cassette._load" flag. vcrpy's ``Cassette._load`` +# replays each *stored* interaction through ``Cassette.append``, which runs +# ``before_record_request`` on it; a ``None`` return there silently drops the +# stored episode. ``_should_drop_telemetry_record`` must therefore NOT fire +# during load, or it would delete already-recorded telemetry episodes the +# instant a non-telemetry-named test (or the very first test in a worker, whose +# ``_current_test_nodeid`` is still empty) loads them — forcing an endless live +# re-record (a phantom MISS:RECORDED on a cassette that was present in Redis). +# The drop is only ever meant to stop *new* incidental telemetry from being +# recorded, never to filter the existing cassette on read. ``_load`` and its +# ``append`` calls run synchronously in one thread, so a thread-local correctly +# scopes the guard and never masks a concurrent background-flush record. +_vcr_load_guard = threading.local() + + +def _vcr_load_in_progress() -> bool: + return getattr(_vcr_load_guard, "active", False) + + +def patch_vcrpy_cassette_load_guard() -> None: + """Wrap ``Cassette._load`` so ``_should_drop_telemetry_record`` is inert + while stored episodes are being replayed into the in-memory cassette.""" + import vcr.cassette as _cassette_mod + + if getattr(_cassette_mod.Cassette._load, "_litellm_load_guarded", False): + return + _orig_load = _cassette_mod.Cassette._load + + def _guarded_load(self): + _vcr_load_guard.active = True + try: + return _orig_load(self) + finally: + _vcr_load_guard.active = False + + _guarded_load._litellm_load_guarded = True + _cassette_mod.Cassette._load = _guarded_load + + +def _should_drop_telemetry_record(request) -> bool: + """Whether to refuse to record this request into the active cassette. + + Several test modules set ``litellm.success_callback = ["langfuse"]`` (and + similar) at *import* time, which globally enables observability logging for + the whole worker. Unrelated tests then emit telemetry whose async flush + (litellm's background logging worker) lands in a *later* test's VCR window + and gets saved as a spurious episode — a non-deterministic MISS:RECORDED on + whichever test happened to be active (observed on + ``test_lowest_latency_routing_buffer`` carrying a Langfuse batch from an + unrelated completion). Refusing to record telemetry for non-telemetry tests + makes the leak a harmless live fire-and-forget call instead (telemetry hosts + are not in ``_LIVE_CALL_HOST_SUFFIXES``, so the probe doesn't flag it, and + vcrpy treats a ``None`` from ``before_record_request`` as "don't record" and + "can't replay" → the request passes through live and is never stored). + Tests that actually assert on telemetry keep recording it. + + Crucially, this never fires while ``Cassette._load`` is replaying stored + interactions (see ``_vcr_load_in_progress``): dropping there would delete an + already-recorded telemetry episode on read and force a live re-record. + + The async-flush leak also rotates *within* the telemetry test set: litellm's + observability loggers flush on a background thread, so an export POST + scheduled by one telemetry test fires mid-way through a *later* + telemetry-named test (after that test's own ``httpx`` mock has exited) and + is recorded as a phantom episode — a non-deterministic MISS:RECORDED / + PARTIAL that lands on a different telemetry test from run to run. Telemetry + *export* POSTs are fire-and-forget; no test asserts on a recorded export + response except the pass-through proxy test (which forwards to Langfuse + ingestion and replays its 207). So drop incidental export POSTs everywhere + else too — dropping returns ``None`` (live fire-and-forget, never stored), + which can only turn a phantom miss into a harmless live call, never the + reverse. Recorded read-back GETs that telemetry tests assert on are matched + by method and so are left untouched. + """ + if _vcr_load_in_progress(): + return False + if not _is_telemetry_request(request): + return False + if ( + _is_telemetry_export_request(request) + and not _current_test_replays_telemetry_export() + ): + return True + return not _current_test_records_telemetry() + + +def _should_passthrough_credential_exchange(request) -> bool: + """Force the Google OAuth2/STS token mint to run live, never from cassette. + + The mint returns a short-lived ``ya29.*`` access token. Recording it lets a + *stale* token replay on a later run; litellm caches it (the recorded + ``expires_in`` keeps ``credentials.expired`` False, so it is never + refreshed) and sends it to a live Vertex/Gemini endpoint, which rejects it + with ``ACCESS_TOKEN_EXPIRED``. The token body carries nothing a test asserts + on, so always mint it live: returning ``None`` from ``before_record_request`` + makes vcrpy neither store nor replay the call. Inert during + ``Cassette._load`` for the same reason as ``_should_drop_telemetry_record``. + """ + if _vcr_load_in_progress(): + return False + return _is_credential_exchange_request(request) + + +# Google APIs (Vertex AI, Gemini, OAuth2/STS). Auth is a ``ya29.*`` OAuth2 +# access token minted fresh on every run, so the per-request key fingerprint +# rotates and never matches a recording. The logical credential — the GCP +# project — is part of the matched URL path (``/projects//...``), so +# skipping the fingerprint comparison for these hosts keeps cache isolation by +# project while letting the existing recordings replay without a re-record. +# (We also collapse ``ya29.*`` tokens to one marker in ``_stable_key_value`` so +# *new* recordings store a stable fingerprint; this matcher relaxation is what +# rescues the cassettes already recorded under the old per-token fingerprints.) +_GOOGLE_HOST_SUFFIXES = ( + "googleapis.com", + "google.internal", +) + + +def _is_google_host_request(request) -> bool: + host = _request_host(request) + if not host: + return False + return any(host == s or host.endswith("." + s) for s in _GOOGLE_HOST_SUFFIXES) + + def _safe_body_matcher(r1, r2) -> None: """Compare request bodies as bytes; never invokes ``json.loads``. @@ -334,11 +691,24 @@ def _safe_body_matcher(r1, r2) -> None: (e.g. the Bedrock batch S3 PUT) before it can return "no match". This matcher is strictly more conservative — the only equivalence it gives up vs. the default is "JSON key order doesn't matter". + + Two compare-time relaxations layer on top, both symmetric so they can + never hide a response-level discrepancy: + + * Requests to a rotating-credential-exchange host (Google OAuth2/STS + token endpoints) skip the body comparison — the signed-JWT body + changes every run. The host matcher still gates the overall match. + * Volatile cache-buster tokens (UUIDs / epoch timestamps) are + canonicalized away via ``_normalize_volatile_tokens``. """ + if _is_credential_exchange_request(r1) or _is_telemetry_request(r1): + return body1, pre1 = _canonical_body(r1) body2, pre2 = _canonical_body(r2) if body1 == body2: return + if _normalize_volatile_tokens(body1) == _normalize_volatile_tokens(body2): + return _emit_body_mismatch_diagnostic(r1, r2, body1, body2, pre1, pre2) raise AssertionError("request bodies differ") @@ -398,6 +768,10 @@ _AWS_SIGV4_CREDENTIAL_RE = re.compile( r"AWS4-HMAC-SHA256\s+Credential=([^/\s,]+)/", re.IGNORECASE ) +# Google OAuth2 access tokens always start with ``ya29.`` regardless of how +# they were minted (service account, metadata server, impersonation). +_GOOGLE_OAUTH_BEARER_RE = re.compile(r"^Bearer\s+ya29\.", re.IGNORECASE) + def _stable_key_value(header_name: str, raw: str) -> str: """Return a *stable* identifier for a credential header. @@ -414,6 +788,14 @@ def _stable_key_value(header_name: str, raw: str) -> str: match = _AWS_SIGV4_CREDENTIAL_RE.search(raw) if match: return f"aws-sigv4:{match.group(1)}" + # Google OAuth2 access tokens (``ya29.*``) are minted fresh from the + # service-account credentials on every run, so hashing the raw token + # would push every Vertex/Gemini request into a new cassette episode — + # exactly the SigV4 failure mode above. The logical credential (the GCP + # project) is already part of the matched URL path, so collapse all such + # tokens to one stable marker. + if _GOOGLE_OAUTH_BEARER_RE.match(raw): + return "google-oauth2" return raw @@ -560,6 +942,14 @@ def _before_record_request(request): this hook is idempotent. The boundary normalizer is also idempotent for the same reason. """ + # Refuse to record incidental telemetry leaked from a globally-enabled + # observability callback into a non-telemetry test (see + # ``_should_drop_telemetry_record``). Returning ``None`` tells vcrpy not to + # store the interaction; the request passes through live (fire-and-forget). + if _should_drop_telemetry_record(request): + return None + if _should_passthrough_credential_exchange(request): + return None headers = getattr(request, "headers", None) if headers is None: return request @@ -626,6 +1016,12 @@ def _coalesce_chunks_to_bytes(chunks): def _key_fingerprint_matcher(r1, r2) -> None: + # Google OAuth2 access tokens rotate every run; the project in the URL + # path (matched separately) is the stable credential identity, so skip the + # fingerprint comparison for Google hosts. See ``_is_google_host_request``. + if _is_google_host_request(r1): + return + def _fp(req): for value in _iter_header_values( getattr(req, "headers", None), KEY_FINGERPRINT_HEADER @@ -649,6 +1045,20 @@ def _key_fingerprint_matcher(r1, r2) -> None: raise AssertionError("API key fingerprints differ") +def _tolerant_query_matcher(r1, r2) -> None: + """vcrpy's ``query`` matcher, but tolerant of telemetry round-trips. + + Observability backends are queried back with a freshly-generated + ``trace_id`` (e.g. ``GET /observations?traceId=litellm-test-``). + Comparing the query string would miss on every run. For telemetry hosts + we skip the query comparison entirely (the host+path matchers still gate + the match); every other host uses vcrpy's stock query matcher unchanged. + """ + if _is_telemetry_request(r1): + return + _vcr_matchers.query(r1, r2) + + def vcr_config_dict() -> dict: return { "decode_compressed_response": True, @@ -660,7 +1070,7 @@ def vcr_config_dict() -> dict: "host", "port", "path", - "query", + TOLERANT_QUERY_MATCHER_NAME, KEY_FINGERPRINT_MATCHER_NAME, SAFE_BODY_MATCHER_NAME, ), @@ -725,7 +1135,9 @@ def register_persister_if_enabled(vcr) -> None: vcr.register_persister(make_redis_persister()) vcr.register_matcher(SAFE_BODY_MATCHER_NAME, _safe_body_matcher) vcr.register_matcher(KEY_FINGERPRINT_MATCHER_NAME, _key_fingerprint_matcher) + vcr.register_matcher(TOLERANT_QUERY_MATCHER_NAME, _tolerant_query_matcher) patch_vcrpy_aiohttp_record_path() + patch_vcrpy_cassette_load_guard() global _atexit_banner_registered if not _atexit_banner_registered: atexit.register(_print_atexit_banner) @@ -1386,6 +1798,12 @@ def install_live_call_probe(request, vcr) -> None: intercepts above the socket layer, so any "outbound" socket would be a recording cycle, not real spend. """ + # Track the current test for telemetry-leak suppression (applies to every + # test, VCR-marked or not). See ``_should_drop_telemetry_record``. + global _current_test_nodeid + _current_test_nodeid = str( + getattr(getattr(request, "node", None), "nodeid", "") or "" + ) if vcr is not None or vcr_disabled(): return None probe = _LiveCallProbe() diff --git a/tests/_vcr_redis_persister.py b/tests/_vcr_redis_persister.py index 373cb66696a..706d2561aa6 100644 --- a/tests/_vcr_redis_persister.py +++ b/tests/_vcr_redis_persister.py @@ -146,8 +146,9 @@ def make_redis_persister( class _RedisPersister: @staticmethod def load_cassette(cassette_path, serializer): + key = redis_key_for(cassette_path) try: - data = redis_client.get(redis_key_for(cassette_path)) + data = redis_client.get(key) except RedisError as exc: _record_cache_failure("load", exc) msg = ( @@ -162,7 +163,7 @@ def make_redis_persister( try: if isinstance(data, bytes): data = data.decode("utf-8") - return deserialize(data, serializer) + result = deserialize(data, serializer) except Exception as exc: _record_cache_failure("load", exc) msg = ( @@ -173,6 +174,22 @@ def make_redis_persister( _log.warning(msg) warnings.warn(msg, VCRCassetteCacheWarning, stacklevel=2) raise CassetteNotFoundError() from exc + # Slide the expiry forward on every successful read. A plain GET + # does not touch the key's TTL, so a cassette that is only ever + # replayed (HIT/NOOP, never re-recorded) expires exactly + # ``ttl_seconds`` after its last *write* no matter how often it is + # read — and whichever CI run happens to cross that boundary + # re-records it live, surfacing as a spurious VCR MISS that no + # amount of matcher tolerance can prevent. Refreshing the TTL on + # read keeps any cassette used at least once per TTL window alive + # indefinitely, so the second/third run of a day replays cleanly. + # Best-effort: a failed refresh must never turn a successful load + # into a miss. + try: + redis_client.expire(key, ttl_seconds) + except RedisError: + pass + return result @staticmethod def save_cassette(cassette_path, cassette_dict, serializer): diff --git a/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py b/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py index a9268da4c31..95d76ba5804 100644 --- a/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py +++ b/tests/agent_tests/local_only_agent_tests/test_a2a_completion_bridge.py @@ -168,7 +168,7 @@ async def test_a2a_completion_bridge_bedrock_agentcore(): litellm._turn_on_debug() # Bedrock AgentCore ARN (streaming-capable runtime) - agentcore_arn = "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" + agentcore_arn = "arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp" send_message_payload = { "message": { diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index 46870f12272..cb2ca385ffc 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -145,6 +145,37 @@ def test_batch_cost_calculator_func_uses_custom_model_info(): ), f"Expected total cost {expected}, got {cost}" +@pytest.mark.parametrize("data_residency", ["eu", "us"]) +def test_batch_cost_calculator_applies_data_residency_uplift( + data_residency, monkeypatch +): + """batch_cost_calculator should apply the regional uplift multiplier when + data_residency is set and the model carries a configured multiplier.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + prev_model_cost = litellm.model_cost + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base_prompt, base_completion = batch_cost_calculator( + usage=usage, + model="gpt-5", + custom_llm_provider="openai", + ) + regional_prompt, regional_completion = batch_cost_calculator( + usage=usage, + model="gpt-5", + custom_llm_provider="openai", + data_residency=data_residency, + ) + + assert base_prompt > 0 and base_completion > 0 + assert regional_prompt == pytest.approx(base_prompt * 1.10, rel=1e-9) + assert regional_completion == pytest.approx(base_completion * 1.10, rel=1e-9) + finally: + litellm.model_cost = prev_model_cost + + @pytest.mark.asyncio async def test_calculate_batch_cost_and_usage_uses_custom_model_info(): """calculate_batch_cost_and_usage should thread model_info.""" diff --git a/tests/batches_tests/test_bedrock_files_and_batches.py b/tests/batches_tests/test_bedrock_files_and_batches.py index 5148ea4db91..97c0802ec99 100644 --- a/tests/batches_tests/test_bedrock_files_and_batches.py +++ b/tests/batches_tests/test_bedrock_files_and_batches.py @@ -38,7 +38,7 @@ async def test_async_create_file(): file=open(file_path, "rb"), purpose="batch", custom_llm_provider="bedrock", - s3_bucket_name="litellm-proxy", + s3_bucket_name="litellm-proxy-941277531214", ) @@ -55,7 +55,7 @@ async def test_async_file_and_batch(): file=open(file_path, "rb"), purpose="batch", custom_llm_provider="bedrock", - s3_bucket_name="litellm-proxy", + s3_bucket_name="litellm-proxy-941277531214", ) print("CREATED FILE RESPONSE=", file_obj) @@ -70,7 +70,7 @@ async def test_async_file_and_batch(): # bedrock specific params ######################################################### model="us.anthropic.claude-haiku-4-5-20251001-v1:0", - aws_batch_role_arn="arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV", + aws_batch_role_arn="arn:aws:iam::941277531214:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV", ) print("CREATED BATCH RESPONSE=", create_batch_response) @@ -129,7 +129,7 @@ async def test_mock_bedrock_file_url_mapping(): ), purpose="batch", custom_llm_provider="bedrock", - s3_bucket_name="litellm-proxy", + s3_bucket_name="litellm-proxy-941277531214", ) print(f"PUT URL: {captured_put_url}") diff --git a/tests/code_coverage_tests/check_licenses.py b/tests/code_coverage_tests/check_licenses.py index 5fb2b495c24..389e534b1ff 100644 --- a/tests/code_coverage_tests/check_licenses.py +++ b/tests/code_coverage_tests/check_licenses.py @@ -376,8 +376,23 @@ class LicenseChecker: all_compliant = True for req in requirements: + # Prefer a lower-bound/exact version (a real released version) for the + # PyPI license lookup. ``next(iter(req.specifier))`` returns an + # arbitrary clause; for a range like ``>=1.0,<2.0`` that can be the + # upper bound (``2.0``) — a version that may not exist on PyPI and + # would 404 to an "unknown" license. try: - version = next(iter(req.specifier)).version if req.specifier else None + floor_versions = [ + spec.version + for spec in req.specifier + if spec.operator in (">=", "==", "===", "~=", ">") + ] + if floor_versions: + version = floor_versions[0] + else: + version = ( + next(iter(req.specifier)).version if req.specifier else None + ) except StopIteration: version = None diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py new file mode 100644 index 00000000000..4dd5c3d88ca --- /dev/null +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py @@ -0,0 +1,42 @@ +"""Shared fixtures for guardrail apply_guardrail tests.""" + +from contextlib import contextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +@contextmanager +def _mock_proxy_logging(): + """Patch the proxy-server globals that apply_guardrail imports at call time.""" + mock_proxy_logging = MagicMock() + mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=None) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock(return_value=None) + mock_logging_obj.async_failure_handler = AsyncMock(return_value=None) + mock_logging_obj.success_handler = MagicMock(return_value=None) + mock_logging_obj.failure_handler = MagicMock(return_value=None) + mock_logging_obj.model_call_details = {} + + with ( + patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing" + ) as mock_proc_cls, + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "0.0.0"), + ): + mock_proc = MagicMock() + mock_proc.common_processing_pre_call_logic = AsyncMock( + return_value=({}, mock_logging_obj) + ) + mock_proc_cls.return_value = mock_proc + yield mock_proxy_logging + + +@pytest.fixture +def mock_proxy_logging_ctx(): + """Return the proxy-logging context manager factory for use as `with ctx():`.""" + return _mock_proxy_logging diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py index 0d27df50d15..e5074c44210 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py @@ -18,14 +18,19 @@ from litellm.types.guardrails import ApplyGuardrailRequest, ApplyGuardrailRespon @pytest.mark.asyncio -async def test_apply_guardrail_endpoint_returns_correct_response(): +async def test_apply_guardrail_endpoint_returns_correct_response( + mock_proxy_logging_ctx, +): """Test that apply_guardrail endpoint returns ApplyGuardrailResponse object""" from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry - with patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" - ) as mock_registry: + with ( + patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry, + mock_proxy_logging_ctx(), + ): # Create a mock guardrail mock_guardrail = Mock(spec=CustomGuardrail) # Apply guardrail returns GenericGuardrailAPIInputs (dict with texts key) @@ -49,7 +54,9 @@ async def test_apply_guardrail_endpoint_returns_correct_response(): # Call the endpoint response = await apply_guardrail( - request=request, user_api_key_dict=user_api_key_dict + fastapi_request=Mock(), + request=request, + user_api_key_dict=user_api_key_dict, ) # Verify the response is of the correct type @@ -65,15 +72,18 @@ async def test_apply_guardrail_endpoint_returns_correct_response(): @pytest.mark.asyncio -async def test_apply_guardrail_endpoint_guardrail_not_found(): +async def test_apply_guardrail_endpoint_guardrail_not_found(mock_proxy_logging_ctx): """Test that apply_guardrail endpoint raises exception when guardrail not found""" from litellm.proxy._types import ProxyException from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry to return None - with patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" - ) as mock_registry: + with ( + patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry, + mock_proxy_logging_ctx(), + ): mock_registry.get_initialized_guardrail_callback.return_value = None # Create the request @@ -86,26 +96,35 @@ async def test_apply_guardrail_endpoint_guardrail_not_found(): # Verify exception is raised with pytest.raises(ProxyException) as exc_info: - await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict) + await apply_guardrail( + fastapi_request=Mock(), + request=request, + user_api_key_dict=user_api_key_dict, + ) assert "non-existent-guardrail" in exc_info.value.message assert "not found" in exc_info.value.message @pytest.mark.asyncio -async def test_apply_guardrail_endpoint_with_presidio_guardrail(): +async def test_apply_guardrail_endpoint_with_presidio_guardrail(mock_proxy_logging_ctx): """Test apply_guardrail endpoint with a Presidio-like guardrail""" from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry - with patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" - ) as mock_registry: + with ( + patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry, + mock_proxy_logging_ctx(), + ): # Create a mock guardrail that simulates Presidio behavior mock_guardrail = Mock(spec=CustomGuardrail) # Simulate masking PII entities - returns GenericGuardrailAPIInputs (dict with texts key) mock_guardrail.apply_guardrail = AsyncMock( - return_value={"texts": ["My name is [PERSON] and my email is [EMAIL_ADDRESS]"]} + return_value={ + "texts": ["My name is [PERSON] and my email is [EMAIL_ADDRESS]"] + } ) # Configure the registry to return our mock guardrail @@ -124,7 +143,9 @@ async def test_apply_guardrail_endpoint_with_presidio_guardrail(): # Call the endpoint response = await apply_guardrail( - request=request, user_api_key_dict=user_api_key_dict + fastapi_request=Mock(), + request=request, + user_api_key_dict=user_api_key_dict, ) # Verify the response is of the correct type @@ -138,14 +159,17 @@ async def test_apply_guardrail_endpoint_with_presidio_guardrail(): @pytest.mark.asyncio -async def test_apply_guardrail_endpoint_without_optional_params(): +async def test_apply_guardrail_endpoint_without_optional_params(mock_proxy_logging_ctx): """Test apply_guardrail endpoint without optional language and entities parameters""" from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry - with patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" - ) as mock_registry: + with ( + patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry, + mock_proxy_logging_ctx(), + ): # Create a mock guardrail mock_guardrail = Mock(spec=CustomGuardrail) # Returns GenericGuardrailAPIInputs (dict with texts key) @@ -166,7 +190,9 @@ async def test_apply_guardrail_endpoint_without_optional_params(): # Call the endpoint response = await apply_guardrail( - request=request, user_api_key_dict=user_api_key_dict + fastapi_request=Mock(), + request=request, + user_api_key_dict=user_api_key_dict, ) # Verify the response is of the correct type diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index dff444168c2..d1caf398540 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -4,7 +4,7 @@ Test the Bedrock guardrail apply_guardrail functionality import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, Mock, patch import pytest @@ -153,7 +153,7 @@ async def test_bedrock_apply_guardrail_api_failure(): @pytest.mark.asyncio -async def test_bedrock_apply_guardrail_endpoint_integration(): +async def test_bedrock_apply_guardrail_endpoint_integration(mock_proxy_logging_ctx): """Test the full endpoint integration with Bedrock guardrail""" from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail @@ -165,9 +165,12 @@ async def test_bedrock_apply_guardrail_endpoint_integration(): ) # Mock the guardrail registry - with patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" - ) as mock_registry: + with ( + patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry, + mock_proxy_logging_ctx(), + ): # Mock the make_bedrock_api_request method with patch.object( guardrail, "make_bedrock_api_request", new_callable=AsyncMock @@ -194,7 +197,9 @@ async def test_bedrock_apply_guardrail_endpoint_integration(): # Call the endpoint response = await apply_guardrail( - request=request, user_api_key_dict=user_api_key_dict + fastapi_request=Mock(), + request=request, + user_api_key_dict=user_api_key_dict, ) # Verify the response diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 6e78a8c4284..ea50fe08ae0 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -20,7 +20,7 @@ async def test_bedrock_guardrails_pii_masking(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", + guardrailIdentifier="zgkmukebruil", guardrailVersion="DRAFT", ) @@ -60,7 +60,7 @@ async def test_bedrock_guardrails_pii_masking_content_list(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", + guardrailIdentifier="zgkmukebruil", guardrailVersion="DRAFT", ) @@ -115,7 +115,7 @@ async def test_bedrock_guardrails_block_messages_api(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", + guardrailIdentifier="4w3d1di3snt5", guardrailVersion="DRAFT", ) @@ -166,7 +166,7 @@ async def test_bedrock_guardrails_block_responses_api(): mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", + guardrailIdentifier="4w3d1di3snt5", guardrailVersion="DRAFT", ) @@ -211,7 +211,7 @@ async def test_bedrock_guardrails_with_streaming(): ) guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", + guardrailIdentifier="4w3d1di3snt5", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", @@ -255,7 +255,7 @@ async def test_bedrock_guardrails_with_streaming_no_violation(): ) guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", + guardrailIdentifier="4w3d1di3snt5", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", @@ -299,7 +299,7 @@ async def test_bedrock_guardrails_streaming_request_body_mock(): # Create the guardrail guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", + guardrailIdentifier="zgkmukebruil", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", @@ -382,7 +382,7 @@ async def test_bedrock_guardrail_aws_param_persistence(): from litellm.types.guardrails import GuardrailEventHooks guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", + guardrailIdentifier="zgkmukebruil", guardrailVersion="DRAFT", aws_access_key_id="test-access-key", aws_secret_access_key="test-secret-key", diff --git a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py index 36ae9e1df67..181691b730d 100644 --- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py +++ b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py @@ -1,3 +1,4 @@ +import json import logging import os import sys @@ -44,6 +45,9 @@ from litellm.llms.bedrock.image_generation.image_handler import ( ) from litellm.llms.bedrock.common_utils import BedrockError +# Base64 placeholder used for mocked Bedrock image responses (a 1x1 PNG). +_MOCK_BEDROCK_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + @pytest.mark.parametrize( "model,expected", @@ -528,17 +532,34 @@ def test_backward_compatibility_regular_nova_model(): def test_amazon_titan_image_gen(): - """Test Amazon Titan image generation with cost tracking.""" - from litellm import image_generation + """Test Amazon Titan image generation with cost tracking. + + The Bedrock CI account is not entitled to amazon.titan-image-generator, so + the network call is mocked and only the transform + cost-tracking path is + exercised. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler # Use v2 as v1 has reached end of life model_id = "bedrock/amazon.titan-image-generator-v2:0" - response = litellm.image_generation( - model=model_id, - prompt="A serene mountain landscape at sunset with a lake reflection", - aws_region_name="us-east-1", - ) + mock_payload = {"images": [_MOCK_BEDROCK_IMAGE_B64]} + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_payload + mock_response.text = json.dumps(mock_payload) + mock_response.headers = {} + + client = HTTPHandler() + with patch.object(client, "post", return_value=mock_response): + response = litellm.image_generation( + model=model_id, + prompt="A serene mountain landscape at sunset with a lake reflection", + aws_region_name="us-east-1", + aws_access_key_id="fake-access-key-id", + aws_secret_access_key="fake-secret-access-key", + client=client, + ) print(f"response cost: {response._hidden_params['response_cost']}") diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 873777189c9..23a94ef389a 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -7,7 +7,6 @@ import sys import traceback from unittest.mock import AsyncMock, MagicMock, patch - sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path @@ -136,6 +135,51 @@ class TestVertexAIGeminiImageGeneration(BaseImageGenTest): } +# Base64 placeholder used for mocked Bedrock image responses (a 1x1 PNG). +_MOCK_BEDROCK_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + + +async def _assert_mocked_bedrock_image_generation(call_args: dict) -> None: + """Run ``aimage_generation`` with the Bedrock HTTP call mocked. + + The CI account is not entitled to Nova Canvas, so the network call is + replaced with a canned Bedrock response. This keeps the request transform, + response transform, and cost-tracking path under test without live access. + """ + mock_payload = {"images": [_MOCK_BEDROCK_IMAGE_B64]} + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_payload + mock_response.text = json.dumps(mock_payload) + mock_response.headers = {} + + custom_logger = TestCustomLogger() + litellm.logging_callback_manager._reset_all_callbacks() + litellm.callbacks = [custom_logger] + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + return_value=mock_response, + ): + response = await litellm.aimage_generation( + **call_args, + prompt="A image of a otter", + aws_access_key_id="fake-access-key-id", + aws_secret_access_key="fake-secret-access-key", + ) + + await asyncio.sleep(1) + + assert custom_logger.standard_logging_payload is not None + assert custom_logger.standard_logging_payload["response_cost"] is not None + assert custom_logger.standard_logging_payload["response_cost"] > 0 + assert response.data is not None + for d in response.data: + assert isinstance(d, Image) + assert d.b64_json is not None or d.url is not None + + class TestBedrockNovaCanvasTextToImage(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: litellm.in_memory_llm_clients_cache = InMemoryCache() @@ -148,6 +192,12 @@ class TestBedrockNovaCanvasTextToImage(BaseImageGenTest): "aws_region_name": "us-east-1", } + @pytest.mark.asyncio(scope="module") + async def test_basic_image_generation(self): + await _assert_mocked_bedrock_image_generation( + self.get_base_image_generation_call_args() + ) + class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: @@ -162,6 +212,12 @@ class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest): "aws_region_name": "us-east-1", } + @pytest.mark.asyncio(scope="module") + async def test_basic_image_generation(self): + await _assert_mocked_bedrock_image_generation( + self.get_base_image_generation_call_args() + ) + class TestOpenAIGPTImage1(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: diff --git a/tests/litellm_utils_tests/conftest.py b/tests/litellm_utils_tests/conftest.py index 418ee76a399..d20203da3a7 100644 --- a/tests/litellm_utils_tests/conftest.py +++ b/tests/litellm_utils_tests/conftest.py @@ -38,7 +38,22 @@ _VCR_INCOMPATIBLE_FILES = frozenset( } ) -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () +# AWS Secrets Manager resource-lifecycle tests. Each run creates a secret +# under a per-run unique name (``litellm_test_``) and either asserts the +# API response echoes that exact unique name or reads it straight back. The +# name *must* be unique per run because AWS enforces a >=7-day deletion +# recovery window — a fixed name can't be re-created on the daily VCR +# re-record. Deterministic replay returns the previously-recorded (different) +# name, so the unique-name round-trip cannot be reproduced offline. The +# config-parsing tests in the same file (settings / STS endpoint) make no such +# unique-resource calls and stay VCR-cached. +_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ( + "::test_write_and_read_simple_secret", + "::test_write_and_read_json_secret", + "::test_read_nonexistent_secret", + "::test_primary_secret_functionality", + "::test_write_secret_with_description_and_tags", +) @pytest.fixture(scope="function", autouse=True) diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py index 3a428e9d588..60ee849f8eb 100644 --- a/tests/litellm_utils_tests/test_litellm_overhead.py +++ b/tests/litellm_utils_tests/test_litellm_overhead.py @@ -82,7 +82,7 @@ async def _vertex_ai_mocks(): "bedrock/mistral.mistral-7b-instruct-v0:2", "openai/gpt-4o", "openai/self_hosted", - "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "vertex_ai/gemini-1.5-flash", ], ) @@ -147,7 +147,7 @@ async def test_litellm_overhead_non_streaming(model): [ "bedrock/mistral.mistral-7b-instruct-v0:2", "openai/gpt-4o", - "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "openai/self_hosted", ], ) diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index ed5346dad71..2f9735274c6 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -1,7 +1,6 @@ from dataclasses import dataclass, field from typing import Dict, FrozenSet, List, Optional, Tuple - OMIT = object() @@ -22,6 +21,8 @@ class ModelEntry: extra_params: Tuple[Tuple[str, str], ...] = field(default_factory=tuple) required_env: FrozenSet[str] = field(default_factory=frozenset) caps: FrozenSet[str] = field(default_factory=frozenset) + fail_reason: Optional[str] = None + bedrock_effort_ceiling: Optional[str] = None def params(self) -> Dict[str, str]: return dict(self.extra_params) @@ -59,9 +60,31 @@ _ADAPTIVE_EFFORT_LABEL: Dict[str, str] = { "max": "max", } +_EFFORT_RANK: Dict[str, int] = { + "low": 0, + "medium": 1, + "high": 2, + "max": 3, + "xhigh": 4, +} + _BAD_REQUEST_EFFORTS: FrozenSet[str] = frozenset({"disabled", "invalid", ""}) +def _bedrock_clamps_effort(model: "ModelEntry", effort: str) -> bool: + """Whether Bedrock will clamp ``effort`` down to ``bedrock_effort_ceiling``. + + Bedrock chat/messages paths clamp unsupported high tiers (e.g. ``xhigh`` + on Opus 4.6) to the model's ceiling rather than rejecting them, so the + missing native capability is OK — the wire effort just degrades. + """ + if model.bedrock_effort_ceiling is None: + return False + if effort not in _EFFORT_RANK or model.bedrock_effort_ceiling not in _EFFORT_RANK: + return False + return _EFFORT_RANK[effort] > _EFFORT_RANK[model.bedrock_effort_ceiling] + + def expected(model: ModelEntry, effort: str) -> CellExpectation: if effort in ("__omit__", "none"): if model.mode == "budget": @@ -73,14 +96,20 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: if effort in ("xhigh", "max"): cap = f"supports_{effort}_reasoning_effort" - if cap not in model.caps: + if cap not in model.caps and not _bedrock_clamps_effort(model, effort): return CellExpectation(status=400, thinking_type=OMIT) if model.mode == "adaptive": + wire_effort = _ADAPTIVE_EFFORT_LABEL[effort] + if model.bedrock_effort_ceiling is not None: + wire_rank = _EFFORT_RANK[wire_effort] + ceiling_rank = _EFFORT_RANK[model.bedrock_effort_ceiling] + if wire_rank > ceiling_rank: + wire_effort = model.bedrock_effort_ceiling return CellExpectation( status=200, thinking_type="adaptive", - output_config_effort=_ADAPTIVE_EFFORT_LABEL[effort], + output_config_effort=wire_effort, ) return CellExpectation( @@ -205,6 +234,12 @@ BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( extra_params=(("aws_region_name", "us-east-1"),), required_env=_BEDROCK_REQ, caps=_CAPS_OPUS_4_7, + fail_reason=( + "claude-opus-4-7 is not entitled on the Bedrock CI account " + "941277531214 (model access requires an AWS Sales request, not " + "self-serve); this cell fails on purpose so it stays loud in CI — " + "remove this fail_reason once access is granted" + ), ), ModelEntry( alias="bedrock-claude-opus-4-6", @@ -213,6 +248,7 @@ BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( extra_params=(("aws_region_name", "us-east-1"),), required_env=_BEDROCK_REQ, caps=_CAPS_4_6, + bedrock_effort_ceiling="max", ), ModelEntry( alias="bedrock-claude-sonnet-4-6", @@ -241,6 +277,7 @@ BEDROCK_INVOKE_CHAT_MODELS: Tuple[ModelEntry, ...] = ( extra_params=(("aws_region_name", "us-east-1"),), required_env=_BEDROCK_REQ, caps=_CAPS_4_6, + bedrock_effort_ceiling="max", ), ModelEntry( alias="bedrock-invoke-claude-sonnet-4-6", diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py index 28e2e402d67..e0b6290ad77 100644 --- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py +++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py @@ -15,7 +15,6 @@ from .grid_spec import ( all_cells, ) - _PROMPT_MESSAGES: List[Dict[str, str]] = [ {"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."} ] @@ -168,6 +167,9 @@ async def test_reasoning_effort_grid( if skip_reason: pytest.skip(skip_reason) + if model.fail_reason: + pytest.xfail(model.fail_reason) + if route_name == "bedrock_invoke_messages": status, exc = await _call_messages(model, effort) else: diff --git a/tests/llm_translation/test_bedrock_agentcore.py b/tests/llm_translation/test_bedrock_agentcore.py index 40774cf3d60..95a814e97e4 100644 --- a/tests/llm_translation/test_bedrock_agentcore.py +++ b/tests/llm_translation/test_bedrock_agentcore.py @@ -19,8 +19,8 @@ import httpx @pytest.mark.parametrize( "model", [ - "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # non-streaming invocation - "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", # streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_13sf6-4046UzHSwy", # non-streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", # streaming invocation ], ) def test_bedrock_agentcore_basic(model): @@ -44,7 +44,7 @@ def test_bedrock_agentcore_basic(model): @pytest.mark.parametrize( "model", [ - "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_13sf6-4046UzHSwy", # streaming invocation ], ) async def test_bedrock_agentcore_with_streaming(model): @@ -54,7 +54,7 @@ async def test_bedrock_agentcore_with_streaming(model): print("running streming test for model=", model) # litellm._turn_on_debug() response = await litellm.acompletion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -82,7 +82,7 @@ def test_bedrock_agentcore_with_custom_params(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -105,7 +105,7 @@ def test_bedrock_agentcore_with_custom_params(): url = call_kwargs["url"] print(f"URL: {url}") assert ( - "/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_r9jvp-3ySZuRHjLC/invocations" + "/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A941277531214%3Aruntime%2Fhosted_agent_r9jvp-Rq79QFC2fp/invocations" in url ) assert "qualifier=DEFAULT" in url @@ -150,7 +150,7 @@ def test_bedrock_agentcore_with_runtime_user_id(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -189,7 +189,7 @@ def test_bedrock_agentcore_with_session_and_user(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -234,7 +234,7 @@ def test_bedrock_agentcore_with_api_key_bearer_token(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -282,7 +282,7 @@ def test_bedrock_agentcore_with_all_parameters(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -350,7 +350,7 @@ def test_bedrock_agentcore_without_api_key_uses_sigv4(): with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", @@ -625,7 +625,7 @@ def test_agentcore_synchronous_non_streaming_response(): with patch.object(client, "post", return_value=mock_response) as mock_post: # Make a synchronous (non-streaming) completion call response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/hosted_agent_r9jvp-Rq79QFC2fp", messages=[ { "role": "user", diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 15f950224d2..a3b4a010f60 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -115,7 +115,7 @@ def test_completion_bedrock_guardrails(streaming): ], max_tokens=10, guardrailConfig={ - "guardrailIdentifier": "ff6ujrregl1q", + "guardrailIdentifier": "4w3d1di3snt5", "guardrailVersion": "DRAFT", "trace": "enabled", }, @@ -144,7 +144,7 @@ def test_completion_bedrock_guardrails(streaming): stream=True, max_tokens=10, guardrailConfig={ - "guardrailIdentifier": "ff6ujrregl1q", + "guardrailIdentifier": "4w3d1di3snt5", "guardrailVersion": "DRAFT", "trace": "enabled", }, @@ -475,7 +475,7 @@ def test_bedrock_claude_3(image_url): ], } response: ModelResponse = completion( - model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", num_retries=3, **data, ) # type: ignore @@ -498,7 +498,7 @@ def test_bedrock_claude_3(image_url): @pytest.mark.parametrize( "model", [ - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", # "meta.llama3-70b-instruct-v1:0", # "anthropic.claude-v2", # "mistral.mixtral-8x7b-instruct-v0:1", @@ -537,7 +537,7 @@ def test_bedrock_stop_value(stop, model): @pytest.mark.parametrize( "model", [ - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "mistral.mixtral-8x7b-instruct-v0:1", ], ) @@ -602,7 +602,7 @@ def test_bedrock_claude_3_tool_calling(): } ] response: ModelResponse = completion( - model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, tools=tools, tool_choice="auto", @@ -630,7 +630,7 @@ def test_bedrock_claude_3_tool_calling(): ) # In the second response, Claude should deduce answer from tool results second_response = completion( - model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, tools=tools, tool_choice="auto", @@ -737,7 +737,7 @@ def test_bedrock_ptu(): from openai.types.chat import ChatCompletion model_id = ( - "arn:aws:bedrock:us-west-2:888602223428:provisioned-model/8fxff74qyhs3" + "arn:aws:bedrock:us-west-2:941277531214:provisioned-model/8fxff74qyhs3" ) try: response = litellm.completion( @@ -752,7 +752,7 @@ def test_bedrock_ptu(): assert "url" in mock_client_post.call_args.kwargs assert ( mock_client_post.call_args.kwargs["url"] - == "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A888602223428%3Aprovisioned-model%2F8fxff74qyhs3/converse" + == "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A941277531214%3Aprovisioned-model%2F8fxff74qyhs3/converse" ) mock_client_post.assert_called_once() @@ -1062,7 +1062,7 @@ def test_bedrock_tools_pt_invalid_names(): print("bedrock tools after prompt formatting=", result) assert len(result) == 2 - assert result[0]["toolSpec"]["name"] == "a123_invalid_name" + assert result[0]["toolSpec"]["name"] == "a123-invalid_name" assert result[1]["toolSpec"]["name"] == "another_invalid_name" @@ -1171,7 +1171,7 @@ def test_bedrock_tools_transformation_valid_params(): assert isinstance(result, list) assert len(result) == 1 assert "toolSpec" in result[0] - assert result[0]["toolSpec"]["name"] == "a123_invalid_name" + assert result[0]["toolSpec"]["name"] == "a123-invalid_name" assert result[0]["toolSpec"]["description"] == "Invalid name test" assert "inputSchema" in result[0]["toolSpec"] assert "json" in result[0]["toolSpec"]["inputSchema"] @@ -2327,7 +2327,7 @@ def test_bedrock_cross_region_inference(monkeypatch): def test_bedrock_empty_content_real_call(): completion( - model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", diff --git a/tests/llm_translation/test_vcr_conftest_common_banner.py b/tests/llm_translation/test_vcr_conftest_common_banner.py index 70ee39abd39..1c4395ef1a8 100644 --- a/tests/llm_translation/test_vcr_conftest_common_banner.py +++ b/tests/llm_translation/test_vcr_conftest_common_banner.py @@ -9,7 +9,9 @@ import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) from tests._vcr_conftest_common import ( # noqa: E402 + VCR_DIAG_EMIT_MAX_LINES, emit_cassette_cache_session_banner, + emit_vcr_diagnostic_log, ) from tests._vcr_redis_persister import ( # noqa: E402 _cache_health, @@ -165,6 +167,55 @@ def test_banner_silent_when_vcr_disabled( assert reporter.output == "" +# --------------------------------------------------------------------------- +# Diagnostic-log dedup + cap. CircleCI truncates step output to the last +# ~400 KB; an unbounded diagnostic dump pushes the VCR classification summary +# out of the retrievable window, so the dump must dedupe and cap. +# --------------------------------------------------------------------------- + + +def test_diagnostic_log_dedupes_repeated_blocks(tmp_path, monkeypatch): + monkeypatch.setenv("LITELLM_VCR_DIAG_DIR", str(tmp_path)) + (tmp_path / "123.log").write_text( + "\n".join(["[vcr-key-fingerprint-matcher] differ"] * 40 + ["unique line"]), + encoding="utf-8", + ) + reporter = _FakeTerminalReporter() + + emit_vcr_diagnostic_log(reporter) + + out = reporter.output + # The repeated block collapses to a single line with an occurrence count. + assert out.count("[vcr-key-fingerprint-matcher] differ") == 1 + assert "(x40)" in out + assert "unique line" in out + + +def test_diagnostic_log_caps_unique_lines(tmp_path, monkeypatch): + monkeypatch.setenv("LITELLM_VCR_DIAG_DIR", str(tmp_path)) + total = VCR_DIAG_EMIT_MAX_LINES + 50 + (tmp_path / "123.log").write_text( + "\n".join(f"unique-diagnostic-{i}" for i in range(total)), encoding="utf-8" + ) + reporter = _FakeTerminalReporter() + + emit_vcr_diagnostic_log(reporter) + + out = reporter.output + emitted = sum(1 for ln in out.splitlines() if ln.startswith("unique-diagnostic-")) + assert emitted == VCR_DIAG_EMIT_MAX_LINES + assert "more unique diagnostic line(s) suppressed" in out + + +def test_diagnostic_log_silent_when_no_dir(tmp_path, monkeypatch): + monkeypatch.setenv("LITELLM_VCR_DIAG_DIR", str(tmp_path / "does-not-exist")) + reporter = _FakeTerminalReporter() + + emit_vcr_diagnostic_log(reporter) + + assert reporter.output == "" + + def test_banner_silent_on_xdist_worker( monkeypatch, vcr_enabled, health_reset, patch_capacity_snapshot ): @@ -179,3 +230,164 @@ def test_banner_silent_on_xdist_worker( emit_cassette_cache_session_banner(reporter) assert reporter.output == "" + + +# --------------------------------------------------------------------------- +# Telemetry-leak suppression. Several modules set ``litellm.success_callback`` +# at import time, so observability logging is globally enabled and an async +# flush can land in an unrelated test's VCR window and be saved as a spurious +# MISS:RECORDED episode. ``_should_drop_telemetry_record`` refuses to record a +# telemetry call for a non-telemetry test (it passes through live instead), +# while tests that actually assert on telemetry keep recording. +# --------------------------------------------------------------------------- + + +class _FakeRequest: + def __init__( + self, host, scheme="https", method="POST", path="/api/public/ingestion" + ): + self.host = host + self.scheme = scheme + self.uri = f"{scheme}://{host}{path}" + self.headers = {} + self.method = method + self.body = b"{}" + + +@pytest.fixture +def current_test(monkeypatch): + """Set the module-global current-test nodeid the suppressor reads.""" + import tests._vcr_conftest_common as common + + def _set(nodeid): + monkeypatch.setattr(common, "_current_test_nodeid", nodeid) + + return _set + + +@pytest.mark.parametrize( + "nodeid,host,method,expected_drop", + [ + # Non-telemetry test: incidental telemetry leak is dropped (not recorded). + ( + "tests/local_testing/test_lowest_latency_routing.py::test_lowest_latency_routing_buffer[1]", + "us.cloud.langfuse.com", + "POST", + True, + ), + ( + "tests/local_testing/test_function_call_parsing.py::test_parse", + "us.cloud.langfuse.com", + "POST", + True, + ), + ( + "tests/llm_translation/test_x.py::test_y", + "otlp.arize.com", + "POST", + True, + ), + # Non-telemetry host on a non-telemetry test: never dropped. + ( + "tests/local_testing/test_lowest_latency_routing.py::test_lowest_latency_routing_buffer[1]", + "api.openai.com", + "POST", + False, + ), + # Telemetry EXPORT POSTs are fire-and-forget and dropped even for + # telemetry-named tests: litellm's background flush makes them rotate + # into a later telemetry test's window as a phantom MISS:RECORDED. The + # e2e suite mocks the export client and asserts on the mock; read-back + # tests assert on a GET — neither needs the recorded export POST. + ( + "tests/local_testing/test_alangfuse.py::test_langfuse_logging", + "us.cloud.langfuse.com", + "POST", + True, + ), + ( + "tests/logging_callback_tests/test_langfuse_e2e_test.py::test_e2e", + "us.cloud.langfuse.com", + "POST", + True, + ), + ( + "tests/logging_callback_tests/test_dynamic_otel_keys.py::test_keys", + "otlp.arize.com", + "POST", + True, + ), + # Read-back GETs that telemetry tests assert on are kept (matched by + # method, so the export-POST drop does not touch them). + ( + "tests/local_testing/test_alangfuse.py::test_langfuse_logging", + "us.cloud.langfuse.com", + "GET", + False, + ), + # ...but a read-back GET on a NON-telemetry test is still incidental. + ( + "tests/local_testing/test_function_call_parsing.py::test_parse", + "us.cloud.langfuse.com", + "GET", + True, + ), + # The pass-through proxy test forwards a client POST to Langfuse + # ingestion and asserts the replayed 207 — its export POST is kept. + ( + "tests/local_testing/test_pass_through_endpoints.py::test_aaapass_through_endpoint_pass_through_keys_langfuse[False-0-207]", + "us.cloud.langfuse.com", + "POST", + False, + ), + ], +) +def test_should_drop_telemetry_record( + current_test, nodeid, host, method, expected_drop +): + import tests._vcr_conftest_common as common + + current_test(nodeid) + req = _FakeRequest(host, method=method) + assert common._should_drop_telemetry_record(req) is expected_drop + + +def test_drop_is_suppressed_while_loading_stored_episodes(current_test): + """During ``Cassette._load`` the drop MUST be inert. + + vcrpy replays each stored interaction through ``Cassette.append`` → + ``before_record_request``; a ``None`` there silently drops the stored + episode. If the telemetry drop fired on load, an already-recorded + telemetry episode would be deleted the instant a non-telemetry-named + test loaded it, forcing an endless live re-record (a phantom + MISS:RECORDED on a cassette that was present in Redis). The drop must + only stop *new* incidental recordings, never filter the cassette on read. + """ + import tests._vcr_conftest_common as common + + # A non-telemetry test loading a stored Langfuse episode: dropped on + # record, but must be KEPT while loading. + current_test("tests/local_testing/test_lowest_latency_routing.py::test_buf") + req = _FakeRequest("us.cloud.langfuse.com") + + assert common._should_drop_telemetry_record(req) is True # record path + + common._vcr_load_guard.active = True + try: + assert common._vcr_load_in_progress() is True + assert common._should_drop_telemetry_record(req) is False # load path + finally: + common._vcr_load_guard.active = False + assert common._should_drop_telemetry_record(req) is True + + +def test_load_guard_patch_is_idempotent(): + import vcr.cassette as cassette_mod + + import tests._vcr_conftest_common as common + + common.patch_vcrpy_cassette_load_guard() + first = cassette_mod.Cassette._load + common.patch_vcrpy_cassette_load_guard() + assert cassette_mod.Cassette._load is first + assert getattr(cassette_mod.Cassette._load, "_litellm_load_guarded", False) diff --git a/tests/llm_translation/test_vcr_filters.py b/tests/llm_translation/test_vcr_filters.py index 03891682781..2b5a6b32a72 100644 --- a/tests/llm_translation/test_vcr_filters.py +++ b/tests/llm_translation/test_vcr_filters.py @@ -21,11 +21,13 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", from tests._vcr_conftest_common import ( # noqa: E402 VCR_FIXED_MULTIPART_BOUNDARY, VCR_IMAGE_B64_PLACEHOLDER, + _before_record_request, _normalize_multipart_boundary, + _should_passthrough_credential_exchange, _strip_image_b64_payloads, + _vcr_load_guard, ) - # --------------------------------------------------------------------------- # Image b64 stripper # --------------------------------------------------------------------------- @@ -218,3 +220,55 @@ def test_normalize_multipart_handles_quoted_boundary(): _normalize_multipart_boundary(req) assert b"quoted-boundary" not in req.body assert VCR_FIXED_MULTIPART_BOUNDARY.encode("utf-8") in req.body + + +# --------------------------------------------------------------------------- +# Credential-exchange passthrough (Google OAuth2/STS token mint must run live) +# --------------------------------------------------------------------------- + + +def _oauth_token_request() -> Request: + return Request( + method="POST", + uri="https://oauth2.googleapis.com/token", + body=b"assertion=eyJhbGciOiJSUzI1NiJ9.signed-jwt&grant_type=urn", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + + +def test_before_record_request_drops_oauth_token_mint(): + # The token mint must never be stored or replayed, else a stale ya29.* token + # gets sent to a live Vertex/Gemini endpoint -> ACCESS_TOKEN_EXPIRED. + assert _before_record_request(_oauth_token_request()) is None + + +def test_before_record_request_keeps_normal_request(): + req = Request( + method="POST", + uri="https://api.openai.com/v1/chat/completions", + body=b'{"model":"gpt-4o"}', + headers={"content-type": "application/json"}, + ) + assert _before_record_request(req) is req + + +def test_credential_exchange_passthrough_inert_during_cassette_load(): + # During Cassette._load stored episodes are replayed through this hook; + # dropping there would mutate the cassette on read. The guard makes it inert. + _vcr_load_guard.active = True + try: + assert _should_passthrough_credential_exchange(_oauth_token_request()) is False + assert _before_record_request(_oauth_token_request()) is not None + finally: + _vcr_load_guard.active = False + + +def test_credential_exchange_passthrough_covers_sts_and_metadata_hosts(): + for host in ("sts.googleapis.com", "metadata.google.internal", "169.254.169.254"): + req = Request( + method="POST", + uri=f"https://{host}/token", + body=b"grant_type=urn", + headers={}, + ) + assert _should_passthrough_credential_exchange(req) is True diff --git a/tests/llm_translation/test_vcr_redis_persister.py b/tests/llm_translation/test_vcr_redis_persister.py index ec86ee73597..7262c09c9b4 100644 --- a/tests/llm_translation/test_vcr_redis_persister.py +++ b/tests/llm_translation/test_vcr_redis_persister.py @@ -79,6 +79,59 @@ def test_load_missing_key_raises_cassette_not_found(): persister.load_cassette("never/recorded", yamlserializer) +def test_load_refreshes_ttl_so_replayed_cassettes_do_not_expire(): + """A successful read must slide the cassette's expiry forward. + + Regression: ``load_cassette`` used a plain ``GET``, which does not + touch the key's TTL. A cassette that is only ever replayed (HIT/NOOP, + never re-recorded) therefore expired exactly ``CASSETTE_TTL_SECONDS`` + after its last *write* no matter how often it was read, and whichever + CI run crossed that 24h boundary re-recorded it live — a spurious VCR + MISS on otherwise-deterministic cassettes. Reading must refresh the + TTL so an actively-used cassette never expires. + """ + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_ttl_refresh" + key = redis_key_for(cassette_id) + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + # Simulate a cassette written ~most-of-a-day ago: only a little TTL left. + fake.expire(key, 60) + assert fake.ttl(key) <= 60 + + persister.load_cassette(cassette_id, yamlserializer) + + refreshed = fake.ttl(key) + assert CASSETTE_TTL_SECONDS - 5 <= refreshed <= CASSETTE_TTL_SECONDS + + +def test_load_ttl_refresh_failure_does_not_break_load(): + """A failed TTL refresh must never turn a successful load into a miss.""" + + class _RefreshFailsRedis: + def __init__(self, inner): + self._inner = inner + + def get(self, *args, **kwargs): + return self._inner.get(*args, **kwargs) + + def set(self, *args, **kwargs): + return self._inner.set(*args, **kwargs) + + def expire(self, *args, **kwargs): + raise RedisConnectionError("simulated outage") + + client = _RefreshFailsRedis(fakeredis.FakeStrictRedis()) + persister = make_redis_persister(client=client) + cassette_id = "tests/llm_translation/test_x/test_ttl_refresh_fail" + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + requests, responses = persister.load_cassette(cassette_id, yamlserializer) + + assert len(requests) == 1 + assert len(responses) == 1 + + def test_redis_key_normalizes_path_passed_by_pytest_recording(): raw = "tests/llm_translation/cassettes/test_anthropic/test_streaming.yaml" assert ( diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index acb79a7577d..abb871789c3 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -68,7 +68,19 @@ _VCR_INCOMPATIBLE_FILES = frozenset( } ) -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () +# Individual tests (vs. whole files above) that VCR replay can't model: +# - ``test_router_text_completion_client``: a concurrency test that fires 300 +# identical requests to verify the async OpenAI client is *reused* across +# calls (per its own comment, it "fails when we create a new Async OpenAI +# client per request"). vcrpy patches the HTTP transport, so replay never +# opens real connections and cannot exercise the client pool the test exists +# to validate. Recording instead stores ~300 near-identical episodes, which +# blows past MAX_EPISODES_PER_CASSETTE (50) so the cassette is refused on +# every run (MISS:OVERFLOW). The endpoint is a free mock, so the live calls +# carry no real provider cost. +_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ( + "test_router.py::test_router_text_completion_client", +) _verbose_state = VerboseReporterState() diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index cce6d33e799..c7abdb5f493 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -299,7 +299,10 @@ def test_completion_claude_3(): @pytest.mark.parametrize( "model", - ["anthropic/claude-sonnet-4-5-20250929", "anthropic.claude-3-sonnet-20240229-v1:0"], + [ + "anthropic/claude-sonnet-4-5-20250929", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + ], ) def test_completion_claude_3_function_call(model): litellm.set_verbose = True @@ -385,7 +388,7 @@ def test_completion_claude_3_function_call(model): [ ("gpt-3.5-turbo", None, None), ("claude-sonnet-4-5-20250929", None, None), - ("anthropic.claude-3-sonnet-20240229-v1:0", None, None), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", None, None), # ( # "azure_ai/command-r-plus", # os.getenv("AZURE_COHERE_API_KEY"), @@ -1578,7 +1581,7 @@ def test_completion_openai(): [ # ("gpt-4o-2024-08-06", None), # ("azure/gpt-4.1-mini", None), - ("bedrock/anthropic.claude-3-sonnet-20240229-v1:0", None), + ("bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", None), # ("azure/gpt-4o-new-test", "2024-08-01-preview"), ], ) @@ -1666,15 +1669,13 @@ def custom_callback( ################################################# - print( - f""" + print(f""" Model: {model}, Messages: {messages}, User: {user}, Seed: {kwargs["seed"]}, temperature: {kwargs["temperature"]}, - """ - ) + """) assert kwargs["user"] == "ishaans app" assert kwargs["model"] == "gpt-3.5-turbo-1106" @@ -2699,7 +2700,7 @@ def test_bedrock_deepseek_custom_prompt_dict(): def test_bedrock_deepseek_known_tokenizer_config(monkeypatch): model = ( - "deepseek_r1/arn:aws:bedrock:us-west-2:888602223428:imported-model/bnnr6463ejgf" + "deepseek_r1/arn:aws:bedrock:us-west-2:941277531214:imported-model/bnnr6463ejgf" ) from litellm.llms.custom_httpx.http_handler import HTTPHandler from unittest.mock import Mock @@ -2914,8 +2915,8 @@ def response_format_tests(response: litellm.ModelResponse): "model", [ "bedrock/mistral.mistral-large-2407-v1:0", - "bedrock/cohere.command-r-plus-v1:0", - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "mistral.mistral-7b-instruct-v0:2", "meta.llama3-8b-instruct-v1:0", ], diff --git a/tests/local_testing/test_function_call_parsing.py b/tests/local_testing/test_function_call_parsing.py index f9582fcc574..2453571f1c4 100644 --- a/tests/local_testing/test_function_call_parsing.py +++ b/tests/local_testing/test_function_call_parsing.py @@ -142,7 +142,8 @@ def trade(model_name: str) -> List[Trade]: # type: ignore @pytest.mark.parametrize( - "model", ["claude-haiku-4-5-20251001", "anthropic.claude-3-haiku-20240307-v1:0"] + "model", + ["claude-haiku-4-5-20251001", "us.anthropic.claude-haiku-4-5-20251001-v1:0"], ) @pytest.mark.flaky(retries=6, delay=10) def test_function_call_parsing(model): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 3c7e004b62e..1cad7d1421e 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -49,7 +49,7 @@ def get_current_weather(location, unit="fahrenheit"): "mistral/mistral-large-latest", "claude-haiku-4-5-20251001", "gemini/gemini-2.5-flash-lite", - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", ], ) @pytest.mark.flaky(retries=3, delay=1) @@ -267,7 +267,6 @@ def test_aaparallel_function_call_with_anthropic_thinking(model): from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message - _PARALLEL_TOOL_HISTORY_MESSAGES = [ { "role": "user", @@ -303,7 +302,7 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ [ # Bedrock Converse still requires modify_params to inject the dummy tool. ( - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", _PARALLEL_TOOL_HISTORY_MESSAGES, True, ), @@ -314,7 +313,7 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ False, ), ( - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", [ { "role": "user", @@ -579,7 +578,7 @@ def test_groq_parallel_function_call(): @pytest.mark.parametrize( "model", [ - "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", ], ) def test_passing_tool_result_as_list(model): diff --git a/tests/local_testing/test_router_max_parallel_requests.py b/tests/local_testing/test_router_max_parallel_requests.py index ab827b057e3..1b81b9eb999 100644 --- a/tests/local_testing/test_router_max_parallel_requests.py +++ b/tests/local_testing/test_router_max_parallel_requests.py @@ -123,8 +123,6 @@ def test_setting_mpr_limits_per_model( async def _handle_router_calls(router): - import random - pre_fill = """ Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nunc ut finibus massa. Quisque a magna magna. Quisque neque diam, varius sit amet tellus eu, elementum fermentum sapien. Integer ut erat eget arcu rutrum blandit. Morbi a metus purus. Nulla porta, urna at finibus malesuada, velit ante suscipit orci, vitae laoreet dui ligula ut augue. Cras elementum pretium dui, nec luctus nulla aliquet ut. Nam faucibus, diam nec semper interdum, nisl nisi viverra nulla, vitae sodales elit ex a purus. Donec tristique malesuada lobortis. Donec posuere iaculis nisl, vitae accumsan libero dignissim dignissim. Suspendisse finibus leo et ex mattis tempor. Praesent at nisl vitae quam egestas lacinia. Donec in justo non erat aliquam accumsan sed vitae ex. Vivamus gravida diam vel ipsum tincidunt dignissim. @@ -141,7 +139,11 @@ async def _handle_router_calls(router): [ { "role": "user", - "content": f"{pre_fill * 3}\n\nRecite the Declaration of independence at a speed of {random.random() * 100} words per minute.", + # Fixed speed (was random.random()*100) so the request body is + # deterministic and the VCR cassette replays instead of + # appending a new episode every run. This is a rate-limiting + # test; the prompt content is irrelevant to what it asserts. + "content": f"{pre_fill * 3}\n\nRecite the Declaration of independence at a speed of 50.0 words per minute.", } ], stream=True, diff --git a/tests/local_testing/test_sagemaker.py b/tests/local_testing/test_sagemaker.py index d4c5a5a857f..fdc8347c36a 100644 --- a/tests/local_testing/test_sagemaker.py +++ b/tests/local_testing/test_sagemaker.py @@ -57,7 +57,7 @@ async def test_completion_sagemaker(sync_mode): print("testing sagemaker") if sync_mode is True: response = litellm.completion( - model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + model="sagemaker/litellm-ci-textgen", messages=[ {"role": "user", "content": "hi"}, ], @@ -67,7 +67,7 @@ async def test_completion_sagemaker(sync_mode): ) else: response = await litellm.acompletion( - model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + model="sagemaker/litellm-ci-textgen", messages=[ {"role": "user", "content": "hi"}, ], @@ -158,7 +158,7 @@ async def test_completion_sagemaker_messages_api(sync_mode): "model", [ # "sagemaker_chat/huggingface-pytorch-tgi-inference-2024-08-23-15-48-59-245", - "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + "sagemaker/litellm-ci-textgen", ], ) # @pytest.mark.flaky(retries=3, delay=1) @@ -218,7 +218,7 @@ async def test_completion_sagemaker_stream(sync_mode, model): "model", [ # "sagemaker_chat/huggingface-pytorch-tgi-inference-2024-08-23-15-48-59-245", - "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + "sagemaker/litellm-ci-textgen", ], ) async def test_completion_sagemaker_streaming_bad_request(sync_mode, model): @@ -256,7 +256,7 @@ async def test_acompletion_sagemaker_non_stream(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + "model": "sagemaker/litellm-ci-textgen", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -282,7 +282,7 @@ async def test_acompletion_sagemaker_non_stream(): ) as mock_post: # Act: Call the litellm.acompletion function response = await litellm.acompletion( - model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + model="sagemaker/litellm-ci-textgen", messages=[ {"role": "user", "content": "hi"}, ], @@ -302,7 +302,7 @@ async def test_acompletion_sagemaker_non_stream(): assert args_to_sagemaker == expected_payload assert ( kwargs["url"] - == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations" + == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/litellm-ci-textgen/invocations" ) @@ -316,7 +316,7 @@ async def test_completion_sagemaker_non_stream(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + "model": "sagemaker/litellm-ci-textgen", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -342,7 +342,7 @@ async def test_completion_sagemaker_non_stream(): ) as mock_post: # Act: Call the litellm.acompletion function response = litellm.completion( - model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + model="sagemaker/litellm-ci-textgen", messages=[ {"role": "user", "content": "hi"}, ], @@ -362,7 +362,7 @@ async def test_completion_sagemaker_non_stream(): assert args_to_sagemaker == expected_payload assert ( kwargs["url"] - == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations" + == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/litellm-ci-textgen/invocations" ) @@ -377,7 +377,7 @@ async def test_completion_sagemaker_prompt_template_non_stream(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + "model": "sagemaker/litellm-ci-textgen", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -433,7 +433,7 @@ async def test_completion_sagemaker_non_stream_with_aws_params(): "id": "cmpl-mockid", "object": "text_completion", "created": 1629800000, - "model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + "model": "sagemaker/litellm-ci-textgen", "choices": [ { "text": "This is a mock response from SageMaker.", @@ -459,7 +459,7 @@ async def test_completion_sagemaker_non_stream_with_aws_params(): ) as mock_post: # Act: Call the litellm.acompletion function response = litellm.completion( - model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614", + model="sagemaker/litellm-ci-textgen", messages=[ {"role": "user", "content": "hi"}, ], @@ -482,5 +482,5 @@ async def test_completion_sagemaker_non_stream_with_aws_params(): assert args_to_sagemaker == expected_payload assert ( kwargs["url"] - == "https://runtime.sagemaker.us-west-5.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations" + == "https://runtime.sagemaker.us-west-5.amazonaws.com/endpoints/litellm-ci-textgen/invocations" ) diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 10f351714e1..eb153404a44 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1174,7 +1174,7 @@ async def test_completion_replicate_llama3_streaming(sync_mode): [ # ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"], # ["bedrock/cohere.command-r-plus-v1:0", None], - ["anthropic.claude-3-sonnet-20240229-v1:0", None], + ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None], # ["mistral.mistral-7b-instruct-v0:2", None], # ["meta.llama3-8b-instruct-v1:0", None], ], @@ -1246,7 +1246,7 @@ def test_bedrock_claude_3_streaming(): try: litellm.set_verbose = True response: ModelResponse = completion( # type: ignore - model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, max_tokens=10, # type: ignore stream=True, @@ -1276,7 +1276,7 @@ def test_bedrock_claude_3_streaming(): "model", [ "claude-haiku-4-5-20251001", - "cohere.command-r-plus-v1:0", # bedrock + "us.anthropic.claude-haiku-4-5-20251001-v1:0", # bedrock "gpt-3.5-turbo", ], ) @@ -3500,7 +3500,7 @@ def test_unit_test_perplexity_citations_chunk(): [ "gpt-3.5-turbo", "claude-sonnet-4-5-20250929", - "anthropic.claude-3-sonnet-20240229-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", # "vertex_ai/claude-3-5-sonnet@20240620", ], ) diff --git a/tests/logging_callback_tests/test_amazing_s3_logs.py b/tests/logging_callback_tests/test_amazing_s3_logs.py index dab2a0cc0b9..e6291a94049 100644 --- a/tests/logging_callback_tests/test_amazing_s3_logs.py +++ b/tests/logging_callback_tests/test_amazing_s3_logs.py @@ -27,7 +27,7 @@ async def test_basic_s3_logging(sync_mode, streaming): verbose_logger.setLevel(level=logging.DEBUG) litellm.success_callback = ["s3"] litellm.s3_callback_params = { - "s3_bucket_name": "load-testing-oct", + "s3_bucket_name": "load-testing-oct-941277531214", "s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY", "s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID", "s3_region_name": "us-west-2", @@ -64,14 +64,14 @@ async def test_basic_s3_logging(sync_mode, streaming): await asyncio.sleep(2) print(f"response: {response}") - total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct") + total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct-941277531214") # assert that atlest one key has response.id in it assert any(response_id in key for key in all_s3_keys) s3 = boto3.client("s3") # delete all objects for key in all_s3_keys: - s3.delete_object(Bucket="load-testing-oct", Key=key) + s3.delete_object(Bucket="load-testing-oct-941277531214", Key=key) @pytest.mark.asyncio @@ -82,7 +82,7 @@ async def test_basic_s3_v2_logging(streaming): from litellm.integrations.s3_v2 import S3Logger litellm.s3_callback_params = { - "s3_bucket_name": "load-testing-oct", + "s3_bucket_name": "load-testing-oct-941277531214", "s3_aws_secret_access_key": "test-secret", "s3_aws_access_key_id": "test-key", "s3_region_name": "us-west-2", diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index d6d0652ed77..0d4405094b5 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -2,7 +2,6 @@ import io import os import sys - sys.path.insert(0, os.path.abspath("../..")) import asyncio @@ -67,7 +66,7 @@ def setup_vector_store_registry(): litellm.vector_store_registry = VectorStoreRegistry( vector_stores=[ LiteLLM_ManagedVectorStore( - vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock" + vector_store_id="LCYXFBR2TU", custom_llm_provider="bedrock" ) ] ) @@ -111,7 +110,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_completion( response = await litellm.acompletion( model="anthropic/claude-3.5-sonnet", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], + vector_store_ids=["LCYXFBR2TU"], client=client, ) except Exception as e: @@ -152,7 +151,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call( response = await litellm.acompletion( model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], + vector_store_ids=["LCYXFBR2TU"], client=async_client, ) print("OPENAI RESPONSE:", json.dumps(dict(response), indent=4, default=str)) @@ -196,7 +195,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_streaming( response = await litellm.acompletion( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], + vector_store_ids=["LCYXFBR2TU"], stream=True, client=async_client, ) @@ -255,7 +254,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools( model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}", messages=[{"role": "user", "content": "what is litellm?"}], max_tokens=10, - tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], + tools=[{"type": "file_search", "vector_store_ids": ["LCYXFBR2TU"]}], ) assert response is not None @@ -279,7 +278,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_ tools=[ { "type": "file_search", - "vector_store_ids": ["T37J8R4WTM"], + "vector_store_ids": ["LCYXFBR2TU"], "filters": { "key": "user_id", "value": "fake-user-id", @@ -387,7 +386,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( tools=[ { "type": "file_search", - "vector_store_ids": ["T37J8R4WTM"], + "vector_store_ids": ["LCYXFBR2TU"], "filters": { "key": "user_id", "value": "fake-user-id", @@ -461,7 +460,7 @@ async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registr await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], + vector_store_ids=["LCYXFBR2TU"], client=client, ) except Exception as e: @@ -537,7 +536,7 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai( await litellm.acompletion( model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], - tools=[{"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}], + tools=[{"type": "file_search", "vector_store_ids": ["LCYXFBR2TU"]}], client=client, ) except Exception as e: @@ -611,7 +610,7 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist model="gpt-5.5", messages=[{"role": "user", "content": "what is litellm?"}], tools=[ - {"type": "file_search", "vector_store_ids": ["T37J8R4WTM"]}, + {"type": "file_search", "vector_store_ids": ["LCYXFBR2TU"]}, {"type": "file_search", "vector_store_ids": ["unknownVS"]}, ], client=client, @@ -645,7 +644,7 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist # model="gpt-5.5", # messages=[{"role": "user", "content": "what is litellm?"}], # vector_store_ids = [ -# "T37J8R4WTM" +# "LCYXFBR2TU" # ], # ) @@ -667,7 +666,7 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist # # expect the vector store request metadata object to have the correct values # vector_store_request_metadata = standard_logging_vector_store_request_metadata[0] -# assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM" +# assert vector_store_request_metadata.get("vector_store_id") == "LCYXFBR2TU" # assert vector_store_request_metadata.get("query") == "what is litellm?" # assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock" @@ -723,7 +722,7 @@ async def test_e2e_bedrock_knowledgebase_retrieval_without_vector_store_registry response = await litellm.acompletion( model="anthropic/claude-3.5-sonnet", messages=[{"role": "user", "content": "what is litellm?"}], - vector_store_ids=["T37J8R4WTM"], + vector_store_ids=["LCYXFBR2TU"], client=client, ) except Exception as e: diff --git a/tests/ocr_tests/conftest.py b/tests/ocr_tests/conftest.py index 94790bd7aa3..7a74dde3e41 100644 --- a/tests/ocr_tests/conftest.py +++ b/tests/ocr_tests/conftest.py @@ -26,6 +26,28 @@ from tests._vcr_conftest_common import ( # noqa: E402,F401 vcr_config_dict, ) +# Vertex AI MaaS Mistral OCR tests that cannot be VCR-cached in CI. +# +# ``vertex_ai/mistral-ocr-2505`` is a Model-as-a-Service partner model that +# must be explicitly enabled in the GCP project's Model Garden. It is not +# provisioned in the CI project (``litellm-ci-cd``), so the live +# ``:rawPredict`` call fails on every run and ``BaseOCRTest`` catches the +# provider error and skips. Because the doomed live call is recorded but the +# test then skips, the persister refuses to save it (skipped tests don't +# persist) and the cassette is never seeded — so the test re-records live and +# is classified MISS:NOT_PERSISTED on every single run, forever. No cassette +# can be recorded until the model is provisioned. Mark the tests VCR- +# incompatible so they are honestly accounted as live calls (UNMARKED:LIVE_CALL) +# rather than phantom cache misses; behaviour is unchanged (they still run and +# still skip on the provider error). The sibling direct-Mistral and Azure OCR +# tests replay from cache normally and are unaffected. Remove these entries if +# the MaaS model is enabled in the CI project. +_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ( + "test_ocr_vertex_ai.py::TestVertexAIMistralOCR::test_ocr_response_structure", + "test_ocr_vertex_ai.py::TestVertexAIMistralOCR::test_basic_ocr_with_url[True]", + "test_ocr_vertex_ai.py::TestVertexAIMistralOCR::test_basic_ocr_with_url[False]", +) + _verbose_state = VerboseReporterState() @@ -62,7 +84,10 @@ def pytest_runtest_logreport(report): def pytest_collection_modifyitems(config, items): - apply_vcr_auto_marker_to_items(items) + apply_vcr_auto_marker_to_items( + items, + skip_nodeid_suffixes=_VCR_INCOMPATIBLE_NODEID_SUFFIXES, + ) def pytest_terminal_summary(terminalreporter, exitstatus, config): diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/proxy_unit_tests/test_proxy_routes.py index 34123e992c2..db41bd65409 100644 --- a/tests/proxy_unit_tests/test_proxy_routes.py +++ b/tests/proxy_unit_tests/test_proxy_routes.py @@ -217,9 +217,120 @@ def _create_request_with_host_header(path: str, host_header: str) -> Request: ], ) def test_get_request_route_not_bypassed_by_malformed_host(host_header: str): - for protected_path in ["/health", "/user/new", "/key/generate", "/get/internal_user_settings"]: - request = _create_request_with_host_header(path=protected_path, host_header=host_header) - result = get_request_route(request) - assert result == protected_path, ( - f"Host: {host_header!r} caused route {protected_path!r} to resolve as {result!r}" + for protected_path in [ + "/health", + "/user/new", + "/key/generate", + "/get/internal_user_settings", + ]: + request = _create_request_with_host_header( + path=protected_path, host_header=host_header ) + result = get_request_route(request) + assert ( + result == protected_path + ), f"Host: {host_header!r} caused route {protected_path!r} to resolve as {result!r}" + + +# --------------------------------------------------------------------------- +# Regression tests for variant call sites that previously read request.url.path +# (Host-derived) instead of the ASGI scope path. Each test sends a Host header +# crafted to collapse url.path to a substring the call site's decision logic +# would match on, while scope["path"] is the real (unmatching) route. +# --------------------------------------------------------------------------- + +_BYPASS_HOSTS = [ + "localhost/?x=1", + "localhost:4000/?x=1", + "localhost/#test", + "localhost:4000/#test", +] + + +def _is_assistants(req): + return RouteChecks._is_assistants_api_request(req) + + +def _metadata_var_name(req): + from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name + + return _get_metadata_variable_name(req) + + +def _vector_store_id_in_path(req): + from litellm.proxy.common_utils.http_parsing_utils import ( + _add_vector_store_id_from_path, + ) + + data: dict = {} + _add_vector_store_id_from_path(request_data=data, request=req) + return "vector_store_id" in data + + +# (label, scope_path, host_suffix_template, predicate, expected) — host_suffix_template +# receives the host_header via %s substitution. The predicate is invoked on a Request +# whose scope["path"] is scope_path and whose Host header is the formatted suffix. +# +# The MCP entries (well_known_mcp_bypass, pkce_token_suffix) call +# get_request_route directly rather than the surrounding production handler +# (MCPRequestHandler.process_mcp_request / _mcp_oauth_user_api_key_auth) — +# those handlers require an ASGI scope plus MCP state to invoke, and the call +# sites do nothing with the path except feed it to this helper. The helper- +# level assertion is the relevant signal. +_CALL_SITES = [ + ("assistants_classification", "/key/generate", "%s/thread", _is_assistants, False), + ( + "metadata_variable_name", + "/chat/completions", + "%s/thread", + _metadata_var_name, + "metadata", + ), + ( + "vector_store_id_extraction", + "/key/generate", + "%s/vector_stores/x/files", + _vector_store_id_in_path, + False, + ), + ( + "well_known_mcp_bypass", + "/mcp/tools/call", + "/.well-known/%s", + lambda r: get_request_route(r).startswith("/.well-known/"), + False, + ), + ( + "pkce_token_suffix", + "/mcp/server-id/token", + "%s", + lambda r: get_request_route(r).rstrip("/").lower().endswith("/token"), + True, + ), + ( + "spend_logs_v2_classification", + "/spend/logs", + "%s/spend/logs/v2", + lambda r: "/spend/logs/v2" in get_request_route(r), + False, + ), + ("health_route_echo", "/test", "%s", lambda r: get_request_route(r), "/test"), +] + + +@pytest.mark.parametrize("host_header", _BYPASS_HOSTS) +@pytest.mark.parametrize( + "label,scope_path,host_suffix_template,predicate,expected", + _CALL_SITES, + ids=[c[0] for c in _CALL_SITES], +) +def test_call_site_uses_scope_path( + label, scope_path, host_suffix_template, predicate, expected, host_header +): + """Each call site that previously read request.url.path must now make its + decision against scope["path"]. The Host header is crafted so url.path + would resolve to a value that flips the decision under the old code.""" + request = _create_request_with_host_header( + path=scope_path, host_header=host_suffix_template % host_header + ) + assert predicate(request) == expected diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index c295805bdb3..176405bb9ca 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -1,3 +1,4 @@ +import json import sys from types import SimpleNamespace from unittest.mock import AsyncMock @@ -91,8 +92,9 @@ async def test_should_not_mutate_dict_container_response_when_recording_owner( assert returned == {"id": "cntr_provider", "object": "container"} data = table.create.await_args.kwargs["data"] - assert data["file_object"]["custom_llm_provider"] == "openai" - assert data["file_object"]["provider_container_id"] == "cntr_provider" + file_obj = json.loads(data["file_object"]) + assert file_obj["custom_llm_provider"] == "openai" + assert file_obj["provider_container_id"] == "cntr_provider" @pytest.mark.asyncio @@ -913,3 +915,195 @@ async def test_admin_with_identity_records_container_ownership(monkeypatch): table.create.assert_awaited_once() created_data = table.create.await_args.kwargs["data"] assert created_data["created_by"] == "proxy-admin" + + +@pytest.mark.asyncio +async def test_should_record_containers_from_responses_output_for_service_account( + monkeypatch, +): + table = AsyncMock() + table.find_unique.return_value = None + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + auth = UserAPIKeyAuth(team_id="team-1") + encoded_container_id = ( + "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR" + "lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" + ) + responses_payload = { + "output": [ + { + "type": "message", + "content": [ + { + "type": "output_text", + "annotations": [ + { + "type": "container_file_citation", + "container_id": encoded_container_id, + "file_id": "cfile_abc", + } + ], + } + ], + } + ], + "_hidden_params": {"custom_llm_provider": "azure"}, + } + + await ownership.record_container_owners_from_responses_response( + response=responses_payload, + user_api_key_dict=auth, + ) + + table.create.assert_awaited_once() + created_data = table.create.await_args.kwargs["data"] + assert created_data["created_by"] == "team:team-1" + assert created_data["unified_object_id"] == encoded_container_id + + +@pytest.mark.asyncio +async def test_service_account_can_access_container_after_responses_tracking( + monkeypatch, +): + encoded_container_id = ( + "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR" + "lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" + ) + table = AsyncMock() + table.find_unique.return_value = None + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + auth = UserAPIKeyAuth(team_id="team-1") + + await ownership.record_container_owners_from_responses_response( + response={ + "output": [ + { + "type": "code_interpreter_call", + "container_id": encoded_container_id, + } + ], + "_hidden_params": {"custom_llm_provider": "azure"}, + }, + user_api_key_dict=auth, + ) + + original_id, provider = await ownership.assert_user_can_access_container( + container_id=encoded_container_id, + user_api_key_dict=auth, + custom_llm_provider="azure", + ) + assert original_id == "cntr_native" + assert provider == "azure" + + +@pytest.mark.asyncio +async def test_should_record_container_ownership_after_streaming_responses_finish( + monkeypatch, +): + """Streaming /v1/responses calls return through the + ``select_data_generator`` branch and never reach the non-streaming + container-ownership tail. The wrapper must read + ``completed_response`` off the upstream iterator once iteration + finishes and write the row, otherwise code-interpreter containers + created during the stream stay unregistered and follow-up file API + calls 403. + """ + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + encoded_container_id = ( + "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR" + "lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" + ) + response_body = SimpleNamespace( + output=[ + SimpleNamespace( + type="code_interpreter_call", + container_id=encoded_container_id, + code_interpreter_call=None, + ) + ] + ) + stream_response = SimpleNamespace( + completed_response=SimpleNamespace(response=response_body), + _hidden_params={"custom_llm_provider": "azure"}, + ) + + async def fake_sse_generator(): + yield "data: chunk-1\n\n" + yield "data: chunk-2\n\n" + + table = AsyncMock() + table.find_unique.return_value = None + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + auth = UserAPIKeyAuth(team_id="team-1") + + wrapped = ( + ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=stream_response, + wrapped_generator=fake_sse_generator(), + user_api_key_dict=auth, + ) + ) + + chunks = [chunk async for chunk in wrapped] + assert chunks == ["data: chunk-1\n\n", "data: chunk-2\n\n"] + + table.create.assert_awaited_once() + created_data = table.create.await_args.kwargs["data"] + assert created_data["created_by"] == "team:team-1" + assert created_data["unified_object_id"] == encoded_container_id + + +@pytest.mark.asyncio +async def test_streaming_ownership_wrap_no_op_when_stream_did_not_complete( + monkeypatch, +): + """If the stream errored before ``response.completed``, + ``completed_response`` is ``None`` — we must skip the ownership + write rather than crash the response generator.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + stream_response = SimpleNamespace(completed_response=None) + + async def fake_sse_generator(): + yield "data: chunk-1\n\n" + + record = AsyncMock() + monkeypatch.setattr( + ownership, + "record_container_owners_from_responses_response", + record, + ) + + wrapped = ( + ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=stream_response, + wrapped_generator=fake_sse_generator(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + ) + chunks = [chunk async for chunk in wrapped] + + assert chunks == ["data: chunk-1\n\n"] + record.assert_not_awaited() diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/test_litellm/integrations/arize/test_arize_phoenix.py index 4a2eab29e8e..afd83f81ce0 100644 --- a/tests/test_litellm/integrations/arize/test_arize_phoenix.py +++ b/tests/test_litellm/integrations/arize/test_arize_phoenix.py @@ -7,7 +7,6 @@ from litellm.integrations.arize.arize_phoenix import ( ArizePhoenixConfig, ArizePhoenixLogger, ) -from litellm.integrations.arize._utils import ArizeOTELAttributes class TestArizePhoenixConfig(unittest.TestCase): @@ -217,44 +216,147 @@ def test_get_arize_phoenix_config_expection_on_missing_api_key(monkeypatch, env_ # --------------------------------------------------------------------------- -# Dynamic project naming from metadata +# Per-project routing via Resource (not span attributes) # --------------------------------------------------------------------------- -class TestGetDynamicProjectName: - """Tests for _get_dynamic_project_name extraction logic.""" +class TestResolveProjectName: + """Tests for _resolve_project_name priority chain.""" - def test_extracts_from_standard_logging_object_metadata(self): + def test_extracts_phoenix_name_from_standard_logging_object_metadata(self): kwargs = { "standard_logging_object": { "metadata": {"phoenix_project_name": "my-project"}, } } - assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "my-project" + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "my-project" - def test_extracts_from_litellm_params_metadata(self): + def test_extracts_phoenix_name_from_litellm_params_metadata(self): kwargs = { "litellm_params": { "metadata": {"phoenix_project_name": "sdk-project"}, } } - assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "sdk-project" + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "sdk-project" - def test_returns_none_when_no_metadata(self): - assert ArizePhoenixLogger._get_dynamic_project_name({}) is None + @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-project"}, clear=False) + def test_falls_back_to_phoenix_env_when_no_metadata(self): + assert ArizePhoenixLogger._resolve_project_name({}) == "env-project" + + @patch.dict( + "os.environ", + {"ARIZE_PROJECT_NAME": "arize-env", "PHOENIX_PROJECT_NAME": ""}, + clear=False, + ) + def test_falls_back_to_arize_env_when_phoenix_unset(self): + assert ArizePhoenixLogger._resolve_project_name({}) == "arize-env" + + @patch.dict("os.environ", {}, clear=True) + def test_falls_back_to_default_when_no_metadata_or_env(self): + assert ArizePhoenixLogger._resolve_project_name({}) == "default" + + def test_phoenix_override_beats_phoenix_metadata(self): + kwargs = { + "standard_logging_object": { + "metadata": { + "phoenix_project_name_override": "override-proj", + "phoenix_project_name": "phoenix-proj", + }, + } + } + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "override-proj" + + def test_whitespace_only_metadata_falls_through_to_default(self): + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name_override": " "}, + } + } + with patch.dict("os.environ", {}, clear=True): + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "default" + + def test_strips_whitespace_from_project_name(self): + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name": " trimmed "}, + } + } + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "trimmed" def test_non_dict_standard_logging_object_does_not_raise(self): - """isinstance(dict) guard prevents AttributeError on non-dict payloads.""" kwargs = {"standard_logging_object": "not-a-dict"} - assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) is None + with patch.dict("os.environ", {}, clear=True): + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "default" + + def test_resolves_override_from_user_api_key_auth_metadata(self): + kwargs = { + "litellm_params": { + "metadata": { + "user_api_key_auth_metadata": { + "phoenix_project_name_override": "claude-code", + }, + }, + }, + } + with patch.dict("os.environ", {}, clear=True): + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "claude-code" + + def test_resolves_phoenix_name_from_user_api_key_auth_metadata(self): + kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_auth_metadata": { + "phoenix_project_name": "team-project", + }, + }, + }, + } + with patch.dict("os.environ", {}, clear=True): + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "team-project" + + def test_proxy_ignores_client_metadata_when_auth_metadata_set(self): + kwargs = { + "litellm_params": { + "proxy_server_request": { + "url": "/v1/chat/completions", + "method": "POST", + "headers": {}, + }, + "metadata": { + "phoenix_project_name_override": "attacker-project", + "user_api_key_auth_metadata": { + "phoenix_project_name_override": "team-project", + }, + }, + }, + } + with patch.dict("os.environ", {}, clear=True): + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "team-project" + + def test_proxy_without_auth_metadata_falls_back_to_env(self): + kwargs = { + "litellm_params": { + "proxy_server_request": { + "url": "/v1/chat/completions", + "method": "POST", + "headers": {}, + }, + "metadata": {"phoenix_project_name": "attacker-project"}, + }, + } + with patch.dict( + "os.environ", {"PHOENIX_PROJECT_NAME": "env-project"}, clear=True + ): + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "env-project" -class TestDynamicProjectNameOnSpan: - """set_arize_phoenix_attributes sets openinference.project.name on the span.""" +class TestProjectNameNotOnSpan: + """Project routing uses Resource on TracerProvider, not span attributes.""" - @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-fallback"}, clear=False) @patch("litellm.integrations.arize._utils.set_attributes") - def test_dynamic_name_sets_span_attribute(self, _mock_set_attrs): + def test_set_arize_phoenix_attributes_does_not_set_project_on_span( + self, _mock_set_attrs + ): span = MagicMock() kwargs = { "standard_logging_object": { @@ -263,20 +365,468 @@ class TestDynamicProjectNameOnSpan: } ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj=None) - span.set_attribute.assert_called_once_with( - "openinference.project.name", "dynamic-proj" + for call in span.set_attribute.call_args_list: + assert call[0][0] != "openinference.project.name" + + +class TestPerProjectTracerProviderCache: + """Spans for different projects use different Resources on export.""" + + def test_different_metadata_routes_to_different_resource(self): + from datetime import datetime + + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, ) - @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-project"}, clear=False) - @patch("litellm.integrations.arize._utils.set_attributes") - def test_falls_back_to_env_var_when_no_dynamic_name(self, _mock_set_attrs): - span = MagicMock() - ArizePhoenixLogger.set_arize_phoenix_attributes(span, {}, response_obj=None) + from litellm.integrations.opentelemetry import OpenTelemetryConfig - span.set_attribute.assert_called_once_with( - "openinference.project.name", "env-project" + exporter = InMemorySpanExporter() + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=exporter), + callback_name="arize_phoenix", ) + start = datetime(2024, 1, 1, 12, 0, 0) + end = datetime(2024, 1, 1, 12, 0, 1) + + logger._handle_success( + { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "project-a"}, + }, + }, + response_obj={}, + start_time=start, + end_time=end, + ) + logger._handle_success( + { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "project-b"}, + }, + }, + response_obj={}, + start_time=start, + end_time=end, + ) + + spans = exporter.get_finished_spans() + project_names = { + s.resource.attributes.get("openinference.project.name") for s in spans + } + assert "project-a" in project_names + assert "project-b" in project_names + + def test_shared_span_processor_created_once_at_init(self): + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + mock_processor = MagicMock() + with patch.object( + OpenTelemetry, "_get_span_processor", return_value=mock_processor + ) as mock_get_processor: + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + assert mock_get_processor.call_count == 1 + assert logger._shared_span_processor is mock_processor + + logger._project_providers.clear() + logger._get_tracer_for("project-a") + logger._get_tracer_for("project-b") + assert mock_get_processor.call_count == 1 + + def test_lru_eviction_does_not_shutdown_provider(self): + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + logger._project_providers.clear() + + logger._get_tracer_for("project-0") + evicted_provider = logger._project_providers["project-0"] + shutdown_mock = MagicMock() + evicted_provider.shutdown = shutdown_mock # type: ignore[method-assign] + + for i in range(1, 65): + logger._get_tracer_for(f"project-{i}") + + assert len(logger._project_providers) == 64 + assert "project-0" not in logger._project_providers + assert "project-64" in logger._project_providers + shutdown_mock.assert_not_called() + + def test_flush_tracer_providers_force_flushes_shared_processor(self): + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + mock_processor = MagicMock() + logger._shared_span_processor = mock_processor + mock_provider = MagicMock() + logger._project_providers["proj"] = mock_provider + + logger.flush_tracer_providers() + + mock_processor.force_flush.assert_called_once() + mock_provider.force_flush.assert_called_once() + + +class TestGetLitellmResourceForProject: + """Resource attrs used by Phoenix OSS and Arize AX for project routing.""" + + def test_project_attrs_win_over_otel_resource_attributes_env(self): + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + + with patch.dict( + "os.environ", + { + "OTEL_RESOURCE_ATTRIBUTES": "openinference.project.name=env-pinned,model_id=env-model" + }, + clear=False, + ): + resource = logger._get_litellm_resource_for_project("dynamic-proj") + + assert resource.attributes["openinference.project.name"] == "dynamic-proj" + assert resource.attributes["model_id"] == "dynamic-proj" + assert resource.attributes["service.name"] == "dynamic-proj" + + @patch.dict("os.environ", {"OTEL_DEPLOYMENT_ENVIRONMENT": "staging"}, clear=False) + def test_preserves_deployment_environment_from_config(self): + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig( + exporter=MagicMock(), deployment_environment="staging" + ), + callback_name="arize_phoenix", + ) + resource = logger._get_litellm_resource_for_project("my-proj") + assert resource.attributes.get("deployment.environment") == "staging" + + +class TestTracerResolutionAndCache: + """_resolve_tracer_for_kwargs, get_tracer_to_use_for_request, provider cache.""" + + def test_get_tracer_to_use_for_request_matches_resolve_tracer(self): + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "same-proj"}, + } + } + project_name, _ = logger._resolve_tracer_for_kwargs(kwargs) + tracer_from_request = logger.get_tracer_to_use_for_request(kwargs) + assert project_name == "same-proj" + assert "same-proj" in logger._project_providers + assert logger._resolve_project_name(kwargs) == project_name + assert tracer_from_request is not None + + def test_cache_reuses_provider_for_same_project(self): + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + logger._project_providers.clear() + + logger._get_tracer_for("cached-proj") + provider_first = logger._project_providers["cached-proj"] + + logger._get_tracer_for("cached-proj") + provider_second = logger._project_providers["cached-proj"] + + assert provider_first is provider_second + assert len(logger._project_providers) == 1 + + def test_parallel_cache_miss_for_same_project_inserts_once(self): + import threading + + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=MagicMock()), + callback_name="arize_phoenix", + ) + logger._project_providers.clear() + + build_calls: list[str] = [] + real_build = logger._build_tracer_provider_for_project + + def tracking_build(project_name: str): + build_calls.append(project_name) + return real_build(project_name) + + barrier = threading.Barrier(10) + errors: list[Exception] = [] + + def worker() -> None: + try: + barrier.wait() + logger._get_tracer_for("race-proj") + except Exception as exc: + errors.append(exc) + + with patch.object( + logger, + "_build_tracer_provider_for_project", + side_effect=tracking_build, + ): + threads = [threading.Thread(target=worker) for _ in range(10)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert len(logger._project_providers) == 1 + assert "race-proj" in logger._project_providers + assert len(build_calls) >= 1 + + def test_injected_tracer_provider_bypasses_project_cache(self): + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=exporter), + callback_name="arize_phoenix", + tracer_provider=provider, + ) + + assert getattr(logger, "_use_injected_tracer_provider", False) is True + assert not hasattr(logger, "_project_providers") or not getattr( + logger, "_project_providers", None + ) + + tracer_a = logger._get_tracer_for("any-project") + tracer_b = logger.get_tracer_to_use_for_request( + {"standard_logging_object": {"metadata": {"phoenix_project_name": "x"}}} + ) + assert tracer_a is logger.tracer + assert tracer_b is logger.tracer + + def test_flush_tracer_providers_noop_for_injected_provider(self): + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=exporter), + callback_name="arize_phoenix", + tracer_provider=provider, + ) + logger.flush_tracer_providers() + exporter.shutdown() + + def test_standard_logging_metadata_wins_over_litellm_params(self): + kwargs = { + "standard_logging_object": { + "metadata": {"phoenix_project_name_override": "from-logging"}, + }, + "litellm_params": { + "metadata": {"phoenix_project_name_override": "from-params"}, + }, + } + assert ArizePhoenixLogger._resolve_project_name(kwargs) == "from-logging" + + +class TestPhoenixTraceHandling: + """_handle_success / _handle_failure span export behavior.""" + + def test_handle_failure_sets_error_status_on_request_span(self): + from datetime import datetime + + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + from opentelemetry.trace import StatusCode + + from litellm.integrations.opentelemetry import ( + LITELLM_REQUEST_SPAN_NAME, + OpenTelemetryConfig, + ) + + exporter = InMemorySpanExporter() + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=exporter), + callback_name="arize_phoenix", + ) + + start = datetime(2024, 1, 1, 12, 0, 0) + end = datetime(2024, 1, 1, 12, 0, 1) + + logger._handle_failure( + { + "standard_logging_object": { + "metadata": {"phoenix_project_name": "fail-proj"}, + }, + "exception": Exception("boom"), + }, + response_obj=None, + start_time=start, + end_time=end, + ) + + spans = exporter.get_finished_spans() + request_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME] + assert len(request_spans) == 1 + assert request_spans[0].status.status_code == StatusCode.ERROR + assert ( + request_spans[0].resource.attributes.get("openinference.project.name") + == "fail-proj" + ) + + def test_proxy_mode_parent_and_child_share_trace_id(self): + from datetime import datetime + + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + from litellm.integrations.opentelemetry import ( + LITELLM_REQUEST_SPAN_NAME, + OpenTelemetryConfig, + ) + + exporter = InMemorySpanExporter() + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=exporter), + callback_name="arize_phoenix", + ) + + start = datetime(2024, 1, 1, 12, 0, 0) + end = datetime(2024, 1, 1, 12, 0, 1) + + logger._handle_success( + { + "litellm_params": { + "proxy_server_request": { + "url": "/chat/completions", + "method": "POST", + "headers": {}, + }, + "metadata": { + "user_api_key_auth_metadata": { + "phoenix_project_name_override": "proxy-proj", + }, + }, + }, + }, + response_obj={}, + start_time=start, + end_time=end, + ) + + spans = exporter.get_finished_spans() + span_names = {s.name for s in spans} + assert "litellm_proxy_request" in span_names + assert LITELLM_REQUEST_SPAN_NAME in span_names + + trace_ids = {s.context.trace_id for s in spans} + assert len(trace_ids) == 1 + for span in spans: + assert ( + span.resource.attributes.get("openinference.project.name") + == "proxy-proj" + ) + + def test_override_routes_all_spans_to_one_project_in_single_request(self): + from datetime import datetime + + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + exporter = InMemorySpanExporter() + logger = ArizePhoenixLogger( + config=OpenTelemetryConfig(exporter=exporter), + callback_name="arize_phoenix", + ) + + start = datetime(2024, 1, 1, 12, 0, 0) + end = datetime(2024, 1, 1, 12, 0, 1) + + logger._handle_success( + { + "standard_logging_object": { + "metadata": { + "user_api_key_auth_metadata": { + "phoenix_project_name_override": "unified-proj", + }, + }, + }, + "litellm_params": { + "proxy_server_request": { + "url": "/v1/chat/completions", + "method": "POST", + "headers": {}, + }, + }, + }, + response_obj={"id": "resp-1"}, + start_time=start, + end_time=end, + ) + + for span in exporter.get_finished_spans(): + assert ( + span.resource.attributes.get("openinference.project.name") + == "unified-proj" + ) + assert span.resource.attributes.get("model_id") == "unified-proj" + + +class TestGetArizePhoenixConfigProjectName: + @patch.dict( + "os.environ", {"PHOENIX_PROJECT_NAME": "phoenix-config-proj"}, clear=True + ) + def test_project_name_from_phoenix_env(self): + config = ArizePhoenixLogger.get_arize_phoenix_config() + assert config.project_name == "phoenix-config-proj" + + @patch.dict("os.environ", {}, clear=True) + def test_project_name_defaults_when_env_unset(self): + config = ArizePhoenixLogger.get_arize_phoenix_config() + assert config.project_name == "default" + if __name__ == "__main__": unittest.main() diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py index be2084969a5..cb786d9c292 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py @@ -3,7 +3,7 @@ import time from unittest.mock import AsyncMock import pytest -from httpx import Response +from httpx import Request, Response from litellm.integrations.datadog.datadog_cost_management import ( DatadogCostManagementLogger, @@ -167,3 +167,230 @@ async def test_async_send_batch(clean_env): content = json.loads(call_args[1]["content"]) assert content[0]["ProviderName"] == "openai" assert content[0]["BilledCost"] == 0.01 + + +_PUT_REQUEST = Request("PUT", "https://api.test.datadoghq.com/api/v2/cost/custom_costs") + + +@pytest.mark.asyncio +async def test_async_send_batch_clears_queue_on_success(clean_env): + """Bug 1 regression: log_queue must be empty after a successful upload.""" + logger = DatadogCostManagementLogger() + logger.async_client = AsyncMock() + logger.async_client.put.return_value = Response( + 202, json={"status": "ok"}, request=_PUT_REQUEST + ) + logger.log_queue = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + ) + ] + await logger.async_send_batch() + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_async_send_batch_preserves_events_added_during_upload(clean_env): + """Events appended while the upload is in flight survive (land on the cleared queue).""" + logger = DatadogCostManagementLogger() + + later_event = StandardLoggingPayload( + custom_llm_provider="anthropic", + model="claude-3", + response_cost=0.02, + startTime=time.time(), + ) + + async def slow_put(*args, **kwargs): + logger.log_queue.append(later_event) + return Response(202, json={"status": "ok"}, request=_PUT_REQUEST) + + logger.async_client = AsyncMock() + logger.async_client.put.side_effect = slow_put + logger.log_queue = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + ) + ] + await logger.async_send_batch() + assert logger.log_queue == [later_event] + + +@pytest.mark.asyncio +async def test_async_send_batch_requeues_on_upload_failure(clean_env): + """Failed upload requeues the original batch (no data loss).""" + logger = DatadogCostManagementLogger() + logger.async_client = AsyncMock() + logger.async_client.put.side_effect = Exception("boom") + original = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + ) + logger.log_queue = [original] + await logger.async_send_batch() + assert logger.log_queue == [original] + + +@pytest.mark.asyncio +async def test_extract_tags_emits_canonical_focus_dimensions(clean_env): + """provider, model, model_id always emitted regardless of cost_tag_keys.""" + logger = DatadogCostManagementLogger() + log = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4o", + model_id="router-id-123", + response_cost=0.01, + startTime=time.time(), + ) + tags = logger._extract_tags(log) + assert tags["provider"] == "openai" + assert tags["model"] == "gpt-4o" + assert tags["model_id"] == "router-id-123" + + +@pytest.mark.asyncio +async def test_extract_tags_allowlist_filters_request_tags(clean_env): + """Only request_tags whose key is in cost_tag_keys reach the Tags dict.""" + logger = DatadogCostManagementLogger(cost_tag_keys=["capability", "tier"]) + log = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + request_tags=["capability:chat", "tier:gold", "secret:disallowed"], + ) + tags = logger._extract_tags(log) + assert tags["capability"] == "chat" + assert tags["tier"] == "gold" + assert "secret" not in tags + + +@pytest.mark.asyncio +async def test_extract_tags_allowlist_filters_metadata(clean_env): + """Only metadata keys in cost_tag_keys flow through; others (and dict/list values) are dropped.""" + logger = DatadogCostManagementLogger(cost_tag_keys=["capability", "owner"]) + log = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + metadata={ + "capability": "chat", + "owner": "team-x", + "secret_field": "sensitive", + "nested_obj": {"a": 1}, + }, + ) + tags = logger._extract_tags(log) + assert tags["capability"] == "chat" + assert tags["owner"] == "team-x" + assert "secret_field" not in tags + assert "nested_obj" not in tags + + +@pytest.mark.asyncio +async def test_extract_tags_empty_allowlist_default(clean_env): + """With no cost_tag_keys, request_tags and arbitrary metadata.* do NOT leak into Tags.""" + logger = DatadogCostManagementLogger() + log = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + request_tags=["capability:chat"], + metadata={"capability": "chat", "user_api_key_alias": "alice"}, + ) + tags = logger._extract_tags(log) + assert "capability" not in tags + # Backwards-compat keys still flow: + assert tags["user"] == "alice" + + +@pytest.mark.asyncio +async def test_extract_tags_nested_metadata_allowlisted(clean_env): + """spend_logs_metadata and requester_metadata get spread one level under the allowlist.""" + logger = DatadogCostManagementLogger(cost_tag_keys=["env", "platform"]) + log = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + metadata={ + "spend_logs_metadata": {"platform": "web", "ignored": "x"}, + "requester_metadata": {"env": "prod"}, + }, + ) + tags = logger._extract_tags(log) + assert tags["platform"] == "web" + # "env" is a reserved trusted dimension — requester_metadata.env must NOT + # overwrite the value sourced from get_datadog_env(). + assert tags["env"] != "prod" + assert "ignored" not in tags + + +@pytest.mark.asyncio +async def test_extract_tags_allowlist_cannot_override_reserved_dimensions(clean_env): + """ + Reserved tag keys (env, service, host, pod_name, provider, model, model_id, + team, user, model_group) must not be overwritten by user-controlled + request_tags or metadata, even when listed in cost_tag_keys. + """ + reserved = [ + "env", + "service", + "host", + "pod_name", + "provider", + "model", + "model_id", + "team", + "user", + "model_group", + ] + logger = DatadogCostManagementLogger(cost_tag_keys=reserved) + + metadata_attack = {k: f"attacker-meta-{k}" for k in reserved} + metadata_attack["user_api_key_alias"] = "trusted-user" + metadata_attack["user_api_key_team_alias"] = "trusted-team" + metadata_attack["model_group"] = "trusted-group" + metadata_attack["spend_logs_metadata"] = { + k: f"attacker-spend-{k}" for k in reserved + } + metadata_attack["requester_metadata"] = {k: f"attacker-req-{k}" for k in reserved} + + log = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + model_id="router-id-123", + response_cost=0.01, + startTime=time.time(), + request_tags=[f"{k}:attacker-rt-{k}" for k in reserved], + metadata=metadata_attack, + ) + + tags = logger._extract_tags(log) + + # Canonical FOCUS dims keep their trusted (top-level payload) values. + assert tags["provider"] == "openai" + assert tags["model"] == "gpt-4" + assert tags["model_id"] == "router-id-123" + + # Backwards-compat trusted dims keep their proxy-controlled metadata values. + assert tags["user"] == "trusted-user" + assert tags["team"] == "trusted-team" + assert tags["model_group"] == "trusted-group" + + # No reserved key carries an attacker-supplied prefix from any path. + for k in reserved: + assert not tags[k].startswith("attacker-"), ( + f"reserved key {k!r} was overwritten by user-controlled input: " + f"{tags[k]!r}" + ) diff --git a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py index 757c558c298..2a26b7fade8 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py @@ -104,6 +104,7 @@ async def test_add_metrics_from_log(clean_env): logger._add_metrics_from_log(log=payload, kwargs=kwargs, status_code="200") # Should have 3 series: total_latency, llm_api_latency, request_count + # (no overhead metric because payload has no hidden_params litellm_overhead_time_ms) assert len(logger.log_queue) == 3 metrics = {s["metric"]: s for s in logger.log_queue} @@ -125,6 +126,72 @@ async def test_add_metrics_from_log(clean_env): assert "status_code:200" in count["tags"] +@pytest.mark.asyncio +async def test_overhead_latency_metric_emitted(clean_env): + """Test that litellm.overhead.latency is emitted when hidden_params contains litellm_overhead_time_ms.""" + logger = DatadogMetricsLogger(batch_size=100, start_periodic_flush=False) + + now = datetime.now() + start_time = now - timedelta(seconds=2) + api_call_start_time = now - timedelta(seconds=1) + + payload = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4o", + hidden_params={ + "litellm_overhead_time_ms": 250.0, # 250 ms of overhead + }, + ) + + kwargs = { + "start_time": start_time, + "api_call_start_time": api_call_start_time, + "end_time": now, + } + + logger._add_metrics_from_log(log=payload, kwargs=kwargs, status_code="200") + + metrics = {s["metric"]: s for s in logger.log_queue} + + # Overhead metric must be present + assert ( + "litellm.overhead.latency" in metrics + ), f"Expected 'litellm.overhead.latency' in emitted metrics, got: {list(metrics.keys())}" + overhead = metrics["litellm.overhead.latency"] + assert overhead["type"] == 3 # gauge + # 250 ms → 0.25 s + assert abs(overhead["points"][0]["value"] - 0.25) < 1e-6 + # status_code should NOT be in overhead tags (it is a latency metric, not a request count) + assert not any(tag.startswith("status_code:") for tag in overhead["tags"]) + + +@pytest.mark.asyncio +async def test_overhead_latency_metric_absent_when_no_hidden_params(clean_env): + """Test that litellm.overhead.latency is NOT emitted when hidden_params has no overhead value.""" + logger = DatadogMetricsLogger(batch_size=100, start_periodic_flush=False) + + now = datetime.now() + start_time = now - timedelta(seconds=2) + api_call_start_time = now - timedelta(seconds=1) + + payload = StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4o", + # No hidden_params / no litellm_overhead_time_ms + ) + + kwargs = { + "start_time": start_time, + "api_call_start_time": api_call_start_time, + "end_time": now, + } + + logger._add_metrics_from_log(log=payload, kwargs=kwargs, status_code="200") + + metrics = {s["metric"]: s for s in logger.log_queue} + assert "litellm.overhead.latency" not in metrics + + @pytest.mark.asyncio async def test_async_log_success_event(clean_env): """Test that success events are added to the queue.""" diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py b/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py index b1a3b834c3d..34103449dad 100644 --- a/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py +++ b/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py @@ -3,6 +3,7 @@ async_management_endpoint_{success,failure}_hook integration points.""" import asyncio from datetime import datetime +from unittest.mock import MagicMock import pytest @@ -14,6 +15,7 @@ from litellm.proxy._types import ( from ._helpers import ( HttpStatusException, assert_server_span_attrs, + get_server_span, make_fastapi_http_exception, make_httpx_status_error, ) @@ -28,6 +30,10 @@ def _real_user_api_key_dict(parent_span): ) +async def _noop_alert(*args, **kwargs): + return None + + async def _drive_admin_failure(*, otel, exception, parent_span, route): payload = ManagementEndpointLoggingPayload( route=route, @@ -180,3 +186,173 @@ def test_admin_endpoint_failure_stamps_server_span( expected_url_path=path, where=f"{path} {expected_status}", ) + + +def test_management_wrapper_success_ends_server_span_without_http_request( + server_span_factory, otel_with_exporter, monkeypatch +): + """Regression: management endpoints whose handler does not declare an + ``http_request`` parameter (``/key/generate``, ``/user/new``, ``/mcp/*``, + ...) must still get their parent SERVER span stamped + ended on success. + + The success hook itself stamps 200 and ``end()``s the parent, but the + wrapper only invoked it when ``http_request`` was present — so on success + the span (created in auth) was never ended and never exported. This drives + the real wrapper around an ``http_request``-less handler and asserts the + SERVER span reaches the exporter with status 200. + """ + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def fake_generate_key_fn(data=None, user_api_key_dict=None): + # No ``http_request`` parameter — mirrors generate_key_fn et al. + return {"key": "sk-xyz", "key_name": "k"} + + asyncio.run( + fake_generate_key_fn( + data={}, + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) + + assert_server_span_attrs( + exporter, + expected_status=200, + expected_url_path=KEY_GENERATE_PATH, + where="management wrapper success without http_request", + ) + + +def test_management_wrapper_failure_ends_server_span( + server_span_factory, otel_with_exporter, monkeypatch +): + """When the handler raises, the wrapper must route through the failure hook + and stamp + end the parent SERVER span with the error status — even for an + ``http_request``-less handler (route falls back to ``func.__name__``).""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def failing_fn(data=None, user_api_key_dict=None): + raise HttpStatusException(500, "boom") + + with pytest.raises(HttpStatusException): + asyncio.run( + failing_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert_server_span_attrs( + exporter, + expected_status=500, + expected_url_path=KEY_GENERATE_PATH, + where="management wrapper failure", + ) + + +def test_management_wrapper_success_with_http_request( + server_span_factory, otel_with_exporter, monkeypatch +): + """Cover the branch where the handler DOES declare ``http_request``: the + route comes from ``http_request.url.path`` and the body is read from it.""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) + + async def _fake_body(request=None): + return {"team_alias": "t"} + + monkeypatch.setattr(mgmt_utils, "_read_request_body", _fake_body) + + server_span = server_span_factory("/team/new") + http_request = MagicMock() + http_request.url.path = "/team/new" + + @mgmt_utils.management_endpoint_wrapper + async def fake_new_team(data=None, http_request=None, user_api_key_dict=None): + return {"team_id": "t-1"} + + asyncio.run( + fake_new_team( + data={}, + http_request=http_request, + user_api_key_dict=_real_user_api_key_dict(server_span), + ) + ) + + assert_server_span_attrs( + exporter, + expected_status=200, + expected_url_path="/team/new", + where="management wrapper success with http_request", + ) + + +def test_management_wrapper_noop_when_otel_logger_absent( + server_span_factory, otel_with_exporter, monkeypatch +): + """When no OTEL logger is registered, the helper early-returns and no SERVER + span is exported — and the handler result is still returned unchanged.""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + _otel, exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", None, raising=False) + monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def fake_fn(data=None, user_api_key_dict=None): + return {"ok": True} + + result = asyncio.run( + fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert result == {"ok": True} + assert get_server_span(exporter) is None + + +def test_management_wrapper_swallows_post_success_errors( + server_span_factory, otel_with_exporter, monkeypatch +): + """A failure in post-success bookkeeping (cache invalidation, alerting) must + not propagate — the handler result is returned regardless (non-blocking).""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_helpers import utils as mgmt_utils + + otel, _exporter = otel_with_exporter + monkeypatch.setattr(proxy_server, "open_telemetry_logger", otel, raising=False) + monkeypatch.setattr(mgmt_utils, "send_management_endpoint_alert", _noop_alert) + + def _boom(*args, **kwargs): + raise RuntimeError("cache backend down") + + monkeypatch.setattr(mgmt_utils, "_delete_api_key_from_cache", _boom) + + server_span = server_span_factory(KEY_GENERATE_PATH) + + @mgmt_utils.management_endpoint_wrapper + async def fake_fn(data=None, user_api_key_dict=None): + return {"ok": True} + + result = asyncio.run( + fake_fn(data={}, user_api_key_dict=_real_user_api_key_dict(server_span)) + ) + + assert result == {"ok": True} diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/test_litellm/integrations/test_galileo.py new file mode 100644 index 00000000000..aab220f46be --- /dev/null +++ b/tests/test_litellm/integrations/test_galileo.py @@ -0,0 +1,397 @@ +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.integrations.galileo import GalileoObserve +from litellm.types.utils import ( + Choices, + EmbeddingResponse, + ImageObject, + ImageResponse, + Message, + ModelResponse, + TextCompletionResponse, +) + + +@pytest.fixture +def galileo_v2_env(monkeypatch): + monkeypatch.setenv("GALILEO_API_KEY", "test-api-key") + monkeypatch.setenv("GALILEO_PROJECT_ID", "86ff8ebe-a297-4134-b167-748bdd8d2c20") + monkeypatch.setenv("GALILEO_LOG_STREAM_ID", "76c4ea50-8aa3-4771-a0d7-8567b112210f") + monkeypatch.setenv("GALILEO_BASE_URL", "https://api.galileo.ai") + + +@pytest.mark.asyncio +async def test_galileo_v2_ingest_url_and_headers(galileo_v2_env): + logger = GalileoObserve() + logger.in_memory_records = [ + { + "latency_ms": 100, + "status_code": 200, + "input_text": "hi", + "output_text": "hello", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 1, + "num_output_tokens": 2, + "created_at": "2026-05-25T12:00:00", + } + ] + + url, payload = logger._get_ingest_request() + assert ( + url + == "https://api.galileo.ai/v2/projects/86ff8ebe-a297-4134-b167-748bdd8d2c20/spans" + ) + assert payload["log_stream_id"] == "76c4ea50-8aa3-4771-a0d7-8567b112210f" + assert payload["spans"][0]["type"] == "llm" + assert payload["spans"][0]["output"]["content"] == "hello" + + assert await logger._ensure_headers() is True + assert logger.headers["Galileo-API-Key"] == "test-api-key" + + +def test_galileo_v2_span_preserves_message_roles(galileo_v2_env): + record = { + "latency_ms": 1, + "status_code": 200, + "input_text": "fallback", + "output_text": "ok", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": "2026-05-25T12:00:00", + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hello"}, + ], + } + span = GalileoObserve._record_to_v2_span(record) + assert span["input"] == [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hello"}, + ] + + +def test_galileo_output_text_from_model_response(galileo_v2_env): + logger = GalileoObserve() + response = ModelResponse( + choices=[ + Choices( + message=Message( + content="assistant reply", + role="assistant", + annotations=[], + ) + ) + ] + ) + + output = logger.get_output_str_from_response(response, {"call_type": "acompletion"}) + assert output == "assistant reply" + + +@pytest.mark.asyncio +async def test_galileo_flush_swallows_http_errors(galileo_v2_env): + logger = GalileoObserve() + logger.in_memory_records = [ + { + "latency_ms": 1, + "status_code": 200, + "input_text": "a", + "output_text": "b", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": "2026-05-25T12:00:00", + } + ] + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = Exception("404 Not Found") + await logger.flush_in_memory_records() + + assert len(logger.in_memory_records) == 1 + + +@pytest.mark.asyncio +async def test_galileo_flush_clears_records_on_201(galileo_v2_env): + logger = GalileoObserve() + logger.in_memory_records = [ + { + "latency_ms": 1, + "status_code": 200, + "input_text": "a", + "output_text": "b", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 0, + "num_output_tokens": 0, + "created_at": "2026-05-25T12:00:00", + } + ] + + mock_response = AsyncMock() + mock_response.is_success = True + mock_response.status_code = 201 + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_response + await logger.flush_in_memory_records() + + assert logger.in_memory_records == [] + + +def test_galileo_normalize_base_url_none(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + assert logger.base_url is None + assert logger._normalize_base_url(None) is None + assert logger._normalize_base_url("https://x.example/") == "https://x.example" + + +def test_galileo_is_configured_branches(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + monkeypatch.delenv("GALILEO_USERNAME", raising=False) + monkeypatch.delenv("GALILEO_PASSWORD", raising=False) + + no_env = GalileoObserve() + assert no_env._is_configured() is False + + monkeypatch.setenv("GALILEO_API_KEY", "k") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + v2 = GalileoObserve() + assert v2._is_configured() is True + + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + legacy = GalileoObserve() + assert legacy._is_configured() is True + + monkeypatch.delenv("GALILEO_PASSWORD", raising=False) + no_pw = GalileoObserve() + assert no_pw._is_configured() is False + + +def test_galileo_input_messages_fallbacks(): + assert GalileoObserve._galileo_input_messages(None, "hi") == [ + {"role": "user", "content": "hi"} + ] + assert GalileoObserve._galileo_input_messages( + ["not-a-dict", {"content": "no role"}], "fallback" + ) == [{"role": "user", "content": "fallback"}] + + +def test_galileo_record_to_v2_span_with_tags_and_offset(): + span = GalileoObserve._record_to_v2_span( + { + "latency_ms": 5, + "status_code": 200, + "input_text": "in", + "output_text": "out", + "node_type": "acompletion", + "model": "gpt-5.2", + "num_input_tokens": 1, + "num_output_tokens": 2, + "created_at": "2026-05-25T12:00:00", + "tags": ["t1"], + } + ) + assert span["tags"] == ["t1"] + assert span["created_at"].endswith("Z") + + offset = GalileoObserve._record_to_v2_span( + {"created_at": "2026-05-25T12:00:00-05:00"} + ) + assert offset["created_at"] == "2026-05-25T12:00:00-05:00" + + +def test_galileo_get_output_str_variants(galileo_v2_env): + logger = GalileoObserve() + assert logger.get_output_str_from_response(None, {}) is None + assert ( + logger.get_output_str_from_response( + EmbeddingResponse(), {"call_type": "embedding"} + ) + is None + ) + + text_resp = TextCompletionResponse() + text_resp.choices = [MagicMock(text="text-completion-output")] + assert ( + logger.get_output_str_from_response(text_resp, {"call_type": "text_completion"}) + == "text-completion-output" + ) + + image_resp = ImageResponse(data=[ImageObject(url="https://x/y.png")]) + assert "y.png" in logger.get_output_str_from_response(image_resp, {}) + + assert logger.get_output_str_from_response("not-a-supported-type", {}) is None + + +def test_galileo_get_ingest_request_unconfigured(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + assert logger._get_ingest_request() is None + + +def test_galileo_get_ingest_request_legacy(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example/") + monkeypatch.setenv("GALILEO_PROJECT_ID", "proj") + logger = GalileoObserve() + logger.in_memory_records = [{"foo": "bar"}] + url, payload = logger._get_ingest_request() + assert url == "https://galileo.example/projects/proj/observe/ingest" + assert payload == {"records": [{"foo": "bar"}]} + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_v2_missing_key(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_BASE_URL", "https://x") + logger = GalileoObserve() + logger.use_v2_api = True + logger.api_key = None + assert await logger._ensure_headers() is False + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_cached(galileo_v2_env): + logger = GalileoObserve() + logger.headers = {"Galileo-API-Key": "already-set"} + assert await logger._ensure_headers() is True + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_legacy_login(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + logger = GalileoObserve() + + login_resp = MagicMock() + login_resp.raise_for_status = MagicMock() + login_resp.json = MagicMock(return_value={"access_token": "tok"}) + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = login_resp + assert await logger._ensure_headers() is True + + assert logger.headers["Authorization"] == "Bearer tok" + + +@pytest.mark.asyncio +async def test_galileo_ensure_headers_legacy_login_failure(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + logger = GalileoObserve() + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = Exception("boom") + assert await logger._ensure_headers() is False + + +@pytest.mark.asyncio +async def test_galileo_flush_noop_when_unconfigured(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + logger.in_memory_records = [{"foo": "bar"}] + await logger.flush_in_memory_records() + assert logger.in_memory_records == [{"foo": "bar"}] + + +@pytest.mark.asyncio +async def test_galileo_flush_resets_headers_on_401(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + logger = GalileoObserve() + logger.headers = {"Authorization": "Bearer stale"} + logger.in_memory_records = [{"records": "x"}] + + mock_response = MagicMock() + mock_response.is_success = False + mock_response.status_code = 401 + mock_response.text = "unauthorized" + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.return_value = mock_response + await logger.flush_in_memory_records() + + assert logger.headers is None + assert logger.in_memory_records == [{"records": "x"}] + + +@pytest.mark.asyncio +async def test_galileo_async_log_success_appends_and_flushes(galileo_v2_env): + import datetime + + logger = GalileoObserve() + response = ModelResponse( + choices=[ + Choices(message=Message(content="reply", role="assistant", annotations=[])) + ], + usage={"prompt_tokens": 1, "completion_tokens": 2}, + ) + + flushed_url: dict = {} + mock_response = MagicMock() + mock_response.is_success = True + mock_response.status_code = 200 + + async def fake_post(**kwargs): + flushed_url["url"] = kwargs.get("url") + return mock_response + + with patch.object(logger.async_httpx_handler, "post", side_effect=fake_post): + await logger.async_log_success_event( + kwargs={ + "call_type": "acompletion", + "model": "gpt", + "messages": [{"role": "user", "content": "hi"}], + }, + response_obj=response, + start_time=datetime.datetime(2026, 5, 25, 12, 0, 0), + end_time=datetime.datetime(2026, 5, 25, 12, 0, 1), + ) + + assert "/v2/projects/" in flushed_url["url"] + assert logger.in_memory_records == [] diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index a7a2b7720d7..2b47a232262 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1418,3 +1418,123 @@ def test_image_count_prevents_text_tokens_fallback(): f"got {prompt_cost}. text_tokens fallback may be double-charging." ) assert completion_cost == 0.0 + + +# --------------------------------------------------------------------------- +# Data-residency (OpenAI regional processing) tests +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _local_model_cost_map(): + prev_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + prev_model_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + yield + finally: + litellm.model_cost = prev_model_cost + if prev_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = prev_env + + +@pytest.mark.parametrize("data_residency", ["eu", "us"]) +def test_data_residency_applies_uplift(data_residency, _local_model_cost_map): + """gpt-5 should apply the regional processing uplift multiplier when + data_residency is set.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token( + model="gpt-5", + usage=usage, + custom_llm_provider="openai", + ) + regional = generic_cost_per_token( + model="gpt-5", + usage=usage, + custom_llm_provider="openai", + data_residency=data_residency, + ) + + base_total = base[0] + base[1] + regional_total = regional[0] + regional[1] + + assert base_total > 0 + assert regional_total == pytest.approx(base_total * 1.10, rel=1e-9) + assert regional[0] == pytest.approx(base[0] * 1.10, rel=1e-9) + assert regional[1] == pytest.approx(base[1] * 1.10, rel=1e-9) + + +def test_data_residency_no_uplift_for_unmarked_model(_local_model_cost_map): + """A model without a regional_processing_uplift_multiplier_* entry should + fall back to base pricing, not error.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token( + model="gpt-3.5-turbo", + usage=usage, + custom_llm_provider="openai", + ) + with_residency = generic_cost_per_token( + model="gpt-3.5-turbo", + usage=usage, + custom_llm_provider="openai", + data_residency="eu", + ) + + assert base == with_residency + + +def test_data_residency_none_no_uplift(_local_model_cost_map): + """data_residency=None should be a no-op even for models with a multiplier.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token( + model="gpt-5", + usage=usage, + custom_llm_provider="openai", + ) + explicit_none = generic_cost_per_token( + model="gpt-5", + usage=usage, + custom_llm_provider="openai", + data_residency=None, + ) + + assert base == explicit_none + + +def test_data_residency_composes_with_service_tier(_local_model_cost_map): + """The uplift multiplies the priority-tier cost, not the standard one.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + priority_base = generic_cost_per_token( + model="gpt-5", + usage=usage, + custom_llm_provider="openai", + service_tier="priority", + ) + priority_eu = generic_cost_per_token( + model="gpt-5", + usage=usage, + custom_llm_provider="openai", + service_tier="priority", + data_residency="eu", + ) + + priority_base_total = priority_base[0] + priority_base[1] + priority_eu_total = priority_eu[0] + priority_eu[1] + + assert priority_base_total > 0 + assert priority_eu_total == pytest.approx(priority_base_total * 1.10, rel=1e-9) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 6c8f7585abb..129ea237efe 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -10,10 +10,12 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, BedrockImageProcessor, _bedrock_converse_messages_pt, + _bedrock_tools_pt, _convert_to_bedrock_tool_call_invoke, _convert_to_bedrock_tool_call_result, anthropic_messages_pt, convert_to_gemini_tool_call_result, + make_valid_bedrock_tool_name, ollama_pt, sanitize_messages_for_tool_calling, ) @@ -2082,6 +2084,90 @@ def test_bedrock_tool_call_invoke_non_dict_arguments(): assert result[0]["toolUse"]["input"] == {} +def test_make_valid_bedrock_tool_name_preserves_hyphens(): + assert make_valid_bedrock_tool_name("my-tool") == "my-tool" + assert ( + make_valid_bedrock_tool_name( + "CreateCaseKnowledgeArticle_foTWsqR6yDt-OnSsvR5e6Q" + ) + == "CreateCaseKnowledgeArticle_foTWsqR6yDt-OnSsvR5e6Q" + ) + + +def test_bedrock_tool_name_sanitized_consistently_in_tools_and_tool_use(): + """toolSpec and toolUse names must match after sanitization (issue #5007).""" + raw_name = "foo@bar" + tools = [ + { + "type": "function", + "function": { + "name": raw_name, + "description": "test", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + tool_spec_name = _bedrock_tools_pt(tools)[0]["toolSpec"]["name"] + + tool_calls = [ + { + "id": "call_1", + "type": "function", + "function": {"name": raw_name, "arguments": "{}"}, + } + ] + tool_use_name = _convert_to_bedrock_tool_call_invoke(tool_calls)[0]["toolUse"][ + "name" + ] + + assert tool_spec_name == "foo_bar" + assert tool_use_name == tool_spec_name + + +def test_bedrock_converse_messages_pt_tool_use_matches_tool_spec_hyphen_name(): + """Hyphenated tool names are preserved and consistent in multi-turn history.""" + tool_name = "my-tool" + messages = [ + {"role": "user", "content": "call the tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_hyphen", + "type": "function", + "function": {"name": tool_name, "arguments": "{}"}, + } + ], + }, + ] + translated = _bedrock_converse_messages_pt( + messages=messages, model="", llm_provider="" + ) + tool_use_blocks = [ + block + for msg in translated + for block in msg.get("content", []) + if "toolUse" in block + ] + assert len(tool_use_blocks) == 1 + assert tool_use_blocks[0]["toolUse"]["name"] == tool_name + + tool_spec_name = _bedrock_tools_pt( + [ + { + "type": "function", + "function": { + "name": tool_name, + "description": "test", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + )[0]["toolSpec"]["name"] + assert tool_spec_name == tool_name + + def test_bedrock_tool_call_invoke_multiple_normal_tools(): """Multiple separate tool calls (normal parallel calling) work correctly.""" tool_calls = [ diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index dbcb048c250..55db31efd2c 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -125,3 +125,40 @@ class TestGetLitellmParamsExplicitFields: def test_no_log_from_explicit_param(self): result = get_litellm_params(no_log=True) assert result["no-log"] is True + + +class TestGetLitellmParamsDataResidency: + """Verify that data_residency is inferred from OpenAI regional api_base.""" + + def test_eu_host_resolves_to_eu(self): + result = get_litellm_params( + custom_llm_provider="openai", + api_base="https://eu.api.openai.com/v1", + ) + assert result["data_residency"] == "eu" + + def test_us_host_resolves_to_us(self): + result = get_litellm_params( + custom_llm_provider="openai", + api_base="https://us.api.openai.com/v1", + ) + assert result["data_residency"] == "us" + + def test_global_host_resolves_to_none(self): + result = get_litellm_params( + custom_llm_provider="openai", + api_base="https://api.openai.com/v1", + ) + assert result["data_residency"] is None + + def test_no_api_base_is_none(self): + result = get_litellm_params(custom_llm_provider="openai") + assert result["data_residency"] is None + + def test_non_openai_provider_does_not_resolve(self): + """Regional OpenAI host doesn't apply to other providers.""" + result = get_litellm_params( + custom_llm_provider="anthropic", + api_base="https://eu.api.openai.com/v1", + ) + assert result["data_residency"] is None diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 2b238b0cdf7..7913efe8294 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -133,7 +133,9 @@ def test_make_disable_auto_response_message_produces_ga_shape(): "turn_detection" not in session ), "turn_detection must not be at the top-level session (beta shape); use audio.input" # turn_detection must be nested under audio.input - assert session["audio"]["input"]["turn_detection"]["create_response"] is False + td = session["audio"]["input"]["turn_detection"] + assert td["type"] == "server_vad" + assert td["create_response"] is False def test_make_disable_auto_response_message_produces_beta_shape_for_beta_clients(): @@ -148,7 +150,55 @@ def test_make_disable_auto_response_message_produces_beta_shape_for_beta_clients assert msg["type"] == "session.update" session = msg["session"] - assert session == {"turn_detection": {"create_response": False}} + assert session == { + "turn_detection": {"type": "server_vad", "create_response": False} + } + + +@pytest.mark.asyncio +async def test_backend_to_client_send_text_receives_str_not_bytes(): + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "session.created", "session": {}}).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.called + sent = client_ws.send_text.call_args_list[0].args[0] + assert isinstance(sent, str) + + +@pytest.mark.asyncio +async def test_backend_to_client_skips_non_utf8_binary_frames(): + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + b"\xff\xfe", + json.dumps({"type": "session.created", "session": {}}).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.call_count == 1 + assert isinstance(client_ws.send_text.call_args_list[0].args[0], str) @pytest.mark.asyncio @@ -426,6 +476,50 @@ async def test_transcription_captured_in_backend_to_client(): assert logging_obj.model_call_details["messages"] == streaming.input_messages +@pytest.mark.asyncio +async def test_client_ack_caches_setup_to_prevent_duplicate_session_update_setup(): + websocket = MagicMock() + backend_ws = MagicMock() + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + + # Two session.update messages arrive before setupComplete round-trip. + websocket.receive_text = AsyncMock( + side_effect=[ + json.dumps({"type": "session.update", "session": {"tools": []}}), + json.dumps({"type": "session.update", "session": {"tools": []}}), + Exception("client done"), + ] + ) + + provider_config = MagicMock() + + def _transform(message: str, model: str, session_configuration_request=None): + if session_configuration_request is None: + return [json.dumps({"setup": {"model": "models/gemini-2.5-flash"}})] + return [] + + provider_config.transform_realtime_request = MagicMock(side_effect=_transform) + + backend_ws.send = AsyncMock() + + streaming = RealTimeStreaming( + websocket=websocket, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + + await streaming.client_ack_messages() + + # Setup should be forwarded exactly once even with repeated session.update. + assert backend_ws.send.await_count == 1 + assert streaming.session_configuration_request is not None + sent_payload = json.loads(backend_ws.send.await_args_list[0].args[0]) + assert "setup" in sent_payload + + def test_collect_session_tools_from_session_update(): """ Test that tools from session.update events are collected. @@ -829,6 +923,169 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(): litellm.callbacks = [] # cleanup +@pytest.mark.asyncio +async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(): + """ + Test that a client-supplied function_call_output whose content triggers a + guardrail is blocked: it is not forwarded to the backend, and an error + event is sent to the client. + """ + from fastapi import HTTPException + + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class BlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ): + texts = inputs.get("texts", []) + for text in texts: + if "@" in text: + raise HTTPException( + status_code=403, + detail={"error": "email address detected"}, + ) + return inputs + + guardrail = BlockingGuardrail( + guardrail_name="email-blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + litellm.callbacks = [guardrail] + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + item_create_msg = json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": "call_123", + "output": "Tool says: my email is test@example.com", + }, + } + ) + + client_ws.receive_text = AsyncMock( + side_effect=[ + item_create_msg, + Exception("connection closed"), + ] + ) + + await streaming.client_ack_messages() + + sent_texts = [json.loads(c.args[0]) for c in client_ws.send_text.call_args_list] + error_events = [e for e in sent_texts if e.get("type") == "error"] + assert len(error_events) == 1, f"Expected one error event, got: {sent_texts}" + assert error_events[0]["error"]["type"] == "guardrail_violation" + + sent_to_backend = [c.args[0] for c in backend_ws.send.call_args_list if c.args] + forwarded_tool_outputs = [ + json.loads(m) + for m in sent_to_backend + if isinstance(m, str) + and json.loads(m).get("type") == "conversation.item.create" + and json.loads(m).get("item", {}).get("type") == "function_call_output" + ] + # A sanitized placeholder must reach the backend so providers that pair + # every toolCall with a toolResponse (Gemini/Vertex Live) exit their + # pending-tool-call state instead of stalling. The placeholder must NOT + # contain any of the blocked content. + assert len(forwarded_tool_outputs) == 1, ( + f"Sanitized function_call_output should be forwarded, got: " + f"{forwarded_tool_outputs}" + ) + sanitized_item = forwarded_tool_outputs[0]["item"] + assert sanitized_item["call_id"] == "call_123" + assert "test@example.com" not in sanitized_item["output"] + + litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_realtime_function_call_output_guardrail_allows_clean_output(): + """ + Test that a clean function_call_output passes through and reaches the backend + when guardrails are configured. + """ + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class BlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ): + return inputs + + guardrail = BlockingGuardrail( + guardrail_name="noop", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + litellm.callbacks = [guardrail] + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + item_create_msg = json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": "call_456", + "output": '{"temperature": 72, "unit": "F"}', + }, + } + ) + + client_ws.receive_text = AsyncMock( + side_effect=[ + item_create_msg, + Exception("connection closed"), + ] + ) + + await streaming.client_ack_messages() + + sent_to_backend = [c.args[0] for c in backend_ws.send.call_args_list if c.args] + forwarded = [ + json.loads(m) + for m in sent_to_backend + if isinstance(m, str) + and json.loads(m).get("type") == "conversation.item.create" + and json.loads(m).get("item", {}).get("type") == "function_call_output" + ] + assert ( + len(forwarded) == 1 + ), f"Clean function_call_output should be forwarded, got: {forwarded}" + + litellm.callbacks = [] # cleanup + + @pytest.mark.asyncio async def test_realtime_text_input_guardrail_uses_pre_call_mode(): """ @@ -1110,3 +1367,406 @@ async def test_on_violation_end_session_closes_on_first_fail(): assert streaming._violation_count == 1 litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_provider_path_suppresses_duplicate_session_created_after_synthetic(): + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[b'{"setupComplete": {}}', ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + + provider_config = MagicMock() + provider_config.transform_realtime_response = MagicMock( + return_value={ + "response": [ + { + "type": "session.created", + "event_id": "event_1", + "session": {"id": "sess_1", "modalities": ["audio"]}, + } + ], + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": [], + "current_conversation_id": None, + "current_item_chunks": [], + "current_delta_type": None, + "session_configuration_request": None, + } + ) + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + # Simulate synthetic session.created already sent by llm_http_handler. + streaming._session_created_sent_to_client = True + + await streaming.backend_to_client_send_messages() + + sent_payloads = [json.loads(c.args[0]) for c in client_ws.send_text.call_args_list] + assert not any( + payload.get("type") == "session.created" for payload in sent_payloads + ), f"Expected duplicate session.created to be suppressed, got: {sent_payloads}" + + +@pytest.mark.asyncio +async def test_duplicate_session_created_still_triggers_guardrail_turn_detection_update(): + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[b'{"setupComplete": {}}', ConnectionClosed(None, None)] + ) + backend_ws.send = AsyncMock() + + provider_config = MagicMock() + provider_config.transform_realtime_response = MagicMock( + return_value={ + "response": [ + { + "type": "session.created", + "event_id": "event_1", + "session": {"id": "sess_1", "modalities": ["audio"]}, + } + ], + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": [], + "current_conversation_id": None, + "current_item_chunks": [], + "current_delta_type": None, + "session_configuration_request": None, + } + ) + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + # Synthetic session.created already sent by llm_http_handler. + streaming._session_created_sent_to_client = True + streaming._has_audio_transcription_guardrails = MagicMock(return_value=True) # type: ignore[method-assign] + streaming._send_to_backend = AsyncMock() # type: ignore[method-assign] + + await streaming.backend_to_client_send_messages() + + # Duplicate session.created should still cause the one-time guardrail + # turn_detection update to be sent to backend. + assert streaming._send_to_backend.await_count == 1 + sent_update = json.loads(streaming._send_to_backend.await_args_list[0].args[0]) + assert sent_update["type"] == "session.update" + injected_session = sent_update["session"] + assert injected_session["type"] == "realtime" + assert ( + injected_session["audio"]["input"]["turn_detection"]["create_response"] is False + ) + + +@pytest.mark.asyncio +async def test_guardrail_update_respects_idempotency_flag(): + """Verify guardrail turn-detection update uses idempotency flag correctly.""" + client_ws = AsyncMock() + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + provider_config = MagicMock() + provider_config.transform_realtime_request = MagicMock( + side_effect=lambda msg, model, session_config: [msg] + ) + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + streaming._has_audio_transcription_guardrails = MagicMock(return_value=True) # type: ignore[method-assign] + + # First call should send the update + assert streaming._guardrail_turn_detection_update_sent is False + await streaming._maybe_send_guardrail_turn_detection_update() + assert streaming._guardrail_turn_detection_update_sent is True + assert backend_ws.send.await_count == 1 + + # Second call should be a no-op (idempotent) + await streaming._maybe_send_guardrail_turn_detection_update() + assert backend_ws.send.await_count == 1 # Still 1, not 2 + + +@pytest.mark.asyncio +async def test_guardrail_turn_detection_injected_into_first_session_update_deferred_mode(): + """Verify turn_detection is injected into first session.update in deferred mode.""" + client_ws = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": { + "modalities": ["text", "audio"], + "tools": [{"type": "function", "name": "get_weather"}], + }, + } + ), + ConnectionClosed(None, None), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + provider_config = MagicMock() + transformed_messages = [] + + def mock_transform(msg, model, session_config): + transformed_messages.append((msg, session_config)) + return [msg] # Pass through for simplicity + + provider_config.transform_realtime_request = MagicMock(side_effect=mock_transform) + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + streaming._has_audio_transcription_guardrails = MagicMock(return_value=True) # type: ignore[method-assign] + + # Simulate first session.update in deferred mode + await streaming.client_ack_messages() + + # Verify turn_detection was injected into the session.update. The + # injection runs before the GA remap, so the create_response flag ends + # up nested under audio.input.turn_detection in the GA-shaped payload. + assert len(transformed_messages) == 1 + transformed_msg, session_config = transformed_messages[0] + msg_obj = json.loads(transformed_msg) + assert msg_obj["type"] == "session.update" + session_obj = msg_obj["session"] + injected_turn_detection = session_obj.get("turn_detection") or session_obj.get( + "audio", {} + ).get("input", {}).get("turn_detection") + assert injected_turn_detection is not None + assert injected_turn_detection["create_response"] is False + assert streaming._guardrail_turn_detection_update_sent is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("existing_turn_detection", [None, "auto", 42, ["server_vad"]]) +async def test_guardrail_turn_detection_injection_tolerates_non_dict_value( + existing_turn_detection, +): + """Client-supplied non-dict turn_detection must not crash client_ack_messages.""" + client_ws = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": { + "modalities": ["text", "audio"], + "turn_detection": existing_turn_detection, + }, + } + ), + ConnectionClosed(None, None), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + provider_config = MagicMock() + transformed_messages = [] + + def mock_transform(msg, model, session_config): + transformed_messages.append((msg, session_config)) + return [msg] + + provider_config.transform_realtime_request = MagicMock(side_effect=mock_transform) + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + streaming._has_audio_transcription_guardrails = MagicMock(return_value=True) # type: ignore[method-assign] + + await streaming.client_ack_messages() + + assert len(transformed_messages) == 1 + transformed_msg, _ = transformed_messages[0] + msg_obj = json.loads(transformed_msg) + session_obj = msg_obj["session"] + injected_turn_detection = session_obj.get("turn_detection") or session_obj.get( + "audio", {} + ).get("input", {}).get("turn_detection") + assert isinstance(injected_turn_detection, dict) + assert injected_turn_detection["create_response"] is False + assert streaming._guardrail_turn_detection_update_sent is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_session", + [ + {"turn_detection": {"type": "server_vad", "create_response": True}}, + { + "audio": { + "input": { + "turn_detection": {"type": "server_vad", "create_response": True} + } + } + }, + ], +) +async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_active( + client_session, +): + """A subsequent client session.update must not be allowed to flip + ``create_response`` back to True once audio transcription guardrails have + disabled VAD auto-response. Covers both the flat beta shape and the + nested GA ``audio.input.turn_detection`` shape. + """ + client_ws = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps({"type": "session.update", "session": client_session}), + ConnectionClosed(None, None), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + provider_config = MagicMock() + transformed_messages = [] + + def mock_transform(msg, model, session_config): + transformed_messages.append((msg, session_config)) + return [msg] + + provider_config.transform_realtime_request = MagicMock(side_effect=mock_transform) + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + streaming._has_audio_transcription_guardrails = MagicMock(return_value=True) # type: ignore[method-assign] + # Simulate that initial setup + guardrail disable have already happened. + streaming.session_configuration_request = json.dumps({"setup": {"model": "x"}}) + streaming._guardrail_turn_detection_update_sent = True + + await streaming.client_ack_messages() + + assert len(transformed_messages) == 1 + forwarded_msg, _ = transformed_messages[0] + msg_obj = json.loads(forwarded_msg) + session_obj = msg_obj["session"] + forwarded_turn_detection = session_obj.get("turn_detection") or session_obj.get( + "audio", {} + ).get("input", {}).get("turn_detection") + assert isinstance(forwarded_turn_detection, dict) + assert forwarded_turn_detection["create_response"] is False + + +@pytest.mark.asyncio +async def test_follow_up_setup_updates_cached_session_configuration_request(): + """A follow-up setup produced by a subsequent session.update must replace + the cached ``session_configuration_request`` so downstream readers + (e.g. modality lookup in ``response.created``) see the latest config.""" + client_ws = AsyncMock() + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps({"type": "session.update", "session": {"tools": []}}), + ConnectionClosed(None, None), + ] + ) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + provider_config = MagicMock() + follow_up_setup = json.dumps( + { + "setup": { + "model": "models/gemini-2.5-flash", + "generationConfig": {"responseModalities": ["TEXT"]}, + "tools": [{"function_declarations": []}], + } + } + ) + provider_config.transform_realtime_request = MagicMock( + return_value=[follow_up_setup] + ) + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + provider_config=provider_config, + model="gemini-2.5-flash", + ) + # Simulate that the original auto-setup was already cached. + streaming.session_configuration_request = json.dumps( + { + "setup": { + "model": "models/gemini-2.5-flash", + "generationConfig": {"responseModalities": ["AUDIO"]}, + } + } + ) + + await streaming.client_ack_messages() + + assert streaming.session_configuration_request == follow_up_setup diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_overhead.py b/tests/test_litellm/litellm_core_utils/test_streaming_overhead.py new file mode 100644 index 00000000000..8fb0659ab5a --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_overhead.py @@ -0,0 +1,508 @@ +""" +Tests for CustomStreamWrapper per-chunk behavior across Anthropic, +Bedrock Invoke, and Bedrock Converse: text passthrough, usage stripping, +hidden_params propagation, finish_reason, sync/async parity, and the +per-stream caches (_GCHUNK_FIELDS, _post_streaming_hooks). +""" + +import asyncio +import time +from typing import List, Optional +from unittest.mock import MagicMock, patch + +import litellm +from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper, + _GCHUNK_FIELDS, + generic_chunk_has_all_required_fields, +) +from litellm.types.utils import ( + Delta, + GenericStreamingChunk as GChunk, + ModelResponseStream, + StreamingChoices, + Usage, +) + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_logging_obj(provider: str = "anthropic") -> MagicMock: + logging_obj = MagicMock() + logging_obj.model_call_details = { + "custom_llm_provider": provider, + "litellm_params": {}, + } + logging_obj.call_type = "completion" + logging_obj.stream_options = None + logging_obj.messages = [{"role": "user", "content": "hi"}] + logging_obj.completion_start_time = None + logging_obj._llm_caching_handler = None + return logging_obj + + +def _make_generic_chunk( + text: str, + is_finished: bool = False, + finish_reason: str = "", + usage: Optional[dict] = None, +) -> GChunk: + return GChunk( + text=text, + is_finished=is_finished, + finish_reason=finish_reason, + usage=usage, + index=0, + tool_use=None, + ) + + +def _make_bedrock_converse_chunk( + text: str = "", + finish_reason: str = "", + usage: Optional[Usage] = None, +) -> ModelResponseStream: + """Simulate what AWSEventStreamDecoder.converse_chunk_parser returns.""" + return ModelResponseStream( + choices=[ + StreamingChoices( + finish_reason=finish_reason or None, + index=0, + delta=Delta(content=text, role="assistant"), + ) + ], + id="msg-test", + model="anthropic.claude-3-5-sonnet", + usage=usage, + ) + + +async def _async_iter(chunks: list): + """Wrap a list as a proper async iterator for use in __anext__ async branch.""" + for chunk in chunks: + yield chunk + + +def _make_wrapper( + chunks: list, + provider: str = "anthropic", + async_stream: bool = False, +) -> CustomStreamWrapper: + logging_obj = _make_logging_obj(provider) + stream = _async_iter(chunks) if async_stream else iter(chunks) + wrapper = CustomStreamWrapper( + completion_stream=stream, + model="claude-3-5-sonnet", + logging_obj=logging_obj, + custom_llm_provider=provider, + ) + return wrapper + + +def _drain_sync(wrapper: CustomStreamWrapper) -> List[ModelResponseStream]: + results = [] + for chunk in wrapper: + results.append(chunk) + return results + + +async def _drain_async(wrapper: CustomStreamWrapper) -> List[ModelResponseStream]: + results = [] + async for chunk in wrapper: + results.append(chunk) + return results + + +# --------------------------------------------------------------------------- +# 1. Module-level _GCHUNK_FIELDS constant +# --------------------------------------------------------------------------- + + +def test_gchunk_fields_is_frozenset(): + """_GCHUNK_FIELDS must be a frozenset built from GChunk.__annotations__.""" + assert isinstance(_GCHUNK_FIELDS, frozenset) + assert _GCHUNK_FIELDS == frozenset(GChunk.__annotations__) + + +def test_generic_chunk_has_all_required_fields_uses_module_constant(monkeypatch): + """generic_chunk_has_all_required_fields must use _GCHUNK_FIELDS, not __annotations__. + + The check semantics: every key in `chunk` must be a known GChunk field. + This identifies GChunk-shaped dicts (all keys are valid GChunk fields). + """ + valid_chunk = _make_generic_chunk("hello") + assert generic_chunk_has_all_required_fields(valid_chunk) is True + + # A dict with an extra unknown key should return False — the unknown key + # is not a GChunk field, so the chunk is not a pure GChunk. + extra_key_chunk = dict(valid_chunk) + extra_key_chunk["unknown_extra_key"] = "value" + assert generic_chunk_has_all_required_fields(extra_key_chunk) is False + + # A dict with only known GChunk fields but fewer keys still passes because + # all its keys are valid (subset of GChunk fields). + partial_chunk = {"text": "hi", "is_finished": False} + assert generic_chunk_has_all_required_fields(partial_chunk) is True + + +# --------------------------------------------------------------------------- +# 2. Cached model name and provider at init time +# --------------------------------------------------------------------------- + + +def test_cached_model_name_simple(): + """For non-openai providers the cached model name must match the model arg.""" + wrapper = _make_wrapper([], provider="anthropic") + assert wrapper._cached_model_name == "claude-3-5-sonnet" + assert wrapper._cached_logging_llm_provider == "anthropic" + + +def test_cached_model_name_openai_prefix(): + """For openai provider when logging provider differs, model name is prefixed.""" + logging_obj = _make_logging_obj(provider="azure") + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + assert wrapper._cached_model_name == "azure/gpt-4o" + assert wrapper._cached_logging_llm_provider == "azure" + + +def test_base_hidden_params_precomputed(): + """_base_hidden_params must be pre-built from _hidden_params at init.""" + wrapper = _make_wrapper([], provider="anthropic") + assert "response_cost" in wrapper._base_hidden_params + assert wrapper._base_hidden_params["response_cost"] is None + # Must include all keys from _hidden_params + for k in wrapper._hidden_params: + assert k in wrapper._base_hidden_params + + +# --------------------------------------------------------------------------- +# 3. Sync path: model_dump() is NOT called on non-usage chunks +# --------------------------------------------------------------------------- + + +def test_sync_path_no_model_dump_on_text_chunks(): + """ + The sync __next__ must NOT call model_dump() on chunks that have no usage. + + ModelResponseStream declares `usage` as a field, so a `hasattr` check + would always succeed and trigger the model_dump()+recreate path on every + chunk. The wrapper must check `is not None` instead. + """ + chunks = [ + _make_generic_chunk("Hello"), + _make_generic_chunk(" world"), + _make_generic_chunk("", is_finished=True, finish_reason="stop"), + ] + wrapper = _make_wrapper(chunks) + + model_dump_call_count = 0 + original_model_dump = ModelResponseStream.model_dump + + def counting_model_dump(self, **kwargs): + nonlocal model_dump_call_count + model_dump_call_count += 1 + return original_model_dump(self, **kwargs) + + with patch.object(ModelResponseStream, "model_dump", counting_model_dump): + results = _drain_sync(wrapper) + + text_chunks = [r for r in results if r.choices and r.choices[0].delta.content] + assert len(text_chunks) >= 2, "Expected at least 2 text chunks" + assert model_dump_call_count <= 1, ( + f"model_dump() called {model_dump_call_count} times — " + "usage check is firing on every chunk" + ) + + +# --------------------------------------------------------------------------- +# 4. Sync path: usage chunk is stripped from body but preserved in hidden_params +# --------------------------------------------------------------------------- + + +def test_sync_path_usage_stripped_from_body_preserved_in_hidden_params(): + """Usage data must be removed from the returned chunk but added to _hidden_params.""" + usage_dict = {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + chunks = [ + _make_generic_chunk("Hello"), + _make_generic_chunk( + "", is_finished=True, finish_reason="stop", usage=usage_dict + ), + ] + wrapper = _make_wrapper(chunks) + results = _drain_sync(wrapper) + + # The usage chunk must be returned (not silently dropped) + finish_chunks = [ + r for r in results if r.choices and r.choices[0].finish_reason == "stop" + ] + assert finish_chunks, "Finish-reason chunk was not returned" + + # The final chunk must carry usage in _hidden_params + final = results[-1] + assert "usage" in final._hidden_params, "usage missing from _hidden_params" + hidden_usage = final._hidden_params["usage"] + assert hidden_usage is not None + + +# --------------------------------------------------------------------------- +# 5. Async path: usage chunk is stripped from body but preserved in hidden_params +# --------------------------------------------------------------------------- + + +def test_async_path_usage_stripped_from_body_preserved_in_hidden_params(): + """Async path mirrors sync path for usage handling.""" + usage_dict = {"prompt_tokens": 5, "completion_tokens": 15, "total_tokens": 20} + chunks = [ + _make_generic_chunk("Hi"), + _make_generic_chunk( + "", is_finished=True, finish_reason="stop", usage=usage_dict + ), + ] + + async def _run(): + # async_stream=True forces the real async-for branch of __anext__ + wrapper = _make_wrapper(chunks, async_stream=True) + return await _drain_async(wrapper) + + results = asyncio.run(_run()) + final = results[-1] + assert "usage" in final._hidden_params + assert final._hidden_params["usage"] is not None + + +# --------------------------------------------------------------------------- +# 6. Bedrock Converse: ModelResponseStream chunks pass through correctly +# --------------------------------------------------------------------------- + + +def test_bedrock_converse_text_chunks_pass_through(): + """ + Bedrock Converse returns ModelResponseStream objects directly. + They should pass through chunk_creator and appear in output unchanged. + """ + chunks = [ + _make_bedrock_converse_chunk("Hello"), + _make_bedrock_converse_chunk(" world"), + _make_bedrock_converse_chunk("", finish_reason="end_turn"), + ] + wrapper = _make_wrapper(chunks, provider="bedrock") + results = _drain_sync(wrapper) + + texts = [ + r.choices[0].delta.content + for r in results + if r.choices and r.choices[0].delta.content + ] + assert "Hello" in texts or any("Hello" in (t or "") for t in texts) + + +def test_bedrock_converse_usage_chunk_stripped_and_in_hidden_params(): + """Usage in a Bedrock Converse ModelResponseStream chunk is handled correctly.""" + usage = Usage(prompt_tokens=8, completion_tokens=12, total_tokens=20) + chunks = [ + _make_bedrock_converse_chunk("Hi"), + _make_bedrock_converse_chunk("", finish_reason="end_turn", usage=usage), + ] + wrapper = _make_wrapper(chunks, provider="bedrock") + results = _drain_sync(wrapper) + + final = results[-1] + assert "usage" in final._hidden_params + assert final._hidden_params["usage"] is not None + + +# --------------------------------------------------------------------------- +# 7. Anthropic generic chunk (GChunk) path +# --------------------------------------------------------------------------- + + +def test_anthropic_generic_chunks_text_pass_through(): + """GChunk text chunks must arrive in the output with correct content.""" + chunks = [ + _make_generic_chunk("The"), + _make_generic_chunk(" answer"), + _make_generic_chunk("", is_finished=True, finish_reason="stop"), + ] + wrapper = _make_wrapper(chunks, provider="anthropic") + results = _drain_sync(wrapper) + + texts = [ + r.choices[0].delta.content + for r in results + if r.choices and r.choices[0].delta.content + ] + assert len(texts) >= 2 + + +def test_anthropic_finish_reason_propagated(): + """finish_reason must be set on the final streaming chunk.""" + chunks = [ + _make_generic_chunk("Hi"), + _make_generic_chunk("", is_finished=True, finish_reason="stop"), + ] + wrapper = _make_wrapper(chunks, provider="anthropic") + results = _drain_sync(wrapper) + + finish_reasons = [ + r.choices[0].finish_reason + for r in results + if r.choices and r.choices[0].finish_reason + ] + assert "stop" in finish_reasons + + +# --------------------------------------------------------------------------- +# 8. Callback caching: _post_streaming_hooks resolved once per stream +# --------------------------------------------------------------------------- + + +def test_post_streaming_hooks_cached_after_first_call(): + """ + _post_streaming_hooks must be None before the first hook call and a list after. + The same list object must be reused on subsequent calls (not re-built). + """ + wrapper = _make_wrapper([], provider="anthropic") + assert wrapper._post_streaming_hooks is None, "Must be None before first call" + + async def _run(): + # Simulate hook resolution with an empty callback list + with patch.object(litellm, "callbacks", []): + await wrapper._call_post_streaming_deployment_hook( + MagicMock(spec=ModelResponseStream) + ) + first_list = wrapper._post_streaming_hooks + assert isinstance(first_list, list) + + # Second call must reuse the same list object + with patch.object(litellm, "callbacks", []): + await wrapper._call_post_streaming_deployment_hook( + MagicMock(spec=ModelResponseStream) + ) + assert ( + wrapper._post_streaming_hooks is first_list + ), "_post_streaming_hooks was rebuilt on second call — caching broken" + + asyncio.run(_run()) + + +def test_post_streaming_hooks_filters_correctly(): + """ + Only CustomLogger instances must be included; plain callables are excluded. + + Note: CustomLogger's base class already defines + async_post_call_streaming_deployment_hook, so ALL CustomLogger subclasses + pass the hasattr() check regardless of whether they override the method. + The filter therefore keeps any CustomLogger instance and drops anything else. + """ + from litellm.integrations.custom_logger import CustomLogger + + class MyLogger(CustomLogger): + pass + + plain_callable = MagicMock() + + wrapper = _make_wrapper([], provider="anthropic") + + async def _run(): + with patch.object(litellm, "callbacks", [MyLogger(), plain_callable]): + await wrapper._call_post_streaming_deployment_hook( + MagicMock(spec=ModelResponseStream) + ) + + # plain_callable must be excluded; MyLogger (CustomLogger subclass) included + assert len(wrapper._post_streaming_hooks) == 1 + assert isinstance(wrapper._post_streaming_hooks[0], MyLogger) + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# 9. model_response_creator: hidden_params built correctly +# --------------------------------------------------------------------------- + + +def test_model_response_creator_hidden_params_no_chunk(): + """model_response_creator() with no args must include all _base_hidden_params.""" + wrapper = _make_wrapper([], provider="anthropic") + response = wrapper.model_response_creator() + + assert response._hidden_params.get("response_cost") is None + assert response._hidden_params.get("custom_llm_provider") == "anthropic" + assert "created_at" in response._hidden_params + + +def test_model_response_creator_hidden_params_caller_merged(): + """When hidden_params are passed by caller, they must be included in result.""" + wrapper = _make_wrapper([], provider="anthropic") + caller_params = {"some_key": "some_value"} + response = wrapper.model_response_creator(hidden_params=caller_params) + + assert response._hidden_params.get("some_key") == "some_value" + assert response._hidden_params.get("response_cost") is None + + +def test_model_response_creator_stream_key_stripped(): + """The 'stream' key must be removed from chunk before constructing ModelResponseStream.""" + wrapper = _make_wrapper([], provider="anthropic") + chunk = {"stream": True, "choices": []} + # Should not raise even if 'stream' would be an invalid ModelResponseStream field + response = wrapper.model_response_creator(chunk=chunk) + assert response is not None + + +# --------------------------------------------------------------------------- +# 10. Per-chunk overhead regression: sync path must not regress +# --------------------------------------------------------------------------- + + +def test_sync_streaming_overhead_not_regressed(): + """ + Micro-benchmark: the sync hot path must process 200 text chunks in < 2 s. + + This test acts as a canary for gross per-chunk overhead regressions. + It is intentionally generous (2 s) to avoid flakiness on slow CI runners. + """ + n_chunks = 200 + chunks = [_make_generic_chunk(f"token-{i}") for i in range(n_chunks)] + chunks.append(_make_generic_chunk("", is_finished=True, finish_reason="stop")) + + wrapper = _make_wrapper(chunks, provider="anthropic") + + start = time.monotonic() + results = _drain_sync(wrapper) + elapsed = time.monotonic() - start + + assert len(results) > 0, "No chunks returned" + assert elapsed < 2.0, ( + f"Sync streaming of {n_chunks} chunks took {elapsed:.3f}s — " + "per-chunk overhead regression detected" + ) + + +def test_async_streaming_overhead_not_regressed(): + """ + Micro-benchmark for the async path: 200 text chunks in < 2 s. + """ + n_chunks = 200 + chunks = [_make_generic_chunk(f"token-{i}") for i in range(n_chunks)] + chunks.append(_make_generic_chunk("", is_finished=True, finish_reason="stop")) + + async def _run(): + wrapper = _make_wrapper(chunks, provider="anthropic") + start = time.monotonic() + results = await _drain_async(wrapper) + return results, time.monotonic() - start + + results, elapsed = asyncio.run(_run()) + assert len(results) > 0 + assert elapsed < 2.0, ( + f"Async streaming of {n_chunks} chunks took {elapsed:.3f}s — " + "per-chunk overhead regression detected" + ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 7d9e4768303..687c5a2e733 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1622,6 +1622,29 @@ def test_effort_output_config_preservation(): assert result["output_config"]["effort"] == "medium" +def test_output_config_format_preservation_and_beta_header(): + """Test that output_config.format is preserved and treated as structured output.""" + config = AnthropicConfig() + output_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + } + optional_params = {"output_config": {"format": output_format, "effort": "xhigh"}} + + result = config.transform_request( + model="claude-opus-4-7", + messages=[{"role": "user", "content": "Test"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + headers = config.update_headers_with_optional_anthropic_beta({}, optional_params) + + assert result["output_config"]["format"] == output_format + assert result["output_config"]["effort"] == "xhigh" + assert "structured-outputs-2025-11-13" in headers["anthropic-beta"] + + def test_effort_beta_header_injection(): """Test that effort beta header is automatically added when output_config is detected.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -1648,7 +1671,7 @@ def test_effort_validation(): messages = [{"role": "user", "content": "Test"}] - # Valid values should work + # Valid values should work (xhigh is Opus 4.7+ only, not 4.5) for effort in ["high", "medium", "low"]: optional_params = {"output_config": {"effort": effort}} result = config.transform_request( @@ -2513,14 +2536,14 @@ def test_reasoning_effort_accepts_dict_shape_for_adaptive_model(reasoning_effort ) # thinking must be set (adaptive for 4.6+) - assert "thinking" in result, ( - f"thinking missing for reasoning_effort={reasoning_effort_value!r}" - ) + assert ( + "thinking" in result + ), f"thinking missing for reasoning_effort={reasoning_effort_value!r}" assert result["thinking"]["type"] == "adaptive" # output_config must carry the mapped effort - assert "output_config" in result, ( - f"output_config missing for reasoning_effort={reasoning_effort_value!r}" - ) + assert ( + "output_config" in result + ), f"output_config missing for reasoning_effort={reasoning_effort_value!r}" assert result["output_config"]["effort"] == "low" @@ -2532,7 +2555,9 @@ def test_reasoning_effort_accepts_dict_shape_for_adaptive_model(reasoning_effort {"effort": "low", "summary": "concise"}, ], ) -def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model(reasoning_effort_value): +def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( + reasoning_effort_value, +): """ Non-adaptive (pre-4.6) branch: dict-shape reasoning_effort must still map to ``thinking.type='enabled'`` + ``budget_tokens``. ``output_config`` must @@ -2547,9 +2572,9 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model(reasoning_ef drop_params=False, ) - assert "thinking" in result, ( - f"thinking missing for reasoning_effort={reasoning_effort_value!r}" - ) + assert ( + "thinking" in result + ), f"thinking missing for reasoning_effort={reasoning_effort_value!r}" assert result["thinking"]["type"] == "enabled" assert "budget_tokens" in result["thinking"] assert result["thinking"]["budget_tokens"] > 0 @@ -2582,12 +2607,12 @@ def test_reasoning_effort_unparseable_dict_is_dropped(bad_value): model="claude-sonnet-4-6-20260219", drop_params=False, ) - assert "thinking" not in result, ( - f"thinking should not be set for bad value {bad_value!r}" - ) - assert "output_config" not in result, ( - f"output_config should not be set for bad value {bad_value!r}" - ) + assert ( + "thinking" not in result + ), f"thinking should not be set for bad value {bad_value!r}" + assert ( + "output_config" not in result + ), f"output_config should not be set for bad value {bad_value!r}" @pytest.mark.parametrize( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py index 3c81bfaa0f9..e6d5c6f4ee1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py @@ -48,6 +48,39 @@ def test_output_format_supported_and_transforms_correctly(): assert "structured-outputs-2025-11-13" in headers["anthropic-beta"] +def test_output_config_format_supported_and_transforms_correctly(): + """Test that output_config.format is preserved and adds the structured-output beta.""" + config = AnthropicMessagesConfig() + + supported_params = config.get_supported_anthropic_messages_params("claude-opus-4-7") + assert "output_config" in supported_params + + output_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + } + optional_params = { + "max_tokens": 1024, + "output_config": {"format": output_format, "effort": "xhigh"}, + } + headers = {} + + result = config.transform_anthropic_messages_request( + model="claude-opus-4-7", + messages=[{"role": "user", "content": "test"}], + anthropic_messages_optional_request_params=optional_params.copy(), + litellm_params={}, + headers=headers, + ) + + headers = config._update_headers_with_anthropic_beta(headers, optional_params) + + assert result["output_config"]["format"] == output_format + assert result["output_config"]["effort"] == "xhigh" + assert "anthropic-beta" in headers + assert "structured-outputs-2025-11-13" in headers["anthropic-beta"] + + def test_output_format_works_with_bedrock_and_azure(): """Test that output_format works with Bedrock and Azure Foundry models.""" config = AnthropicMessagesConfig() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py index 83716b8c8d3..54bf0c4ac0f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -6,6 +6,9 @@ from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) @pytest.mark.parametrize( @@ -102,7 +105,6 @@ def test_invalid_reasoning_effort_raises_400(bad_effort): "model,bad_effort", [ ("claude-opus-4-6", "xhigh"), - ("bedrock/invoke/us.anthropic.claude-opus-4-6-v1", "xhigh"), ("claude-sonnet-4-6", "xhigh"), ], ) @@ -123,6 +125,56 @@ def test_reasoning_effort_unsupported_tier_raises_400_messages(model, bad_effort assert "not supported by this model" in str(exc_info.value) +@pytest.mark.parametrize( + "model,effort,expected_effort", + [ + ("invoke/us.anthropic.claude-opus-4-6-v1", "xhigh", "max"), + ("invoke/us.anthropic.claude-opus-4-6-v1", "max", "max"), + ("invoke/us.anthropic.claude-opus-4-6-v1", "high", "high"), + ("invoke/us.anthropic.claude-opus-4-7", "xhigh", "xhigh"), + ], +) +def test_bedrock_invoke_messages_clamps_effort_to_ceiling( + model, effort, expected_effort +): + """Bedrock Invoke /v1/messages degrades effort to the model's ceiling. + + Claude Code "goal mode" sends ``xhigh``; Opus 4.6 must clamp to ``max`` + instead of raising, while Opus 4.7 (ceiling ``xhigh``) keeps ``xhigh``. + """ + config = AmazonAnthropicClaudeMessagesConfig() + optional_params = {"max_tokens": 1024, "reasoning_effort": effort} + + result = config.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["output_config"]["effort"] == expected_effort + assert result["thinking"]["type"] == "adaptive" + + +def test_bedrock_invoke_messages_rejects_xhigh_without_ceiling(): + """Sonnet 4.6 on Bedrock has no effort ceiling, so xhigh is still rejected.""" + config = AmazonAnthropicClaudeMessagesConfig() + optional_params = {"max_tokens": 1024, "reasoning_effort": "xhigh"} + + with pytest.raises(AnthropicError) as exc_info: + config.transform_anthropic_messages_request( + model="invoke/us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert exc_info.value.status_code == 400 + assert "not supported by this model" in str(exc_info.value) + + @pytest.mark.parametrize( "model", [ diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py index d42d109f21b..08fef8c6a24 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py @@ -63,7 +63,13 @@ class TestGetModelInfoReasoningEffortFields: class TestModelRegistryReasoningEffortFields: """Verify specific models have the expected reasoning effort capability - values in the JSON registry file.""" + values in the JSON registry file. + + Claude models intentionally OMIT ``supports_minimal_reasoning_effort``: + ``minimal`` is not a real Anthropic effort level (the API accepts only + low/medium/high/xhigh/max), so LiteLLM degrades ``minimal`` to ``low`` + regardless of the flag. These tests guard against the flag being + re-added to the Claude fleet.""" @pytest.fixture(autouse=True) def _load_registry(self): @@ -77,41 +83,41 @@ class TestModelRegistryReasoningEffortFields: entry = self.registry["claude-opus-4-6"] assert entry.get("supports_max_reasoning_effort") is True - def test_opus_4_7_supports_minimal(self): + def test_opus_4_7_omits_minimal(self): entry = self.registry["claude-opus-4-7"] - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry - def test_opus_4_6_supports_minimal(self): + def test_opus_4_6_omits_minimal(self): entry = self.registry["claude-opus-4-6"] - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry - def test_sonnet_4_6_supports_minimal(self): + def test_sonnet_4_6_omits_minimal(self): entry = self.registry["anthropic.claude-sonnet-4-6"] - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry def test_bedrock_opus_4_7_supports_max(self): entry = self.registry["anthropic.claude-opus-4-7"] assert entry.get("supports_max_reasoning_effort") is True - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry def test_vertex_opus_4_7_supports_max(self): entry = self.registry["vertex_ai/claude-opus-4-7"] assert entry.get("supports_max_reasoning_effort") is True - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry def test_vertex_opus_4_6_supports_max(self): entry = self.registry["vertex_ai/claude-opus-4-6"] assert entry.get("supports_max_reasoning_effort") is True - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry - def test_azure_ai_opus_4_6_supports_minimal(self): + def test_azure_ai_opus_4_6_omits_minimal(self): entry = self.registry["azure_ai/claude-opus-4-6"] - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry def test_azure_ai_opus_4_7_supports_max(self): entry = self.registry["azure_ai/claude-opus-4-7"] assert entry.get("supports_max_reasoning_effort") is True - assert entry.get("supports_minimal_reasoning_effort") is True + assert "supports_minimal_reasoning_effort" not in entry # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 3fa794375e7..413241adf37 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -1646,6 +1646,336 @@ def test_azure_v1_api_uses_openai_client(api_version): ), f"base_url should contain /openai/v1/, got {async_client.base_url}" +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_with_azure_ad_token_provider(api_version): + """ + The v1 OpenAI client path must forward `azure_ad_token_provider` so Azure AD + auth works for `api_version` in {"v1", "latest", "preview"}. + + Regression: https://github.com/BerriAI/litellm/issues/27945 — before the fix + the v1 branch only forwarded `api_key`, so AD-only configs raised + "The api_key client option must be set" on every request. + + The OpenAI SDK accepts a callable for `api_key` and re-invokes it on every + request, so passing the provider directly preserves token refresh. + """ + from openai import AsyncOpenAI, OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + token_value = "mock-azure-ad-token-from-provider" + + def token_provider(): + return token_value + + init_return = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": token_provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + + client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + # The SDK stores callables as `_api_key_provider` and refreshes + # `self.api_key` before each request. + assert client._api_key_provider is token_provider + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + + async_client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + assert isinstance(async_client, AsyncOpenAI) + # Async client requires an async provider; we wrap the sync provider + # so the SDK can `await` it. + assert async_client._api_key_provider is not None + assert async_client._api_key_provider is not token_provider + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_async_token_provider_resolves_to_current_token(api_version): + """ + The async wrapper must call the underlying sync provider on each invocation + (not cache its first return value), so token rotation is honored. + """ + import asyncio + + from openai import AsyncOpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + tokens = iter(["token-1", "token-2", "token-3"]) + + def rotating_provider(): + return next(tokens) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": rotating_provider, + } + + async_client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + assert isinstance(async_client, AsyncOpenAI) + loop = asyncio.new_event_loop() + try: + first = loop.run_until_complete(async_client._api_key_provider()) + second = loop.run_until_complete(async_client._api_key_provider()) + finally: + loop.close() + + assert first == "token-1" + assert second == "token-2" + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_with_static_azure_ad_token(api_version): + """ + When only `azure_ad_token` (a static string) is set, the v1 client should + receive it as `api_key`. + """ + from openai import OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + token_value = "static-azure-ad-token" + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": token_value, + "azure_ad_token_provider": None, + } + + client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + assert client.api_key == token_value + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_key_wins_over_ad_token(api_version): + """ + Explicit `api_key` takes precedence over `azure_ad_token_provider` / + `azure_ad_token`, matching the priority documented in + `initialize_azure_sdk_client`. + """ + from openai import OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "explicit-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": "should-be-ignored", + "azure_ad_token_provider": lambda: "also-ignored", + } + + client = base_llm.get_azure_openai_client( + api_key="explicit-key", + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + assert client.api_key == "explicit-key" + assert client._api_key_provider is None + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_separates_distinct_ad_providers(api_version): + """ + Two configs sharing api_base/api_version but with different AD token + providers must not share a cached OpenAI client, otherwise requests for + one config would be sent with another config's AD credentials. + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def provider_a(): + return "token-a" + + def provider_b(): + return "token-b" + + def _init_for(provider): + return { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_for(provider_a) + client_a = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider_a}, + _is_async=True, + ) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_for(provider_b) + client_b = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider_b}, + _is_async=True, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert isinstance(client_b, AsyncOpenAI) + assert client_a is not client_b + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_separates_distinct_entra_credentials(api_version): + """ + Configs that synthesize an AD provider from tenant_id/client_id/client_secret + must not share a cached client when those inputs differ. + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def synth_provider(): + return "synthesized-token" + + def _init_synth(): + return { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": synth_provider, + } + + common = { + "api_key": None, + "api_base": api_base, + "api_version": api_version, + "_is_async": True, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_synth() + client_a = base_llm.get_azure_openai_client( + litellm_params={ + "tenant_id": "tenant-a", + "client_id": "client-a", + "client_secret": "secret-a", + }, + **common, + ) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_synth() + client_b = base_llm.get_azure_openai_client( + litellm_params={ + "tenant_id": "tenant-b", + "client_id": "client-b", + "client_secret": "secret-b", + }, + **common, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert isinstance(client_b, AsyncOpenAI) + assert client_a is not client_b + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_reuses_for_identical_ad_config(api_version): + """ + Identical AD configs should still share a cached client (regression guard + so the cache-key change doesn't accidentally disable caching). + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def provider(): + return "tok" + + init_return = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + client_a = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider}, + _is_async=True, + ) + client_b = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider}, + _is_async=True, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert client_a is client_b + + def test_azure_traditional_api_uses_azure_openai_client(): """ Test that traditional Azure API versions still use AzureOpenAI client. diff --git a/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py b/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py new file mode 100644 index 00000000000..3cecb7fa963 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py @@ -0,0 +1,134 @@ +""" +Tests for `litellm.llms.base_llm.managed_resources.utils.extract_model_id_from_unified_id`. + +The regex inside this helper is shared by both the vector-store unified-ID +format (`...;model_id,;...`) and the file-ID format (`...;llm_output_file_model_id,`). +A naive regex (`r"model_id,([^;]+)"`) substring-matches the latter and +returns the deployment UUID, which then gets fed as a model candidate +into the team-access check and 403s every team-BYOK file attach +(LIT-3244 patch/1.86.0 second-order finding). These tests pin the +field-boundary anchor that prevents that. +""" + +import pytest + +from litellm.llms.base_llm.managed_resources.utils import ( + encode_unified_id, + extract_model_id_from_unified_id, +) + +# --------------------------------------------------------------------------- +# Vector-store unified-ID shape — has a top-level `model_id,` field. +# Existing behavior must be preserved: returns the value. +# --------------------------------------------------------------------------- + + +def test_extract_model_id_returns_value_for_vector_store_unified_id(): + unified_id = ( + "litellm_proxy:vector_store" + ";unified_id,abc-123" + ";target_model_names,gpt-4,gemini" + ";resource_id,vs_xyz" + ";model_id,deployment-uuid-456" + ) + assert extract_model_id_from_unified_id(unified_id) == "deployment-uuid-456" + + +def test_extract_model_id_returns_value_when_field_is_first(): + """`model_id` is the very first field after the prefix (anchor must accept start-of-string).""" + unified_id = "litellm_proxy:vector_store;model_id,first-field-value;unified_id,abc" + # First field after the prefix is preceded by `;`, so it matches via the + # `;model_id,` branch. Pin that the anchor isn't accidentally too strict. + assert extract_model_id_from_unified_id(unified_id) == "first-field-value" + + +# --------------------------------------------------------------------------- +# File-ID shape — has `llm_output_file_model_id,` but no top-level +# `model_id,` field. Must return None (the previous regex would have +# substring-matched and returned the deployment UUID). +# --------------------------------------------------------------------------- + + +def test_extract_model_id_returns_none_for_file_id_without_model_id_field(): + """Regression pin for LIT-3244 patch/1.86.0. + + File-IDs constructed via `LITELLM_MANAGED_FILE_COMPLETE_STR` have + `llm_output_file_model_id,` but no top-level + `model_id,` field. The previous regex matched the substring and + returned the UUID, which then 403'd team-BYOK file attaches with + `Tried to access `. + """ + file_id = ( + "litellm_proxy:text/plain" + ";unified_id,file-uuid-123" + ";target_model_names,openai/gpt-4o" + ";llm_output_file_id,file-OpenAIReturnedId" + ";llm_output_file_model_id,813bf25f-e5a7-4658-8253-a6f677be8eb5" + ) + assert extract_model_id_from_unified_id(file_id) is None, ( + "File-ID has no top-level `model_id,` field — the deployment UUID " + "in `llm_output_file_model_id,` must NOT be returned. Returning it " + "feeds the UUID as a model candidate into the team-access check " + "and 403s every team-BYOK file attach (LIT-3244 patch/1.86.0)." + ) + + +def test_extract_model_id_returns_none_for_file_id_with_model_id_value_null(): + """The current file-ID builder writes `llm_output_file_model_id,None` + (the Python `None` stringified) when the upstream model_id isn't known. + Still no top-level `model_id,` field → must return None. + """ + file_id = ( + "litellm_proxy:text/plain" + ";unified_id,uuid" + ";target_model_names,openai/gpt-4o" + ";llm_output_file_id,file-Y" + ";llm_output_file_model_id,None" + ) + assert extract_model_id_from_unified_id(file_id) is None + + +# --------------------------------------------------------------------------- +# Base64-encoded inputs must decode and apply the same anchor. +# --------------------------------------------------------------------------- + + +def test_extract_model_id_decodes_base64_then_anchors(): + file_id_plain = ( + "litellm_proxy:text/plain" + ";unified_id,uuid" + ";target_model_names,openai/gpt-4o" + ";llm_output_file_id,file-Y" + ";llm_output_file_model_id,813bf25f-e5a7-4658-8253-a6f677be8eb5" + ) + encoded = encode_unified_id(file_id_plain) + assert extract_model_id_from_unified_id(encoded) is None + + vector_store_plain = ( + "litellm_proxy:vector_store" + ";unified_id,abc" + ";target_model_names,gpt-4" + ";resource_id,vs_xyz" + ";model_id,real-model-id" + ) + encoded_vs = encode_unified_id(vector_store_plain) + assert extract_model_id_from_unified_id(encoded_vs) == "real-model-id" + + +# --------------------------------------------------------------------------- +# Defensive: malformed / non-string inputs must not raise. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bad_input", [None, 42, b"bytes-not-str", []]) +def test_extract_model_id_returns_none_for_non_string_input(bad_input): + assert extract_model_id_from_unified_id(bad_input) is None # type: ignore[arg-type] + + +def test_extract_model_id_returns_none_when_field_absent(): + assert ( + extract_model_id_from_unified_id( + "litellm_proxy:other;unified_id,abc;some_field,whatever" + ) + is None + ) diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 64b43b15dcd..3287061d37e 100644 --- a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -76,7 +76,7 @@ class TestAgentCoreAcceptHeader: with patch.object(client, "post", return_value=MagicMock()) as mock_post: try: litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_runtime", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_runtime", messages=[{"role": "user", "content": "test"}], api_key="test-jwt-token", client=client, @@ -281,7 +281,7 @@ class TestAgentCoreStreamingJsonFallback: with patch.object(client, "post", return_value=mock_response): response = litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, @@ -318,7 +318,7 @@ class TestAgentCoreStreamingJsonFallback: client, "post", new_callable=AsyncMock, return_value=mock_response ): response = await litellm.acompletion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, @@ -353,7 +353,7 @@ class TestAgentCoreStreamingJsonFallback: Exception, match="Failed to read/parse JSON response body" ): litellm.completion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, @@ -383,7 +383,7 @@ class TestAgentCoreStreamingJsonFallback: Exception, match="Failed to read/parse JSON response body" ): await litellm.acompletion( - model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/test_agent", + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:941277531214:runtime/test_agent", messages=[{"role": "user", "content": "test"}], stream=True, client=client, 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 b2e254901f4..4c4c0e17a38 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 @@ -430,6 +430,61 @@ def test_output_config_forwarded_for_bedrock_chat_invoke_request(): assert result["max_tokens"] == 100 +def test_output_config_format_converted_for_bedrock_chat_invoke_request(): + """Bedrock Invoke chat path consumes ``output_config.format`` before forwarding.""" + config = AmazonAnthropicClaudeConfig() + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + + result = config.transform_request( + model="anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "test"}], + optional_params={ + "max_tokens": 100, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"effort": "xhigh"} + last_content = result["messages"][0]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +@pytest.mark.parametrize( + "model,expected_effort", + [ + ("anthropic.claude-opus-4-5-20251101-v1:0", "high"), + ("anthropic.claude-opus-4-6-v1", "max"), + ("anthropic.claude-opus-4-7", "xhigh"), + ], +) +def test_output_config_effort_normalized_for_bedrock_chat_invoke_request( + model, expected_effort +): + """Bedrock Invoke chat path accepts ``xhigh`` and forwards the provider-safe effort.""" + config = AmazonAnthropicClaudeConfig() + + result = config.transform_request( + model=model, + messages=[{"role": "user", "content": "test"}], + optional_params={ + "max_tokens": 100, + "output_config": {"effort": "xhigh"}, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"effort": expected_effort} + + def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider(): config = AmazonAnthropicClaudeConfig() messages = [{"role": "user", "content": "test"}] diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 5f2ed3dc00f..c8e72b7ac5b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -318,6 +318,7 @@ def test_reasoning_effort_none_omits_thinking_for_anthropic_converse(model): ("bedrock/converse/us.anthropic.claude-opus-4-7", "high", "high"), ("bedrock/converse/us.anthropic.claude-opus-4-7", "xhigh", "xhigh"), ("bedrock/converse/us.anthropic.claude-opus-4-7", "max", "max"), + ("bedrock/converse/us.anthropic.claude-opus-4-6-v1", "xhigh", "max"), ("bedrock/converse/us.anthropic.claude-opus-4-6-v1", "max", "max"), ("bedrock/converse/us.anthropic.claude-sonnet-4-6", "high", "high"), ("bedrock/converse/us.anthropic.claude-sonnet-4-6", "minimal", "low"), @@ -369,6 +370,132 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +def test_output_config_format_translated_to_native_output_config_converse(): + """``output_config.format`` becomes Bedrock ``outputConfig`` and is not forwarded raw.""" + config = AmazonConverseConfig() + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive"}, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + }, + litellm_params={}, + headers={}, + ) + + additional = result.get("additionalModelRequestFields", {}) + assert additional.get("output_config") == {"effort": "xhigh"} + assert "format" not in additional["output_config"] + assert result["outputConfig"]["textFormat"]["type"] == "json_schema" + parsed_schema = json.loads( + result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"] + ) + assert parsed_schema == {**schema, "additionalProperties": False} + + +def test_output_config_format_dropped_on_unsupported_converse_model_warns(caplog): + """When Converse model lacks native structured-output support, the silently + dropped ``output_config.format`` must surface as a warning so callers can + diagnose plain-text responses.""" + from unittest.mock import patch + + config = AmazonConverseConfig() + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + + with patch.object( + AmazonConverseConfig, + "_supports_native_structured_outputs", + return_value=False, + ): + with caplog.at_level("WARNING"): + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "output_config": { + "format": {"type": "json_schema", "schema": schema}, + }, + }, + litellm_params={}, + headers={}, + ) + + assert "outputConfig" not in result + assert any( + "dropping `output_config.format`" in record.getMessage() + for record in caplog.records + ) + + +def test_output_config_normalized_marker_does_not_leak_into_optional_params(): + """The internal ``_output_config_normalized`` marker set by + ``_handle_reasoning_effort_parameter`` must be consumed during request + preparation so it does not linger on the caller's ``optional_params``.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-opus-4-6-v1", + drop_params=False, + ) + assert optional_params.get("_output_config_normalized") is True + + config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-6-v1", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert "_output_config_normalized" not in optional_params + + +@pytest.mark.parametrize( + "model,expected_effort", + [ + ("bedrock/converse/us.anthropic.claude-opus-4-5-20251101-v1:0", "high"), + ("bedrock/converse/us.anthropic.claude-opus-4-6-v1", "max"), + ("bedrock/converse/us.anthropic.claude-opus-4-7", "xhigh"), + ], +) +def test_output_config_effort_normalized_for_bedrock_converse_opus( + model, expected_effort +): + """Bedrock Converse accepts ``xhigh`` and forwards the provider-safe effort.""" + config = AmazonConverseConfig() + + result = config._transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "xhigh"}, + }, + litellm_params={}, + headers={}, + ) + + additional = result.get("additionalModelRequestFields", {}) + assert additional.get("output_config") == {"effort": expected_effort} + + @pytest.mark.parametrize( "effort", ["disabled", "invalid", ""], diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 2e315a535f0..c92a9905229 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -767,6 +767,163 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): assert "output_format" not in result +def test_bedrock_messages_converts_output_config_format_to_inline_schema(): + """``output_config.format`` is consumed so Bedrock does not see an unknown nested key.""" + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"effort": "xhigh"} + assert "output_format" not in result + last_content = result["messages"][0]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +@pytest.mark.parametrize( + "model,expected_effort", + [ + ("anthropic.claude-opus-4-5-20251101-v1:0", "high"), + ("anthropic.claude-opus-4-6-v1", "max"), + ("anthropic.claude-opus-4-7", "xhigh"), + ], +) +def test_bedrock_messages_normalizes_output_config_effort_for_opus( + model, expected_effort +): + """Bedrock /v1/messages accepts ``xhigh`` and forwards the provider-safe effort.""" + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 4096, + "output_config": {"effort": "xhigh"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"effort": expected_effort} + + +def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema(): + """Inline-schema embedding must not mutate the caller's ``messages`` list, + message dicts, or content list.""" + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + caller_content = [{"type": "text", "text": "Hello"}] + caller_message = {"role": "user", "content": caller_content} + caller_messages = [caller_message] + schema = {"type": "object", "properties": {"answer": {"type": "string"}}} + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=caller_messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert caller_messages == [ + {"role": "user", "content": [{"type": "text", "text": "Hello"}]} + ] + assert caller_message == { + "role": "user", + "content": [{"type": "text", "text": "Hello"}], + } + assert caller_content == [{"type": "text", "text": "Hello"}] + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_bedrock_messages_does_not_mutate_callers_output_config(): + """`pop_bedrock_invoke_output_config_format` / effort normalization must not + leak into the caller's ``optional_params`` dict.""" + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + caller_output_config = { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + } + optional_params = { + "max_tokens": 4096, + "output_config": caller_output_config, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20251101-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert caller_output_config == { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + } + + def test_bedrock_messages_strips_output_config_with_output_format(): """ When both output_config and output_format are present, output_format @@ -1071,9 +1228,7 @@ def test_bedrock_messages_preserves_compact_context_management_and_adds_beta(): messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}] optional_params = { "max_tokens": 4096, - "context_management": { - "edits": [{"type": "compact_20260112"}] - }, + "context_management": {"edits": [{"type": "compact_20260112"}]}, } result = cfg.transform_anthropic_messages_request( @@ -1084,9 +1239,7 @@ def test_bedrock_messages_preserves_compact_context_management_and_adds_beta(): headers={}, ) - assert result.get("context_management") == { - "edits": [{"type": "compact_20260112"}] - } + assert result.get("context_management") == {"edits": [{"type": "compact_20260112"}]} assert "compact-2026-01-12" in result.get("anthropic_beta", []) assert result["max_tokens"] == 4096 @@ -1118,9 +1271,7 @@ def test_bedrock_messages_filters_unsupported_context_management_edits(): headers={}, ) - assert result.get("context_management") == { - "edits": [{"type": "compact_20260112"}] - } + assert result.get("context_management") == {"edits": [{"type": "compact_20260112"}]} assert "compact-2026-01-12" in result.get("anthropic_beta", []) 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 c39fb427a01..6298eeb25e9 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -1,9 +1,7 @@ -import json import os import sys import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -12,7 +10,6 @@ sys.path.insert( from litellm.llms.bedrock.common_utils import BedrockModelInfo - # --------------------------------------------------------------------------- # # get_bedrock_response_stream_shape lazy-load tests # # --------------------------------------------------------------------------- # @@ -24,8 +21,10 @@ def _reset_bedrock_response_stream_shape_cache(): import litellm.llms.bedrock.common_utils as mod mod.get_bedrock_response_stream_shape.cache_clear() + mod._get_local_model_cost_map.cache_clear() yield mod.get_bedrock_response_stream_shape.cache_clear() + mod._get_local_model_cost_map.cache_clear() def test_bedrock_response_stream_shape_lazy_loads_once(): @@ -222,3 +221,45 @@ def test_context_window_suffix_stripped_for_cost_lookup(): get_bedrock_base_model("anthropic.claude-3-5-sonnet-20241022-v2:0:51k") == "anthropic.claude-3-5-sonnet-20241022-v2:0" ) + + +def test_output_config_effort_normalization_uses_model_info_ceiling(monkeypatch): + import litellm.llms.bedrock.common_utils as mod + + calls = [] + + def fake_get_model_info(model, custom_llm_provider=None): + calls.append((model, custom_llm_provider)) + return {"bedrock_output_config_effort_ceiling": "max"} + + monkeypatch.setattr(mod, "_get_model_info", fake_get_model_info) + output_config = {"effort": "xhigh"} + + mod.normalize_bedrock_opus_output_config_effort( + model="custom-bedrock-alias-without-opus-pattern", + output_config=output_config, + ) + + assert output_config == {"effort": "max"} + assert calls == [("custom-bedrock-alias-without-opus-pattern", "bedrock")] + + +@pytest.mark.parametrize( + "model,expected_ceiling", + [ + ("anthropic.claude-opus-4-5-20251101-v1:0", "high"), + ("anthropic.claude-opus-4-6-v1", "max"), + ("anthropic.claude-opus-4-7", "xhigh"), + ("us.anthropic.claude-opus-4-5-20251101-v1:0", "high"), + ("us.anthropic.claude-opus-4-6-v1", "max"), + ("us.anthropic.claude-opus-4-7", "xhigh"), + ], +) +def test_bundled_bedrock_opus_model_info_declares_output_config_effort_ceiling( + model, expected_ceiling +): + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + + model_info = GetModelCostMap.load_local_model_cost_map()[model] + + assert model_info["bedrock_output_config_effort_ceiling"] == expected_ceiling diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index cc0a32d2ce6..cc0adc4277c 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -20,8 +20,10 @@ def test_gemini_realtime_transformation_session_created(): assert config is not None session_configuration_request = { - "model": "gemini-1.5-flash", - "generationConfig": {"responseModalities": ["TEXT"]}, + "setup": { + "model": "gemini-1.5-flash", + "generationConfig": {"responseModalities": ["TEXT"]}, + } } session_configuration_request_str = json.dumps(session_configuration_request) session_created_message = {"setupComplete": {}} @@ -45,8 +47,54 @@ def test_gemini_realtime_transformation_session_created(): }, ) - print(transformed_message) - assert transformed_message["response"][0]["type"] == "session.created" + session_created = transformed_message["response"][0] + assert session_created["type"] == "session.created" + # Verify the setup-wrapped configuration reaches the modality lookup so + # the synthetic session.created reflects the cached responseModalities. + assert session_created["session"]["modalities"] == ["text"] + + +def test_session_created_does_not_overwrite_session_configuration_request(): + config = GeminiRealtimeConfig() + + session_configuration_request_str = json.dumps( + { + "setup": { + "model": "models/gemini-2.5-flash-native-audio", + "generationConfig": {"responseModalities": ["AUDIO"]}, + } + } + ) + setup_complete_message = {"setupComplete": {}} + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + transformed = config.transform_realtime_response( + json.dumps(setup_complete_message), + "gemini-2.5-flash-native-audio", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": session_configuration_request_str, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + # Must keep original setup payload (with "setup"), not overwrite with session.created event. + assert ( + transformed["session_configuration_request"] + == session_configuration_request_str + ) + + # Also verify emitted session.created reflects audio modality from setup payload. + session_created = transformed["response"][0] + assert session_created["type"] == "session.created" + assert "audio" in session_created["session"]["modalities"] def test_gemini_realtime_transformation_content_delta(): @@ -54,8 +102,10 @@ def test_gemini_realtime_transformation_content_delta(): assert config is not None session_configuration_request = { - "model": "gemini-1.5-flash", - "generationConfig": {"responseModalities": ["TEXT"]}, + "setup": { + "model": "gemini-1.5-flash", + "generationConfig": {"responseModalities": ["TEXT"]}, + } } session_configuration_request_str = json.dumps(session_configuration_request) session_created_message = { @@ -147,8 +197,10 @@ def test_gemini_realtime_transformation_audio_delta(): assert config is not None session_configuration_request = { - "model": "gemini-1.5-flash", - "generationConfig": {"responseModalities": ["AUDIO"]}, + "setup": { + "model": "gemini-1.5-flash", + "generationConfig": {"responseModalities": ["AUDIO"]}, + } } session_configuration_request_str = json.dumps(session_configuration_request) @@ -196,8 +248,10 @@ def test_gemini_realtime_transformation_generation_complete(): assert config is not None session_configuration_request = { - "model": "gemini-1.5-flash", - "generationConfig": {"responseModalities": ["AUDIO"]}, + "setup": { + "model": "gemini-1.5-flash", + "generationConfig": {"responseModalities": ["AUDIO"]}, + } } session_configuration_request_str = json.dumps(session_configuration_request) @@ -225,9 +279,9 @@ def test_gemini_realtime_transformation_generation_complete(): contains_audio_done_event = False for response in responses: if response["type"] == OpenAIRealtimeEventTypes.RESPONSE_AUDIO_DONE.value: - contains_audio_delta = True + contains_audio_done_event = True break - assert contains_audio_delta, "Expected audio delta event" + assert contains_audio_done_event, "Expected audio done event" def test_gemini_3_1_flash_live_preview_model_cost_map_entry(): @@ -242,3 +296,1211 @@ def test_gemini_3_1_flash_live_preview_model_cost_map_entry(): assert info.get("max_output_tokens") == 65536 assert "video" in info.get("supported_modalities", []) assert info.get("supports_function_calling") is True + + +def test_gemini_realtime_tool_call_transformation(): + """Test transformation of Gemini toolCall to OpenAI function_call_arguments.done format.""" + config = GeminiRealtimeConfig() + + # Gemini toolCall message format + gemini_tool_call = { + "toolCall": { + "functionCalls": [ + { + "id": "call_123", + "name": "get_weather", + "args": {"location": "San Francisco", "unit": "fahrenheit"}, + } + ] + } + } + + gemini_tool_call_str = json.dumps(gemini_tool_call) + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "test-trace-123" + + # Transform the toolCall message + result = config.transform_realtime_response( + gemini_tool_call_str, + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": "item_123", + "current_response_id": "resp_123", + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + print("Tool call transformation result:", json.dumps(result, indent=2)) + + # Verify the transformation + responses = result["response"] + assert len(responses) > 0, "Expected at least one response event" + + # Find the function_call_arguments.done event + function_call_event = None + for event in responses: + if event.get("type") == "response.function_call_arguments.done": + function_call_event = event + break + + assert ( + function_call_event is not None + ), "Expected function_call_arguments.done event" + assert function_call_event["call_id"] == "call_123" + assert function_call_event["name"] == "get_weather" + assert function_call_event["response_id"] == "resp_123" + assert function_call_event["item_id"] == "item_123_tool_0" + assert function_call_event["output_index"] == 0 + + # Verify arguments are properly serialized as JSON string + args = json.loads(function_call_event["arguments"]) + assert args["location"] == "San Francisco" + assert args["unit"] == "fahrenheit" + + +def test_gemini_realtime_session_update_with_tools(): + """Test transformation of OpenAI session.update with tools to Gemini setup format.""" + config = GeminiRealtimeConfig() + + # OpenAI format session update with tools + session_update = { + "type": "session.update", + "session": { + "instructions": "You are a helpful assistant with weather tools.", + "temperature": 0.7, + "max_response_output_tokens": 1024, + "modalities": ["audio"], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city name", + }, + "unit": { + "type": "string", + "enum": ["fahrenheit", "celsius"], + }, + }, + "required": ["location"], + }, + }, + } + ], + }, + } + + # Transform to Gemini format (first session.update, so setup should be sent) + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash", + session_configuration_request=None, + ) + + assert len(messages) == 1, "Expected one setup message" + + gemini_setup = json.loads(messages[0]) + assert "setup" in gemini_setup + + setup_config = gemini_setup["setup"] + + # Verify tools are at top level, not in generationConfig + assert "tools" in setup_config + assert "tools" not in setup_config.get("generationConfig", {}) + + # Verify tool structure matches Gemini format + tools = setup_config["tools"] + assert len(tools) == 1 + assert "function_declarations" in tools[0] + + function_decl = tools[0]["function_declarations"][0] + assert function_decl["name"] == "get_weather" + assert "Get the current weather" in function_decl["description"] + assert "parameters" in function_decl + + +def test_gemini_session_update_defaults_to_audio_modality(): + config = GeminiRealtimeConfig() + + session_update = { + "type": "session.update", + "session": { + "instructions": "You are a helpful assistant.", + # No modalities on purpose + }, + } + + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash", + session_configuration_request=None, + ) + + assert len(messages) == 1 + setup_payload = json.loads(messages[0])["setup"] + assert setup_payload["generationConfig"]["responseModalities"] == ["AUDIO"] + + +def test_gemini_requires_session_configuration_feature_flag(monkeypatch): + config = GeminiRealtimeConfig() + + # Default behavior remains backwards-compatible (auto setup on connect) + monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False) + assert config.requires_session_configuration() is True + + # Opt-in behavior: defer setup until client sends session.update + monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False) + assert config.requires_session_configuration() is False + + +def test_gemini_realtime_function_call_output_transformation(): + """Test transformation of OpenAI function_call_output to Gemini toolResponse format. + + Exercises the full production round-trip: a Gemini toolCall arrives first + and populates the call_id -> name mapping, then the OpenAI + function_call_output is transformed and must carry the function name back + to Gemini in functionResponses. + """ + config = GeminiRealtimeConfig() + + # Receive a toolCall from Gemini first to populate the call_id -> name mapping. + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_func_output" + config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": "call_123", + "name": "get_weather", + "args": {"location": "San Francisco"}, + } + ] + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + assert config._tool_call_id_to_name.get("call_123") == "get_weather" + + # OpenAI format function call output + function_output = { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": "call_123", + "output": json.dumps( + { + "location": "San Francisco", + "temperature": 72, + "unit": "fahrenheit", + "conditions": "sunny", + } + ), + }, + } + + # Transform to Gemini format + messages = config.transform_realtime_request( + json.dumps(function_output), + "gemini-2.5-flash", + session_configuration_request="existing", + ) + + assert len(messages) == 1, "Expected one toolResponse message" + + gemini_response = json.loads(messages[0]) + assert "toolResponse" in gemini_response + + tool_response = gemini_response["toolResponse"] + assert "functionResponses" in tool_response + assert len(tool_response["functionResponses"]) == 1 + + func_response = tool_response["functionResponses"][0] + assert func_response["id"] == "call_123" + assert func_response["name"] == "get_weather" + assert "response" in func_response + assert func_response["response"]["temperature"] == 72 + assert func_response["response"]["conditions"] == "sunny" + + # A retry of the same function_call_output (e.g. a client SDK that + # re-sends the result) must still produce a functionResponses payload + # carrying ``name`` — the call_id → name mapping must not be evicted + # after the first lookup. + retry_messages = config.transform_realtime_request( + json.dumps(function_output), + "gemini-2.5-flash", + session_configuration_request="existing", + ) + retry_response = json.loads(retry_messages[0])["toolResponse"]["functionResponses"][ + 0 + ] + assert retry_response["name"] == "get_weather" + + +def test_gemini_realtime_user_text_transformation(): + """Test transformation of OpenAI user message to Gemini clientContent format.""" + config = GeminiRealtimeConfig() + + # OpenAI format user message + user_message = { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "What's the weather in London?"} + ], + }, + } + + # Transform to Gemini format + messages = config.transform_realtime_request( + json.dumps(user_message), + "gemini-2.5-flash", + session_configuration_request="existing", + ) + + assert len(messages) == 1, "Expected one clientContent message" + + gemini_message = json.loads(messages[0]) + assert "clientContent" in gemini_message + + client_content = gemini_message["clientContent"] + assert "turns" in client_content + assert len(client_content["turns"]) == 1 + + turn = client_content["turns"][0] + assert turn["role"] == "user" + assert len(turn["parts"]) == 1 + assert turn["parts"][0]["text"] == "What's the weather in London?" + assert client_content["turnComplete"] is True + + +def test_return_new_content_delta_events_without_session_config_does_not_error(): + config = GeminiRealtimeConfig() + + events = config.return_new_content_delta_events( + response_id="resp_1", + output_item_id="item_1", + conversation_id="conv_1", + delta_type="text", + session_configuration_request=None, + ) + + assert len(events) >= 1 + assert events[0]["type"] == "response.created" + + +def test_gemini_realtime_multi_tool_calls_have_unique_item_ids(): + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "test-trace-123" + + gemini_tool_call = { + "toolCall": { + "functionCalls": [ + { + "id": "call_1", + "name": "get_weather", + "args": {"location": "SF"}, + }, + { + "id": "call_2", + "name": "get_weather", + "args": {"location": "NYC"}, + }, + ] + } + } + + result = config.transform_realtime_response( + json.dumps(gemini_tool_call), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": "item_123", + "current_response_id": "resp_123", + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + responses = [ + ev + for ev in result["response"] + if ev.get("type") == "response.function_call_arguments.done" + ] + assert len(responses) == 2 + assert responses[0]["response_id"] == "resp_123" + assert responses[1]["response_id"] == "resp_123" + assert responses[0]["item_id"] == "item_123_tool_0" + assert responses[1]["item_id"] == "item_123_tool_1" + assert responses[0]["item_id"] != responses[1]["item_id"] + assert responses[0]["output_index"] == 0 + assert responses[1]["output_index"] == 1 + + +def test_gemini_session_update_includes_input_audio_transcription_default(): + """Verify _handle_session_update includes inputAudioTranscription default.""" + config = GeminiRealtimeConfig() + session_update = { + "type": "session.update", + "session": { + "modalities": ["text", "audio"], + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + }, + } + ], + }, + } + + result = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash", + session_configuration_request=None, + ) + + assert len(result) == 1 + setup = json.loads(result[0]) + assert "setup" in setup + assert "inputAudioTranscription" in setup["setup"] + assert setup["setup"]["inputAudioTranscription"] == {} + + +def test_gemini_tool_call_emits_response_created_preamble(): + """Verify response.created is emitted before tool call events when response_id is None.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + gemini_tool_call = { + "toolCall": { + "functionCalls": [ + { + "id": "call_123", + "name": "get_weather", + "args": {"location": "San Francisco", "unit": "fahrenheit"}, + } + ] + } + } + + # Transform with current_response_id=None to trigger preamble emission + result = config.transform_realtime_response( + json.dumps(gemini_tool_call), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + responses = result["response"] + # Should have: response.created, output_item.added, function_call_arguments.delta, function_call_arguments.done, output_item.done, conversation.item.created, response.done + assert len(responses) >= 7 + assert responses[0]["type"] == "response.created" + assert "response" in responses[0] + assert responses[0]["response"]["status"] == "in_progress" + # response.created on the tool-call path mirrors the audio/text preamble: + # modalities/temperature/max_output_tokens are present so spec-compliant + # clients see consistent response metadata regardless of payload type. + assert "modalities" in responses[0]["response"] + assert "temperature" in responses[0]["response"] + assert "max_output_tokens" in responses[0]["response"] + assert responses[1]["type"] == "response.output_item.added" + assert responses[1]["item"]["type"] == "function_call" + assert responses[1]["item"]["status"] == "in_progress" + assert responses[2]["type"] == "response.function_call_arguments.delta" + assert responses[2]["call_id"] == "call_123" + assert responses[2]["delta"] == responses[3]["arguments"] + assert responses[3]["type"] == "response.function_call_arguments.done" + assert responses[4]["type"] == "response.output_item.done" + assert responses[4]["item"]["type"] == "function_call" + assert responses[4]["item"]["status"] == "completed" + assert responses[5]["type"] == "conversation.item.created" + assert responses[5]["item"]["type"] == "function_call" + assert responses[5]["item"]["status"] == "completed" + assert responses[6]["type"] == "response.done" + assert responses[6]["response"]["status"] == "completed" + assert len(responses[6]["response"]["output"]) == 1 + assert responses[6]["response"]["output"][0]["type"] == "function_call" + assert result["current_output_item_id"] is None + assert result["current_response_id"] is None + + +def test_gemini_tool_call_resets_ids_for_post_tool_model_turn(): + """After tool-call response.done, a subsequent modelTurn must emit response.created.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + session_configuration_request = json.dumps( + { + "setup": { + "model": "gemini-1.5-flash", + "generationConfig": {"responseModalities": ["TEXT"]}, + } + } + ) + + tool_result = config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": "call_123", + "name": "get_weather", + "args": {"location": "San Francisco"}, + } + ] + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": session_configuration_request, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + tool_response_id = tool_result["response"][0]["response"]["id"] + assert tool_result["current_output_item_id"] is None + assert tool_result["current_response_id"] is None + + post_tool_result = config.transform_realtime_response( + json.dumps( + { + "serverContent": { + "modelTurn": {"parts": [{"text": "The weather is sunny."}]} + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": session_configuration_request, + "current_output_item_id": tool_result["current_output_item_id"], + "current_response_id": tool_result["current_response_id"], + "current_conversation_id": tool_result["current_conversation_id"], + "current_delta_chunks": tool_result["current_delta_chunks"], + "current_item_chunks": tool_result["current_item_chunks"], + "current_delta_type": tool_result["current_delta_type"], + }, + ) + + post_tool_events = post_tool_result["response"] + assert post_tool_events[0]["type"] == "response.created" + assert post_tool_events[0]["response"]["id"] != tool_response_id + assert ( + post_tool_result["current_response_id"] == post_tool_events[0]["response"]["id"] + ) + + +def test_gemini_empty_tool_call_does_not_crash_websocket(): + """A toolCall payload with no functionCalls must not raise the + 'Unknown message type' guard — that would terminate the WebSocket session + on what is at worst a benign no-op from Gemini.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_empty_tool_call" + + result = config.transform_realtime_response( + json.dumps({"toolCall": {"functionCalls": []}}), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + assert result["response"] == [] + assert result["current_response_id"] is None + assert result["current_output_item_id"] is None + + +def test_gemini_empty_tool_call_with_sibling_usage_metadata_does_not_crash(): + """A toolCall with empty functionCalls alongside a sibling key (e.g. + ``usageMetadata``) must still be handled as a benign no-op: the empty + toolCall is consumed and the metadata sibling is skipped, without + raising ``Unknown message type``.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_empty_tool_call_with_sibling" + + result = config.transform_realtime_response( + json.dumps( + { + "toolCall": {"functionCalls": []}, + "usageMetadata": {"totalTokenCount": 7}, + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": "item_existing", + "current_response_id": "resp_existing", + "current_conversation_id": "conv_existing", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + assert result["response"] == [] + # In-flight response IDs must survive the benign no-op. + assert result["current_response_id"] == "resp_existing" + assert result["current_output_item_id"] == "item_existing" + + +def test_gemini_tool_call_response_done_includes_usage_from_sibling_metadata(): + """A ``toolCall`` frame with a sibling ``usageMetadata`` must propagate the + real token counts onto the emitted ``response.done`` so spend/budget + accounting records tokens consumed by the tool-call turn — otherwise an + authenticated client can repeatedly drive tool calls with zero spend.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_tool_call_usage" + + result = config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": "call_usage", + "name": "get_weather", + "args": {"location": "NYC"}, + } + ] + }, + "usageMetadata": { + "promptTokenCount": 17, + "responseTokenCount": 4, + "totalTokenCount": 21, + }, + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + response_done = next( + ev for ev in result["response"] if ev.get("type") == "response.done" + ) + usage = response_done["response"]["usage"] + assert usage["input_tokens"] == 17 + assert usage["output_tokens"] == 4 + assert usage["total_tokens"] == 21 + + +def test_gemini_tool_call_response_done_falls_back_to_empty_usage(): + """Without sibling ``usageMetadata`` the tool-call ``response.done`` still + carries a valid empty usage block so OpenAI-compatible clients (which + expect ``usage`` on every ``response.done``) don't break.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_tool_call_no_usage" + + result = config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": "call_no_usage", + "name": "get_weather", + "args": {"location": "NYC"}, + } + ] + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + response_done = next( + ev for ev in result["response"] if ev.get("type") == "response.done" + ) + usage = response_done["response"]["usage"] + assert usage["input_tokens"] == 0 + assert usage["output_tokens"] == 0 + assert usage["total_tokens"] == 0 + + +def test_gemini_function_call_output_includes_name(): + """Verify function_call_output includes name field from stored mapping.""" + config = GeminiRealtimeConfig() + + # First, receive a toolCall from Gemini (this stores the call_id → name mapping) + gemini_tool_call = { + "toolCall": { + "functionCalls": [ + { + "id": "call_123", + "name": "get_weather", + "args": {"location": "San Francisco"}, + } + ] + } + } + + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + config.transform_realtime_response( + json.dumps(gemini_tool_call), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + # Verify mapping was stored + assert "call_123" in config._tool_call_id_to_name + assert config._tool_call_id_to_name["call_123"] == "get_weather" + + # Now send a function_call_output back (this should include the name) + function_output = { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": "call_123", + "output": json.dumps({"result": "72 degrees"}), + }, + } + + result = config.transform_realtime_request( + json.dumps(function_output), + "gemini-2.5-flash", + session_configuration_request="{}", + ) + + assert len(result) == 1 + tool_response = json.loads(result[0]) + assert "toolResponse" in tool_response + assert "functionResponses" in tool_response["toolResponse"] + assert len(tool_response["toolResponse"]["functionResponses"]) == 1 + + function_response = tool_response["toolResponse"]["functionResponses"][0] + assert function_response["id"] == "call_123" + assert function_response["name"] == "get_weather" # ✅ Name is included + assert "response" in function_response + + +def test_gemini_subsequent_session_update_forwards_tools_merged_with_original_setup(): + """A client session.update sent after the auto-setup must forward tools/ + instructions as a follow-up setup, merged with the original setup so we + don't drop the pre-existing config (model, generationConfig, etc.).""" + config = GeminiRealtimeConfig() + + original_setup = { + "setup": { + "model": "models/gemini-2.5-flash-native-audio", + "generationConfig": {"responseModalities": ["AUDIO"]}, + "inputAudioTranscription": {}, + "systemInstruction": {"role": "user", "parts": [{"text": "original"}]}, + } + } + + session_update = { + "type": "session.update", + "session": { + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather.", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + "instructions": "Be concise.", + }, + } + + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash-native-audio", + session_configuration_request=json.dumps(original_setup), + ) + + assert len(messages) == 1 + follow_up = json.loads(messages[0])["setup"] + assert "tools" in follow_up + assert follow_up["tools"][0]["function_declarations"][0]["name"] == "get_weather" + # systemInstruction overwritten by client's instructions + assert follow_up["systemInstruction"]["parts"][0]["text"] == "Be concise." + # Original generationConfig / model / inputAudioTranscription preserved + assert follow_up["generationConfig"]["responseModalities"] == ["AUDIO"] + assert follow_up["model"] == "models/gemini-2.5-flash-native-audio" + assert follow_up["inputAudioTranscription"] == {} + + +def test_gemini_subsequent_session_update_with_turn_detection_only_preserves_original_tools(): + """A subsequent session.update carrying only turn_detection (the + guardrail-injected disable) must keep the original tools/generationConfig.""" + config = GeminiRealtimeConfig() + + original_setup = { + "setup": { + "model": "models/gemini-2.5-flash-native-audio", + "generationConfig": {"responseModalities": ["AUDIO"]}, + "inputAudioTranscription": {}, + "tools": [ + { + "function_declarations": [ + {"name": "lookup", "description": "x", "parameters": {}} + ] + } + ], + } + } + + session_update = { + "type": "session.update", + "session": {"turn_detection": {"create_response": False}}, + } + + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash-native-audio", + session_configuration_request=json.dumps(original_setup), + ) + + assert len(messages) == 1 + follow_up = json.loads(messages[0])["setup"] + assert follow_up["tools"] == original_setup["setup"]["tools"] + assert ( + follow_up["realtimeInputConfig"]["automaticActivityDetection"]["disabled"] + is True + ) + + +def test_gemini_follow_up_session_update_preserves_response_modalities_on_partial_generation_config(): + """A follow-up session.update that only sets `temperature` (or any other + generationConfig sub-field) must not wipe `responseModalities` from the + original setup.""" + config = GeminiRealtimeConfig() + + original_setup = { + "setup": { + "model": "models/gemini-2.5-flash-native-audio", + "generationConfig": { + "responseModalities": ["AUDIO"], + "maxOutputTokens": 2048, + }, + "inputAudioTranscription": {}, + } + } + + session_update = { + "type": "session.update", + "session": {"temperature": 0.7}, + } + + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash-native-audio", + session_configuration_request=json.dumps(original_setup), + ) + + follow_up = json.loads(messages[0])["setup"] + assert follow_up["generationConfig"]["responseModalities"] == ["AUDIO"] + assert follow_up["generationConfig"]["maxOutputTokens"] == 2048 + assert follow_up["generationConfig"]["temperature"] == 0.7 + + +def test_gemini_subsequent_session_update_preserves_automatic_activity_detection_subfields(): + config = GeminiRealtimeConfig() + + original_setup = { + "setup": { + "model": "models/gemini-2.5-flash-native-audio", + "generationConfig": {"responseModalities": ["AUDIO"]}, + "realtimeInputConfig": { + "automaticActivityDetection": { + "disabled": False, + "silenceDurationMs": 500, + "prefixPaddingMs": 100, + } + }, + } + } + + session_update = { + "type": "session.update", + "session": {"turn_detection": {"create_response": False}}, + } + + messages = config.transform_realtime_request( + json.dumps(session_update), + "gemini-2.5-flash-native-audio", + session_configuration_request=json.dumps(original_setup), + ) + + automatic_activity_detection = json.loads(messages[0])["setup"][ + "realtimeInputConfig" + ]["automaticActivityDetection"] + assert automatic_activity_detection["disabled"] is True + assert automatic_activity_detection["silenceDurationMs"] == 500 + assert automatic_activity_detection["prefixPaddingMs"] == 100 + + +def test_gemini_tool_call_id_to_name_evicts_oldest_when_capped(): + """The call_id → name LRU must evict the oldest entry once the cap is + reached so long sessions with many tool calls don't grow unboundedly, + while keeping recently-seen call_ids resolvable for retried + function_call_output messages.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_lru" + + config._TOOL_CALL_ID_TO_NAME_MAX = 4 + + for idx in range(8): + config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": f"call_{idx}", + "name": f"fn_{idx}", + "args": {}, + } + ] + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + assert len(config._tool_call_id_to_name) == 4 + # Most recent 4 retained; oldest 4 evicted. + assert list(config._tool_call_id_to_name) == [ + "call_4", + "call_5", + "call_6", + "call_7", + ] + + +def test_gemini_standalone_usage_metadata_does_not_crash_websocket(): + """A Gemini frame containing only sibling metadata (e.g. a standalone + ``usageMetadata`` block emitted between turns) must not trip the + ``Unknown message type`` guard — that would terminate the WebSocket + session on a benign no-op frame.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_usage_only" + + result = config.transform_realtime_response( + json.dumps( + { + "usageMetadata": { + "promptTokenCount": 12, + "responseTokenCount": 34, + "totalTokenCount": 46, + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": "item_existing", + "current_response_id": "resp_existing", + "current_conversation_id": "conv_existing", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + assert result["response"] == [] + # State must be returned unchanged so subsequent frames continue the + # in-flight response correctly. + assert result["current_output_item_id"] == "item_existing" + assert result["current_response_id"] == "resp_existing" + assert result["current_conversation_id"] == "conv_existing" + + +def test_gemini_standalone_usage_metadata_is_attributed_to_next_tool_call_response_done(): + """A standalone ``usageMetadata`` frame emitted between turns must not + silently drop the consumed tokens. The next tool-call ``response.done`` + must carry those token counts so an authenticated client cannot drive + tool-call turns whose token usage is recorded as zero, bypassing + spend/budget accounting.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_standalone_usage_then_tool_call" + + standalone_result = config.transform_realtime_response( + json.dumps( + { + "usageMetadata": { + "promptTokenCount": 31, + "responseTokenCount": 9, + "totalTokenCount": 40, + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + assert standalone_result["response"] == [] + + tool_call_result = config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": "call_buffered", + "name": "get_weather", + "args": {"location": "NYC"}, + } + ] + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + response_done = next( + ev for ev in tool_call_result["response"] if ev.get("type") == "response.done" + ) + usage = response_done["response"]["usage"] + assert usage["input_tokens"] == 31 + assert usage["output_tokens"] == 9 + assert usage["total_tokens"] == 40 + # Buffer must be cleared after attribution so a subsequent tool-call + # turn without its own usage does not double-count the previous frame. + assert config._pending_usage_metadata is None + + +def test_gemini_standalone_usage_metadata_is_attributed_to_next_response_done(): + """A standalone ``usageMetadata`` frame must also flow into the normal + (non-tool-call) ``response.done`` path so audio/text turns whose usage + arrives in a separate frame are still billed correctly.""" + config = GeminiRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_standalone_usage_then_turn_complete" + + config.transform_realtime_response( + json.dumps( + { + "usageMetadata": { + "promptTokenCount": 5, + "responseTokenCount": 11, + "totalTokenCount": 16, + } + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + turn_complete_result = config.transform_realtime_response( + json.dumps({"serverContent": {"turnComplete": True}}), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + response_done = next( + ev + for ev in turn_complete_result["response"] + if ev.get("type") == "response.done" + ) + usage = response_done["response"]["usage"] + assert usage["input_tokens"] == 5 + assert usage["output_tokens"] == 11 + assert usage["total_tokens"] == 16 + assert config._pending_usage_metadata is None + + +def test_gemini_in_frame_usage_metadata_clears_pending_buffer(): + """When ``usageMetadata`` arrives in the same frame as the closing + ``toolCall`` / ``turnComplete``, the in-frame counts are authoritative + and any buffered standalone metadata must be discarded so a later + turn's ``response.done`` does not double-count tokens.""" + config = GeminiRealtimeConfig() + config._pending_usage_metadata = { + "promptTokenCount": 99, + "responseTokenCount": 99, + "totalTokenCount": 198, + } + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_in_frame_clears_buffer" + + result = config.transform_realtime_response( + json.dumps( + { + "toolCall": { + "functionCalls": [ + { + "id": "call_in_frame", + "name": "get_weather", + "args": {"location": "NYC"}, + } + ] + }, + "usageMetadata": { + "promptTokenCount": 3, + "responseTokenCount": 2, + "totalTokenCount": 5, + }, + } + ), + "gemini-2.5-flash", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + response_done = next( + ev for ev in result["response"] if ev.get("type") == "response.done" + ) + usage = response_done["response"]["usage"] + assert usage["input_tokens"] == 3 + assert usage["output_tokens"] == 2 + assert usage["total_tokens"] == 5 + assert config._pending_usage_metadata is None diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py new file mode 100644 index 00000000000..ac89428617d --- /dev/null +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py @@ -0,0 +1,134 @@ +""" +Tests that data_residency is correctly populated on the litellm logging +object's litellm_params for OpenAI Responses paths, even when +custom_llm_provider is resolved from the model string inside responses() +rather than passed explicitly. +""" + +import json +from unittest.mock import MagicMock, patch + +import litellm + + +def _make_responses_api_response_body() -> dict: + return { + "id": "resp-test", + "object": "response", + "created_at": 1234567890, + "model": "gpt-4.1", + "output": [ + { + "type": "message", + "id": "msg-test", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "ok", + "annotations": [], + } + ], + } + ], + "status": "completed", + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + }, + } + + +def _make_mock_http_client(response_body: dict) -> MagicMock: + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = response_body + mock_response.text = json.dumps(response_body) + mock_client.post.return_value = mock_response + return mock_client + + +def _capture_logging_obj(): + captured = {} + + real_init = litellm.Logging.__init__ + + def init_spy(self, *args, **kwargs): + real_init(self, *args, **kwargs) + captured["logging_obj"] = self + + return captured, init_spy + + +def test_responses_eu_api_base_sets_data_residency(): + """When api_base is a regional OpenAI host and custom_llm_provider is + inferred from the model (not passed explicitly), data_residency must end + up on the logging object's litellm_params so the cost calculator can apply + the regional uplift.""" + mock_client = _make_mock_http_client(_make_responses_api_response_body()) + captured, init_spy = _capture_logging_obj() + + with ( + patch( + "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", + return_value=mock_client, + ), + patch.object(litellm.Logging, "__init__", init_spy), + ): + litellm.responses( + model="gpt-4.1", + input="hi", + api_base="https://eu.api.openai.com/v1", + api_key="test-key", + ) + + logging_obj = captured["logging_obj"] + assert logging_obj.litellm_params.get("data_residency") == "eu" + + +def test_responses_us_api_base_sets_data_residency(): + mock_client = _make_mock_http_client(_make_responses_api_response_body()) + captured, init_spy = _capture_logging_obj() + + with ( + patch( + "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", + return_value=mock_client, + ), + patch.object(litellm.Logging, "__init__", init_spy), + ): + litellm.responses( + model="gpt-4.1", + input="hi", + api_base="https://us.api.openai.com/v1", + api_key="test-key", + ) + + logging_obj = captured["logging_obj"] + assert logging_obj.litellm_params.get("data_residency") == "us" + + +def test_responses_global_api_base_leaves_data_residency_none(): + mock_client = _make_mock_http_client(_make_responses_api_response_body()) + captured, init_spy = _capture_logging_obj() + + with ( + patch( + "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", + return_value=mock_client, + ), + patch.object(litellm.Logging, "__init__", init_spy), + ): + litellm.responses( + model="gpt-4.1", + input="hi", + api_base="https://api.openai.com/v1", + api_key="test-key", + ) + + logging_obj = captured["logging_obj"] + assert logging_obj.litellm_params.get("data_residency") is None diff --git a/tests/test_litellm/llms/openai/test_data_residency.py b/tests/test_litellm/llms/openai/test_data_residency.py new file mode 100644 index 00000000000..ecb5739133c --- /dev/null +++ b/tests/test_litellm/llms/openai/test_data_residency.py @@ -0,0 +1,34 @@ +"""Tests for the OpenAI data-residency inference helper.""" + +import pytest + +from litellm.llms.openai.data_residency import infer_openai_data_residency + + +@pytest.mark.parametrize( + "api_base, expected", + [ + ("https://eu.api.openai.com/v1", "eu"), + ("https://eu.api.openai.com", "eu"), + ("https://us.api.openai.com/v1", "us"), + ("https://us.api.openai.com", "us"), + ("https://EU.api.openai.com/v1", "eu"), + ("https://api.openai.com/v1", None), + ("https://api.openai.com", None), + ("https://example.com/v1", None), + ("https://my-azure-endpoint.openai.azure.com/openai/deployments/foo", None), + ("", None), + (None, None), + ("not a url", None), + ], +) +def test_infer_openai_data_residency(api_base, expected): + assert infer_openai_data_residency("openai", api_base) == expected + + +@pytest.mark.parametrize("custom_llm_provider", [None, "anthropic", "azure", "bedrock"]) +def test_infer_openai_data_residency_non_openai_provider(custom_llm_provider): + assert ( + infer_openai_data_residency(custom_llm_provider, "https://eu.api.openai.com/v1") + is None + ) diff --git a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index 1baaf912568..0ad614099de 100644 --- a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -19,6 +19,7 @@ import websockets.exceptions # registers websockets.exceptions on the websocket sys.path.insert(0, os.path.abspath("../../../../..")) +import litellm from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig # --------------------------------------------------------------------------- @@ -82,6 +83,85 @@ def test_session_configuration_request_model_format(): ) +def test_vertex_requires_session_configuration_feature_flag(monkeypatch): + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + # Default remains backwards-compatible (auto setup on connect) + monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False) + assert cfg.requires_session_configuration() is True + + # Opt-in deferred setup for tool-injection flow + monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False) + assert cfg.requires_session_configuration() is False + + +def test_vertex_session_update_defaults_to_audio_modality(): + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + session_update = { + "type": "session.update", + "session": { + "instructions": "You are a helpful assistant.", + # No modalities provided on purpose + }, + } + + messages = cfg.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=None, + ) + assert len(messages) == 1 + setup_payload = json.loads(messages[0])["setup"] + assert setup_payload["generationConfig"]["responseModalities"] == ["AUDIO"] + + +def test_vertex_session_update_normalizes_ga_remapped_fields(): + """GA-format clients send ``output_modalities`` and nested + ``audio.input.transcription`` / ``audio.input.turn_detection``. These must + be normalised back to the flat beta keys before ``map_openai_params`` + runs so client preferences aren't silently dropped. + """ + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + session_update = { + "type": "session.update", + "session": { + "instructions": "Be concise.", + "output_modalities": ["text"], + "audio": { + "input": { + "transcription": {}, + "turn_detection": {"silence_duration_ms": 1500}, + }, + }, + }, + } + + messages = cfg.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=None, + ) + assert len(messages) == 1 + setup_payload = json.loads(messages[0])["setup"] + + assert setup_payload["generationConfig"]["responseModalities"] == ["TEXT"] + assert setup_payload["inputAudioTranscription"] == {} + assert ( + setup_payload["realtimeInputConfig"]["automaticActivityDetection"][ + "silenceDurationMs" + ] + == 1500 + ) + + # --------------------------------------------------------------------------- # Round-trip test: text-in / text-out via RealTimeStreaming # --------------------------------------------------------------------------- @@ -208,3 +288,61 @@ async def test_vertex_realtime_text_in_text_out(): # response.done should have been forwarded done_msgs = [m for m in sent_to_client if '"response.done"' in m] assert done_msgs, "Expected response.done to be sent to client" + + +def test_vertex_warns_when_dropping_guardrail_turn_detection_update(caplog): + """A subsequent session.update carrying the guardrail's + ``create_response: False`` cannot be forwarded as a follow-up setup on + Vertex AI (1007). Surface a warning so operators know the auto-response + suppression is being silently dropped.""" + import logging + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + session_update = { + "type": "session.update", + "session": {"turn_detection": {"create_response": False}}, + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result = cfg.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=json.dumps({"setup": {"model": "x"}}), + ) + + assert result == [] + assert any( + "Vertex AI Realtime" in record.message + and "create_response=False" in record.message + for record in caplog.records + ) + + +def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog): + """A subsequent session.update without ``create_response: False`` is a + routine drop and should stay at debug level (no warning).""" + import logging + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + session_update = { + "type": "session.update", + "session": {"instructions": "Be concise."}, + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cfg.transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=json.dumps({"setup": {"model": "x"}}), + ) + + assert not any( + "Vertex AI Realtime" in record.message and "session.update" in record.message + for record in caplog.records + ) diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 70583cfe61b..55197d3165c 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -456,6 +456,127 @@ class TestVertexAIVideoConfig: raw_response=mock_response, logging_obj=self.mock_logging_obj ) + def test_get_video_edit_prefetch_params(self): + """Test that prefetch params returns the fetchPredictOperation URL and body.""" + operation_name = "projects/test-project/locations/us-central1/publishers/google/models/veo-3.1-generate-001/operations/op-123" + api_base = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models" + + fetch_url, fetch_body = self.config.get_video_edit_prefetch_params( + video_id=operation_name, + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "fetchPredictOperation" in fetch_url + assert "veo-3.1-generate-001" in fetch_url + assert fetch_body == {"operationName": operation_name} + + def test_transform_video_edit_request_with_bytes(self): + """Test video edit request builds predictLongRunning body from pre-fetched bytes.""" + operation_name = "projects/test-project/locations/us-central1/publishers/google/models/veo-3.1-generate-001/operations/op-123" + api_base = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models" + fake_bytes = base64.b64encode(b"fake_video").decode() + + prefetched = { + "done": True, + "response": { + "videos": [{"bytesBase64Encoded": fake_bytes, "mimeType": "video/mp4"}] + }, + } + + url, data = self.config.transform_video_edit_request( + prompt="Make it brighter", + video_id=operation_name, + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={"Authorization": "Bearer token"}, + prefetched_source_data=prefetched, + ) + + assert url.endswith(":predictLongRunning") + assert "veo-3.1-generate-001" in url + instance = data["instances"][0] + assert instance["prompt"] == "Make it brighter" + assert instance["video"]["bytesBase64Encoded"] == fake_bytes + assert instance["video"]["mimeType"] == "video/mp4" + + def test_transform_video_edit_request_with_gcs_uri(self): + """Test that gcsUri is used when present in source video.""" + operation_name = "projects/test-project/locations/us-central1/publishers/google/models/veo-3.1-generate-001/operations/op-456" + api_base = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models" + + prefetched = { + "done": True, + "response": { + "videos": [{"gcsUri": "gs://bucket/video.mp4", "mimeType": "video/mp4"}] + }, + } + + _, data = self.config.transform_video_edit_request( + prompt="Make it darker", + video_id=operation_name, + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + prefetched_source_data=prefetched, + ) + + assert data["instances"][0]["video"] == {"gcsUri": "gs://bucket/video.mp4"} + + def test_transform_video_edit_request_source_not_done_raises(self): + """Test that editing an in-progress video raises a clear error.""" + operation_name = "projects/test-project/locations/us-central1/publishers/google/models/veo-3.1-generate-001/operations/op-789" + api_base = "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models" + + with pytest.raises(ValueError, match="not complete yet"): + self.config.transform_video_edit_request( + prompt="Make it brighter", + video_id=operation_name, + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + prefetched_source_data={"done": False}, + ) + + def test_transform_video_edit_response(self): + """Test that edit response returns a processing VideoObject with encoded ID.""" + operation_name = "projects/test-project/locations/us-central1/publishers/google/models/veo-3.1-generate-001/operations/new-op-123" + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = {"name": operation_name} + + video_obj = self.config.transform_video_edit_response( + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + custom_llm_provider="vertex_ai", + ) + + assert isinstance(video_obj, VideoObject) + assert video_obj.status == "processing" + assert video_obj.id + assert video_obj.model == "veo-3.1-generate-001" + + def test_transform_video_edit_response_includes_usage_for_cost(self): + """Edit responses include duration/resolution usage for spend accounting.""" + operation_name = "projects/test-project/locations/us-central1/publishers/google/models/veo-3.1-generate-001/operations/new-op-123" + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = {"name": operation_name} + request_data = { + "instances": [{"prompt": "Make it brighter", "video": {}}], + "parameters": {"durationSeconds": 8, "resolution": "1080p"}, + } + + video_obj = self.config.transform_video_edit_response( + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + custom_llm_provider="vertex_ai", + request_data=request_data, + ) + + assert video_obj.usage is not None + assert video_obj.usage["duration_seconds"] == 8.0 + assert video_obj.usage["video_resolution"] == "1080p" + def test_transform_video_remix_request_not_supported(self): """Test that video remix raises NotImplementedError.""" with pytest.raises(NotImplementedError, match="Video remix is not supported"): 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 88742c67a86..1499f7e474f 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 @@ -880,7 +880,7 @@ class TestMCPPublicRouteGuard: with patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", ) as mock_auth: - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) mock_auth.assert_not_called() assert isinstance(auth_result, UserAPIKeyAuth) @@ -997,7 +997,7 @@ class TestMCPOAuth2FallbackTargetGating: mock_mgr.get_mcp_server_by_name.return_value = ( TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.oauth2) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) async def test_fallback_blocked_when_any_target_in_header_is_not_oauth2(self): @@ -1157,7 +1157,7 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=True, ) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) mock_auth.assert_not_called() @@ -1400,7 +1400,7 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=True, ) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) assert auth_result.user_id == "real-user" mock_auth.assert_called_once() @@ -1437,7 +1437,7 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=True, ) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) assert auth_result.user_id == "real-user" mock_auth.assert_called_once() @@ -2444,13 +2444,14 @@ async def test_get_team_object_permission_with_core_auth_auto_loading(): @pytest.mark.asyncio async def test_get_allowed_mcp_servers_for_team_uses_helper(): """ - Test that _get_allowed_mcp_servers_for_team properly uses _get_team_object_permission - helper which handles both loaded and unloaded object_permission cases. + Test that _get_allowed_mcp_servers_for_team resolves both legacy + object_permission fields (mcp_servers, mcp_access_groups) and the unified + team.access_group_ids → access_mcp_server_ids path. """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -2464,53 +2465,51 @@ async def test_get_allowed_mcp_servers_for_team_uses_helper(): transport=MCPTransport.http, ) try: - # Create mock object permission with servers and access groups mock_object_permission = LiteLLM_ObjectPermissionTable( object_permission_id="perm-789", mcp_servers=["direct-server1", "direct-server2"], mcp_access_groups=["dev-group"], vector_stores=[], ) + mock_team = LiteLLM_TeamTable( + team_id="team-789", + access_group_ids=[], + object_permission_id="perm-789", + ) + mock_team.object_permission = mock_object_permission - # Create mock user auth mock_user_auth = UserAPIKeyAuth( api_key="test-key", user_id="test-user", team_id="team-789", ) - # Mock the helper methods - with patch.object( - MCPRequestHandler, "_get_team_object_permission" - ) as mock_get_team_perm: - with patch.object( - MCPRequestHandler, "_get_mcp_servers_from_access_groups" - ) as mock_get_access_group_servers: - # Configure mocks - mock_get_team_perm.return_value = mock_object_permission - mock_get_access_group_servers.return_value = [ - "group-server1", - "group-server2", - ] + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ), + patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=["group-server1", "group-server2"], + ) as mock_get_access_group_servers, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( + mock_user_auth + ) - # Call the method - result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( - mock_user_auth - ) + assert set(result) == { + "direct-server1", + "direct-server2", + "group-server1", + "group-server2", + } - # Assert the result contains both direct and access group servers - assert set(result) == { - "direct-server1", - "direct-server2", - "group-server1", - "group-server2", - } - - # Verify _get_team_object_permission was called (the helper we fixed) - mock_get_team_perm.assert_called_once_with(mock_user_auth) - - # Verify access groups were resolved - mock_get_access_group_servers.assert_called_once_with(["dev-group"]) + mock_get_access_group_servers.assert_called_once_with(["dev-group"]) finally: for sid in ("direct-server1", "direct-server2"): global_mcp_server_manager.registry.pop(sid, None) @@ -2520,32 +2519,36 @@ async def test_get_allowed_mcp_servers_for_team_uses_helper(): async def test_get_allowed_mcp_servers_for_team_with_no_object_permission(): """ Test that _get_allowed_mcp_servers_for_team returns empty list when - team has no object_permission. + the team has no object_permission and no access_group_ids. """ - # Create mock user auth + from litellm.proxy._types import LiteLLM_TeamTable + + mock_team = LiteLLM_TeamTable( + team_id="team-no-perm", + access_group_ids=[], + object_permission_id=None, + ) + mock_user_auth = UserAPIKeyAuth( api_key="test-key", user_id="test-user", team_id="team-no-perm", ) - # Mock the helper to return None (no object permission) - with patch.object( - MCPRequestHandler, "_get_team_object_permission" - ) as mock_get_team_perm: - mock_get_team_perm.return_value = None - - # Call the method + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ), + ): result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( mock_user_auth ) - # Assert empty list is returned assert result == [] - # Verify the helper was called - mock_get_team_perm.assert_called_once_with(mock_user_auth) - @pytest.mark.asyncio async def test_get_allowed_mcp_servers_for_team_without_user_auth_returns_empty(): @@ -3160,3 +3163,481 @@ class TestOrgMCPPermissions: user_api_key_auth=auth, ) assert sorted(result) == ["tool_a", "tool_b"] + + +# --------------------------------------------------------------------------- +# LIT-3189: key unified access_group_ids extend team MCP scope +# --------------------------------------------------------------------------- + + +def _patch_proxy_server_globals_for_mcp(): + """Non-None mocks so the helper's None-guard doesn't short-circuit.""" + return [ + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ] + + +def _fake_mcp_access_group( + access_group_id, + access_mcp_server_ids=None, + assigned_team_ids=None, + assigned_key_ids=None, +): + from litellm.proxy._types import LiteLLM_AccessGroupTable + + return LiteLLM_AccessGroupTable( + access_group_id=access_group_id, + access_group_name=access_group_id, + access_mcp_server_ids=access_mcp_server_ids or [], + assigned_team_ids=assigned_team_ids or [], + assigned_key_ids=assigned_key_ids or [], + ) + + +def _start_patches(patches): + for p in patches: + p.start() + + +def _stop_patches(patches): + for p in patches: + p.stop() + + +@pytest.mark.asyncio +async def test_mcp_key_access_group_extras_when_team_authorized(): + """Group's assigned_team_ids includes key's team and grants an MCP server → server returned.""" + valid_token = UserAPIKeyAuth( + token="test-token", + access_group_ids=["mcp-premium"], + team_id="team-a", + ) + fake_ag = _fake_mcp_access_group( + access_group_id="mcp-premium", + access_mcp_server_ids=["srv-stripe"], + assigned_team_ids=["team-a"], + ) + + mock_mgr = MagicMock() + mock_mgr.expand_permission_list.side_effect = lambda x: list(x) + + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_mgr, + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + valid_token + ) + assert result == ["srv-stripe"] + finally: + _stop_patches(patches) + + +@pytest.mark.asyncio +async def test_mcp_key_access_group_extras_when_key_directly_authorized(): + """Group's assigned_key_ids includes the key's token → server returned (per-key auth).""" + valid_token = UserAPIKeyAuth( + token="test-token-hashed", + access_group_ids=["mcp-per-key"], + team_id="team-a", + ) + fake_ag = _fake_mcp_access_group( + access_group_id="mcp-per-key", + access_mcp_server_ids=["srv-stripe"], + assigned_team_ids=[], + assigned_key_ids=["test-token-hashed"], + ) + + mock_mgr = MagicMock() + mock_mgr.expand_permission_list.side_effect = lambda x: list(x) + + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_mgr, + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + valid_token + ) + assert result == ["srv-stripe"] + finally: + _stop_patches(patches) + + +@pytest.mark.asyncio +async def test_mcp_key_access_group_extras_when_key_has_no_groups(): + """Empty access_group_ids → no extras, no DB read.""" + valid_token = UserAPIKeyAuth( + token="test-token", + access_group_ids=[], + team_id="team-a", + ) + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + valid_token + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_mcp_key_access_group_extras_when_group_has_no_servers(): + """Group authorizes the team but its access_mcp_server_ids is empty → no extras.""" + valid_token = UserAPIKeyAuth( + token="test-token", + access_group_ids=["mcp-empty"], + team_id="team-a", + ) + fake_ag = _fake_mcp_access_group( + access_group_id="mcp-empty", + access_mcp_server_ids=[], + assigned_team_ids=["team-a"], + ) + + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + valid_token + ) + assert result == [] + finally: + _stop_patches(patches) + + +@pytest.mark.asyncio +async def test_mcp_key_access_group_extras_when_group_authorizes_neither(): + """ + Escalation regression: team member attaches a foreign access group to their key. + Group grants servers BUT assigned_team_ids/assigned_key_ids exclude this caller. + No extras contributed. + """ + valid_token = UserAPIKeyAuth( + token="team-a-token", + access_group_ids=["team-b-mcp-group"], + team_id="team-a", + ) + fake_ag = _fake_mcp_access_group( + access_group_id="team-b-mcp-group", + access_mcp_server_ids=["srv-finance-only"], + assigned_team_ids=["team-b"], + assigned_key_ids=["team-b-token"], + ) + + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + valid_token + ) + assert result == [] + finally: + _stop_patches(patches) + + +@pytest.mark.asyncio +async def test_mcp_key_access_group_extras_when_get_access_object_raises(): + """Group lookup failure is treated as no authorization (does not crash).""" + valid_token = UserAPIKeyAuth( + token="test-token", + access_group_ids=["missing-mcp-group"], + team_id="team-a", + ) + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + side_effect=Exception("not found"), + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler._get_key_access_group_mcp_server_extras( + valid_token + ) + assert result == [] + finally: + _stop_patches(patches) + + +@pytest.mark.asyncio +async def test_get_allowed_mcp_servers_unions_key_access_group_extras(): + """End-to-end: team has [srv-team], key access group grants [srv-extra] → both in final list. + + Without this fix [srv-extra] would be intersected away because the team doesn't list it. + """ + auth = UserAPIKeyAuth( + token="test-token", + api_key="test-key", + team_id="team-a", + access_group_ids=["mcp-extra-group"], + ) + + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["srv-team"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=["srv-extra"], + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert sorted(result) == ["srv-extra", "srv-team"] + + +@pytest.mark.asyncio +async def test_get_allowed_mcp_servers_no_union_when_no_authorized_extras(): + """End-to-end: no authorized extras → behavior identical to today (team ceiling enforced).""" + auth = UserAPIKeyAuth( + token="test-token", + api_key="test-key", + team_id="team-a", + access_group_ids=["mcp-foreign-group"], + ) + + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=["srv-key-only"], + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["srv-team"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=[], + ), + ): + # key ∩ team = {} (no overlap), extras = [] → final = [] + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert result == [] + + +# --------------------------------------------------------------------------- +# Issue #27657: team unified access_group_ids resolve to MCP servers +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_team_access_group_ids_resolve_to_mcp_servers(): + """A virtual key with empty access_group_ids inherits MCP servers from + its team's access_group_ids (mirror of the model-side resolution). + + Reproduction of https://github.com/BerriAI/litellm/issues/27657: + the runtime used to ignore team.access_group_ids when computing the + MCP scope, so virtual keys saw empty server lists even when their + team had an MCP-granting access group attached. + """ + from litellm.proxy._types import LiteLLM_TeamTable + + mock_team = LiteLLM_TeamTable( + team_id="team-a", + access_group_ids=["mcp-premium"], + object_permission_id=None, + ) + + auth = UserAPIKeyAuth( + token="test-token-hash", + api_key="sk-test", + team_id="team-a", + access_group_ids=[], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=["srv-stripe"], + ) as mock_resolver, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + + assert result == ["srv-stripe"] + mock_resolver.assert_called_once() + assert mock_resolver.call_args.kwargs["access_group_ids"] == ["mcp-premium"] + + +@pytest.mark.asyncio +async def test_team_access_group_ids_union_with_object_permission(): + """When both legacy object_permission and unified team.access_group_ids + grant MCP servers, the final list is their union.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + for sid in ("srv-direct",): + global_mcp_server_manager.registry[sid] = MCPServer( + server_id=sid, + name=sid, + server_name=sid, + url=f"https://{sid}.example.com", + transport=MCPTransport.http, + ) + try: + mock_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-1", + mcp_servers=["srv-direct"], + mcp_access_groups=[], + vector_stores=[], + ) + mock_team = LiteLLM_TeamTable( + team_id="team-a", + access_group_ids=["mcp-premium"], + object_permission_id="perm-1", + ) + mock_team.object_permission = mock_object_permission + + auth = UserAPIKeyAuth( + token="test-token-hash", + api_key="sk-test", + team_id="team-a", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=["srv-stripe"], + ), + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + + assert set(result) == {"srv-direct", "srv-stripe"} + finally: + global_mcp_server_manager.registry.pop("srv-direct", None) + + +@pytest.mark.asyncio +async def test_team_access_group_ids_empty_returns_no_extras(): + """Empty team.access_group_ids → resolver called with [], short-circuits + without DB access, no extras added.""" + from litellm.proxy._types import LiteLLM_TeamTable + + mock_team = LiteLLM_TeamTable( + team_id="team-a", + access_group_ids=[], + object_permission_id=None, + ) + + auth = UserAPIKeyAuth( + token="test-token-hash", + api_key="sk-test", + team_id="team-a", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ) as mock_resolver, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + + assert result == [] + mock_resolver.assert_called_once() + assert mock_resolver.call_args.kwargs["access_group_ids"] == [] + + +@pytest.mark.asyncio +async def test_get_allowed_mcp_servers_includes_team_access_group_extras_end_to_end(): + """End-to-end: virtual key has nothing of its own, team has an MCP + access group → key sees the granted server through get_allowed_mcp_servers.""" + auth = UserAPIKeyAuth( + token="test-token", + api_key="sk-test", + team_id="team-a", + access_group_ids=[], + ) + + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["srv-stripe"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=[], + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert result == ["srv-stripe"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_callback_oauth_error_responses.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_callback_oauth_error_responses.py new file mode 100644 index 00000000000..11ef40b9961 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_callback_oauth_error_responses.py @@ -0,0 +1,210 @@ +"""Regression tests for LIT-2750. + +The MCP OAuth ``/callback`` endpoint must handle IdP error responses +(e.g. ``?error=access_denied``) gracefully instead of returning a 422 +because ``code`` and ``state`` were declared as required FastAPI query +params. Per RFC 6749 §4.1.2.1 the IdP redirects to the configured +redirect URI with ``error`` / ``error_description`` / ``error_uri`` +query params and no ``code`` when the user denies access. + +These tests cover both the propagate-to-client path (when state decodes +to a trusted ``redirect_uri``) and the in-page fallback (when state is +missing, undecryptable, or carries an untrusted redirect_uri). They also +pin the success path (``code`` + ``state``) against accidental +regressions. +""" + +import pytest + + +@pytest.fixture(autouse=True) +def _mock_mcp_client_ip(): + """Bypass IP-based access control for the in-process TestClient. + + Mirrors the autouse fixture in ``test_discoverable_endpoints.py`` so + these tests don't require a real client IP context. + """ + from unittest.mock import patch + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", + return_value=None, + ): + yield + + +@pytest.fixture +def callback_test_client(monkeypatch): + """FastAPI TestClient mounted with the MCP discoverable router. + + Sets a deterministic ``LITELLM_SALT_KEY`` so encoded states minted + in-test can be decrypted by the handler. + """ + from fastapi import FastAPI + from fastapi.testclient import TestClient + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-LIT-2750") + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + router, + ) + + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +class TestCallbackOAuthErrorResponses: + """LIT-2750: IdP error responses to ``/callback`` must not 422.""" + + def test_idp_error_with_no_state_returns_400_html(self, callback_test_client): + """Pre-fix: 422 Pydantic. Post-fix: 400 HTML with the IdP's error.""" + resp = callback_test_client.get( + "/callback", + params={ + "error": "access_denied", + "error_description": "User declined access", + }, + follow_redirects=False, + ) + assert resp.status_code == 400 + assert "text/html" in resp.headers["content-type"] + body = resp.text + assert "access_denied" in body + assert "User declined access" in body + # Sanity: must not leak the Pydantic validation error. + assert "Field required" not in body + + def test_idp_error_html_escapes_user_controlled_fields( + self, callback_test_client + ): + """A malicious IdP must not be able to inject HTML/JS via error params.""" + resp = callback_test_client.get( + "/callback", + params={ + "error": "", + "error_description": "", + }, + follow_redirects=False, + ) + assert resp.status_code == 400 + body = resp.text + # Raw tags must be escaped, not present verbatim. + assert "" not in body + assert "" not in body + assert "<script>alert(1)</script>" in body + + def test_idp_error_with_trusted_state_propagates_to_client_redirect_uri( + self, callback_test_client + ): + """When state decodes to a trusted (loopback) redirect_uri, propagate + the error back so the MCP client's OAuth library can surface it + instead of timing out waiting on the loopback.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + encode_state_with_base_url, + ) + + state = encode_state_with_base_url( + base_url="http://localhost:3000/", + original_state="client-original-state-xyz", + client_redirect_uri="http://127.0.0.1:60108/callback", + ) + + resp = callback_test_client.get( + "/callback", + params={ + "error": "access_denied", + "error_description": "User declined access", + "state": state, + }, + follow_redirects=False, + ) + assert resp.status_code == 302 + location = resp.headers["location"] + assert location.startswith("http://127.0.0.1:60108/callback?") + assert "error=access_denied" in location + # Original client state must be round-tripped, not our wrapped state. + assert "state=client-original-state-xyz" in location + # error_description percent-encoded but present. + assert "error_description=User" in location + # Wrapped/encrypted state must NOT leak to the client. + assert state not in location + + def test_idp_error_with_untrusted_redirect_uri_does_not_open_redirect( + self, callback_test_client + ): + """If the state minted earlier carries a redirect_uri that the proxy + no longer trusts, we must surface the error inline rather than + 302-ing to an attacker-controlled URL (open-redirect).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + encode_state_with_base_url, + ) + + state = encode_state_with_base_url( + base_url="http://localhost:3000/", + original_state="x", + client_redirect_uri="https://attacker.example.com/steal", + ) + + resp = callback_test_client.get( + "/callback", + params={"error": "access_denied", "state": state}, + follow_redirects=False, + ) + # Must not 3xx — open redirect would defeat the redirect_uri allowlist. + assert resp.status_code == 400 + assert "attacker.example.com" not in resp.headers.get("location", "") + assert "access_denied" in resp.text + + def test_idp_error_with_undecryptable_state_falls_back_to_html( + self, callback_test_client + ): + resp = callback_test_client.get( + "/callback", + params={ + "error": "server_error", + "error_description": "boom", + "state": "not-a-valid-encrypted-state", + }, + follow_redirects=False, + ) + assert resp.status_code == 400 + assert "server_error" in resp.text + assert "boom" in resp.text + + def test_bare_callback_with_no_params_returns_400_not_422( + self, callback_test_client + ): + """An SSO redirect chain that drops the original /authorize query + params should land on a human-readable 400, not a Pydantic 422.""" + resp = callback_test_client.get("/callback", follow_redirects=False) + assert resp.status_code == 400 + assert "invalid_request" in resp.text + assert "Field required" not in resp.text + + def test_success_path_still_redirects_with_code_and_state( + self, callback_test_client + ): + """Regression: the successful (``code``+``state``) flow must still + redirect back to the trusted client redirect_uri with the original + state preserved.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + encode_state_with_base_url, + ) + + state = encode_state_with_base_url( + base_url="http://localhost:3000/", + original_state="orig-state-success", + client_redirect_uri="http://127.0.0.1:60108/callback", + ) + + resp = callback_test_client.get( + "/callback", + params={"code": "auth-code-abc", "state": state}, + follow_redirects=False, + ) + assert resp.status_code == 302 + location = resp.headers["location"] + assert location.startswith("http://127.0.0.1:60108/callback?") + assert "code=auth-code-abc" in location + assert "state=orig-state-success" in location diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py index c2e42d2f592..d8e4a342e52 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py @@ -462,6 +462,9 @@ async def test_e2e_jwt_team_mcp_key_intersection(monkeypatch): monkeypatch.setattr( "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object ) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_team_object", mock_get_team_object + ) jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") @@ -495,28 +498,25 @@ async def test_e2e_jwt_team_mcp_key_intersection(monkeypatch): object_permission=key_object_permission, # Key has its own permissions ) - # Mock the helper methods to return our test data - with patch.object( - MCPRequestHandler, "_get_team_object_permission" - ) as mock_team_perm: - mock_team_perm.return_value = team_object_permission + with ( + patch.object( + MCPRequestHandler, + "_get_key_object_permission", + return_value=key_object_permission, + ), + patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + ): + allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth + ) - with patch.object( - MCPRequestHandler, "_get_key_object_permission" - ) as mock_key_perm: - mock_key_perm.return_value = key_object_permission - - with patch.object( - MCPRequestHandler, "_get_mcp_servers_from_access_groups" - ) as mock_access_groups: - mock_access_groups.return_value = [] - - allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers( - user_api_key_auth - ) - - # Should be intersection: only server-2 is in both - expected = ["server-2"] - assert sorted(allowed_servers) == sorted( - expected - ), f"Expected intersection {expected}, got {allowed_servers}" + # Should be intersection: only server-2 is in both + expected = ["server-2"] + assert sorted(allowed_servers) == sorted( + expected + ), f"Expected intersection {expected}, got {allowed_servers}" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py index 2ae575b6d99..052231b562a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py @@ -41,37 +41,44 @@ async def test_simple_jwt_mcp_permissions_enforced(): object_permission_id="perm-123", mcp_servers=team_mcp_servers, ) + team_obj = LiteLLM_TeamTable( + team_id="my-team", + access_group_ids=[], + object_permission_id="perm-123", + ) + team_obj.object_permission = team_object_permission - # 3. Mock the team permission lookup - with patch.object( - MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock - ) as mock_team_perm: - mock_team_perm.return_value = team_object_permission + # 3. Mock the team object lookup (object_permission attached) and prisma_client + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=team_obj, + ) as mock_get_team, + patch.object( + MCPRequestHandler, + "_get_key_object_permission", + new_callable=AsyncMock, + return_value=None, + ), + patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + ): + # 4. Call get_allowed_mcp_servers - this is what MCP routes use + allowed = await MCPRequestHandler.get_allowed_mcp_servers(user_auth) - # Mock key permissions (empty - user has no key-level MCP permissions) - with patch.object( - MCPRequestHandler, "_get_key_object_permission", new_callable=AsyncMock - ) as mock_key_perm: - mock_key_perm.return_value = None + # 5. Verify only team's MCP servers are returned + assert sorted(allowed) == sorted( + team_mcp_servers + ), f"Expected {team_mcp_servers}, got {allowed}" - # Mock access groups (empty) - with patch.object( - MCPRequestHandler, - "_get_mcp_servers_from_access_groups", - new_callable=AsyncMock, - ) as mock_access_groups: - mock_access_groups.return_value = [] - - # 4. Call get_allowed_mcp_servers - this is what MCP routes use - allowed = await MCPRequestHandler.get_allowed_mcp_servers(user_auth) - - # 5. Verify only team's MCP servers are returned - assert sorted(allowed) == sorted( - team_mcp_servers - ), f"Expected {team_mcp_servers}, got {allowed}" - - # Verify team permission was looked up - mock_team_perm.assert_called_once_with(user_auth) + # Verify team was looked up + mock_get_team.assert_called() @pytest.mark.asyncio @@ -120,25 +127,33 @@ async def test_simple_jwt_team_id_required_for_mcp_permissions(): object_permission_id="perm-1", mcp_servers=team_mcp_servers, ) + team_obj = LiteLLM_TeamTable( + team_id="team-abc", + access_group_ids=[], + object_permission_id="perm-1", + ) + team_obj.object_permission = team_perm - with patch.object( - MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock - ) as mock_perm: - mock_perm.return_value = team_perm - - with patch.object( + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=team_obj, + ) as mock_get_team, + patch.object( MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock, - ) as mock_groups: - mock_groups.return_value = [] + return_value=[], + ), + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( + user_with_team + ) - result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( - user_with_team - ) - - assert sorted(result) == sorted(team_mcp_servers) - mock_perm.assert_called_once() # Permission WAS checked + assert sorted(result) == sorted(team_mcp_servers) + mock_get_team.assert_called() # Team WAS looked up # Case 2: team_id is None -> team permissions NOT checked user_without_team = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 116ba83f42e..52e7a0ff373 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1625,6 +1625,50 @@ async def test_reject_clientside_metadata_tags_non_llm_route(): assert result is True +@pytest.mark.asyncio +async def test_reject_clientside_metadata_tags_allows_key_tags_without_client_tags(): + """Key metadata.tags are injected after the reject check; requests without + client metadata.tags must not be blocked when reject_clientside_metadata_tags is on.""" + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "test"}], + } + + general_settings = {"reject_clientside_metadata_tags": True} + mock_request = MagicMock(spec=Request) + valid_token = UserAPIKeyAuth( + token="test-token", + models=["gpt-3.5-turbo"], + metadata={"tags": ["engineering"]}, + ) + + with patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={}, + ): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=mock_request, + ) + + assert result is True + assert request_body["metadata"]["tags"] == ["engineering"] + + @pytest.mark.asyncio async def test_virtual_key_soft_budget_check_with_user_obj(): """Test _virtual_key_soft_budget_check includes user_email when user_obj is provided""" @@ -3370,3 +3414,102 @@ async def test_resolve_end_user_reraises_budget_exceeded( prisma_client=MagicMock(), user_api_key_cache=cache, ) + + +@pytest.mark.asyncio +async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): + """ + Regression pin for LIT-3244 patch/1.86.0 follow-up. + + `_cache_team_object` is the canonical "refresh this team" primitive. + Two cache keys are in play: + - "team_id:" — used by `get_team_object(team_id=...)`, + i.e. API-key auth and JWT-with-team_id_jwt_field + - "team_alias:" — used by `get_team_object_by_alias(team_alias=...)`, + i.e. JWT-with-team_alias_jwt_field + + Invariants this test pins: + 1. Writes the team_id-keyed entry with the refreshed object (team_id + is the table PK — guaranteed unique, safe to write). + 2. DELETES (does NOT write) the team_alias-keyed entry. `team_alias` + has no UNIQUE constraint in schema.prisma, so writing it from + this generic refresh path would let a team admin who renames + their team to collide with another team's alias silently + overwrite the cached team for JWT-by-alias auth (veria-ai + review on #28739). Deleting forces the next JWT-by-alias + reader through `get_team_object_by_alias`, which enforces + len(teams)==1 before populating the cache. + 3. When team_alias is None, NO alias-key operation happens (no + delete of an empty-keyed entry, no spurious write). + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object + + base_team_row = { + "team_id": "team-1234", + "team_alias": "H-Capacity", + "models": ["openai/*", "bedrock-claude-sonnet-4"], + } + + # ===== team_alias is set ===== + team_table = LiteLLM_TeamTableCachedObj(**base_team_row) + cache = MagicMock() + cache.async_set_cache = AsyncMock() + cache.delete_cache = MagicMock() + logging_obj = MagicMock() + logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + + await _cache_team_object( + team_id="team-1234", + team_table=team_table, + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + ) + + # (1) team_id-keyed write fires with the refreshed object + written_keys = [ + (c.kwargs.get("key") or c.args[0]) + for c in cache.async_set_cache.await_args_list + ] + assert written_keys == ["team_id:team-1234"], ( + "Only the team_id-keyed write should fire; the alias key must be " + "deleted, NOT written. " + f"Got writes: {written_keys}" + ) + written_value = ( + cache.async_set_cache.await_args.kwargs.get("value") + or cache.async_set_cache.await_args.args[1] + ) + assert written_value is team_table + + # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache + # and the Redis dual cache (mirrors _delete_cache_key_object pattern). + cache.delete_cache.assert_called_once_with(key="team_alias:H-Capacity") + logging_obj.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( + key="team_alias:H-Capacity" + ) + + # ===== team_alias is None: no alias-key operation ===== + aliasless = LiteLLM_TeamTableCachedObj(**{**base_team_row, "team_alias": None}) + cache2 = MagicMock() + cache2.async_set_cache = AsyncMock() + cache2.delete_cache = MagicMock() + logging_obj2 = MagicMock() + logging_obj2.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + + await _cache_team_object( + team_id="team-no-alias", + team_table=aliasless, + user_api_key_cache=cache2, + proxy_logging_obj=logging_obj2, + ) + + cache2.delete_cache.assert_not_called() + logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_not_awaited() + written_keys_aliasless = [ + (c.kwargs.get("key") or c.args[0]) + for c in cache2.async_set_cache.await_args_list + ] + assert written_keys_aliasless == ["team_id:team-no-alias"] diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 68e1636d380..2d40db9017e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1644,6 +1644,7 @@ class TestObservabilityCallbackBans: "braintrust_api_key", "braintrust_project", "phoenix_project_name", + "phoenix_project_name_override", "wandb_api_key", "weave_project_id", "gcs_bucket_name", @@ -1675,6 +1676,7 @@ class TestObservabilityCallbackBans: "posthog_api_url", "braintrust_project", "phoenix_project_name", + "phoenix_project_name_override", ], ) def test_observability_field_in_metadata_dict_is_rejected( diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 4acf42996e0..b308a665062 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -262,12 +262,166 @@ def test_virtual_key_mcp_routes_allows_v1_mcp_server_subpaths(route): ) def test_mcp_management_routes_classified_as_management_not_llm_api(route): """MCP server CRUD must be management routes, not llm_api routes, so - DISABLE_LLM_API_ENDPOINTS on admin nodes does not block the Admin UI.""" + DISABLE_LLM_API_ENDPOINTS on admin nodes does not block the Admin UI. + + Note: virtual keys with allowed_routes=["llm_api_routes"] can still call + *GET* `/v1/mcp/server` and *GET* `/v1/mcp/server/{server_id}` — that + carve-out is enforced method-aware inside + `is_virtual_key_allowed_to_call_route`, not by adding the paths to + `llm_api_routes`. So `is_llm_api_route()` still returns False here and + `DISABLE_LLM_API_ENDPOINTS` still does not block these paths. + """ assert RouteChecks.is_llm_api_route(route=route) is False assert RouteChecks.is_management_route(route=route) is True +def _mock_request(method: str) -> Request: + request = MagicMock(spec=Request) + request.method = method + return request + + +@pytest.mark.parametrize( + "route", + [ + "/v1/mcp/server", + "/v1/mcp/server/abc-123", + ], +) +def test_virtual_key_llm_api_routes_allows_get_mcp_server_discovery(route): + """ + Regression test: virtual keys with allowed_routes=["llm_api_routes"] must + be able to list/inspect MCP servers via GET /v1/mcp/server[/{server_id}]. + + The handlers strip credential-bearing fields via + `_sanitize_mcp_server_list_for_virtual_key` when the caller is a + restricted virtual key, so GET is safe to expose. The carve-out is + method-aware (see below) — non-GET requests to the same paths are + rejected at this layer, so admin-only writes remain gated. + """ + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + result = RouteChecks.is_virtual_key_allowed_to_call_route( + route=route, + valid_token=valid_token, + request=_mock_request("GET"), + ) + + assert result is True + + +@pytest.mark.parametrize( + "route", + [ + "/v1/mcp/server", + "/v1/mcp/server/abc-123", + ], +) +@pytest.mark.parametrize("method", ["POST", "PUT", "PATCH", "DELETE"]) +def test_virtual_key_llm_api_routes_rejects_non_get_mcp_server_discovery(route, method): + """Method-aware: the MCP server discovery carve-out is GET-only. + + POST/PUT/PATCH/DELETE on `/v1/mcp/server[/{server_id}]` are admin-only + management writes and must not be reachable via llm_api_routes. + """ + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route=route, + valid_token=valid_token, + request=_mock_request(method), + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.parametrize( + "route", + [ + # Multi-segment admin-only sub-paths must NOT be reachable via + # llm_api_routes, even on GET. + "/v1/mcp/server/abc-123/approve", + "/v1/mcp/server/abc-123/reject", + "/v1/mcp/server/oauth/session", + "/v1/mcp/server/abc-123/user-credential", + ], +) +def test_virtual_key_llm_api_routes_rejects_mcp_multi_segment_admin_subpaths( + route, +): + """Multi-segment admin-only MCP sub-paths are not reachable via llm_api_routes. + + The discovery carve-out only matches `/v1/mcp/server` and + `/v1/mcp/server/{server_id}` (single segment after `/server/`), so any + path with additional segments is rejected even when the request is GET. + """ + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route=route, + valid_token=valid_token, + request=_mock_request("GET"), + ) + + assert exc_info.value.status_code == 403 + + +def test_spend_logs_v2_classified_as_management_not_llm_api(): + """Paginated spend logs are a management/spend read route, not an LLM API.""" + + assert RouteChecks.is_llm_api_route(route="/spend/logs/v2") is False + assert RouteChecks.is_management_route(route="/spend/logs/v2") is True + + +def test_virtual_key_management_routes_allows_spend_logs_v2(): + """Management virtual keys should be allowed to call the v2 spend logs endpoint.""" + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["management_routes"], + ) + + result = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/spend/logs/v2", + valid_token=valid_token, + ) + + assert result is True + + +def test_virtual_key_llm_api_routes_denies_spend_logs_v2(): + """AI API virtual keys should not gain spend-log access.""" + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route="/spend/logs/v2", + valid_token=valid_token, + ) + + assert exc_info.value.status_code == 403 + assert "Virtual key is not allowed to call this route" in str(exc_info.value.detail) + + @pytest.mark.parametrize( "route", [ @@ -1322,6 +1476,7 @@ ADMIN_VIEWER_LOGS_PAGE_ROUTES = [ "/cost/estimate", # Public spend logs / spend tracking routes that admin viewer should read "/spend/logs", + "/spend/logs/v2", "/spend/keys", "/spend/users", "/spend/tags", diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index cb30970a34e..d328d68dcd4 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -8,11 +8,14 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.proxy.common_utils.callback_utils import ( + add_policy_to_applied_policies_header, decrypt_callback_vars, encrypt_callback_vars, + get_logging_caching_headers, initialize_callbacks_on_proxy, get_remaining_tokens_and_requests_from_request_data, normalize_callback_names, + sanitize_openai_provider_metadata, ) import litellm @@ -92,6 +95,50 @@ def test_normalize_callback_names_lowercases_strings(): ] +def test_add_policy_to_applied_policies_header_uses_litellm_metadata_bucket(): + request_data = { + "input_file_id": "file-abc123", + "litellm_metadata": {}, + } + + add_policy_to_applied_policies_header( + request_data=request_data, policy_name="global-baseline" + ) + + assert request_data["litellm_metadata"]["applied_policies"] == ["global-baseline"] + assert "applied_policies" not in request_data.get("metadata", {}) + + +def test_sanitize_openai_provider_metadata_strips_internal_tracking_fields(): + metadata = { + "customer_id": "cust-123", + "applied_policies": ["global-baseline"], + "applied_guardrails": ["pii_blocker"], + "note": 42, + } + + sanitized = sanitize_openai_provider_metadata(metadata) + + assert sanitized == {"customer_id": "cust-123"} + + +def test_get_logging_caching_headers_merges_metadata_and_litellm_metadata(): + request_data = { + "metadata": {"customer_id": "cust-123"}, + "litellm_metadata": { + "applied_policies": ["global-baseline"], + "applied_guardrails": ["pii_blocker"], + "policy_sources": {"global-baseline": "team_default"}, + }, + } + + headers = get_logging_caching_headers(request_data) + + assert headers["x-litellm-applied-policies"] == "global-baseline" + assert headers["x-litellm-applied-guardrails"] == "pii_blocker" + assert headers["x-litellm-policy-sources"] == "global-baseline=team_default" + + def test_initialize_callbacks_on_proxy_instantiates_compression_interception( monkeypatch, ): diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 0d7becd3e2e..ce8f0802ae1 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1149,6 +1149,13 @@ async def test_apply_guardrail_not_found(mocker): "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry ) + mock_proxy_logging = mocker.Mock() + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + mocker.patch("litellm.proxy.proxy_server.general_settings", {}) + mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock()) + mocker.patch("litellm.proxy.proxy_server.version", "test") + # Create request request = ApplyGuardrailRequest( guardrail_name="non-existent-guardrail", text="Test input text" @@ -1159,7 +1166,11 @@ async def test_apply_guardrail_not_found(mocker): # Call endpoint and expect ProxyException with pytest.raises(ProxyException) as exc_info: - await apply_guardrail(request=request, user_api_key_dict=mock_user_auth) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=mock_user_auth, + ) # Verify error details assert str(exc_info.value.code) == "404" @@ -1186,6 +1197,25 @@ async def test_apply_guardrail_execution_error(mocker): "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry ) + mock_logging_obj = mocker.Mock() + mock_logging_obj.async_failure_handler = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_processor = mocker.Mock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj) + ) + mocker.patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) + mock_proxy_logging = mocker.Mock() + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + mocker.patch("litellm.proxy.proxy_server.general_settings", {}) + mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock()) + mocker.patch("litellm.proxy.proxy_server.version", "test") + mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor") + # Create request request = ApplyGuardrailRequest( guardrail_name="test-guardrail", text="Test input text with forbidden content" @@ -1196,12 +1226,70 @@ async def test_apply_guardrail_execution_error(mocker): # Call endpoint and expect ProxyException with pytest.raises(ProxyException) as exc_info: - await apply_guardrail(request=request, user_api_key_dict=mock_user_auth) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=mock_user_auth, + ) # Verify error is properly handled assert "Bedrock guardrail failed" in str(exc_info.value.message) +@pytest.mark.asyncio +async def test_apply_guardrail_invokes_logging_pipeline(mocker): + mock_guardrail = mocker.Mock() + mock_guardrail.apply_guardrail = AsyncMock(return_value={"texts": ["masked"]}) + + mock_registry = mocker.Mock() + mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry + ) + + mock_logging_obj = mocker.Mock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_processor = mocker.Mock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj) + ) + mocker.patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) + + mock_proxy_logging = mocker.Mock() + mock_proxy_logging.post_call_success_hook = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + mocker.patch("litellm.proxy.proxy_server.general_settings", {}) + mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock()) + mocker.patch("litellm.proxy.proxy_server.version", "test") + mock_executor = mocker.Mock() + mocker.patch( + "litellm.litellm_core_utils.thread_pool_executor.executor", mock_executor + ) + + request = ApplyGuardrailRequest( + guardrail_name="test-guardrail", text="hello@example.com" + ) + response = await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response.response_text == "masked" + mock_processor.common_processing_pre_call_logic.assert_awaited_once() + mock_proxy_logging.post_call_success_hook.assert_awaited_once() + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_logging_obj.call_type == "pass_through_endpoint" + mock_executor.submit.assert_called_once() + assert mock_logging_obj.async_success_handler.await_args.kwargs["result"] == { + "response": {"response_text": "masked"} + } + + @pytest.mark.asyncio async def test_get_guardrail_info_endpoint_config_guardrail(mocker): """ 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 b65f6305b77..85c7c130b36 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 @@ -1497,6 +1497,305 @@ class TestUpdateDBModelBlocked: assert "blocked" not in result +def _build_db_model_with_pricing(): + """Wildcard deployment with custom pricing in litellm_params; Deployment.__init__ + mirrors SPECIAL_MODEL_INFO_PARAMS into model_info, so both blobs hold the rate.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + input_cost_per_token=0.000001, + output_cost_per_token=0.000002, + ), + model_info=ModelInfo(id="dep-pricing-0"), + ) + + +class TestUpdateDBModelClearPricing: + """Sending an explicit `null` for a pricing field must remove it from both + `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored + between the two by Deployment.__init__). + + Restricted to SPECIAL_MODEL_INFO_PARAMS so non-pricing fields (e.g. team_id) + cannot be cleared via this path. + """ + + def test_clear_input_cost_removes_from_both_blobs(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import updateLiteLLMParams + + result = update_db_model( + db_model=_build_db_model_with_pricing(), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(input_cost_per_token=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + # Other pricing untouched + assert params.get("output_cost_per_token") == 0.000002 + assert info.get("output_cost_per_token") == 0.000002 + + def test_clear_output_cost_removes_from_both_blobs(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import updateLiteLLMParams + + result = update_db_model( + db_model=_build_db_model_with_pricing(), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(output_cost_per_token=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "output_cost_per_token" not in params + assert "output_cost_per_token" not in info + + def test_non_null_pricing_update_still_works(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import updateLiteLLMParams + + result = update_db_model( + db_model=_build_db_model_with_pricing(), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(input_cost_per_token=0.000005) + ), + ) + + params = json.loads(result["litellm_params"]) + assert params["input_cost_per_token"] == 0.000005 + + def test_omitted_pricing_field_is_preserved(self): + """PATCH semantics: fields not in the patch keep their existing value.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import updateLiteLLMParams + + result = update_db_model( + db_model=_build_db_model_with_pricing(), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(output_cost_per_token=0.000007) + ), + ) + + params = json.loads(result["litellm_params"]) + assert params["input_cost_per_token"] == 0.000001 + assert params["output_cost_per_token"] == 0.000007 + + def test_null_on_non_pricing_field_does_not_clear(self): + """Security guard: only SPECIAL_MODEL_INFO_PARAMS can be cleared via null. + Privileged or unrelated model_info fields (e.g. team_id) must be unaffected + by the null-clearing path so a team admin can't ungate a team-scoped model. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateLiteLLMParams, + ) + + db_model = Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + input_cost_per_token=0.000001, + ), + model_info=ModelInfo(id="dep-pricing-1", team_id="team-keep-me"), + ) + + # Patch sends a null for api_base (non-SPECIAL field). Must NOT clear team_id + # or any other non-pricing field from the merged dict. + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(api_base=None) + ), + ) + + info = json.loads(result["model_info"]) + # Pricing still present (not part of this patch) + assert "input_cost_per_token" in info + # team_id must survive + assert info.get("team_id") == "team-keep-me" + + def test_clear_survives_model_info_passthrough_with_old_pricing(self): + """Realistic UI submit shape: the patch carries BOTH blobs. The + model_info portion still has the old pricing because the form + re-serializes the source blob. The litellm_params null must beat the + model_info merge — i.e. the clear runs after both merges, not between. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ModelInfo, updateLiteLLMParams + + result = update_db_model( + db_model=_build_db_model_with_pricing(), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(input_cost_per_token=None), + # The UI passes the OLD model_info blob through unchanged. + model_info=ModelInfo( + id="dep-pricing-0", + input_cost_per_token=0.000001, # stale value from the page state + ), + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "input_cost_per_token" not in params + assert ( + "input_cost_per_token" not in info + ), "model_info passthrough must not resurrect the cleared override" + + def test_clear_via_model_info_clears_both_blobs(self): + """The mirror works in the reverse direction too: nulling a pricing field + via the model_info patch should clear it from litellm_params as well.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ModelInfo + + result = update_db_model( + db_model=_build_db_model_with_pricing(), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-pricing-0", input_cost_per_token=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + + def test_clear_cache_read_cost_removes_from_both_blobs(self): + """cache_read_input_token_cost was added to SPECIAL_MODEL_INFO_PARAMS so + the same null-clear path works for cache-read overrides.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateLiteLLMParams, + ) + + db_model = Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + cache_read_input_token_cost=0.0000005, + ), + model_info=ModelInfo(id="dep-cache-read-0"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(cache_read_input_token_cost=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "cache_read_input_token_cost" not in params + assert "cache_read_input_token_cost" not in info + + def test_clear_cache_write_cost_removes_from_both_blobs(self): + """cache_creation_input_token_cost was added to SPECIAL_MODEL_INFO_PARAMS so + the same null-clear path works for cache-write overrides.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateLiteLLMParams, + ) + + db_model = Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + cache_creation_input_token_cost=0.000003, + ), + model_info=ModelInfo(id="dep-cache-write-0"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(cache_creation_input_token_cost=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "cache_creation_input_token_cost" not in params + assert "cache_creation_input_token_cost" not in info + + def test_clear_cache_read_preserves_other_pricing(self): + """Clearing cache_read must not touch input/output cost overrides.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateLiteLLMParams, + ) + + db_model = Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + input_cost_per_token=0.000001, + output_cost_per_token=0.000002, + cache_read_input_token_cost=0.0000005, + cache_creation_input_token_cost=0.000003, + ), + model_info=ModelInfo(id="dep-cache-mixed-0"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(cache_read_input_token_cost=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "cache_read_input_token_cost" not in params + assert "cache_read_input_token_cost" not in info + # Other pricing untouched in both blobs + assert params["input_cost_per_token"] == 0.000001 + assert params["output_cost_per_token"] == 0.000002 + assert params["cache_creation_input_token_cost"] == 0.000003 + assert info["input_cost_per_token"] == 0.000001 + assert info["output_cost_per_token"] == 0.000002 + assert info["cache_creation_input_token_cost"] == 0.000003 + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 41a5b891ad3..13bb39c35c9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1540,6 +1540,137 @@ def test_add_new_models_to_team_with_existing_models(): assert updated_models.sort() == ["model1", "model2", "model3", "model4"].sort() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "endpoint_name", + ["team_model_add", "team_model_delete"], +) +async def test_team_model_add_delete_refresh_team_cache(endpoint_name): + """ + Regression pin for LIT-3244 vector-store BYOK 403. + + `team_model_add` and `team_model_delete` mutate `team.models` in the + DB. Without a cache refresh, the in-memory `LiteLLM_TeamTableCachedObj` + used by `common_checks` stays stale and team members 403 on a model + the DB has just granted (or, symmetrically, keep using a model the DB + has just revoked). + + Pin: after the DB update, the endpoint must call `_cache_team_object` + with the updated team row so the cached team stays in sync. + """ + from unittest.mock import AsyncMock, MagicMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import ( + LitellmUserRoles, + TeamModelAddRequest, + TeamModelDeleteRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import ( + team_model_add, + team_model_delete, + ) + + mock_request = Mock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" + ) + + existing_team = MagicMock() + existing_team.model_dump.return_value = { + "team_id": "team-1234", + "models": ["bedrock-claude-sonnet-4", "openai/*"], + "object_permission_id": "op-1234", + "object_permission": { + "object_permission_id": "op-1234", + "search_tools": ["allowed-tool-A"], + }, + } + + updated_team = MagicMock() + updated_team.team_id = "team-1234" + updated_team.model_dump.return_value = { + "team_id": "team-1234", + "models": ["bedrock-claude-sonnet-4", "openai/*", "team-byok-1"], + # The Prisma update must come back with `object_permission` populated + # (via `include={"object_permission": True}`), otherwise the cache + # write below would null it out — see LIT-3244 follow-up. + "object_permission_id": "op-1234", + "object_permission": { + "object_permission_id": "op-1234", + "search_tools": ["allowed-tool-A"], + }, + } + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, + ): + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=updated_team + ) + mock_cache_team.return_value = None + + if endpoint_name == "team_model_add": + await team_model_add( + data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]), + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + else: + await team_model_delete( + data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]), + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + + # The pin: cache refresh must run with the updated team row. + assert mock_cache_team.await_count == 1, ( + f"{endpoint_name} must call _cache_team_object exactly once " + f"after the DB update (LIT-3244 regression pin); " + f"got await_count={mock_cache_team.await_count}" + ) + call_kwargs = mock_cache_team.await_args.kwargs + assert call_kwargs["team_id"] == "team-1234" + # The cached object must be built from the *updated* row, not the + # pre-mutation `existing_team` — that's the whole point. Both rows + # share team_id, so the only assertion that actually pins this is + # against the field that differs between them: `models`. + assert call_kwargs["team_table"].team_id == "team-1234" + assert call_kwargs["team_table"].models == [ + "bedrock-claude-sonnet-4", + "openai/*", + "team-byok-1", + ] + # And the cached object MUST carry the `object_permission` relation + # (LIT-3244 follow-up). If the Prisma update were missing + # `include={"object_permission": True}`, the cached team would have + # object_permission=None, and downstream consumers like + # `validate_key_search_tools_against_team` would treat that as + # "no team-level restriction" and stop enforcing the team's + # search-tool allowlist on key issuance. + assert call_kwargs["team_table"].object_permission is not None + assert call_kwargs["team_table"].object_permission.search_tools == [ + "allowed-tool-A" + ] + # Pin the Prisma call shape too — the regression is in *what the + # update returns*, so the contract that the update asks for + # `object_permission` belongs in this test. + update_call_kwargs = ( + mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs + ) + assert update_call_kwargs.get("include", {}).get("object_permission") is True + + @pytest.mark.asyncio async def test_update_team_team_member_budget_not_passed_to_db(): """ @@ -1568,7 +1699,9 @@ async def test_update_team_team_member_budget_not_passed_to_db(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget, @@ -1999,7 +2132,9 @@ async def test_update_team_with_team_member_budget_duration(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 97a21136198..344742ffe89 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -20,6 +20,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, pass_through_request, ) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, +) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -2153,7 +2156,12 @@ async def test_create_pass_through_route_custom_body_url_target(): endpoint_func = create_pass_through_route( endpoint=unique_path, target="https://bedrock-agent-runtime.us-east-1.amazonaws.com", - custom_headers={"Content-Type": "application/json"}, + custom_headers=Headers( + { + "Authorization": "AWS4-HMAC-SHA256 signed", + "Content-Type": "application/json", + } + ), _forward_headers=True, ) @@ -2213,6 +2221,147 @@ async def test_create_pass_through_route_custom_body_url_target(): # The critical assertion: custom_body takes precedence over # the body parsed from the raw request assert call_kwargs["custom_body"] == bedrock_body + # HeadersDict-like custom_headers (e.g. botocore SigV4) must be coerced + # to a plain dict so signed headers actually reach the upstream. + assert call_kwargs["custom_headers"] == { + "authorization": "AWS4-HMAC-SHA256 signed", + "content-type": "application/json", + } + + +@pytest.mark.asyncio +async def test_pass_through_request_non_streaming_uses_content_for_state_raw_body(): + """ + Bedrock SigV4 path: exact signed bytes live on request.state; upstream must receive + content=... even if pre_call_hook mutates the parsed dict (would change json=). + """ + # Bytes that were signed (simulated); parsed body + hook will diverge on purpose. + raw_signed = b'{"retrievalQuery":{"text":"signed"},"sig":"intact"}' + parsed_from_wire = {"retrievalQuery": {"text": "signed"}, "sig": "intact"} + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.query_params = QueryParams({}) + mock_request.headers = Headers({"Content-Type": "application/json"}) + mock_request.state = SimpleNamespace() + setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed) + mock_request.body = AsyncMock( + return_value=json.dumps(parsed_from_wire).encode("utf-8") + ) + + mock_user = MagicMock() + mock_user.api_key = "sk-test" + + upstream = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"ok": true}', + request=httpx.Request( + "POST", + "https://bedrock-agent-runtime.us-east-1.amazonaws.com/knowledgebases/KB/retrieve", + ), + ) + + mock_async_client = AsyncMock() + mock_async_client.request = AsyncMock(return_value=upstream) + mock_client_obj = MagicMock() + mock_client_obj.client = mock_async_client + + async def _hook_mutates_body(**kwargs): + data = kwargs["data"] + if isinstance(data, dict): + data["hook_mutated"] = True + return data + + with ( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client", + return_value=mock_client_obj, + ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook", + new=AsyncMock(side_effect=_hook_mutates_body), + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ), + ): + await pass_through_request( + request=mock_request, + target="https://bedrock-agent-runtime.us-east-1.amazonaws.com/knowledgebases/KB/retrieve", + custom_headers={"content-type": "application/json"}, + user_api_key_dict=mock_user, + stream=False, + ) + + mock_async_client.request.assert_called_once() + req_kw = mock_async_client.request.call_args[1] + assert req_kw.get("content") == raw_signed + assert "json" not in req_kw + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_uses_content_for_state_raw_body(): + """Streaming pass-through with state raw body must use build_request(..., content=...).""" + raw_signed = b'{"model":"m","stream":true}' + parsed_from_wire = {"model": "m", "stream": True} + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.query_params = QueryParams({}) + mock_request.headers = Headers({"Content-Type": "application/json"}) + mock_request.state = SimpleNamespace() + setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed) + mock_request.body = AsyncMock( + return_value=json.dumps(parsed_from_wire).encode("utf-8") + ) + + mock_user = MagicMock() + mock_user.api_key = "sk-test" + + mock_built = MagicMock() + mock_async_client = AsyncMock() + mock_async_client.build_request = MagicMock(return_value=mock_built) + stream_resp = httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + content=b"data: {}\n\n", + request=httpx.Request("POST", "https://example.com/v1/messages"), + ) + mock_async_client.send = AsyncMock(return_value=stream_resp) + mock_client_obj = MagicMock() + mock_client_obj.client = mock_async_client + + with ( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client", + return_value=mock_client_obj, + ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook", + new=AsyncMock(side_effect=lambda **kw: kw["data"]), + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ), + ): + response = await pass_through_request( + request=mock_request, + target="https://example.com/v1/messages", + custom_headers={"Authorization": "Bearer x"}, + user_api_key_dict=mock_user, + stream=None, + ) + + from fastapi.responses import StreamingResponse + + assert isinstance(response, StreamingResponse) + mock_async_client.build_request.assert_called_once() + br_kw = mock_async_client.build_request.call_args[1] + assert br_kw.get("content") == raw_signed + assert "json" not in br_kw @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 7176cf455c8..aaf1dad4910 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -538,6 +538,36 @@ def test_forward_headers_from_request_protected_headers_not_overwritten(): assert "Anthropic-Beta" not in result +def test_forward_headers_custom_wins_case_insensitive_over_request_authorization(): + """ + When forwarding request headers, provider-signed/custom headers must win + even if the incoming request uses a different case for the same header name. + """ + from litellm.passthrough.utils import BasePassthroughUtils + + request_headers = { + "authorization": "Bearer sk-litellm-key", + "content-type": "application/json", + "x-request-id": "req-123", + } + signed_headers = { + "Authorization": "AWS4-HMAC-SHA256 signed", + "Content-Type": "application/json", + } + + result = BasePassthroughUtils.forward_headers_from_request( + request_headers=request_headers, + headers=signed_headers.copy(), + forward_headers=True, + ) + + assert result["Authorization"] == "AWS4-HMAC-SHA256 signed" + assert "authorization" not in result + assert result["Content-Type"] == "application/json" + assert "content-type" not in result + assert result["x-request-id"] == "req-123" + + @pytest.mark.asyncio async def test_vertex_passthrough_custom_model_name_replaced_in_url(): """ diff --git a/tests/test_litellm/proxy/proxy_server/.coverage_baseline b/tests/test_litellm/proxy/proxy_server/.coverage_baseline new file mode 100644 index 00000000000..287ff5be9f5 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/.coverage_baseline @@ -0,0 +1 @@ +line:0.0 branch:0.0 diff --git a/tests/test_litellm/proxy/proxy_server/__init__.py b/tests/test_litellm/proxy/proxy_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/proxy_server/_coverage_check.py b/tests/test_litellm/proxy/proxy_server/_coverage_check.py new file mode 100644 index 00000000000..5db1045eca4 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/_coverage_check.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +"""Coverage gate for the proxy_server.py behavior-pinning project. + +Reads a coverage XML report (produced by ``pytest --cov-branch +--cov-report=xml:``) and asserts that line + branch coverage on +``litellm/proxy/proxy_server.py`` meets the per-PR target. + +Target selection: + --pr-target {1|2|3} explicit target + (none) self-selected by inspecting which placeholder + test files have been filled (PR1 fills before + PR2, PR2 before PR3). With nothing filled, the + target is "PR0" (baseline, no minimum). + +Exits 0 on PASS, non-zero on FAIL. +""" + +from __future__ import annotations + +import argparse +import ast +import sys +import xml.etree.ElementTree as ET +from pathlib import Path +from typing import Dict, List, Tuple + +HERE = Path(__file__).resolve().parent +SOURCE_FILE = "litellm/proxy/proxy_server.py" + +# PR target gates: (line%, branch%) +TARGETS: Dict[str, Tuple[float, float]] = { + "PR0": (0.0, 0.0), + "PR1": (25.0, 18.0), + "PR2": (50.0, 38.0), + "PR3": (70.0, 55.0), +} + +# Which placeholder files each PR is expected to fill (see Notion plan). +PR1_FILES: List[str] = [ + "test_lifecycle.py", + "test_proxy_config.py", + "test_spend_counters.py", + "test_background_health.py", + "test_openapi_customization.py", + "test_exception_handlers.py", + "test_streaming_helpers.py", +] +PR2_FILES: List[str] = [ + "test_routes_models.py", + "test_routes_chat_completions.py", + "test_routes_completions.py", + "test_routes_embeddings.py", + "test_routes_moderations.py", + "test_routes_audio.py", + "test_routes_assistants.py", + "test_routes_threads.py", + "test_routes_utils.py", + "test_routes_model_info.py", + "test_routes_model_metrics.py", + "test_routes_queue.py", +] +PR3_FILES: List[str] = [ + "test_routes_login_sso.py", + "test_routes_onboarding.py", + "test_routes_invitation.py", + "test_routes_config.py", + "test_routes_model_cost_map.py", + "test_routes_anthropic_beta.py", + "test_routes_misc.py", +] + + +def file_has_tests(path: Path) -> bool: + """A test file is considered filled if it defines at least one ``test_*``.""" + if not path.is_file(): + return False + try: + tree = ast.parse(path.read_text()) + except SyntaxError: + return False + for node in ast.walk(tree): + if isinstance( + node, (ast.FunctionDef, ast.AsyncFunctionDef) + ) and node.name.startswith("test_"): + return True + return False + + +def detect_pr_target(dir_path: Path) -> str: + """Pick the strictest PR whose files are fully filled in this directory.""" + pr3_filled = all(file_has_tests(dir_path / f) for f in PR3_FILES) + pr2_filled = all(file_has_tests(dir_path / f) for f in PR2_FILES) + pr1_filled = all(file_has_tests(dir_path / f) for f in PR1_FILES) + if pr3_filled and pr2_filled and pr1_filled: + return "PR3" + if pr2_filled and pr1_filled: + return "PR2" + if pr1_filled: + return "PR1" + return "PR0" + + +def parse_coverage_xml(xml_path: Path) -> Tuple[float, float]: + """Extract (line%, branch%) for proxy_server.py from a coverage XML report. + + Returns (0.0, 0.0) if the file isn't found in the report. + """ + if not xml_path.is_file(): + raise FileNotFoundError(f"Coverage XML not found at {xml_path}") + tree = ET.parse(xml_path) + root = tree.getroot() + for class_elem in root.iter("class"): + filename = class_elem.get("filename", "") + # Coverage tools emit either a repo-relative path or just the basename + # depending on configuration. Match by suffix. + if filename.endswith("proxy/proxy_server.py") or filename.endswith( + "proxy_server.py" + ): + line_rate = float(class_elem.get("line-rate", "0")) + branch_rate = float(class_elem.get("branch-rate", "0")) + return line_rate * 100.0, branch_rate * 100.0 + return 0.0, 0.0 + + +def parse_baseline(baseline_path: Path) -> Tuple[float, float]: + """Parse ``line: branch:`` baseline; missing file -> (0, 0).""" + if not baseline_path.is_file(): + return 0.0, 0.0 + line_pct = 0.0 + branch_pct = 0.0 + for token in baseline_path.read_text().split(): + if ":" not in token: + continue + key, _, value = token.partition(":") + try: + num = float(value) + except ValueError: + continue + if key == "line": + line_pct = num + elif key == "branch": + branch_pct = num + return line_pct, branch_pct + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--pr-target", + choices=["1", "2", "3"], + default=None, + help="Explicit PR target (1, 2, or 3). If omitted, self-selected.", + ) + parser.add_argument( + "--coverage-xml", + default=str(HERE.parent.parent.parent.parent / ".cov_new.xml"), + help="Path to coverage XML (default: /.cov_new.xml)", + ) + args = parser.parse_args() + + if args.pr_target: + target = f"PR{args.pr_target}" + else: + target = detect_pr_target(HERE) + target_line, target_branch = TARGETS[target] + + # The effective floor is the max of the PR target and the committed + # baseline. The baseline is updated as each PR lands so a future + # regression (e.g. a test deletion) trips this gate even if the + # static PR target is already met. + baseline_line, baseline_branch = parse_baseline(HERE / ".coverage_baseline") + line_min = max(target_line, baseline_line) + branch_min = max(target_branch, baseline_branch) + + xml_path = Path(args.coverage_xml) + try: + line_pct, branch_pct = parse_coverage_xml(xml_path) + except FileNotFoundError as exc: + print(f"FAIL: {exc}", file=sys.stderr) + return 2 + + line_ok = line_pct >= line_min + branch_ok = branch_pct >= branch_min + status = "PASS" if (line_ok and branch_ok) else "FAIL" + + print( + f"target={target} baseline=(line:{baseline_line:.2f} branch:{baseline_branch:.2f})" + ) + print( + f"line: {line_pct:6.2f}% / {line_min:6.2f}% " f"{'OK' if line_ok else 'MISS'}" + ) + print( + f"branch: {branch_pct:6.2f}% / {branch_min:6.2f}% " + f"{'OK' if branch_ok else 'MISS'}" + ) + print(status) + return 0 if status == "PASS" else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_litellm/proxy/proxy_server/_pin_check.py b/tests/test_litellm/proxy/proxy_server/_pin_check.py new file mode 100644 index 00000000000..3a3cdfccac7 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/_pin_check.py @@ -0,0 +1,249 @@ +#!/usr/bin/env python3 +"""Pin-list gate for the proxy_server.py behavior-pinning project. + +For each identifier in a pin list, asserts that the test directory contains: + 1. At least one happy-path test that references the identifier and uses + a real assertion (normalize(response.json()) == {...}, .model_validate, + or a dict-equality with >= 3 keys). + 2. At least one error-path test (name hints at error OR asserts a 4xx/5xx + status OR uses pytest.raises). + 3. No test that is "status-only" (its sole assert is on response.status_code). + +``test_harness_smoke.py`` is ignored (harness self-tests don't count toward +behavior pinning). + +Exits 0 on PASS, non-zero on FAIL. +""" + +from __future__ import annotations + +import argparse +import ast +import re +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Set, Tuple + +HERE = Path(__file__).resolve().parent + +PIN_LINE_RE = re.compile(r"^- `([^`]+)`\s*$") +ERROR_NAME_HINTS = ( + "error", + "fail", + "invalid", + "unauthorized", + "forbidden", + "missing", + "denied", + "rejected", + "bad", + "raises", + "exception", + "404", + "401", + "403", + "422", + "500", +) +ERROR_STATUS_CODES = frozenset({400, 401, 402, 403, 404, 405, 409, 422, 500, 502, 503}) + + +@dataclass +class TestFunction: + name: str + file: Path + source: str + asserts: List[ast.Assert] = field(default_factory=list) + raises_calls: int = 0 + status_code_asserts: List[int] = field(default_factory=list) + has_strong_assertion: bool = ( + False # normalize() or .model_validate() or large dict-eq + ) + + +def parse_pin_list(path: Path) -> List[str]: + items: List[str] = [] + for line in path.read_text().splitlines(): + m = PIN_LINE_RE.match(line) + if m: + items.append(m.group(1).strip()) + return items + + +def _has_strong_assertion(node: ast.AST) -> bool: + """True if an assert subtree contains normalize(), .model_validate(), or dict-eq with >=3 keys.""" + for sub in ast.walk(node): + if isinstance(sub, ast.Call): + func = sub.func + if isinstance(func, ast.Name) and func.id == "normalize": + return True + if isinstance(func, ast.Attribute) and func.attr == "model_validate": + return True + if ( + isinstance(sub, ast.Compare) + and len(sub.ops) == 1 + and isinstance(sub.ops[0], ast.Eq) + ): + # response.json() == {= 3 keys>} + rhs = sub.comparators[0] + if isinstance(rhs, ast.Dict) and len(rhs.keys) >= 3: + return True + return False + + +def _extract_status_code(node: ast.Assert) -> Optional[int]: + """If this assert is exactly ``X.status_code == ``, return the int.""" + test = node.test + if not isinstance(test, ast.Compare): + return None + if len(test.ops) != 1 or not isinstance(test.ops[0], ast.Eq): + return None + left = test.left + if not (isinstance(left, ast.Attribute) and left.attr == "status_code"): + return None + right = test.comparators[0] + if isinstance(right, ast.Constant) and isinstance(right.value, int): + return right.value + return None + + +def collect_test_functions(test_dir: Path) -> List[TestFunction]: + funcs: List[TestFunction] = [] + for path in sorted(test_dir.glob("test_*.py")): + # Skip the harness's own smoke tests — they don't count toward + # behavior pinning. + if path.name == "test_harness_smoke.py": + continue + source = path.read_text() + try: + tree = ast.parse(source) + except SyntaxError: + continue + for node in ast.walk(tree): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + if not node.name.startswith("test_"): + continue + tf = TestFunction(name=node.name, file=path, source=source) + for sub in ast.walk(node): + if isinstance(sub, ast.Assert): + tf.asserts.append(sub) + sc = _extract_status_code(sub) + if sc is not None: + tf.status_code_asserts.append(sc) + if _has_strong_assertion(sub): + tf.has_strong_assertion = True + if isinstance(sub, ast.With): + for item in sub.items: + ctx = item.context_expr + if isinstance(ctx, ast.Call) and isinstance( + ctx.func, ast.Attribute + ): + if ctx.func.attr == "raises": + tf.raises_calls += 1 + funcs.append(tf) + return funcs + + +def _is_status_only(tf: TestFunction) -> bool: + """A test that has >=1 status_code assert and ALL its asserts are status_code.""" + return len(tf.asserts) >= 1 and len(tf.status_code_asserts) == len(tf.asserts) + + +def _looks_like_error_test(tf: TestFunction) -> bool: + name_lower = tf.name.lower() + if any(hint in name_lower for hint in ERROR_NAME_HINTS): + return True + if tf.raises_calls > 0: + return True + if any(sc in ERROR_STATUS_CODES for sc in tf.status_code_asserts): + return True + return False + + +def _references_pin(tf: TestFunction, pin: str) -> bool: + """Cheap string-contains check against the test function's source. + + This is intentionally permissive — if the pin identifier (e.g. + ``update_cache`` or ``POST /chat/completions``) appears anywhere in + the test file we count it. Aliased route paths or parametrize + cases trigger the same reference. + """ + return pin in tf.source + + +def check(pin_list: List[str], funcs: List[TestFunction]) -> Tuple[bool, List[str]]: + failures: List[str] = [] + + status_only = [tf for tf in funcs if _is_status_only(tf)] + for tf in status_only: + failures.append( + f"status-only test (only asserts response.status_code): " + f"{tf.file.name}::{tf.name}" + ) + + by_pin: Dict[str, List[TestFunction]] = {pin: [] for pin in pin_list} + for tf in funcs: + for pin in pin_list: + if _references_pin(tf, pin): + by_pin[pin].append(tf) + + for pin, matches in by_pin.items(): + if not matches: + failures.append(f"no tests reference pin: {pin}") + continue + has_happy = any( + tf.has_strong_assertion and not _looks_like_error_test(tf) for tf in matches + ) + has_error = any(_looks_like_error_test(tf) for tf in matches) + if not has_happy: + failures.append( + f"no happy-path test with strong assertion (normalize/model_validate/dict-eq>=3) " + f"for pin: {pin}" + ) + if not has_error: + failures.append(f"no error-path test for pin: {pin}") + + return (not failures), failures + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--list", + required=True, + help="Path to pin list file (markdown bullets in `- ` + backtick + symbol + backtick format)", + ) + parser.add_argument( + "--test-dir", + default=str(HERE), + help="Test directory to scan (default: this directory)", + ) + args = parser.parse_args() + + pin_path = Path(args.list) + if not pin_path.is_file(): + print(f"FAIL: pin list not found at {pin_path}", file=sys.stderr) + return 2 + + pin_list = parse_pin_list(pin_path) + if not pin_list: + print(f"FAIL: pin list at {pin_path} contained zero items", file=sys.stderr) + return 2 + + test_dir = Path(args.test_dir) + funcs = collect_test_functions(test_dir) + + ok, failures = check(pin_list, funcs) + print(f"pins: {len(pin_list)}") + print(f"tests: {len(funcs)}") + if failures: + for f in failures: + print(f" - {f}") + print("PASS" if ok else "FAIL") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_litellm/proxy/proxy_server/conftest.py b/tests/test_litellm/proxy/proxy_server/conftest.py new file mode 100644 index 00000000000..c545965f9a9 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/conftest.py @@ -0,0 +1,513 @@ +"""Shared fixtures for tests/test_litellm/proxy/proxy_server/. + +All fixtures and helpers used by PR1/PR2/PR3 test files live here. Do NOT +add fixtures inside individual test files. If a fixture is missing, add it +here and update the Notion plan. +""" + +from __future__ import annotations + +import contextlib +import os +import sys +from pathlib import Path +from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest + +# Repo root, anchored to this file (not CWD) so the path is correct no +# matter where pytest is invoked from. With the project installed via +# uv this is defensive — `litellm` already resolves through site-packages +# — but it lets the harness work in editable-source layouts too. +sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + + +# --------------------------------------------------------------------------- +# normalize() — used by every dict-equality assertion to scrub volatile fields +# --------------------------------------------------------------------------- + +VOLATILE_KEYS = frozenset( + { + "created_at", + "updated_at", + "key", + "token", + "id", + "request_id", + "expires", + "expires_at", + "litellm_call_id", + "key_alias", + "created", + } +) + + +def normalize(data: Any, volatile: frozenset[str] = VOLATILE_KEYS) -> Any: + """Replace volatile field values with "" so dict equality works. + + Recursive over dicts and lists. Pass an explicit ``volatile`` set to + extend or override the default. + """ + if isinstance(data, dict): + return { + k: ("" if k in volatile else normalize(v, volatile)) + for k, v in data.items() + } + if isinstance(data, list): + return [normalize(v, volatile) for v in data] + return data + + +# --------------------------------------------------------------------------- +# app + client — session-scoped so app import + TestClient setup amortize +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="session") +def app(): + """Return the proxy_server FastAPI app with lifespan effectively disabled. + + TestClient used WITHOUT the ``with`` context manager skips the lifespan, + so the startup event (DB connect, Router init, OTEL setup) never fires. + Module import still runs once; module-level globals are harmless. + """ + os.environ.setdefault("LITELLM_LOG", "ERROR") + from litellm.proxy.proxy_server import app as _app + + return _app + + +@pytest.fixture(scope="session") +def client(app): + """TestClient wrapping the session app. + + NOT entered as a context manager — lifespan does not fire. Tests that + require a real lifespan should use a function-scoped TestClient with + a ``with`` block locally and accept the per-test cost. + """ + from fastapi.testclient import TestClient + + return TestClient(app, raise_server_exceptions=False) + + +# --------------------------------------------------------------------------- +# mock_prisma — function-scoped MagicMock with the common table methods stubbed +# --------------------------------------------------------------------------- + +# Tables most-touched by proxy_server.py routes. Add to this list if a +# test discovers a missing table. +_PRISMA_TABLES: List[str] = [ + "litellm_verificationtoken", + "litellm_teamtable", + "litellm_usertable", + "litellm_endusertable", + "litellm_organizationtable", + "litellm_organizationmembership", + "litellm_proxymodeltable", + "litellm_modeltable", + "litellm_budgettable", + "litellm_spendlogs", + "litellm_invitationlink", + "litellm_credentialstable", + "litellm_mcpservertable", + "litellm_objectpermissiontable", + "litellm_configtable", + "litellm_audit_log", + "litellm_dailyuserspend", + "litellm_dailyteamspend", + "litellm_dailytagspend", + "litellm_managed_object_table", + "litellm_managed_vector_stores_table", + "litellm_promptstable", + "litellm_guardrailstable", + "litellm_managed_files", + "litellm_session_token_table", + "litellm_passthrough_endpoint_table", + "litellm_cron_job", + "litellm_passthrough_logs", + "litellm_health_check_table", + "litellm_mcpusercredentials", +] + + +def _make_table_mock() -> MagicMock: + table = MagicMock() + table.find_unique = AsyncMock(return_value=None) + table.find_many = AsyncMock(return_value=[]) + table.find_first = AsyncMock(return_value=None) + table.create = AsyncMock() + table.create_many = AsyncMock() + table.update = AsyncMock() + table.update_many = AsyncMock() + table.upsert = AsyncMock() + table.delete = AsyncMock() + table.delete_many = AsyncMock() + table.count = AsyncMock(return_value=0) + table.group_by = AsyncMock(return_value=[]) + table.aggregate = AsyncMock(return_value={}) + return table + + +@pytest.fixture +def mock_prisma() -> MagicMock: + """MagicMock prisma_client with .db. methods stubbed. + + Default returns: find_unique/find_first -> None, find_many/group_by -> [], + count -> 0. Override in a test with:: + + mock_prisma.db.litellm_teamtable.find_unique.return_value = ... + """ + client_mock = MagicMock() + client_mock.db = MagicMock() + client_mock.connect = AsyncMock() + client_mock.disconnect = AsyncMock() + client_mock.health_check = AsyncMock(return_value=True) + for table_name in _PRISMA_TABLES: + setattr(client_mock.db, table_name, _make_table_mock()) + return client_mock + + +# --------------------------------------------------------------------------- +# auth_as — context manager that overrides user_api_key_auth dependency +# --------------------------------------------------------------------------- + + +@pytest.fixture +def auth_as(app) -> Callable[..., contextlib.AbstractContextManager]: + """Context manager that overrides ``user_api_key_auth`` for a role. + + Usage:: + + def test_admin_only(client, auth_as): + from litellm.proxy._types import LitellmUserRoles + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/some/admin/route") + assert response.status_code == 200 + + Outside the ``with`` block the override is removed so other tests see + the real dependency. + """ + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + @contextlib.contextmanager + def _auth_as( + role: Any = None, + user_id: str = "test-user-id", + team_id: Optional[str] = None, + api_key: str = "sk-test-key", + **kwargs: Any, + ) -> Iterator[Any]: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + if role is None: + role = LitellmUserRoles.PROXY_ADMIN + + fake_auth = UserAPIKeyAuth( + api_key=api_key, + user_id=user_id, + team_id=team_id, + user_role=role, + **kwargs, + ) + + async def _override() -> UserAPIKeyAuth: + return fake_auth + + previous = app.dependency_overrides.get(user_api_key_auth) + app.dependency_overrides[user_api_key_auth] = _override + try: + yield fake_auth + finally: + if previous is None: + app.dependency_overrides.pop(user_api_key_auth, None) + else: + app.dependency_overrides[user_api_key_auth] = previous + + return _auth_as + + +# --------------------------------------------------------------------------- +# Response builders — used by mock_router for parametrized responses +# --------------------------------------------------------------------------- + + +def make_acompletion_response( + model: str = "gpt-4", + messages: Optional[List[Dict[str, Any]]] = None, + stream: bool = False, + tools: Optional[List[Dict[str, Any]]] = None, + content: str = "Hello from mock", + **kwargs: Any, +) -> Any: + """Build a deterministic chat-completion response. + + Returns: + - An async generator when ``stream=True`` + - A tool-call shape when ``tools`` is non-empty + - A plain text response otherwise + """ + from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Function, + Message, + ModelResponse, + Usage, + ) + + if stream: + return _stream_chunks(model=model, content=content) + + if tools: + tool_name = tools[0].get("function", {}).get("name", "fake_tool") + message = Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_test", + type="function", + function=Function(name=tool_name, arguments="{}"), + ) + ], + ) + else: + message = Message(role="assistant", content=content) + + return ModelResponse( + id="chatcmpl-test", + choices=[Choices(finish_reason="stop", index=0, message=message)], + created=0, + model=model, + object="chat.completion", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + +async def _stream_chunks( + model: str = "gpt-4", content: str = "Hi" +) -> AsyncIterator[Any]: + from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + ) + + for piece in [content, ""]: + yield ModelResponseStream( + id="chatcmpl-test", + choices=[ + StreamingChoices( + finish_reason=None if piece else "stop", + index=0, + delta=Delta(content=piece or None, role="assistant"), + ) + ], + created=0, + model=model, + object="chat.completion.chunk", + ) + + +def make_embedding_response( + model: str = "text-embedding-ada-002", + input: Any = None, + dimensions: int = 8, + **kwargs: Any, +) -> Any: + from litellm.types.utils import EmbeddingResponse + + if isinstance(input, list): + n = len(input) + elif input is None: + n = 1 + else: + n = 1 + return EmbeddingResponse( + model=model, + data=[ + {"embedding": [0.0] * dimensions, "index": i, "object": "embedding"} + for i in range(n) + ], + object="list", + usage={"prompt_tokens": n, "total_tokens": n}, + ) + + +def make_image_response(model: str = "dall-e-3", **kwargs: Any) -> Any: + from litellm.types.utils import ImageResponse + + return ImageResponse( + created=0, + data=[{"url": "https://example.invalid/image.png"}], + ) + + +def make_speech_response(**kwargs: Any) -> bytes: + """Return a fake audio blob. The route serializes bytes to a streaming response.""" + return b"\x00" * 128 + + +def make_transcription_response(**kwargs: Any) -> Any: + from litellm.types.utils import TranscriptionResponse + + return TranscriptionResponse(text="hello world") + + +def make_moderation_response(**kwargs: Any) -> Dict[str, Any]: + return { + "id": "modr-test", + "model": "text-moderation-latest", + "results": [ + { + "flagged": False, + "categories": {}, + "category_scores": {}, + } + ], + } + + +# --------------------------------------------------------------------------- +# mock_router — fake Router with all the *async* call surfaces stubbed +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_router() -> MagicMock: + """A MagicMock standing in for ``llm_router`` with parametrized responses.""" + + async def _acompletion(model: str = "gpt-4", messages=None, **kwargs): + return make_acompletion_response(model=model, messages=messages, **kwargs) + + async def _aembedding(model: str = "text-embedding-ada-002", input=None, **kwargs): + return make_embedding_response(model=model, input=input, **kwargs) + + async def _aimage_generation(**kwargs): + return make_image_response(**kwargs) + + async def _aspeech(**kwargs): + return make_speech_response(**kwargs) + + async def _atranscription(**kwargs): + return make_transcription_response(**kwargs) + + async def _amoderation(**kwargs): + return make_moderation_response(**kwargs) + + router = MagicMock() + router.acompletion = AsyncMock(side_effect=_acompletion) + router.aembedding = AsyncMock(side_effect=_aembedding) + router.aimage_generation = AsyncMock(side_effect=_aimage_generation) + router.aspeech = AsyncMock(side_effect=_aspeech) + router.atranscription = AsyncMock(side_effect=_atranscription) + router.amoderation = AsyncMock(side_effect=_amoderation) + router.model_list = [ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}, + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-latest"}, + }, + { + "model_name": "bedrock-claude", + "litellm_params": {"model": "bedrock/anthropic.claude-3-5-sonnet"}, + }, + ] + router.model_names = ["gpt-4", "claude-sonnet", "bedrock-claude"] + router.get_model_list = MagicMock(return_value=router.model_list) + return router + + +# --------------------------------------------------------------------------- +# mock_callbacks_disabled — autouse: zero out global callbacks per test +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def mock_callbacks_disabled(monkeypatch) -> None: + """Wipe ``litellm.callbacks`` and friends so tests don't leak side effects.""" + import litellm + + for attr in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + "input_callback", + "service_callback", + ): + if hasattr(litellm, attr): + monkeypatch.setattr(litellm, attr, [], raising=False) + + +# --------------------------------------------------------------------------- +# Builders for DB-like objects (used by routes that load from DB) +# --------------------------------------------------------------------------- + + +def make_user( + user_id: str = "user-test", + role: Any = None, + teams: Optional[List[str]] = None, + max_budget: Optional[float] = None, + spend: float = 0.0, + **kwargs: Any, +) -> Any: + from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles + + if role is None: + role = LitellmUserRoles.INTERNAL_USER + + return LiteLLM_UserTable( + user_id=user_id, + user_role=role, + teams=teams or [], + max_budget=max_budget, + spend=spend, + **kwargs, + ) + + +def make_team( + team_id: str = "team-test", + team_alias: str = "Test Team", + max_budget: Optional[float] = None, + spend: float = 0.0, + members_with_roles: Optional[List[Dict[str, Any]]] = None, + **kwargs: Any, +) -> Any: + from litellm.proxy._types import LiteLLM_TeamTable + + return LiteLLM_TeamTable( + team_id=team_id, + team_alias=team_alias, + max_budget=max_budget, + spend=spend, + members_with_roles=members_with_roles or [], + **kwargs, + ) + + +def make_key( + token: str = "hashed-test-key", + key_alias: Optional[str] = None, + team_id: Optional[str] = None, + user_id: str = "user-test", + spend: float = 0.0, + max_budget: Optional[float] = None, + **kwargs: Any, +) -> Any: + from litellm.proxy._types import LiteLLM_VerificationToken + + return LiteLLM_VerificationToken( + token=token, + key_alias=key_alias, + team_id=team_id, + user_id=user_id, + spend=spend, + max_budget=max_budget, + **kwargs, + ) diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_harness_smoke.py b/tests/test_litellm/proxy/proxy_server/test_harness_smoke.py new file mode 100644 index 00000000000..566b040a5e7 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_harness_smoke.py @@ -0,0 +1,283 @@ +"""Smoke tests for the proxy_server/ test harness. + +Validates that fixtures + scripts work end-to-end before PR1/PR2/PR3 depend +on them. ``_pin_check.py`` skips this file explicitly so it doesn't count +toward behavior pinning. +""" + +from __future__ import annotations + +import importlib.util +import sys +import textwrap +from pathlib import Path + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from .conftest import ( # type: ignore[import-not-found] + make_acompletion_response, + make_embedding_response, + normalize, +) + +HERE = Path(__file__).resolve().parent + + +# --------------------------------------------------------------------------- +# Fixture smoke tests +# --------------------------------------------------------------------------- + + +def test_app_fixture_returns_fastapi_app(app): + assert isinstance(app, FastAPI) + assert app.router is not None + + +def test_client_fixture_returns_testclient(client): + assert isinstance(client, TestClient) + assert hasattr(client, "post") + assert hasattr(client, "get") + + +def test_mock_prisma_has_team_table(mock_prisma): + assert hasattr(mock_prisma.db, "litellm_teamtable") + assert callable(mock_prisma.db.litellm_teamtable.find_unique) + assert callable(mock_prisma.db.litellm_teamtable.find_many) + + +def test_mock_prisma_has_key_table(mock_prisma): + assert hasattr(mock_prisma.db, "litellm_verificationtoken") + assert callable(mock_prisma.db.litellm_verificationtoken.find_unique) + + +def test_auth_as_admin_overrides_dependency(app, auth_as): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + assert user_api_key_auth in app.dependency_overrides + + +def test_auth_as_internal_user_overrides_dependency(app, auth_as): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + with auth_as(LitellmUserRoles.INTERNAL_USER) as fake_auth: + assert user_api_key_auth in app.dependency_overrides + assert fake_auth.user_role == LitellmUserRoles.INTERNAL_USER + + +def test_auth_as_cleans_up_on_exit(app, auth_as): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + assert user_api_key_auth not in app.dependency_overrides + with auth_as(LitellmUserRoles.PROXY_ADMIN): + pass + assert user_api_key_auth not in app.dependency_overrides + + +def test_mock_router_acompletion_callable(mock_router): + from unittest.mock import AsyncMock + + assert isinstance(mock_router.acompletion, AsyncMock) + assert isinstance(mock_router.aembedding, AsyncMock) + assert isinstance(mock_router.aimage_generation, AsyncMock) + + +@pytest.mark.asyncio +async def test_make_acompletion_response_stream(): + gen = make_acompletion_response(model="gpt-4", stream=True) + chunks = [chunk async for chunk in gen] + assert len(chunks) >= 1 + # Last chunk should have finish_reason set + assert chunks[-1].choices[0].finish_reason == "stop" + + +def test_make_acompletion_response_tools(): + resp = make_acompletion_response( + model="gpt-4", + tools=[{"type": "function", "function": {"name": "fake_tool"}}], + ) + assert resp.choices[0].message.tool_calls is not None + assert resp.choices[0].message.tool_calls[0].function.name == "fake_tool" + + +def test_make_embedding_response_shape(): + resp = make_embedding_response(input=["a", "b", "c"], dimensions=4) + data = resp.data + assert len(data) == 3 + assert len(data[0]["embedding"]) == 4 + + +def test_normalize_replaces_volatile_keys(): + out = normalize({"key": "abc", "spend": 0, "nested": {"id": "x", "value": 5}}) + assert out == { + "key": "", + "spend": 0, + "nested": {"id": "", "value": 5}, + } + + +def test_normalize_handles_lists(): + out = normalize([{"key": "a"}, {"key": "b"}]) + assert out == [{"key": ""}, {"key": ""}] + + +# --------------------------------------------------------------------------- +# Script smoke tests — _coverage_check.py +# --------------------------------------------------------------------------- + + +def _load_script(name: str): + spec = importlib.util.spec_from_file_location(name, HERE / f"{name}.py") + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + # Register in sys.modules so dataclasses can resolve cls.__module__. + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +def _write_cov_xml(tmp_path: Path, line_rate: float, branch_rate: float) -> Path: + xml = textwrap.dedent(f"""\ + + + + + + + + + + + """) + path = tmp_path / "cov.xml" + path.write_text(xml) + return path + + +def test_coverage_check_pass_on_synthetic_xml(tmp_path): + cov_check = _load_script("_coverage_check") + xml = _write_cov_xml(tmp_path, line_rate=0.75, branch_rate=0.60) + line_pct, branch_pct = cov_check.parse_coverage_xml(xml) + assert line_pct == pytest.approx(75.0) + assert branch_pct == pytest.approx(60.0) + + +def test_coverage_check_fail_on_low_coverage(tmp_path, monkeypatch, capsys): + cov_check = _load_script("_coverage_check") + xml = _write_cov_xml(tmp_path, line_rate=0.10, branch_rate=0.05) + monkeypatch.setattr( + sys, + "argv", + ["_coverage_check.py", "--pr-target", "3", "--coverage-xml", str(xml)], + ) + rc = cov_check.main() + assert rc == 1 + out = capsys.readouterr().out + assert "FAIL" in out + + +def test_coverage_check_pass_on_high_coverage(tmp_path, monkeypatch, capsys): + cov_check = _load_script("_coverage_check") + xml = _write_cov_xml(tmp_path, line_rate=0.75, branch_rate=0.60) + monkeypatch.setattr( + sys, + "argv", + ["_coverage_check.py", "--pr-target", "3", "--coverage-xml", str(xml)], + ) + rc = cov_check.main() + assert rc == 0 + out = capsys.readouterr().out + assert "PASS" in out + + +# --------------------------------------------------------------------------- +# Script smoke tests — _pin_check.py +# --------------------------------------------------------------------------- + + +def _write_pin_list(tmp_path: Path, items: list) -> Path: + path = tmp_path / "pins.txt" + path.write_text("\n".join(f"- `{item}`" for item in items) + "\n") + return path + + +def _write_test_file(tmp_path: Path, name: str, body: str) -> Path: + path = tmp_path / name + path.write_text(textwrap.dedent(body)) + return path + + +def test_pin_check_pass_on_complete_pins(tmp_path): + pin_check = _load_script("_pin_check") + _write_pin_list(tmp_path, ["update_cache"]) + _write_test_file( + tmp_path, + "test_thing.py", + """\ + def test_update_cache_happy(): + data = update_cache(value=1) + assert data == {"key1": 1, "key2": 2, "key3": 3} + + def test_update_cache_error(): + import pytest + with pytest.raises(ValueError): + update_cache(value=None) + """, + ) + pin_list = pin_check.parse_pin_list(tmp_path / "pins.txt") + funcs = pin_check.collect_test_functions(tmp_path) + ok, failures = pin_check.check(pin_list, funcs) + assert ok, failures + + +def test_pin_check_fail_on_missing_pin(tmp_path): + pin_check = _load_script("_pin_check") + _write_pin_list(tmp_path, ["update_cache", "never_referenced_symbol"]) + _write_test_file( + tmp_path, + "test_thing.py", + """\ + def test_update_cache_happy(): + data = update_cache(value=1) + assert data == {"key1": 1, "key2": 2, "key3": 3} + + def test_update_cache_error(): + import pytest + with pytest.raises(ValueError): + update_cache(value=None) + """, + ) + pin_list = pin_check.parse_pin_list(tmp_path / "pins.txt") + funcs = pin_check.collect_test_functions(tmp_path) + ok, failures = pin_check.check(pin_list, funcs) + assert not ok + assert any("never_referenced_symbol" in f for f in failures) + + +def test_pin_check_fail_on_status_only_test(tmp_path): + pin_check = _load_script("_pin_check") + _write_pin_list(tmp_path, ["some_route"]) + _write_test_file( + tmp_path, + "test_thing.py", + """\ + def test_some_route_happy(): + response = client.get("/some_route") + assert response.status_code == 200 + + def test_some_route_error(): + response = client.get("/some_route") + assert response.status_code == 404 + """, + ) + pin_list = pin_check.parse_pin_list(tmp_path / "pins.txt") + funcs = pin_check.collect_test_functions(tmp_path) + ok, failures = pin_check.check(pin_list, funcs) + assert not ok + assert any("status-only" in f for f in failures) diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py b/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_openapi_customization.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py b/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_assistants.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py b/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_chat_completions.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_completions.py b/tests/test_litellm/proxy/proxy_server/test_routes_completions.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_completions.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py b/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_embeddings.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py b/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_invitation.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py b/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_moderations.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_queue.py b/tests/test_litellm/proxy/proxy_server/test_routes_queue.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_queue.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_threads.py b/tests/test_litellm/proxy/proxy_server/test_routes_threads.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_threads.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py new file mode 100644 index 00000000000..ad6b4016461 --- /dev/null +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -0,0 +1 @@ +"""Placeholder. Filled by a follow-up PR per the Notion plan.""" diff --git a/tests/test_litellm/proxy/test_batch_expiry.py b/tests/test_litellm/proxy/test_batch_expiry.py index d63f278e715..38c4a71608d 100644 --- a/tests/test_litellm/proxy/test_batch_expiry.py +++ b/tests/test_litellm/proxy/test_batch_expiry.py @@ -178,6 +178,77 @@ class TestBatchEndpointTeamOverride: assert kwargs["output_expires_after"] == TEAM_EXPIRY +class TestBatchEndpointPolicyMetadata: + """Batch create must not forward LiteLLM policy tracking via OpenAI metadata.""" + + def test_create_batch_does_not_forward_applied_policies_metadata( + self, monkeypatch, llm_router + ): + from litellm.proxy.policy_engine.attachment_registry import ( + get_attachment_registry, + ) + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import ( + Policy, + PolicyAttachment, + PolicyGuardrails, + ) + + policy_registry = get_policy_registry() + policy_registry._policies = { + "global-baseline": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker"]), + ), + } + policy_registry._initialized = True + + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [ + PolicyAttachment(policy="global-baseline", scope="*"), + ] + attachment_registry._initialized = True + + _setup_proxy(monkeypatch, llm_router) + + user_key = UserAPIKeyAuth( + api_key="test-key", + team_alias="batch-team", + key_alias="batch-key", + ) + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + captured_kwargs = {} + + async def mock_acreate_batch(**kwargs): + captured_kwargs.update(kwargs) + return _make_batch_response() + + monkeypatch.setattr(litellm, "acreate_batch", mock_acreate_batch) + + try: + response = client.post( + "/v1/batches", + json={ + "input_file_id": "file-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + headers={"Authorization": "Bearer test-key"}, + ) + assert response.status_code == 200 + finally: + app.dependency_overrides.clear() + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert captured_kwargs.get("metadata") in (None, {}) + assert ( + "global-baseline" in captured_kwargs["litellm_metadata"]["applied_policies"] + ) + + class TestBatchEndpointTeamValidation: """Verify validation errors for malformed team metadata on batch endpoint.""" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index fc9813ba530..f336c632546 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -515,6 +515,59 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_secret_fields( ) +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_request(): + """Regression: the body snapshot used to include the proxy_server_request + key itself, producing the path + ``proxy_server_request.body.proxy_server_request.body == body``. Custom + loggers and audit consumers must not see the self-referencing structure + (independent of redaction — fires on every successful call). + """ + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="test-user", + metadata={}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + snapshot_body = updated["proxy_server_request"]["body"] + assert "proxy_server_request" not in snapshot_body, ( + "proxy_server_request must be excluded from its own body snapshot " + "to prevent the body from self-referencing" + ) + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection(): """Regression: metadata arriving as a JSON string (multipart/form-data or @@ -4182,6 +4235,209 @@ class TestApplyClientTagPolicyPreAuth: assert exc_info.value.max_budget == 0.10 +class TestApplyKeyTagsPreAuth: + def test_merges_key_tags_into_metadata(self): + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["engineering", "production"]}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert data["metadata"]["tags"] == ["engineering", "production"] + + def test_unions_key_tags_with_existing_request_tags(self): + data = { + "model": "gpt-3.5-turbo", + "metadata": {"tags": ["request-tag"]}, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["key-tag", "request-tag"]}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + # request-tag deduplicated; key-tag appended + assert data["metadata"]["tags"] == ["request-tag", "key-tag"] + + def test_no_key_tags_no_mutation(self): + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert "metadata" not in data or "tags" not in data.get("metadata", {}) + + def test_empty_key_metadata_no_mutation(self): + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert "metadata" not in data + + def test_uses_litellm_metadata_when_present(self): + data = { + "model": "gpt-3.5-turbo", + "litellm_metadata": {"foo": "bar"}, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["key-tag"]}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert data["litellm_metadata"]["tags"] == ["key-tag"] + assert "tags" not in data.get("metadata", {}) + + def test_string_metadata_parsed_before_merge(self): + data = { + "model": "gpt-3.5-turbo", + "metadata": '{"tags": ["existing"]}', + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["key-tag"]}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert isinstance(data["metadata"], dict) + assert data["metadata"]["tags"] == ["existing", "key-tag"] + + @pytest.mark.asyncio + async def test_key_tags_visible_to_tag_max_budget_check(self): + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TagTable + from litellm.proxy.auth.auth_checks import _tag_max_budget_check + from litellm.proxy.utils import ProxyLogging + + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["engineering"]}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + tag_object = LiteLLM_TagTable( + tag_name="engineering", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:tag:engineering": + return 0.50 + return fallback_spend + + with ( + patch( + "litellm.proxy.proxy_server.get_current_spend", + mock_get_current_spend, + ), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"engineering": tag_object}, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body=data, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 0.50 + assert exc_info.value.max_budget == 0.10 + + @pytest.mark.asyncio + async def test_key_tags_within_budget_passes_check(self): + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TagTable + from litellm.proxy.auth.auth_checks import _tag_max_budget_check + from litellm.proxy.utils import ProxyLogging + + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["engineering"]}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + tag_object = LiteLLM_TagTable( + tag_name="engineering", + spend=0.05, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:tag:engineering": + return 0.05 + return fallback_spend + + with ( + patch( + "litellm.proxy.proxy_server.get_current_spend", + mock_get_current_spend, + ), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"engineering": tag_object}, + ), + ): + await _tag_max_budget_check( + request_body=data, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + + # ============================================================================ # Tests for #27516: provider hint resolution from deployment when the # user-facing model name has no provider prefix. diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ae0996d16d5..ba9c3b75bae 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -604,6 +604,49 @@ def test_ui_extensionless_route_requires_restructure(tmp_path): assert "login" in response.text +def test_admin_ui_export_serves_nested_extensionless_routes(): + out_dir = ( + Path(litellm.__file__).parent / "proxy" / "_experimental" / "out" + ) + assert out_dir.is_dir(), f"missing UI export at {out_dir}" + + nested_html_offenders = [ + path.relative_to(out_dir).as_posix() + for path in out_dir.rglob("*.html") + if path.parent != out_dir + and path.name != "index.html" + and "_next" not in path.parts + and "litellm-asset-prefix" not in path.parts + ] + assert not nested_html_offenders, ( + "Nested routes must be named index.html. Offenders: " + f"{nested_html_offenders}" + ) + + callback_index = out_dir / "mcp" / "oauth" / "callback" / "index.html" + assert callback_index.is_file(), ( + f"MCP OAuth callback page must exist at {callback_index}; " + "without it /ui/mcp/oauth/callback 404s after Linear redirects back." + ) + + fastapi_app = FastAPI() + fastapi_app.mount( + "/ui", StaticFiles(directory=str(out_dir), html=True), name="ui" + ) + client = TestClient(fastapi_app) + + redirect = client.get( + "/ui/mcp/oauth/callback?code=abc&state=xyz", + follow_redirects=False, + ) + assert redirect.status_code == 307 + assert redirect.headers["location"].endswith("/ui/mcp/oauth/callback/?code=abc&state=xyz") + + landed = client.get("/ui/mcp/oauth/callback?code=abc&state=xyz") + assert landed.status_code == 200 + assert " dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_opus_4_8_model_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_models = { + "claude-opus-4-8": { + "provider": "anthropic", + "max_input_tokens": 1000000, + }, + "anthropic.claude-opus-4-8": { + "provider": "bedrock_converse", + "max_input_tokens": 1000000, + }, + "vertex_ai/claude-opus-4-8": { + "provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + }, + # Microsoft Foundry / Azure caps Opus 4.8 at a 200k context window. + "azure_ai/claude-opus-4-8": { + "provider": "azure_ai", + "max_input_tokens": 200000, + }, + } + + for model_name, config in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == config["provider"] + assert info["mode"] == "chat" + assert info["max_input_tokens"] == config["max_input_tokens"] + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + # Base pricing matches Opus 4.7: $5 / $25 per MTok, with the standard + # 1.25x cache-write and 0.1x cache-read multipliers. + assert info["input_cost_per_token"] == 5e-06 + assert info["output_cost_per_token"] == 2.5e-05 + assert info["cache_creation_input_token_cost"] == 6.25e-06 + assert info["cache_read_input_token_cost"] == 5e-07 + + # Opus 4.x flagships are flat-rate across the full context window. + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + assert info["supports_assistant_prefill"] is False + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + + +def test_opus_4_8_bedrock_regional_model_pricing(): + model_data = _load_root_cost_map() + + # Global endpoints use base pricing; regional endpoints carry a 10% premium. + expected_models = { + "global.anthropic.claude-opus-4-8": { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 2.5e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + }, + "us.anthropic.claude-opus-4-8": { + "input_cost_per_token": 5.5e-06, + "output_cost_per_token": 2.75e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_read_input_token_cost": 5.5e-07, + }, + "eu.anthropic.claude-opus-4-8": { + "input_cost_per_token": 5.5e-06, + "output_cost_per_token": 2.75e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_read_input_token_cost": 5.5e-07, + }, + "au.anthropic.claude-opus-4-8": { + "input_cost_per_token": 5.5e-06, + "output_cost_per_token": 2.75e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_read_input_token_cost": 5.5e-07, + }, + } + + for model_name, expected in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["bedrock_output_config_effort_ceiling"] == "xhigh" + for key, value in expected.items(): + assert info[key] == value + + +def test_opus_4_8_fast_mode_multiplier(): + """Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok); + Opus 4.7 was 6x ($30/$150).""" + model_data = _load_root_cost_map() + entry = model_data["claude-opus-4-8"]["provider_specific_entry"] + assert entry["us"] == 1.1 + assert entry["fast"] == 2.0 + + +def test_opus_4_8_present_in_bundled_backup(): + """The bundled backup is the runtime fallback (and what tests load with + ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as + the root cost map, otherwise the model resolves on one path but not the + other.""" + backup = GetModelCostMap.load_local_model_cost_map() + for model_name in ( + "claude-opus-4-8", + "anthropic.claude-opus-4-8", + "global.anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-4-8", + "eu.anthropic.claude-opus-4-8", + "au.anthropic.claude-opus-4-8", + "vertex_ai/claude-opus-4-8", + "vertex_ai/claude-opus-4-8@default", + "azure_ai/claude-opus-4-8", + ): + assert model_name in backup, f"Missing from backup cost map: {model_name}" + + +def test_opus_4_8_registered_for_bedrock_converse(): + assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS + + +def test_opus_4_8_provider_resolves_via_model_info(local_model_cost_map): + """Regression: ``claude-opus-4-8`` must resolve to provider ``anthropic``. + + Before the cost-map entry existed, the model was unknown to LiteLLM, so it + could not be tied to the ``anthropic`` provider and an ``anthropic/*`` + wildcard deployment would not match it. + """ + info = litellm.get_model_info(model="claude-opus-4-8") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 diff --git a/tests/test_litellm/test_claude_sonnet_4_6_config.py b/tests/test_litellm/test_claude_sonnet_4_6_config.py index 434ef9bdeb1..27023d4ee6d 100644 --- a/tests/test_litellm/test_claude_sonnet_4_6_config.py +++ b/tests/test_litellm/test_claude_sonnet_4_6_config.py @@ -50,7 +50,6 @@ def test_bedrock_sonnet_4_6_region_prefixes(): assert model_info.get("supports_pdf_input") is True assert model_info.get("supports_assistant_prefill") is True assert model_info.get("supports_reasoning") is True - assert model_info.get("tool_use_system_prompt_tokens") == 346 def test_bedrock_sonnet_4_6_jp_matches_other_regional_pricing(): diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 00902890da3..1a9bf5a9428 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -13,6 +13,7 @@ from pydantic import BaseModel import litellm from litellm.cost_calculator import ( completion_cost, + cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator, ) @@ -21,6 +22,55 @@ from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.utils import TranscriptionResponse +def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch): + """ + Router/proxy configs may use deployment ids like openai/openai/. Cost lookup must + resolve to model_prices keys (e.g. gpt-5.5), not fail or multiply prefixes. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + prompt_usd, completion_usd = cost_per_token( + model="openai/openai/gpt-5.5", + prompt_tokens=100, + completion_tokens=50, + custom_llm_provider="openai", + ) + + assert prompt_usd + completion_usd > 0 + + +def test_cost_per_token_non_string_model_does_not_hang(): + """ + The provider-prefix dedup loop must not spin forever when `model` is a + non-string object (e.g. a MagicMock from a mocked transport). It should + return or raise promptly instead of looping on a truthy `.startswith()`. + """ + import threading + from unittest.mock import MagicMock + + result: dict = {} + + def _run(): + try: + cost_per_token( + model=MagicMock(), + prompt_tokens=10, + completion_tokens=5, + custom_llm_provider="anthropic", + ) + result["status"] = "returned" + except Exception: + result["status"] = "raised" + + worker = threading.Thread(target=_run, daemon=True) + worker.start() + worker.join(timeout=10) + + assert not worker.is_alive(), "cost_per_token hung on a non-string model" + assert result.get("status") in ("returned", "raised") + + def test_completion_cost_uses_response_model_for_dynamic_routing(): """ Test that completion_cost uses the model from the response object diff --git a/tests/test_litellm/test_thinking_enabled.py b/tests/test_litellm/test_thinking_enabled.py new file mode 100644 index 00000000000..8ba406c395a --- /dev/null +++ b/tests/test_litellm/test_thinking_enabled.py @@ -0,0 +1,74 @@ +""" +Unit tests for is_thinking_enabled method in BaseConfig. + +Tests the fix for issue #28576: handle None thinking param without crashing. +""" + +import pytest +from litellm.llms.base_llm.chat.transformation import BaseConfig + + +class TestIsThinkingEnabled: + """Test is_thinking_enabled handles various thinking parameter values.""" + + @pytest.fixture + def transformer(self): + """Create a BaseConfig instance for testing.""" + # BaseConfig is abstract, so we create a minimal concrete subclass + class ConcreteConfig(BaseConfig): + def __init__(self): + pass + + def get_complete_url(self, *args, **kwargs): + return "" + + def validate_environment(self, *args, **kwargs): + return {} + + def transform_request(self, *args, **kwargs): + return {}, {} + + def transform_response(self, *args, **kwargs): + return None + + def get_supported_openai_params(self, model: str): + return [] + + def map_openai_params(self, *args, **kwargs): + return {} + + def get_error_class(self, *args, **kwargs): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + return BaseLLMException(500, "test error") + + return ConcreteConfig() + + @pytest.mark.parametrize( + "non_default_params,expected", + [ + # thinking=None should not crash, returns False + ({"thinking": None}, False), + # thinking={'type': 'enabled'} returns True + ({"thinking": {"type": "enabled"}}, True), + # thinking key missing returns False + ({}, False), + # thinking={} returns False + ({"thinking": {}}, False), + # thinking with different type returns False + ({"thinking": {"type": "disabled"}}, False), + # reasoning_effort present returns True + ({"reasoning_effort": "medium"}, True), + # both thinking enabled and reasoning_effort returns True + ({"thinking": {"type": "enabled"}, "reasoning_effort": "high"}, True), + # falsy thinking values should not crash + ({"thinking": False}, False), + ({"thinking": 0}, False), + ({"thinking": ""}, False), + ], + ) + def test_is_thinking_enabled(self, transformer, non_default_params, expected): + """Test is_thinking_enabled with various parameter combinations.""" + result = transformer.is_thinking_enabled(non_default_params) + assert result == expected, ( + f"Expected {expected} for params {non_default_params}, got {result}" + ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 40d9cf3231e..6a78653ec99 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -737,6 +737,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token_priority": {"type": "number"}, "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "regional_processing_uplift_multiplier_eu": {"type": "number"}, + "regional_processing_uplift_multiplier_us": {"type": "number"}, "input_cost_per_pixel": {"type": "number"}, "input_cost_per_query": {"type": "number"}, "input_cost_per_request": {"type": "number"}, @@ -857,8 +859,11 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_adaptive_thinking": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, - "supports_output_config": {"type": "boolean"}, - "tool_use_system_prompt_tokens": {"type": "number"}, + "supports_output_config": {"type": "boolean"}, + "bedrock_output_config_effort_ceiling": { + "type": "string", + "enum": ["low", "medium", "high", "max", "xhigh"], + }, "tpm": {"type": "number"}, "provider_specific_entry": {"type": "object"}, "supported_endpoints": { diff --git a/tests/test_litellm/test_vcr_safe_body_matcher.py b/tests/test_litellm/test_vcr_safe_body_matcher.py index 0ed6ad69e3c..77b5416a15c 100644 --- a/tests/test_litellm/test_vcr_safe_body_matcher.py +++ b/tests/test_litellm/test_vcr_safe_body_matcher.py @@ -14,15 +14,22 @@ from tests._vcr_conftest_common import ( # noqa: E402 KEY_FINGERPRINT_HEADER, KEY_FINGERPRINT_MATCHER_NAME, SAFE_BODY_MATCHER_NAME, + TOLERANT_QUERY_MATCHER_NAME, _before_record_request, + _is_credential_exchange_request, + _is_telemetry_request, _key_fingerprint_matcher, + _normalize_volatile_tokens, _safe_body_matcher, + _tolerant_query_matcher, vcr_config_dict, ) -def _req(body): - return SimpleNamespace(body=body, headers={"Content-Type": "application/json"}) +def _req(body, uri="https://api.openai.com/v1/chat/completions"): + return SimpleNamespace( + body=body, uri=uri, headers={"Content-Type": "application/json"} + ) def _req_with_headers(headers, body=b""): @@ -150,6 +157,132 @@ def test_before_record_request_is_deterministic_across_distinct_requests(): ) +def test_google_oauth_bearer_tokens_collapse_to_one_fingerprint(): + """Rotating ``ya29.*`` access tokens must share one fingerprint so + Vertex/Gemini cassettes match across runs (cf. AWS SigV4 access-key + stabilization).""" + run1 = _before_record_request( + _req_with_headers({"Authorization": "Bearer ya29.FIRST-token-aaaaaaaa"}) + ) + run2 = _before_record_request( + _req_with_headers({"Authorization": "Bearer ya29.SECOND-token-bbbbbbbb"}) + ) + assert run1.headers[KEY_FINGERPRINT_HEADER] == run2.headers[KEY_FINGERPRINT_HEADER] + _key_fingerprint_matcher(run1, run2) + + +def test_non_google_bearer_tokens_still_distinguished(): + """The ya29 collapse must not make every Bearer token identical.""" + google = _before_record_request( + _req_with_headers({"Authorization": "Bearer ya29.something"}) + ) + real = _before_record_request( + _req_with_headers({"Authorization": "Bearer sk-real-openai-key"}) + ) + assert ( + google.headers[KEY_FINGERPRINT_HEADER] != real.headers[KEY_FINGERPRINT_HEADER] + ) + + +def test_normalize_volatile_tokens_collapses_uuid_and_timestamps(): + a = b'{"content": "news today b92ed205-0fa9-4e79-939c-2365023e9cb3"}' + b = b'{"content": "news today 1a4e1afa-2915-4dcf-b043-33b991cae879"}' + assert _normalize_volatile_tokens(a) == _normalize_volatile_tokens(b) + + c = b'{"input": "embed data 1779581429.9713597"}' + d = b'{"input": "embed data 1779583432.6874988"}' + assert _normalize_volatile_tokens(c) == _normalize_volatile_tokens(d) + + e = b'{"timestamp": "2026-05-25T03:40:37.262045Z"}' + f = b'{"timestamp": "2026-05-25T06:10:20.830356Z"}' + assert _normalize_volatile_tokens(e) == _normalize_volatile_tokens(f) + + +def test_normalize_volatile_tokens_leaves_deterministic_bodies_unchanged(): + body = b'{"model":"claude-haiku-4-5-20251001","temperature":0.0,"n":2}' + assert _normalize_volatile_tokens(body) == body + + +def test_safe_body_matcher_matches_bodies_differing_only_by_cachebuster(): + a = _req(b'{"messages":[{"content":"hi 1779579395.5545585"}],"model":"gpt-4.1"}') + b = _req(b'{"messages":[{"content":"hi 1779579663.595344"}],"model":"gpt-4.1"}') + _safe_body_matcher(a, b) # must not raise + + +def test_safe_body_matcher_still_rejects_genuinely_different_bodies(): + a = _req(b'{"messages":[{"content":"hello"}]}') + b = _req(b'{"messages":[{"content":"goodbye"}]}') + with pytest.raises(AssertionError): + _safe_body_matcher(a, b) + + +def test_credential_exchange_request_skips_body_comparison(): + assert _is_credential_exchange_request( + _req(b"assertion=AAA", uri="https://oauth2.googleapis.com/token") + ) + assert not _is_credential_exchange_request( + _req(b"x", uri="https://api.openai.com/v1/chat/completions") + ) + # Freshly-signed JWT assertions differ every run but must still match. + a = _req( + b"grant_type=x&assertion=eyJ0AAAA", uri="https://oauth2.googleapis.com/token" + ) + b = _req( + b"grant_type=x&assertion=eyJ0BBBB", uri="https://oauth2.googleapis.com/token" + ) + _safe_body_matcher(a, b) # must not raise + + +def test_match_on_uses_tolerant_query_not_builtin(): + cfg = vcr_config_dict() + assert TOLERANT_QUERY_MATCHER_NAME in cfg["match_on"] + assert "query" not in cfg["match_on"] + + +def test_telemetry_request_detection(): + assert _is_telemetry_request( + _req(b"x", uri="https://us.cloud.langfuse.com/api/public/ingestion") + ) + assert _is_telemetry_request(_req(b"x", uri="https://otlp.arize.com/v1/traces")) + assert not _is_telemetry_request( + _req(b"x", uri="https://api.openai.com/v1/chat/completions") + ) + + +def test_safe_body_matcher_skips_telemetry_body(): + a = _req( + b'{"batch":[{"id":"aaa","timestamp":"2026-05-25T03:40:37Z"}]}', + uri="https://us.cloud.langfuse.com/api/public/ingestion", + ) + b = _req( + b'{"batch":[{"id":"zzz","timestamp":"2026-05-25T09:99:99Z","extra":1}]}', + uri="https://us.cloud.langfuse.com/api/public/ingestion", + ) + _safe_body_matcher(a, b) # must not raise despite wholly different bodies + + +def test_tolerant_query_skips_telemetry_but_enforces_others(): + from vcr.request import Request + + def _greq(uri): + return Request(method="GET", uri=uri, body=b"", headers={}) + + # Telemetry GET with a fresh trace_id in the query must still match. + a = _greq( + "https://us.cloud.langfuse.com/api/public/observations?traceId=litellm-test-AAA" + ) + b = _greq( + "https://us.cloud.langfuse.com/api/public/observations?traceId=litellm-test-BBB" + ) + _tolerant_query_matcher(a, b) # must not raise + + # Non-telemetry hosts keep vcrpy's strict query comparison. + c = _greq("https://api.openai.com/v1/models?page=1") + d = _greq("https://api.openai.com/v1/models?page=2") + with pytest.raises(AssertionError): + _tolerant_query_matcher(c, d) + + def test_before_record_request_is_idempotent_on_the_same_request_object(): """vcrpy invokes ``before_record_request`` more than once per request. diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index b0eb2438b95..3d0472ef96e 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -398,6 +398,34 @@ class TestVideoGeneration: ) assert abs(cost - 0.8) < 0.001 + def test_completion_cost_video_edit_uses_video_calculator(self): + """video_edit is charged via the same video cost path as create_video.""" + from litellm.cost_calculator import completion_cost + + mock_response = MagicMock() + mock_response.usage = MagicMock() + mock_response.usage.duration_seconds = 10.0 + type(mock_response)._hidden_params = {} + + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "metadata": { + "model_info": { + "output_cost_per_video_per_second": 0.05, + } + } + } + + cost = completion_cost( + completion_response=mock_response, + model="vertex_ai/veo-3.1-generate-001", + call_type="video_edit", + custom_llm_provider="vertex_ai", + custom_pricing=True, + litellm_logging_obj=mock_logging_obj, + ) + assert cost == 0.5 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig() diff --git a/tests/test_openai_endpoints.py b/tests/test_openai_endpoints.py index e898b88a556..29875a04413 100644 --- a/tests/test_openai_endpoints.py +++ b/tests/test_openai_endpoints.py @@ -446,7 +446,7 @@ async def test_chat_completion_anthropic_structured_output(): client = AsyncOpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") res = await client.beta.chat.completions.parse( - model="bedrock/us.anthropic.claude-3-sonnet-20240229-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, response_format=EventsList, timeout=60, diff --git a/tests/vector_store_tests/test_bedrock_vector_store.py b/tests/vector_store_tests/test_bedrock_vector_store.py index d8af1c7188b..47e73e61c59 100644 --- a/tests/vector_store_tests/test_bedrock_vector_store.py +++ b/tests/vector_store_tests/test_bedrock_vector_store.py @@ -22,7 +22,7 @@ class TestBedrockVectorStore(BaseVectorStoreTest): def get_base_request_args(self): return { - "vector_store_id": "T37J8R4WTM", + "vector_store_id": "LCYXFBR2TU", "custom_llm_provider": "bedrock", "query": "what happens after we add a model", } @@ -106,7 +106,7 @@ async def test_bedrock_search_with_router(): _router = Router(model_list=[]) search_response = await _router.avector_store_search( query="what happens after we add a model", - vector_store_id="T37J8R4WTM", + vector_store_id="LCYXFBR2TU", custom_llm_provider="bedrock", ) print(search_response) @@ -150,7 +150,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): # Create vector store with credential reference vector_store = LiteLLM_ManagedVectorStore( - vector_store_id="T37J8R4WTM", + vector_store_id="LCYXFBR2TU", custom_llm_provider="bedrock", created_at=datetime.now(timezone.utc), updated_at=datetime.now(timezone.utc), @@ -162,7 +162,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): litellm.vector_store_registry = registry # Verify credentials can be retrieved from registry - retrieved_credentials = registry.get_credentials_for_vector_store("T37J8R4WTM") + retrieved_credentials = registry.get_credentials_for_vector_store("LCYXFBR2TU") assert retrieved_credentials, "Should retrieve credentials from registry" assert retrieved_credentials.get("aws_access_key_id") == "test_access_key" assert retrieved_credentials.get("aws_secret_access_key") == "test_secret_key" @@ -194,7 +194,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): search_response = await _router.avector_store_search( query="what happens after we add a model", - vector_store_id="T37J8R4WTM", + vector_store_id="LCYXFBR2TU", custom_llm_provider="bedrock", ) @@ -203,7 +203,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): call_kwargs = mock_handler.call_args[1] # Verify that the credential accessor was called with the correct vector store ID - mock_get_creds.assert_called_with("T37J8R4WTM") + mock_get_creds.assert_called_with("LCYXFBR2TU") # Verify the credentials were injected into the search call litellm_params = call_kwargs.get("litellm_params", {}) @@ -224,7 +224,7 @@ async def test_bedrock_search_with_credentials_managed_registry(): assert search_response["data"][0]["id"] == "test_result" print( - f"✅ Test passed: Credential accessor was called with vector store ID: T37J8R4WTM" + f"✅ Test passed: Credential accessor was called with vector store ID: LCYXFBR2TU" ) print(f"✅ Retrieved credentials: {retrieved_credentials}") print(f"✅ Credentials were injected into search call") diff --git a/ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts b/ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts new file mode 100644 index 00000000000..83831f82da0 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/serverRootPath.config.ts @@ -0,0 +1,32 @@ +import { defineConfig, devices } from "@playwright/test"; + +// Minimal config for the SERVER_ROOT_PATH redirect spec. Deliberately does NOT +// reuse the main e2e config because: +// - globalSetup logs in via http://localhost:4000/ui/login, which 404s when +// the proxy is mounted under a non-root path. +// - The redirect spec must run against a clean, unauthenticated session, so +// no storage state should be loaded. +export default defineConfig({ + testDir: "./tests/login", + testMatch: ["serverRootPathRedirect.spec.ts"], + fullyParallel: false, + forbidOnly: !!process.env.CI, + retries: process.env.CI ? 2 : 0, + workers: 1, + reporter: "list", + use: { + trace: "on-first-retry", + actionTimeout: 15 * 1000, + navigationTimeout: 30 * 1000, + }, + projects: [ + { + name: "chromium", + use: { ...devices["Desktop Chrome"] }, + }, + ], + timeout: 60 * 1000, + expect: { + timeout: 10 * 1000, + }, +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts index 5d4b2508444..994d211cc18 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts @@ -10,4 +10,21 @@ test("user can log in", async ({ page }) => { await expect(loginButton).toBeEnabled(); await loginButton.click(); await expect(page.getByText("Virtual Keys")).toBeVisible(); + + // Match the navbar account button by its stable aria-label (UserDropdown.tsx + // emits "Account menu — — signed in as "). Earlier this used + // `hasText: /^User$/`, which never matched the rendered button (text is + // displayName = "Account" for the master-key admin), so the trigger evaluate + // would time out in CI. + const userTrigger = page.locator('button[aria-label^="Account menu"]').first(); + await userTrigger.click(); + + // Filter by the popupRender wrapper class to disambiguate from other + // ant-dropdown popups. + const popup = page.locator(".ant-dropdown:visible").filter({ + has: page.locator(".bg-white.rounded-lg.shadow-lg"), + }).first(); + await expect(popup).toBeVisible({ timeout: 5_000 }); + await expect(popup.getByText("Admin", { exact: true })).toBeVisible({ timeout: 5_000 }); + await expect(popup.getByText("default_user_id", { exact: true })).toBeVisible({ timeout: 5_000 }); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts new file mode 100644 index 00000000000..62a3e913184 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/login/serverRootPathRedirect.spec.ts @@ -0,0 +1,32 @@ +import { expect, test } from "@playwright/test"; + +// Driven by the SERVER_ROOT_PATH env var injected by the workflow; the container +// is booted with the same value, so the asset paths and the runtime config it +// serves at /litellm/.well-known/litellm-ui-config will both reflect it. +const ROOT_PATH = process.env.SERVER_ROOT_PATH ?? ""; + +test.skip(!ROOT_PATH, "Requires SERVER_ROOT_PATH env var"); + +// Contract: an unauthenticated visit must redirect to a login URL that preserves +// the SERVER_ROOT_PATH prefix. The redirect URL is built client-side from +// `proxyBaseUrl`, which is populated by an async fetch of the runtime UI config. +// If the redirect fires before that fetch resolves, the URL is missing the +// prefix and the user lands on a 404. To make the race deterministic across +// runners, the config endpoint is intentionally delayed. +test("unauth redirect preserves SERVER_ROOT_PATH prefix", async ({ page }) => { + // Matches both `/litellm/.well-known/litellm-ui-config` and + // `${SERVER_ROOT_PATH}/.well-known/litellm-ui-config` (the proxy rewrites the + // bundle at boot when a root path is set). + await page.route("**/.well-known/litellm-ui-config", async (route) => { + await new Promise((resolve) => setTimeout(resolve, 500)); + await route.continue(); + }); + + await page.context().clearCookies(); + + await page.goto(`http://localhost:4000${ROOT_PATH}/ui/?page=virtual-keys`); + + await page.waitForURL((url) => url.pathname.endsWith("/ui/login"), { timeout: 15_000 }); + + expect(page.url()).toContain(`${ROOT_PATH}/ui/login`); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts new file mode 100644 index 00000000000..f953a82daaa --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/mcp/mcpServers.spec.ts @@ -0,0 +1,62 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; + +// Coverage scope: only the happy-path Streamable HTTP + None auth create flow. +// See E2E_COVERAGE.md (#29 row) for the full list of uncovered MCP surfaces +// — SSE / stdio / OpenAPI transports, API Key / Bearer / OAuth2 / Basic / Token +// / AWS SigV4 auth, edit/delete, BYOK credentials, tool list/call (needs a real +// or mocked MCP server in the e2e fixture stack), and access-group permissions. +test.describe("MCP Servers", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Add a custom MCP server via the discovery → custom form", async ({ page }) => { + await navigateToPage(page, Page.McpServers); + + // Open the discovery modal, then drop into the custom-server form + await page.getByRole("button", { name: /Add New MCP Server/i }).click(); + const discovery = page.locator(".ant-modal:visible").filter({ hasText: "Add MCP Server" }); + await expect(discovery).toBeVisible({ timeout: 5_000 }); + await discovery.getByRole("button", { name: /Custom Server/i }).click(); + + const formModal = page.locator(".ant-modal:visible").filter({ hasText: "MCP Server Name" }); + await expect(formModal).toBeVisible({ timeout: 5_000 }); + + // Name — no spaces or hyphens per validateMCPServerName + const uniqueName = `e2e_mcp_${Date.now()}`; + await formModal.locator('input[id="server_name"]').fill(uniqueName); + + // Transport: Streamable HTTP — the only value the proxy actually accepts is "http" + const transportField = formModal.locator(".ant-form-item", { hasText: "Transport Type" }); + await transportField.locator(".ant-select").click(); + await page.locator(".ant-select-dropdown:visible").getByText("Streamable HTTP").click(); + + // URL — use a fake URL; the form just persists it, it doesn't have to be reachable + await formModal.locator('input[id="url"]').fill("https://e2e-fake-mcp.test.local/mcp"); + + // Authentication: None + // The auth_type Form.Item has no label prop (create_mcp_server.tsx:795), so + // it can't be anchored by label text. Scope via the enclosing Collapse + // panel ("Authentication") instead — that anchor is stable even if the + // placeholder copy changes. + const authSection = formModal.locator(".ant-collapse-item", { hasText: /^Authentication/ }); + const authField = authSection.locator(".ant-form-item").first(); + await authField.locator(".ant-select").click(); + await page.locator(".ant-select-dropdown:visible").getByText("None", { exact: true }).click(); + + // Submit + await formModal.getByRole("button", { name: /^Add MCP Server$/ }).click(); + + // No teardown needed — the e2e runner spins up a fresh DB per invocation. + + // Success toast and the new row in the table. Scope the row lookup to + // the MCP servers table so the form modal's `server_name` input — which + // still holds the timestamped value during its close animation — can't + // satisfy the assertion before the server actually lands in the list. + await expect(page.getByText("MCP Server created successfully").first()) + .toBeVisible({ timeout: 15_000 }); + await expect(page.locator("table tbody").getByText(uniqueName).first()) + .toBeVisible({ timeout: 10_000 }); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts new file mode 100644 index 00000000000..ada4dfb735e --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/modelHub/modelHub.spec.ts @@ -0,0 +1,79 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; + +test.describe("AI Hub (internal admin view)", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Make models public via the multi-step modal", async ({ page }) => { + await navigateToPage(page, Page.ModelHubTable); + + // Open the "Select Models to Make Public" modal + await page.getByRole("button", { name: /Select Models to Make Public/i }).click(); + + const modal = page.locator(".ant-modal:visible").filter({ hasText: "Make Models Public" }); + await expect(modal).toBeVisible({ timeout: 5_000 }); + + // Guard: the "Select All (N)" label only shows a count when filteredData + // has at least one row. Asserting N>=1 here turns a missing-seed-data + // failure into an immediate diagnostic rather than a downstream timeout + // on the disabled-Next button or the success toast. + await expect(modal.getByText(/Select All \(\d+\)/)).toBeVisible({ timeout: 5_000 }); + + // Step 1: pick the seeded models via "Select All" + await modal.getByText(/Select All/i).click(); + + // Move to confirm step + await modal.getByRole("button", { name: "Next" }).click(); + await expect(modal.getByText("Confirm Making Models Public")).toBeVisible({ timeout: 5_000 }); + + // Submit + await modal.getByRole("button", { name: "Make Public" }).click(); + + await expect(page.getByText(/Successfully made .* model group\(s\) public/i).first()) + .toBeVisible({ timeout: 15_000 }); + }); + + test("AI Hub tab list renders Model Hub, Agent Hub, MCP Hub and Skill Hub", async ({ page }) => { + await navigateToPage(page, Page.ModelHubTable); + + // The tab strip lives in the main view; check each tab is present and clickable. + // (The "Claude Code Plugin Marketplace" tab from the manual-QA checklist was + // renamed to "Skill Hub" — verify the current label here so the test stays + // in sync with the UI.) + // + // Note: unlike the public /ui/model_hub_table view (test below), the admin + // ModelHubTable renders all four tabs unconditionally — there are no `&&` + // guards around Agent Hub or MCP Hub in the source + // (ModelHubTable.tsx ~L436-439). Asserting all four here is intentional: + // this pins the manual-QA contract that the AI Hub tab strip exposes + // exactly these labels regardless of seeded agent/MCP data. + for (const tabName of ["Model Hub", "Agent Hub", "MCP Hub", "Skill Hub"]) { + const tab = page.getByRole("tab", { name: tabName }); + await expect(tab, `${tabName} tab should be present`).toBeVisible({ timeout: 5_000 }); + await tab.click(); + } + }); +}); + +test.describe("Public model hub (/ui/model_hub_table)", () => { + // No storageState — the public page is reached anonymously with a `key` query param. + + test("Public model_hub_table loads and renders the Model Hub tab", async ({ page }) => { + // The page expects the proxy key as the `key` query param. Use the master + // key the e2e runner already exports — this matches what the AI Hub copy + // button hands out. + const masterKey = process.env.LITELLM_MASTER_KEY || "sk-1234"; + await page.goto(`/ui/model_hub_table?key=${masterKey}`); + + // Dismiss the feedback popup before asserting on the tab, so a popup + // race can't briefly mask the tab while we're evaluating visibility. + await dismissFeedbackPopup(page); + + // Page loads (no auth redirect) and the Model Hub tab is always present. + // Agent Hub and MCP Hub tabs are conditionally rendered only when public + // agents/MCP servers exist, so we don't assert on them in a fresh CI run. + await expect(page.getByRole("tab", { name: "Model Hub" })).toBeVisible({ timeout: 10_000 }); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts index c3bd8489027..bb53fb7a23b 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -1,5 +1,5 @@ import { test, expect } from "@playwright/test"; -import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID } from "../../constants"; import { Role, users } from "../../fixtures/users"; import { navigateToPage } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; @@ -150,6 +150,111 @@ test.describe("Add Model", () => { await expect(tableBody.getByText("claude-haiku-4-5").first()).toBeVisible({ timeout: 15_000 }); }); + test("Add team-only model via Team-BYOK toggle and verify it appears with the team", async ({ page, request }) => { + // The Team-BYOK switch is gated on `premiumUser` — without a license set + // for the proxy under test, the toggle is disabled and this manual-QA + // step cannot be exercised. + test.skip( + !process.env.LITELLM_LICENSE, + "LITELLM_LICENSE not set in test env — Team-BYOK switch is disabled", + ); + + // Make the test idempotent across retries and local reruns: delete any + // Cohere model already scoped to the e2e team before we start, and again + // after we finish. The sibling "Add wildcard route" test creates a + // team-less Cohere wildcard, so we only target rows that have BOTH the + // cohere/* model_name AND team_id == e2e-team-crud. + const masterKey = users[Role.ProxyAdmin].password; + const auth = { Authorization: `Bearer ${masterKey}` }; + const deleteTeamScopedCohereModels = async () => { + const res = await request.get("/v2/model/info", { headers: auth }); + if (!res.ok()) return; + const body = await res.json(); + const matches: Array<{ id: string }> = (body?.data ?? []).filter((m: any) => + typeof m?.model_name === "string" && + m.model_name.startsWith("cohere") && + m?.model_info?.team_id === E2E_TEAM_CRUD_ID, + ); + for (const m of matches) { + await request.post("/model/delete", { headers: auth, data: { id: m.id } }); + } + }; + await deleteTeamScopedCohereModels(); + + try { + await navigateToPage(page, Page.Models); + await page.getByRole("tab", { name: "Add Model" }).click(); + + await selectProvider(page, "Cohere"); + + const modelDropdown = page.locator(".ant-select-selection-overflow").first(); + await modelDropdown.click(); + const wildcardOption = page.getByTitle(/All .* Models \(Wildcard\)/); + await wildcardOption.click(); + await page.keyboard.press("Escape"); + + const apiKeyInput = page.locator('input[type="password"]').first(); + await apiKeyInput.fill("sk-any-key-for-team-byok-test"); + + // Flip the Team-BYOK switch on (Form.Item label "Team-BYOK Model") + const teamByokRow = page.locator(".ant-form-item", { hasText: "Team-BYOK Model" }); + await teamByokRow.getByRole("switch").click(); + + // The Team dropdown appears underneath once the switch is on. TeamDropdown + // renders its Select.Option children with custom / markup, so + // the popup items don't carry role="option" — match by text content, + // scoped to the visible dropdown so a stale tag elsewhere in the form + // can't satisfy it. + const teamDropdown = page.getByTestId("team-dropdown"); + await expect(teamDropdown).toBeVisible({ timeout: 5_000 }); + await teamDropdown.click(); + const teamOption = page.locator(".ant-select-dropdown:visible") + .getByText(E2E_TEAM_CRUD_ID) + .first(); + await expect(teamOption).toBeVisible({ timeout: 5_000 }); + await teamOption.click(); + + await page.getByRole("button", { name: "Add Model" }).last().click(); + + // Scope the success toast to antd's notification container so a stale + // success message from an earlier test in the same context can't satisfy + // the assertion. + await expect(page.locator(".ant-notification").getByText("created successfully").last()) + .toBeVisible({ timeout: 15_000 }); + + // Verify the model is now in All Models with the team_id attached. The + // Models table renders team-scoped models with the team id in the row. + await page.getByRole("tab", { name: "All Models" }).click(); + await page.waitForLoadState("networkidle"); + // Match the sibling tests in this file — networkidle fires before the + // table finishes re-rendering, so give it the same 2s settle before + // searching. + await page.waitForTimeout(2000); + + await page.locator('input[placeholder="Search model names..."]').fill("cohere"); + await page.waitForTimeout(1000); + + // Confirm the search returned at least one result — gives a clear + // failure message when the table is empty instead of timing out on a + // row assertion. + await expect(page.getByTestId("models-results-count")).toHaveText( + /Showing \d+ - \d+ of \d+ results/, + { timeout: 15_000 }, + ); + + // Stronger than "alias appears somewhere in tbody" — pin the assertion + // to a single row that has BOTH the cohere model_name AND the seeded + // team alias, so a stale cohere row from "Add wildcard route" (no team) + // can't satisfy the check. + const teamCohereRow = page.locator("table tbody tr") + .filter({ hasText: "cohere/" }) + .filter({ hasText: E2E_TEAM_CRUD_ALIAS }); + await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); + } finally { + await deleteTeamScopedCohereModels(); + } + }); + test("Add wildcard route and verify it appears in All Models", async ({ page }) => { await navigateToPage(page, Page.Models); await page.getByRole("tab", { name: "Add Model" }).click(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts new file mode 100644 index 00000000000..d21192d237d --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/clearCustomPricing.spec.ts @@ -0,0 +1,177 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Role, users } from "../../fixtures/users"; + +/** + * Regression: clearing the Input / Output / Cache Read / Cache Write Cost + * fields on a deployment with a user-set pricing override must actually remove + * the override from both `litellm_params` and `model_info`. + * + * Pre-fix, the UI sent the old pricing back on every save (the spread of + * `values.litellm_params` re-injected it), and the backend's `exclude_none=True` + * stripped any null that did make it through. End-result: the dashboard + * displayed "Saved" but the override remained in the DB. The cache fields had + * the same bug in a parallel code path and are covered here too. + */ +test.describe("Clear custom pricing on a deployment", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + const masterKey = users[Role.ProxyAdmin].password; + const SEED_INPUT_PER_TOKEN = 0.0000777; + const SEED_OUTPUT_PER_TOKEN = 0.0000999; + const SEED_CACHE_READ_PER_TOKEN = 0.0000333; + const SEED_CACHE_WRITE_PER_TOKEN = 0.0000555; + + // Unique-per-run name so concurrent / repeated runs don't collide on the + // shared dashboard DB. Captured here so afterEach can clean it up. + let createdModelId: string | null = null; + let modelName: string; + + test.beforeEach(async ({ page }) => { + modelName = `e2e-clear-pricing-${Date.now()}`; + const res = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey}` }, + data: { + model_name: modelName, + litellm_params: { + model: "openai/gpt-4o", + api_key: "sk-e2e-not-used", + input_cost_per_token: SEED_INPUT_PER_TOKEN, + output_cost_per_token: SEED_OUTPUT_PER_TOKEN, + cache_read_input_token_cost: SEED_CACHE_READ_PER_TOKEN, + cache_creation_input_token_cost: SEED_CACHE_WRITE_PER_TOKEN, + }, + model_info: {}, + }, + }); + expect(res.ok(), `POST /model/new for ${modelName}`).toBe(true); + const body = await res.json(); + createdModelId = body.model_info?.id ?? body.model_id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); + }); + + test.afterEach(async ({ page }) => { + // The dashboard DB persists across this suite (not just per-test), so every + // model created here must be cleaned up regardless of test outcome. + if (createdModelId) { + await page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey}` }, + data: { id: createdModelId }, + }); + createdModelId = null; + } + }); + + test("UI sends null for cleared pricing and backend removes the override", async ({ + page, + }) => { + // Navigate to the model detail view. + await page.goto("/ui"); + await page.getByText("Models + Endpoints").click(); + + const modelRow = page.locator("tr", { hasText: modelName }).first(); + await expect(modelRow).toBeVisible({ timeout: 15_000 }); + await modelRow.click(); + await expect(page.getByText("Back to Models").first()).toBeVisible({ + timeout: 10_000, + }); + + // Sanity: the seeded pricing is shown in the detail view (77.7000 / 99.9000 + // per 1M tokens). The dashboard renders the per-token rate × 1e6. + await expect(page.getByText("77.7000")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByText("99.9000")).toBeVisible({ timeout: 10_000 }); + + // Open the edit form and clear all four pricing fields. + await page.getByRole("button", { name: "Edit Settings" }).click(); + const inputCost = page.getByPlaceholder("Enter input cost"); + const outputCost = page.getByPlaceholder("Enter output cost"); + // Both cache fields share the same placeholder ("Defaults to Input Cost if blank"), + // so disambiguate via the Form.Item id (AntD assigns the `name` prop as input id). + const cacheReadCost = page.locator("#cache_read_cost"); + const cacheWriteCost = page.locator("#cache_write_cost"); + await inputCost.waitFor({ timeout: 15_000 }); + for (const field of [inputCost, outputCost, cacheReadCost, cacheWriteCost]) { + await field.click({ clickCount: 3 }); + await page.keyboard.press("Delete"); + } + + // Capture the outgoing PATCH so we can assert the UI sends explicit nulls. + const patchPromise = page.waitForRequest( + (req) => + req.method() === "PATCH" && + req.url().includes(`/model/${createdModelId}/update`) + ); + await page.getByRole("button", { name: "Save Changes" }).click(); + const patchReq = await patchPromise; + const patchBody = JSON.parse(patchReq.postData() ?? "{}"); + expect( + patchBody.litellm_params.input_cost_per_token, + "UI sends explicit null for cleared input cost" + ).toBeNull(); + expect( + patchBody.litellm_params.output_cost_per_token, + "UI sends explicit null for cleared output cost" + ).toBeNull(); + expect( + patchBody.litellm_params.cache_read_input_token_cost, + "UI sends explicit null for cleared cache_read cost" + ).toBeNull(); + expect( + patchBody.litellm_params.cache_creation_input_token_cost, + "UI sends explicit null for cleared cache_write cost" + ).toBeNull(); + + // Success toast confirms the save was accepted. + await expect( + page.getByText("Model settings updated successfully") + ).toBeVisible({ timeout: 10_000 }); + + // Verify via the management API: the user-set rate is gone from both blobs. + // The cost-map may synthesize a default for known providers in the response, + // so the assertion is "no longer the seeded value" rather than literally + // undefined. + const infoRes = await page.request.get( + `/v2/model/info?include_team_models=true&page=1&size=100&modelId=${createdModelId}`, + { headers: { Authorization: `Bearer ${masterKey}` } } + ); + expect(infoRes.ok()).toBe(true); + const infoBody = await infoRes.json(); + const row = (infoBody.data ?? infoBody).find?.( + (m: any) => m?.model_info?.id === createdModelId + ); + expect(row, "model info row").toBeTruthy(); + + expect( + "input_cost_per_token" in row.litellm_params, + "litellm_params.input_cost_per_token key removed" + ).toBe(false); + expect( + "output_cost_per_token" in row.litellm_params, + "litellm_params.output_cost_per_token key removed" + ).toBe(false); + expect( + "cache_read_input_token_cost" in row.litellm_params, + "litellm_params.cache_read_input_token_cost key removed" + ).toBe(false); + expect( + "cache_creation_input_token_cost" in row.litellm_params, + "litellm_params.cache_creation_input_token_cost key removed" + ).toBe(false); + expect( + row.model_info.input_cost_per_token, + "model_info.input_cost_per_token no longer the seeded override" + ).not.toBe(SEED_INPUT_PER_TOKEN); + expect( + row.model_info.output_cost_per_token, + "model_info.output_cost_per_token no longer the seeded override" + ).not.toBe(SEED_OUTPUT_PER_TOKEN); + expect( + row.model_info.cache_read_input_token_cost, + "model_info.cache_read_input_token_cost no longer the seeded override" + ).not.toBe(SEED_CACHE_READ_PER_TOKEN); + expect( + row.model_info.cache_creation_input_token_cost, + "model_info.cache_creation_input_token_cost no longer the seeded override" + ).not.toBe(SEED_CACHE_WRITE_PER_TOKEN); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts index a1864b22a43..6f6e8373390 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/teams.spec.ts @@ -131,4 +131,49 @@ test.describe("Proxy Admin - Teams", () => { await expect(page.getByText(/updated|success/i).first()).toBeVisible({ timeout: 10_000 }); }); + + test("Edit team model selection", async ({ page, request }) => { + // Restore the seeded models via API in case a prior run (or a CI retry) + // left this team mutated — the assertion below requires fake-anthropic-claude + // to be present. + const masterKey = process.env.LITELLM_MASTER_KEY || "sk-1234"; + const seededModels = ["fake-openai-gpt-4", "fake-anthropic-claude"]; + const restore = async () => { + const res = await request.post("http://localhost:4000/team/update", { + headers: { Authorization: `Bearer ${masterKey}` }, + data: { team_id: E2E_TEAM_CRUD_ID, models: seededModels }, + }); + expect(res.ok(), `restore failed: ${res.status()} ${await res.text()}`).toBeTruthy(); + }; + await restore(); + + try { + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + + await clickTeamId(page, E2E_TEAM_CRUD_ID); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + // Remove the anthropic tag — other tests against this team use "All Team + // Models" so they pick up whatever remains. + const modelsSelect = page.locator("[data-testid='models-select']"); + await expect(modelsSelect).toBeVisible({ timeout: 10_000 }); + + const anthropicTag = modelsSelect + .locator(".ant-select-selection-item") + .filter({ hasText: "fake-anthropic-claude" }); + await expect(anthropicTag).toBeVisible({ timeout: 5_000 }); + await anthropicTag.locator(".ant-select-selection-item-remove").click(); + + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect(page.getByText(/Team settings updated|updated successfully/i).first()) + .toBeVisible({ timeout: 10_000 }); + } finally { + // Leave the team in its seeded state for any subsequent test or rerun. + await restore(); + } + }); }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts new file mode 100644 index 00000000000..8dd5571f7af --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/settings/routerSettings.spec.ts @@ -0,0 +1,102 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { Role, users } from "../../fixtures/users"; + +const PRIMARY = "fake-openai-gpt-4"; +const FALLBACK = "fake-anthropic-claude"; + +/** + * Wipe any fallbacks for the primary model so the test is idempotent across + * retries and local reruns (the proxy persists router_settings to the DB). + */ +async function clearFallbackForPrimary(request: import("@playwright/test").APIRequestContext) { + const masterKey = users[Role.ProxyAdmin].password; + const auth = { Authorization: `Bearer ${masterKey}` }; + + const current = await request.get("http://localhost:4000/get/config/callbacks", { headers: auth }); + if (!current.ok()) return; + const body = await current.json(); + const router = body?.router_settings ?? {}; + const existing: Array> = Array.isArray(router.fallbacks) ? router.fallbacks : []; + const next = existing.filter((entry) => !(entry && PRIMARY in entry)); + if (next.length === existing.length) return; + + await request.post("http://localhost:4000/config/update", { + headers: auth, + data: { router_settings: { ...router, fallbacks: next } }, + }); +} + +test.describe("Router Settings - Fallbacks", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test.beforeEach(async ({ request }) => { + await clearFallbackForPrimary(request); + }); + + test.afterEach(async ({ request }) => { + await clearFallbackForPrimary(request); + }); + + test("Add a fallback and verify it appears in the table", async ({ page }) => { + await navigateToPage(page, Page.RouterSettings); + + // Four tabs: Loadbalancing / Routing Groups / Fallbacks / General — click Fallbacks + await page.getByRole("tab", { name: "Fallbacks" }).click(); + + // The model options come from /model_group/info, which AddFallbacks + // fires only after the modal mounts. Wait for that response so the + // dropdown is populated before we try to pick from it — without this + // the test races on CI (local SLOWMO masks the gap). + const modelsLoaded = page.waitForResponse( + (res) => res.url().includes("/model_group/info") && res.status() === 200, + { timeout: 15_000 }, + ); + await page.getByRole("button", { name: /Add Fallbacks/i }).click(); + await modelsLoaded; + + const modal = page.locator(".ant-modal:visible"); + await expect(modal).toBeVisible({ timeout: 5_000 }); + + // FallbackGroupConfig.tsx renders both selects with `showSearch`. The + // most stable interaction is: click to open + focus, type the model name to + // narrow the listbox to a single highlighted option, then press Enter. + // Verify each selection landed by watching the dialog's own state transition + // (the tab title updates to the picked primary; the fallback chain list + // populates) rather than by asserting on the dropdown popup, which sits in + // a custom getPopupContainer and is awkward to scope reliably. + const primarySelect = modal.locator(".ant-select").filter({ hasText: "Select primary model" }); + await primarySelect.click(); + await page.keyboard.type(PRIMARY); + await page.keyboard.press("Enter"); + await expect(modal.getByRole("tab", { name: PRIMARY })).toBeVisible({ timeout: 10_000 }); + + const fallbackSelect = modal.locator(".ant-select").filter({ hasText: "Select fallback models" }); + await fallbackSelect.click(); + await page.keyboard.type(FALLBACK); + await page.keyboard.press("Enter"); + await page.keyboard.press("Escape"); + // The Fallback Chain helper text reads "(N/10 used)"; once it ticks to 1 the + // selection has been recorded. + await expect(modal.getByText("(1/10 used)")).toBeVisible({ timeout: 10_000 }); + + // Save + await modal.getByRole("button", { name: /Save All Configurations/i }).click(); + + // Success toast + await expect(page.getByText(/fallback configuration\(s\) added successfully/i).first()) + .toBeVisible({ timeout: 10_000 }); + + // Modal closes, and a single row contains BOTH the primary and the fallback + // model — stronger than asserting each name appears somewhere in tbody, + // which could be satisfied by leftover rows from prior runs. + await expect(modal).not.toBeVisible({ timeout: 5_000 }); + + const newRow = page.locator("table tbody tr") + .filter({ hasText: PRIMARY }) + .filter({ hasText: FALLBACK }); + await expect(newRow).toHaveCount(1, { timeout: 10_000 }); + }); +}); diff --git a/ui/litellm-dashboard/next.config.mjs b/ui/litellm-dashboard/next.config.mjs index cfaeb24dc5d..19a2ca298fe 100644 --- a/ui/litellm-dashboard/next.config.mjs +++ b/ui/litellm-dashboard/next.config.mjs @@ -14,6 +14,7 @@ const nextConfig = { }, basePath: "", assetPrefix: "/litellm-asset-prefix", + trailingSlash: true, turbopack: { // Must be absolute; "." is no longer allowed root: __dirname, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/api-playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/api-playground/page.tsx deleted file mode 100644 index 0948b7626db..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/api-playground/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import TransformRequestPanel from "@/components/transform_request"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const APIPlaygroundPage = () => { - const { accessToken } = useAuthorized(); - - return ; -}; - -export default APIPlaygroundPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/budgets/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/budgets/page.tsx deleted file mode 100644 index e49bd342c05..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/budgets/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import BudgetPanel from "@/components/budgets/budget_panel"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const BudgetsPage = () => { - const { accessToken } = useAuthorized(); - - return ; -}; - -export default BudgetsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/caching/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/caching/page.tsx deleted file mode 100644 index 6dcbcdc697c..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/caching/page.tsx +++ /dev/null @@ -1,20 +0,0 @@ -"use client"; - -import CacheDashboard from "@/components/cache_dashboard"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const CachingPage = () => { - const { token, accessToken, userRole, userId, premiumUser } = useAuthorized(); - - return ( - - ); -}; - -export default CachingPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/claude-code-plugins/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/claude-code-plugins/page.tsx deleted file mode 100644 index c92c39639c6..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/claude-code-plugins/page.tsx +++ /dev/null @@ -1,17 +0,0 @@ -"use client"; - -import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const ClaudeCodePluginsPage = () => { - const { accessToken, userRole } = useAuthorized(); - - return ( - - ); -}; - -export default ClaudeCodePluginsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/old-usage/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/old-usage/page.tsx deleted file mode 100644 index 9521f4f69f1..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/old-usage/page.tsx +++ /dev/null @@ -1,23 +0,0 @@ -"use client"; - -import Usage from "@/components/usage"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useState } from "react"; - -const OldUsagePage = () => { - const { accessToken, token, userRole, userId, premiumUser } = useAuthorized(); - const [keys, setKeys] = useState([]); - - return ( - - ); -}; - -export default OldUsagePage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/prompts/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/prompts/page.tsx deleted file mode 100644 index 0836a03b7e7..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/prompts/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import PromptsPanel from "@/components/prompts"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const PromptsPage = () => { - const { accessToken } = useAuthorized(); - - return ; -}; - -export default PromptsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/experimental/tag-management/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/experimental/tag-management/page.tsx deleted file mode 100644 index 0e686387b34..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/experimental/tag-management/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import TagManagement from "@/components/tag_management"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const TagManagementPage = () => { - const { accessToken, userId, userRole } = useAuthorized(); - - return ; -}; - -export default TagManagementPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx deleted file mode 100644 index 50cee215eb9..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import GuardrailsPanel from "@/components/guardrails"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const GuardrailsPage = () => { - const { accessToken } = useAuthorized(); - - return ; -}; - -export default GuardrailsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useHideAgentPlatformBanner.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useHideAgentPlatformBanner.ts new file mode 100644 index 00000000000..15b44a5abc6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useHideAgentPlatformBanner.ts @@ -0,0 +1,36 @@ +// hooks/useHideAgentPlatformBanner.ts +import { useSyncExternalStore } from "react"; +import { getLocalStorageItem, LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; + +export const HIDE_AGENT_PLATFORM_BANNER_KEY = "litellmHideAgentPlatformBanner"; + +function subscribe(callback: () => void) { + const onStorage = (e: StorageEvent) => { + if (e.key === HIDE_AGENT_PLATFORM_BANNER_KEY) { + callback(); + } + }; + + const onCustom = (e: Event) => { + const { key } = (e as CustomEvent).detail; + if (key === HIDE_AGENT_PLATFORM_BANNER_KEY) { + callback(); + } + }; + + window.addEventListener("storage", onStorage); + window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); + + return () => { + window.removeEventListener("storage", onStorage); + window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); + }; +} + +function getSnapshot() { + return getLocalStorageItem(HIDE_AGENT_PLATFORM_BANNER_KEY) === "true"; +} + +export function useHideAgentPlatformBanner() { + return useSyncExternalStore(subscribe, getSnapshot); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 13e93798def..5bb55ee8d10 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -30,14 +30,12 @@ function withBase(path: string): string { * * Key = legacy page id used in leftnav, Value = route segment under (dashboard)/ */ -const MIGRATED_PAGES: Record = { - "api-reference": "api-reference", -}; +const MIGRATED_PAGES: Record = {}; function LayoutContent({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); - const { accessToken, userRole, userId, userEmail, premiumUser } = useAuthorized(); + const { accessToken } = useAuthorized(); const [sidebarCollapsed, setSidebarCollapsed] = React.useState(false); const [page, setPage] = useState(() => { return searchParams.get("page") || "api-keys"; @@ -70,15 +68,9 @@ function LayoutContent({ children }: { children: React.ReactNode }) { isPublicPage={false} sidebarCollapsed={sidebarCollapsed} onToggleSidebar={toggleSidebar} - userID={userId} - userEmail={userEmail} - userRole={userRole} - premiumUser={premiumUser} proxySettings={undefined} setProxySettings={() => { }} accessToken={accessToken} - isDarkMode={false} - toggleDarkMode={() => { }} />
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx deleted file mode 100644 index 43ce427131b..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx +++ /dev/null @@ -1,20 +0,0 @@ -"use client"; - -import SpendLogsTable from "@/components/view_logs"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const LogsPage = () => { - const { accessToken, token, userRole, userId, premiumUser } = useAuthorized(); - - return ( - - ); -}; - -export default LogsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx deleted file mode 100644 index c37a935976b..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import ModelHubTable from "@/components/AIHub/ModelHubTable"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const ModelHubPage = () => { - const { accessToken, premiumUser, userRole } = useAuthorized(); - - return ; -}; - -export default ModelHubPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx deleted file mode 100644 index c1f6ec51d78..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx +++ /dev/null @@ -1,17 +0,0 @@ -"use client"; - -import PoliciesPanel from "@/components/policies"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const PoliciesPage = () => { - const { accessToken, userRole } = useAuthorized(); - - return ( - - ); -}; - -export default PoliciesPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/settings/admin-settings/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/settings/admin-settings/page.tsx deleted file mode 100644 index 8dae33afe7e..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/settings/admin-settings/page.tsx +++ /dev/null @@ -1,13 +0,0 @@ -"use client"; - -import AdminPanel from "@/components/AdminPanel"; - -const AdminSettings = () => { - - return ( - - ); -}; - -export default AdminSettings; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/settings/logging-and-alerts/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/settings/logging-and-alerts/page.tsx deleted file mode 100644 index b13e3c42f9e..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/settings/logging-and-alerts/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import Settings from "@/components/settings"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const LoggingAndAlertsPage = () => { - const { accessToken, userRole, userId, premiumUser } = useAuthorized(); - - return ; -}; - -export default LoggingAndAlertsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/settings/router-settings/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/settings/router-settings/page.tsx deleted file mode 100644 index 2b5463cd81f..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/settings/router-settings/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import GeneralSettings from "@/components/general_settings"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const RouterSettingsPage = () => { - const { accessToken, userRole, userId } = useAuthorized(); - - return ; -}; - -export default RouterSettingsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/settings/ui-theme/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/settings/ui-theme/page.tsx deleted file mode 100644 index c6826cf11df..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/settings/ui-theme/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import UIThemeSettings from "@/components/ui_theme_settings"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const UIThemePage = () => { - const { userId, userRole, accessToken } = useAuthorized(); - - return ; -}; - -export default UIThemePage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx deleted file mode 100644 index 47d2331bed8..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx +++ /dev/null @@ -1,17 +0,0 @@ -"use client"; - -import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const SkillsPage = () => { - const { accessToken, userRole } = useAuthorized(); - - return ( - - ); -}; - -export default SkillsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/TeamsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/TeamsView.tsx deleted file mode 100644 index 94a0e03304e..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/TeamsView.tsx +++ /dev/null @@ -1,370 +0,0 @@ -import React, { useState, useEffect } from "react"; -import { useQueryClient } from "@tanstack/react-query"; -import { organizationKeys } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { teamDeleteCall, Organization } from "@/components/networking"; -import { fetchTeams } from "@/components/common_components/fetch_teams"; -import { Form } from "antd"; -import TeamInfoView from "@/components/team/TeamInfo"; -import TeamSSOSettings from "@/components/TeamSSOSettings"; -import { isAdminRole } from "@/utils/roles"; -import { Card, Button, Col, Text, Grid, TabPanel } from "@tremor/react"; -import AvailableTeamsPanel from "@/components/team/available_teams"; -import type { KeyResponse, Team } from "@/components/key_team_helpers/key_list"; - -import { Member, v2TeamListCall } from "@/components/networking"; -import { updateExistingKeys } from "@/utils/dataUtils"; -import TeamsHeaderTabs from "@/app/(dashboard)/teams/components/TeamsHeaderTabs"; -import TeamsFilters from "@/app/(dashboard)/teams/components/TeamsFilters"; -import useFetchTeams from "@/app/(dashboard)/teams/hooks/useFetchTeams"; -import TeamsTable from "@/app/(dashboard)/teams/components/TeamsTable/TeamsTable"; -import DeleteTeamModal from "@/app/(dashboard)/teams/components/modals/DeleteTeamModal"; -import CreateTeamModal from "@/app/(dashboard)/teams/components/modals/CreateTeamModal"; - -interface TeamProps { - teams: Team[] | null; - accessToken: string | null; - setTeams: React.Dispatch>; - userID: string | null; - userRole: string | null; - organizations: Organization[] | null; - premiumUser?: boolean; -} - -interface FilterState { - team_id: string; - team_alias: string; - organization_id: string; - sort_by: string; - sort_order: "asc" | "desc"; -} - -interface TeamInfo { - members_with_roles: Member[]; -} - -interface PerTeamInfo { - keys: KeyResponse[]; - team_info: TeamInfo; -} - -const TeamsView: React.FC = ({ - teams, - accessToken, - setTeams, - userID, - userRole, - organizations, - premiumUser = false, -}) => { - const queryClient = useQueryClient(); - const [currentOrg, setCurrentOrg] = useState(null); - const [showFilters, setShowFilters] = useState(false); - const [filters, setFilters] = useState({ - team_id: "", - team_alias: "", - organization_id: "", - sort_by: "created_at", - sort_order: "desc", - }); - - const [form] = Form.useForm(); - const [memberForm] = Form.useForm(); - - const [selectedTeamId, setSelectedTeamId] = useState(null); - const [editTeam, setEditTeam] = useState(false); - - const [isTeamModalVisible, setIsTeamModalVisible] = useState(false); - const [isAddMemberModalVisible, setIsAddMemberModalVisible] = useState(false); - const [isEditMemberModalVisible, setIsEditMemberModalVisible] = useState(false); - const [userModels, setUserModels] = useState([]); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [teamToDelete, setTeamToDelete] = useState(null); - const [perTeamInfo, setPerTeamInfo] = useState>({}); - - const [loggingSettings, setLoggingSettings] = useState([]); - const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); - const { lastRefreshed, onRefreshClick: handleRefreshClick } = useFetchTeams({ currentOrg, setTeams }); - - useEffect(() => { - const fetchTeamInfo = () => { - if (!teams) return; - - const newPerTeamInfo = teams.reduce( - (acc, team) => { - acc[team.team_id] = { - keys: team.keys || [], - team_info: { - members_with_roles: team.members_with_roles || [], - }, - }; - return acc; - }, - {} as Record, - ); - - setPerTeamInfo(newPerTeamInfo); - }; - - fetchTeamInfo(); - }, [teams]); - - const handleOk = () => { - setIsTeamModalVisible(false); - form.resetFields(); - setLoggingSettings([]); - setModelAliases({}); - }; - - const handleMemberOk = () => { - setIsAddMemberModalVisible(false); - setIsEditMemberModalVisible(false); - memberForm.resetFields(); - }; - - const handleCancel = () => { - setIsTeamModalVisible(false); - form.resetFields(); - setLoggingSettings([]); - setModelAliases({}); - }; - - const handleDelete = async (team_id: string) => { - // Set the team to delete and open the confirmation modal - setTeamToDelete(team_id); - setIsDeleteModalOpen(true); - }; - - const confirmDelete = async () => { - if (teamToDelete == null || teams == null || accessToken == null) { - return; - } - - try { - await teamDeleteCall(accessToken, teamToDelete); - queryClient.invalidateQueries({ queryKey: organizationKeys.all }); - // Successfully completed the deletion. Update the state to trigger a rerender. - fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); - } catch (error) { - console.error("Error deleting the team:", error); - // Handle any error situations, such as displaying an error message to the user. - } - - // Close the confirmation modal and reset the teamToDelete - setIsDeleteModalOpen(false); - setTeamToDelete(null); - }; - - const cancelDelete = () => { - // Close the confirmation modal and reset the teamToDelete - setIsDeleteModalOpen(false); - setTeamToDelete(null); - }; - - const is_team_admin = (team: any) => { - if (team == null || team.members_with_roles == null) { - return false; - } - for (let i = 0; i < team.members_with_roles.length; i++) { - let member = team.members_with_roles[i]; - if (member.user_id == userID && member.role == "admin") { - return true; - } - } - return false; - }; - - const handleFilterChange = (key: keyof FilterState, value: string) => { - const newFilters = { ...filters, [key]: value }; - setFilters(newFilters); - // Call teamListCall with the new filters - if (accessToken) { - v2TeamListCall( - accessToken, - newFilters.organization_id || null, - null, - newFilters.team_id || null, - newFilters.team_alias || null, - ) - .then((response) => { - if (response && response.teams) { - setTeams(response.teams); - } - }) - .catch((error) => { - console.error("Error fetching teams:", error); - }); - } - }; - - const handleSortChange = (sortBy: string, sortOrder: "asc" | "desc") => { - const newFilters = { - ...filters, - sort_by: sortBy, - sort_order: sortOrder, - }; - setFilters(newFilters); - // Call teamListCall with the new sort parameters - if (accessToken) { - v2TeamListCall( - accessToken, - filters.organization_id || null, - null, - filters.team_id || null, - filters.team_alias || null, - ) - .then((response) => { - if (response && response.teams) { - setTeams(response.teams); - } - }) - .catch((error) => { - console.error("Error fetching teams:", error); - }); - } - }; - - const handleFilterReset = () => { - setFilters({ - team_id: "", - team_alias: "", - organization_id: "", - sort_by: "created_at", - sort_order: "desc", - }); - // Reset teams list - if (accessToken) { - v2TeamListCall(accessToken, null, userID || null, null, null) - .then((response) => { - if (response && response.teams) { - setTeams(response.teams); - } - }) - .catch((error) => { - console.error("Error fetching teams:", error); - }); - } - }; - - return ( -
- -
- {(userRole == "Admin" || userRole == "Org Admin") && ( - - )} - {selectedTeamId ? ( - { - setTeams((teams) => { - if (teams == null) { - return teams; - } - const updated = teams.map((team) => { - if (data.team_id === team.team_id) { - return updateExistingKeys(team, data); - } - return team; - }); - // Minimal fix: refresh the full team list after an update - if (accessToken) { - fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); - } - return updated; - }); - }} - onClose={() => { - setSelectedTeamId(null); - setEditTeam(false); - }} - accessToken={accessToken} - is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} - is_proxy_admin={userRole == "Admin"} - is_org_admin={(() => { - const team = teams?.find((t) => t.team_id === selectedTeamId); - if (!team?.organization_id || !organizations || !userID) return false; - const org = organizations.find((o) => o.organization_id === team.organization_id); - return org?.members?.some((m: any) => m.user_id === userID && m.user_role === "org_admin") ?? false; - })()} - userModels={userModels} - editTeam={editTeam} - premiumUser={premiumUser} - /> - ) : ( - - - - Click on “Team ID” to view team details and manage team members. - - - - -
-
- -
-
- - {isDeleteModalOpen && ( - - )} - - - - - - - - {isAdminRole(userRole || "") && ( - - - - )} - - )} - {(userRole == "Admin" || userRole == "Org Admin") && ( - - )} - - - - ); -}; - -export default TeamsView; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsFilters.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsFilters.test.tsx deleted file mode 100644 index 9a818c27624..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsFilters.test.tsx +++ /dev/null @@ -1,151 +0,0 @@ -import { render, screen, within } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import React from "react"; -import { describe, expect, it, vi } from "vitest"; -import { Organization } from "@/components/networking"; -import TeamsFilters from "./TeamsFilters"; - -type FilterState = { - team_id: string; - team_alias: string; - organization_id: string; - sort_by: string; - sort_order: "asc" | "desc"; -}; - -const emptyFilters: FilterState = { - team_alias: "", - team_id: "", - organization_id: "", - sort_by: "", - sort_order: "asc", -}; - -const mockOrganizations: Organization[] = [ - { organization_id: "org-1", organization_alias: "Acme Corp" } as Organization, - { organization_id: "org-2", organization_alias: "Globex" } as Organization, -]; - -const renderFilters = (overrides: Partial[0]> = {}) => { - const defaults = { - filters: emptyFilters, - organizations: mockOrganizations, - showFilters: false, - onToggleFilters: vi.fn(), - onChange: vi.fn(), - onReset: vi.fn(), - }; - return render(); -}; - -describe("TeamsFilters", () => { - it("should render the team name search input, Filters button, and Reset Filters button", () => { - renderFilters(); - - expect(screen.getByPlaceholderText("Search by Team Name...")).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /^filters$/i })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /reset filters/i })).toBeInTheDocument(); - }); - - it("should reflect the current team_alias filter value in the search input", () => { - renderFilters({ filters: { ...emptyFilters, team_alias: "Platform" } }); - - expect(screen.getByPlaceholderText("Search by Team Name...")).toHaveValue("Platform"); - }); - - it("should call onChange with 'team_alias' key when the search input changes", async () => { - const user = userEvent.setup(); - const onChange = vi.fn(); - renderFilters({ onChange }); - - await user.type(screen.getByPlaceholderText("Search by Team Name..."), "Dev"); - - expect(onChange).toHaveBeenCalledWith("team_alias", expect.stringContaining("D")); - }); - - it("should call onToggleFilters with the inverted boolean when the Filters button is clicked", async () => { - const user = userEvent.setup(); - const onToggleFilters = vi.fn(); - renderFilters({ showFilters: false, onToggleFilters }); - - await user.click(screen.getByRole("button", { name: /^filters$/i })); - - expect(onToggleFilters).toHaveBeenCalledWith(true); - }); - - it("should call onToggleFilters(false) when filters are currently expanded", async () => { - const user = userEvent.setup(); - const onToggleFilters = vi.fn(); - renderFilters({ showFilters: true, onToggleFilters }); - - await user.click(screen.getByRole("button", { name: /^filters$/i })); - - expect(onToggleFilters).toHaveBeenCalledWith(false); - }); - - it("should call onReset when the Reset Filters button is clicked", async () => { - const user = userEvent.setup(); - const onReset = vi.fn(); - renderFilters({ onReset }); - - await user.click(screen.getByRole("button", { name: /reset filters/i })); - - expect(onReset).toHaveBeenCalledTimes(1); - }); - - it("should not show the Team ID input when showFilters is false", () => { - renderFilters({ showFilters: false }); - - expect(screen.queryByPlaceholderText("Enter Team ID")).not.toBeInTheDocument(); - }); - - it("should show the Team ID input when showFilters is true", () => { - renderFilters({ showFilters: true }); - - expect(screen.getByPlaceholderText("Enter Team ID")).toBeInTheDocument(); - }); - - it("should call onChange with 'team_id' key when the Team ID input changes", async () => { - const user = userEvent.setup(); - const onChange = vi.fn(); - renderFilters({ showFilters: true, onChange }); - - await user.type(screen.getByPlaceholderText("Enter Team ID"), "abc"); - - expect(onChange).toHaveBeenCalledWith("team_id", expect.stringContaining("a")); - }); - - it("should reflect the current team_id filter value in the Team ID input", () => { - renderFilters({ showFilters: true, filters: { ...emptyFilters, team_id: "team-xyz" } }); - - expect(screen.getByPlaceholderText("Enter Team ID")).toHaveValue("team-xyz"); - }); - - it("should show the active filter indicator on the Filters button when team_alias is set", () => { - renderFilters({ filters: { ...emptyFilters, team_alias: "Platform" } }); - - const filtersButton = screen.getByRole("button", { name: /^filters$/i }); - expect(within(filtersButton).getByTestId("active-filter-indicator")).toBeInTheDocument(); - }); - - it("should show the active filter indicator on the Filters button when team_id is set", () => { - renderFilters({ filters: { ...emptyFilters, team_id: "team-123" } }); - - const filtersButton = screen.getByRole("button", { name: /^filters$/i }); - expect(within(filtersButton).getByTestId("active-filter-indicator")).toBeInTheDocument(); - }); - - it("should show the active filter indicator on the Filters button when organization_id is set", () => { - renderFilters({ filters: { ...emptyFilters, organization_id: "org-1" } }); - - const filtersButton = screen.getByRole("button", { name: /^filters$/i }); - expect(within(filtersButton).getByTestId("active-filter-indicator")).toBeInTheDocument(); - }); - - it("should not show the active filter indicator when all filters are empty", () => { - renderFilters({ filters: emptyFilters }); - - const filtersButton = screen.getByRole("button", { name: /^filters$/i }); - expect(within(filtersButton).queryByTestId("active-filter-indicator")).not.toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsFilters.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsFilters.tsx deleted file mode 100644 index 04c65ffe268..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsFilters.tsx +++ /dev/null @@ -1,141 +0,0 @@ -import { Select, SelectItem } from "@tremor/react"; -import React from "react"; -import { Organization } from "@/components/networking"; - -interface TeamsFiltersProps { - filters: FilterState; - organizations: Organization[] | null; - showFilters: boolean; - onToggleFilters: (toggle: boolean) => void; - onChange: (key: K, value: FilterState[K]) => void; - onReset: () => void; -} - -type FilterState = { - team_id: string; - team_alias: string; - organization_id: string; - sort_by: string; - sort_order: "asc" | "desc"; -}; - -const TeamsFilters = ({ - filters, - organizations, - showFilters, - onToggleFilters, - onChange, - onReset, -}: TeamsFiltersProps) => { - return ( -
- {/* Search and Filter Controls */} -
- {/* Team Alias Search */} -
- onChange("team_alias", e.target.value)} - /> - - - -
- - {/* Filter Button */} - - - {/* Reset Filters Button */} - -
- - {/* Additional Filters */} - {showFilters && ( -
- {/* Team ID Search */} -
- onChange("team_id", e.target.value)} - /> - - - -
- - {/* Organization Dropdown */} -
- -
-
- )} -
- ); -}; - -export default TeamsFilters; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx deleted file mode 100644 index 50a7f10f047..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx +++ /dev/null @@ -1,54 +0,0 @@ -import { render, screen } from "@testing-library/react"; -import React from "react"; -import { describe, expect, it, vi } from "vitest"; -import TeamsHeaderTabs from "./TeamsHeaderTabs"; - -vi.mock("@tremor/react", () => ({ - TabGroup: ({ children, ...props }: any) =>
{children}
, - TabList: ({ children, ...props }: any) =>
{children}
, - Tab: ({ children, ...props }: any) => , - TabPanels: ({ children, ...props }: any) =>
{children}
, - Text: ({ children, ...props }: any) => {children}, - Icon: ({ onClick, ...props }: any) => {children}, - TableBody: ({ children }: any) => {children}, - TableRow: ({ children }: any) => {children}, - TableHeaderCell: ({ children }: any) => , - TableCell: ({ children, ...props }: any) => , - Text: ({ children }: any) => {children}, -})); - -vi.mock("antd", () => ({ - Tooltip: ({ children }: any) => <>{children}, -})); - -vi.mock("@heroicons/react/outline", () => ({ - PencilAltIcon: () => , - TrashIcon: () => , -})); - -vi.mock("@/utils/dataUtils", () => ({ - formatNumberWithCommas: (val: number, decimals: number) => - val != null ? val.toFixed(decimals) : "N/A", -})); - -vi.mock("@/app/(dashboard)/teams/components/TeamsTable/ModelsCell", () => ({ - default: ({ team }: any) => , -})); - -vi.mock("@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell", () => ({ - default: ({ team }: any) => , -})); - -const makeTeam = (overrides: Partial = {}): Team => ({ - team_id: "team-abc1234", - team_alias: "Platform", - models: ["gpt-4"], - max_budget: 500, - budget_duration: null, - tpm_limit: null, - rpm_limit: null, - organization_id: "org-1", - created_at: "2024-06-01T00:00:00Z", - keys: [], - members_with_roles: [], - spend: 123.4567, - ...overrides, -}); - -const defaultPerTeamInfo = { - "team-abc1234": { - keys: [{ token: "tok-1" } as any, { token: "tok-2" } as any], - team_info: { - members_with_roles: [{ user_id: "u1", role: "admin" } as any], - }, - }, -}; - -const renderTable = (overrides: Partial[0]> = {}) => { - const defaults = { - teams: [makeTeam()], - currentOrg: null, - perTeamInfo: defaultPerTeamInfo, - userRole: "Admin", - userId: "user-1", - setSelectedTeamId: vi.fn(), - setEditTeam: vi.fn(), - onDeleteTeam: vi.fn(), - }; - return render(); -}; - -describe("TeamsTable", () => { - it("should render table headers", () => { - renderTable(); - - expect(screen.getByText("Team Name")).toBeInTheDocument(); - expect(screen.getByText("Team ID")).toBeInTheDocument(); - expect(screen.getByText("Created")).toBeInTheDocument(); - expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); - expect(screen.getByText("Budget (USD)")).toBeInTheDocument(); - expect(screen.getByText("Models")).toBeInTheDocument(); - expect(screen.getByText("Organization")).toBeInTheDocument(); - expect(screen.getByText("Your Role")).toBeInTheDocument(); - expect(screen.getByText("Info")).toBeInTheDocument(); - }); - - it("should render team rows with team data", () => { - renderTable(); - - expect(screen.getByText("Platform")).toBeInTheDocument(); - expect(screen.getByText("team-ab...")).toBeInTheDocument(); - expect(screen.getByText("org-1")).toBeInTheDocument(); - }); - - it("should show edit and delete icons for Admin users", () => { - renderTable({ userRole: "Admin" }); - - expect(screen.getAllByTestId("icon-btn").length).toBeGreaterThanOrEqual(2); - }); - - it("should not show edit and delete icons for non-Admin users", () => { - renderTable({ userRole: "Internal User" }); - - // Only the team ID button should be present, no icon-btn for edit/delete - const iconBtns = screen.queryAllByTestId("icon-btn"); - expect(iconBtns).toHaveLength(0); - }); - - it("should call setSelectedTeamId when team ID button is clicked", async () => { - const user = userEvent.setup(); - const setSelectedTeamId = vi.fn(); - renderTable({ setSelectedTeamId }); - - await user.click(screen.getByText("team-ab...")); - - expect(setSelectedTeamId).toHaveBeenCalledWith("team-abc1234"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.tsx deleted file mode 100644 index f881065d4ab..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.tsx +++ /dev/null @@ -1,166 +0,0 @@ -import { Button, Icon, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text } from "@tremor/react"; -import { Tooltip } from "antd"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import { PencilAltIcon, TrashIcon } from "@heroicons/react/outline"; -import React from "react"; -import { type KeyResponse, Team } from "@/components/key_team_helpers/key_list"; -import { Member, Organization } from "@/components/networking"; -import ModelsCell from "@/app/(dashboard)/teams/components/TeamsTable/ModelsCell"; -import YourRoleCell from "@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell"; - -type TeamsTableProps = { - teams: Team[] | null; - currentOrg: Organization | null; - perTeamInfo: Record; - userRole: string | null; - userId: string | null; - setSelectedTeamId: (teamId: string) => void; - setEditTeam: (editTeam: boolean) => void; - onDeleteTeam: (teamId: string) => void; -}; - -interface TeamInfo { - members_with_roles: Member[]; -} - -interface PerTeamInfo { - keys: KeyResponse[]; - team_info: TeamInfo; -} - -const TeamsTable = ({ - teams, - currentOrg, - setSelectedTeamId, - perTeamInfo, - userRole, - userId, - setEditTeam, - onDeleteTeam, -}: TeamsTableProps) => { - return ( -
from TableCell renders without HTML warnings. -const renderModelsCell = (team: Team) => - render( - - - - - - -
, - ); - -describe("ModelsCell", () => { - it("should show 'All Proxy Models' badge when the models array is empty", () => { - renderModelsCell(makeTeam([])); - - expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); - }); - - it("should show an 'All Proxy Models' badge when the model value is 'all-proxy-models'", () => { - renderModelsCell(makeTeam(["all-proxy-models"])); - - expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); - }); - - it("should display individual model badges for up to 3 models without an accordion", () => { - renderModelsCell(makeTeam(["gpt-4", "gpt-3.5-turbo", "claude-3"])); - - expect(screen.getByText("gpt-4")).toBeInTheDocument(); - expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); - expect(screen.getByText("claude-3")).toBeInTheDocument(); - expect(screen.queryByRole("button", { name: /accordion/i })).not.toBeInTheDocument(); - }); - - it("should truncate model names longer than 30 characters with an ellipsis", () => { - const longName = "a-very-long-model-name-exceeding-thirty-chars"; - renderModelsCell(makeTeam([longName])); - - const badge = screen.getByText((text) => text.endsWith("...")); - expect(badge).toBeInTheDocument(); - expect(badge.textContent!.length).toBeLessThanOrEqual(33); // 30 chars + "..." - }); - - it("should show the first 3 models and a '+N more models' badge when there are more than 3 models", () => { - renderModelsCell(makeTeam(["m1", "m2", "m3", "m4", "m5"])); - - expect(screen.getByText("m1")).toBeInTheDocument(); - expect(screen.getByText("m2")).toBeInTheDocument(); - expect(screen.getByText("m3")).toBeInTheDocument(); - expect(screen.getByText("+2 more models")).toBeInTheDocument(); - expect(screen.queryByText("m4")).not.toBeInTheDocument(); - expect(screen.queryByText("m5")).not.toBeInTheDocument(); - }); - - it("should use singular 'more model' when there is exactly 1 overflow model", () => { - renderModelsCell(makeTeam(["m1", "m2", "m3", "m4"])); - - expect(screen.getByText("+1 more model")).toBeInTheDocument(); - }); - - it("should show the accordion toggle button when there are more than 3 models", () => { - renderModelsCell(makeTeam(["m1", "m2", "m3", "m4"])); - - expect(screen.getByRole("button", { name: /accordion/i })).toBeInTheDocument(); - }); - - it("should expand to show all models when the accordion toggle is clicked", () => { - renderModelsCell(makeTeam(["m1", "m2", "m3", "m4", "m5"])); - - act(() => { - screen.getByRole("button", { name: /accordion/i }).click(); - }); - - expect(screen.getByText("m4")).toBeInTheDocument(); - expect(screen.getByText("m5")).toBeInTheDocument(); - expect(screen.queryByText("+2 more models")).not.toBeInTheDocument(); - }); - - it("should collapse back to show the overflow badge after a second click on the toggle", () => { - renderModelsCell(makeTeam(["m1", "m2", "m3", "m4", "m5"])); - - const toggle = screen.getByRole("button", { name: /accordion/i }); - act(() => { - toggle.click(); - }); - act(() => { - toggle.click(); - }); - - expect(screen.queryByText("m4")).not.toBeInTheDocument(); - expect(screen.getByText("+2 more models")).toBeInTheDocument(); - }); - - it("should collapse to a single 'All Proxy Models' badge when the models list includes 'all-proxy-models'", () => { - renderModelsCell(makeTeam(["m1", "m2", "m3", "all-proxy-models"])); - - // When all-proxy-models is present, all individual models are hidden and no accordion is shown - expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); - expect(screen.queryByText("m1")).not.toBeInTheDocument(); - expect(screen.queryByText("m2")).not.toBeInTheDocument(); - expect(screen.queryByText("m3")).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: /accordion/i })).not.toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx deleted file mode 100644 index 62a7fdb783f..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/ModelsCell.tsx +++ /dev/null @@ -1,107 +0,0 @@ -import { Badge, Icon, TableCell, Text } from "@tremor/react"; -import { ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; -import React, { useMemo, useState } from "react"; -import { Team } from "@/components/key_team_helpers/key_list"; - -interface ModelsCellProps { - team: Team; -} - -interface ModelEntry { - name: string; - source: "direct" | "access_group"; -} - -const ModelsCell = ({ team }: ModelsCellProps) => { - const [expandedAccordion, setExpandedAccordion] = useState(false); - - const isAllModels = !team.models || team.models.length === 0 || team.models.includes("all-proxy-models"); - - const modelEntries: ModelEntry[] = useMemo(() => { - if (isAllModels) return []; - const entries: ModelEntry[] = team.models.map((m) => ({ - name: m, - source: "direct" as const, - })); - for (const m of team.access_group_models || []) { - entries.push({ name: m, source: "access_group" }); - } - return entries; - }, [team.models, team.access_group_models, isAllModels]); - - const renderBadge = (entry: ModelEntry, index: number) => { - if (entry.name === "all-proxy-models") { - return ( - - All Proxy Models - - ); - } - const displayName = getModelDisplayName(entry.name); - const truncated = displayName.length > 30 ? `${displayName.slice(0, 30)}...` : displayName; - return ( - - {truncated} - - ); - }; - - return ( - 3 ? "px-0" : ""} - > -
- {modelEntries.length === 0 ? ( - - All Proxy Models - - ) : ( -
-
- {modelEntries.length > 3 && ( -
- { - setExpandedAccordion((prev) => !prev); - }} - /> -
- )} -
- {modelEntries.slice(0, 3).map((entry, index) => renderBadge(entry, index))} - {modelEntries.length > 3 && !expandedAccordion && ( - - - +{modelEntries.length - 3} {modelEntries.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordion && ( -
- {modelEntries.slice(3).map((entry, index) => renderBadge(entry, index + 3))} -
- )} -
-
-
- )} -
- - ); -}; - -export default ModelsCell; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx deleted file mode 100644 index 6b072ababb6..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx +++ /dev/null @@ -1,129 +0,0 @@ -import { render, screen } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import React from "react"; -import { describe, expect, it, vi } from "vitest"; -import { Team } from "@/components/key_team_helpers/key_list"; -import TeamsTable from "./TeamsTable"; - -vi.mock("@tremor/react", () => ({ - Button: React.forwardRef(({ children, ...props }, ref) => - React.createElement("button", { ...props, ref }, children), - ), - Icon: ({ onClick, ...props }: any) =>
{children}{children}{team.models.join(",")}{team.team_id}
- - - Team Name - Team ID - Created - Spend (USD) - Budget (USD) - Models - Organization - Your Role - Info - - - - - {teams && teams.length > 0 - ? teams - .filter((team) => { - if (!currentOrg) return true; - return team.organization_id === currentOrg.organization_id; - }) - .sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime()) - .map((team: any) => ( - - - {team["team_alias"]} - - -
- - - -
- - - {team.created_at ? new Date(team.created_at).toLocaleDateString() : "N/A"} - - - {formatNumberWithCommas(team["spend"], 4)} - - - {team["max_budget"] !== null && team["max_budget"] !== undefined ? team["max_budget"] : "No limit"} - - - {team.organization_id} - - - - {perTeamInfo && - team.team_id && - perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].keys && - perTeamInfo[team.team_id].keys.length}{" "} - Keys - - - {perTeamInfo && - team.team_id && - perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].team_info && - perTeamInfo[team.team_id].team_info.members_with_roles && - perTeamInfo[team.team_id].team_info.members_with_roles.length}{" "} - Members - - - - {userRole == "Admin" ? ( - <> - { - setSelectedTeamId(team.team_id); - setEditTeam(true); - }} - /> - onDeleteTeam(team.team_id)} icon={TrashIcon} size="sm" /> - - ) : null} - - - )) - : null} - -
- ); -}; - -export default TeamsTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge.test.tsx deleted file mode 100644 index b7e659403cb..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge.test.tsx +++ /dev/null @@ -1,38 +0,0 @@ -import React from "react"; -import { describe, it, expect } from "vitest"; -import { render, screen } from "@testing-library/react"; -import "@testing-library/jest-dom"; -import TeamRoleBadge from "./TeamRoleBadge"; - -const renderBadge = (role: string | null) => render(
{TeamRoleBadge(role)}
); - -describe("TeamRoleBadge", () => { - it("renders admin badge with correct label, base classes, styles, and an icon", () => { - renderBadge("admin"); - const label = screen.getByText("Admin"); - const badge = label.closest("span")!; - expect(badge).toHaveClass("inline-flex", "items-center", "border", "text-xs", "font-medium"); - expect(badge).toHaveStyle({ - backgroundColor: "#EEF2FF", - color: "#3730A3", - borderColor: "#C7D2FE", - }); - expect(badge.querySelector("svg")).toBeInTheDocument(); // ShieldIcon renders as an SVG - }); - - it.each<[string | null]>([["user"], [null], ["viewer" as unknown as string]])( - "renders member badge for non-admin role (%p) with correct styles", - (role) => { - renderBadge(role); - const label = screen.getByText("Member"); - const badge = label.closest("span")!; - expect(badge).toHaveClass("inline-flex", "items-center", "border", "text-xs", "font-medium"); - expect(badge).toHaveStyle({ - backgroundColor: "#F3F4F6", - color: "#4B5563", - borderColor: "#E5E7EB", - }); - expect(badge.querySelector("svg")).toBeInTheDocument(); // UserIcon renders as an SVG - }, - ); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge.tsx deleted file mode 100644 index 394b5d85348..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge.tsx +++ /dev/null @@ -1,47 +0,0 @@ -import { ShieldIcon, UserIcon } from "lucide-react"; - -const MEMBER_BADGE_BG = "#F3F4F6"; // gray-100 -const MEMBER_BADGE_TEXT = "#4B5563"; // gray-600 -const MEMBER_BADGE_BORDER = "#E5E7EB"; // gray-200 - -const ADMIN_BADGE_BG = "#EEF2FF"; // indigo-50 -const ADMIN_BADGE_TEXT = "#3730A3"; // indigo-800 -const ADMIN_BADGE_BORDER = "#C7D2FE"; // indigo-200 - -const TeamRoleBadge = (role: string | null) => { - const base = "inline-flex items-center px-2.5 py-0.5 rounded-md text-xs font-medium border"; - - switch (role) { - case "admin": - return ( - - - Admin - - ); - case "user": - default: - return ( - - - Member - - ); - } -}; - -export default TeamRoleBadge; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell.test.tsx deleted file mode 100644 index 20a4497159d..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell.test.tsx +++ /dev/null @@ -1,43 +0,0 @@ -import React from "react"; -import { describe, it, expect, vi } from "vitest"; -import { render, screen } from "@testing-library/react"; -import "@testing-library/jest-dom"; -import type { Team } from "@/components/key_team_helpers/key_list"; -import YourRoleCell from "./YourRoleCell"; - -// Lightweight mocks for stable, focused tests -vi.mock("@tremor/react", () => ({ - TableCell: ({ children }: { children: React.ReactNode }) =>
{children}
, -})); - -// The component invokes TeamRoleBadge as a function, so mock it as such -vi.mock("@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge", () => ({ - __esModule: true, - default: (role: string | null) => {role === "admin" ? "Admin" : "Member"}, -})); - -const team = (members?: Array<{ user_id: string; role: "admin" | "user" }>): Team => - ({ members_with_roles: members }) as unknown as Team; - -describe("YourRoleCell", () => { - it("renders Admin when the user is an admin of the team", () => { - render(); - expect(screen.getByTestId("cell")).toBeInTheDocument(); - expect(screen.getByTestId("badge")).toHaveTextContent("Admin"); - }); - - it("renders Member when the user is a regular member", () => { - render(); - expect(screen.getByTestId("badge")).toHaveTextContent("Member"); - }); - - it.each<[string, Team, string | null]>([ - ["userId is null", team([{ user_id: "u3", role: "admin" }]), null], - ["user not in team", team([{ user_id: "x", role: "user" }]), "y"], - ["team has no members", team([]), "u4"], - ["members field undefined", team(undefined), "u5"], - ])("falls back to Member when no role can be determined (%s)", (_label, t, uid) => { - render(); - expect(screen.getByTestId("badge")).toHaveTextContent("Member"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell.tsx deleted file mode 100644 index 66592943fd5..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell.tsx +++ /dev/null @@ -1,22 +0,0 @@ -import { TableCell } from "@tremor/react"; -import { Team } from "@/components/key_team_helpers/key_list"; -import TeamRoleBadge from "@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/TeamRoleBadge"; - -interface YourRoleCellProps { - team: Team; - userId: string | null; -} - -const getUserRole = (team: Team, userId: string | null): string | null => { - if (!userId) return null; - const member = team.members_with_roles?.find((m) => m.user_id === userId); - return member?.role ?? null; -}; - -const YourRoleCell = ({ team, userId }: YourRoleCellProps) => { - const roleBadge = TeamRoleBadge(getUserRole(team, userId)); - - return {roleBadge}; -}; - -export default YourRoleCell; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx deleted file mode 100644 index 1a8c6632a03..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx +++ /dev/null @@ -1,822 +0,0 @@ -import { Button as Button2, Form, Input, Modal, Select as Select2, Switch, Tooltip } from "antd"; -import { Accordion, AccordionBody, AccordionHeader, Text, TextInput } from "@tremor/react"; -import { InfoCircleOutlined } from "@ant-design/icons"; -import { - fetchAvailableModelsForTeamOrKey, - getModelDisplayName, - unfurlWildcardModelsInList, -} from "@/components/key_team_helpers/fetch_available_models_team_key"; -import NumericalInput from "@/components/shared/numerical_input"; -import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector"; -import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; -import AgentSelector from "@/components/agent_management/AgentSelector"; -import PremiumLoggingSettings from "@/components/common_components/PremiumLoggingSettings"; -import ModelAliasManager from "@/components/common_components/ModelAliasManager"; -import React, { useEffect, useState } from "react"; -import { useQueryClient } from "@tanstack/react-query"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { - fetchMCPAccessGroups, - getGuardrailsList, - getPoliciesList, - Organization, - Team, - teamCreateCall, -} from "@/components/networking"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { organizationKeys } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; -import SearchToolSelector from "@/components/SearchTools/SearchToolSelector"; - -interface ModelAliases { - [key: string]: string; -} - -interface CreateTeamModalProps { - isTeamModalVisible: boolean; - handleOk: () => void; - handleCancel: () => void; - currentOrg: Organization | null; - organizations: Organization[] | null; - teams: Team[] | null; - setTeams: (teams: Team[] | null) => void; - modelAliases: ModelAliases; - setModelAliases: (modelAliases: ModelAliases) => void; - loggingSettings: any[]; - setLoggingSettings: (loggingSettings: any[]) => void; - setIsTeamModalVisible: (isTeamModalVisible: boolean) => void; -} - -const getOrganizationModels = (organization: Organization | null, userModels: string[]) => { - let tempModelsToPick = []; - - if (organization) { - if (organization.models.length > 0) { - console.log(`organization.models: ${organization.models}`); - tempModelsToPick = organization.models; - } else { - // show all available models if the team has no models set - tempModelsToPick = userModels; - } - } else { - // no team set, show all available models - tempModelsToPick = userModels; - } - - return unfurlWildcardModelsInList(tempModelsToPick, userModels); -}; - -const CreateTeamModal = ({ - isTeamModalVisible, - handleOk, - handleCancel, - currentOrg, - organizations, - teams, - setTeams, - modelAliases, - setModelAliases, - loggingSettings, - setLoggingSettings, - setIsTeamModalVisible, -}: CreateTeamModalProps) => { - const { userId: userID, userRole, accessToken, premiumUser } = useAuthorized(); - const queryClient = useQueryClient(); - const [form] = Form.useForm(); - const [userModels, setUserModels] = useState([]); - const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState(null); - const [modelsToPick, setModelsToPick] = useState([]); - const [guardrailsList, setGuardrailsList] = useState([]); - const [policiesList, setPoliciesList] = useState([]); - const [mcpAccessGroups, setMcpAccessGroups] = useState([]); - const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false); - - useEffect(() => { - const fetchUserModels = async () => { - try { - if (userID === null || userRole === null || accessToken === null) { - return; - } - const models = await fetchAvailableModelsForTeamOrKey(userID, userRole, accessToken); - if (models) { - setUserModels(models); - } - } catch (error) { - console.error("Error fetching user models:", error); - } - }; - - fetchUserModels(); - }, [accessToken, userID, userRole, teams]); - - useEffect(() => { - console.log(`currentOrgForCreateTeam: ${currentOrgForCreateTeam}`); - const models = getOrganizationModels(currentOrgForCreateTeam, userModels); - console.log(`models: ${models}`); - setModelsToPick(models); - form.setFieldValue("models", []); - }, [currentOrgForCreateTeam, userModels, form]); - - const fetchMcpAccessGroups = async () => { - try { - if (accessToken == null) { - return; - } - const groups = await fetchMCPAccessGroups(accessToken); - setMcpAccessGroups(groups); - } catch (error) { - console.error("Failed to fetch MCP access groups:", error); - } - }; - - useEffect(() => { - fetchMcpAccessGroups(); - }, [accessToken, fetchMcpAccessGroups]); - - useEffect(() => { - const fetchGuardrails = async () => { - try { - if (accessToken == null) { - return; - } - - const response = await getGuardrailsList(accessToken); - const guardrailNames = response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); - setGuardrailsList(guardrailNames); - } catch (error) { - console.error("Failed to fetch guardrails:", error); - } - }; - - const fetchPolicies = async () => { - try { - if (accessToken == null) { - return; - } - - const response = await getPoliciesList(accessToken); - const policyNames = response.policies.map((p: { policy_name: string }) => p.policy_name); - setPoliciesList(policyNames); - } catch (error) { - console.error("Failed to fetch policies:", error); - } - }; - - fetchGuardrails(); - fetchPolicies(); - }, [accessToken]); - - const handleCreate = async (formValues: Record) => { - try { - console.log(`formValues: ${JSON.stringify(formValues)}`); - if (accessToken != null) { - const newTeamAlias = formValues?.team_alias; - const existingTeamAliases = teams?.map((t) => t.team_alias) ?? []; - let organizationId = formValues?.organization_id || currentOrg?.organization_id; - if (organizationId === "" || typeof organizationId !== "string") { - formValues.organization_id = null; - } else { - formValues.organization_id = organizationId.trim(); - } - - // Remove guardrails from top level since it's now in metadata - if (existingTeamAliases.includes(newTeamAlias)) { - throw new Error(`Team alias ${newTeamAlias} already exists, please pick another alias`); - } - - NotificationsManager.info("Creating Team"); - - // Handle logging settings in metadata - if (loggingSettings.length > 0) { - let metadata = {}; - if (formValues.metadata) { - try { - metadata = JSON.parse(formValues.metadata); - } catch (e) { - console.warn("Invalid JSON in metadata field, starting with empty object"); - } - } - - // Add logging settings to metadata - metadata = { - ...metadata, - logging: loggingSettings.filter((config) => config.callback_name), // Only include configs with callback_name - }; - - formValues.metadata = JSON.stringify(metadata); - } - - if (formValues.secret_manager_settings) { - if (typeof formValues.secret_manager_settings === "string") { - if (formValues.secret_manager_settings.trim() === "") { - delete formValues.secret_manager_settings; - } else { - try { - formValues.secret_manager_settings = JSON.parse(formValues.secret_manager_settings); - } catch (e) { - throw new Error("Failed to parse secret manager settings: " + e); - } - } - } - } - - // Transform integrations into object_permission (vector stores, MCP, agents, search tools) - const hasAgents = - formValues.allowed_agents_and_groups && - ((formValues.allowed_agents_and_groups.agents?.length ?? 0) > 0 || - (formValues.allowed_agents_and_groups.accessGroups?.length ?? 0) > 0); - const hasSearchTools = - Array.isArray(formValues.object_permission_search_tools) && - formValues.object_permission_search_tools.length > 0; - - if ( - (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) || - (formValues.allowed_mcp_servers_and_groups && - (formValues.allowed_mcp_servers_and_groups.servers?.length > 0 || - formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 || - formValues.allowed_mcp_servers_and_groups.toolPermissions)) || - hasAgents || - hasSearchTools - ) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } - if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { - formValues.object_permission.vector_stores = formValues.allowed_vector_store_ids; - delete formValues.allowed_vector_store_ids; - } - if (formValues.allowed_mcp_servers_and_groups) { - const { servers, accessGroups } = formValues.allowed_mcp_servers_and_groups; - if (servers && servers.length > 0) { - formValues.object_permission.mcp_servers = servers; - } - if (accessGroups && accessGroups.length > 0) { - formValues.object_permission.mcp_access_groups = accessGroups; - } - delete formValues.allowed_mcp_servers_and_groups; - } - - // Add tool permissions separately - if (formValues.mcp_tool_permissions && Object.keys(formValues.mcp_tool_permissions).length > 0) { - formValues.object_permission.mcp_tool_permissions = formValues.mcp_tool_permissions; - delete formValues.mcp_tool_permissions; - } - - // Handle agent permissions - if (formValues.allowed_agents_and_groups) { - const { agents, accessGroups } = formValues.allowed_agents_and_groups; - if (agents && agents.length > 0) { - formValues.object_permission.agents = agents; - } - if (accessGroups && accessGroups.length > 0) { - formValues.object_permission.agent_access_groups = accessGroups; - } - delete formValues.allowed_agents_and_groups; - } - - if (hasSearchTools) { - formValues.object_permission.search_tools = formValues.object_permission_search_tools; - delete formValues.object_permission_search_tools; - } - } - - // Transform allowed_mcp_access_groups into object_permission - if (formValues.allowed_mcp_access_groups && formValues.allowed_mcp_access_groups.length > 0) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } - formValues.object_permission.mcp_access_groups = formValues.allowed_mcp_access_groups; - delete formValues.allowed_mcp_access_groups; - } - - // Add model_aliases if any are defined - if (Object.keys(modelAliases).length > 0) { - formValues.model_aliases = modelAliases; - } - - const response: any = await teamCreateCall(accessToken, formValues); - queryClient.invalidateQueries({ queryKey: organizationKeys.all }); - if (teams !== null) { - setTeams([...teams, response]); - } else { - setTeams([response]); - } - console.log(`response for team create call: ${response}`); - NotificationsManager.success("Team created"); - form.resetFields(); - setLoggingSettings([]); - setModelAliases({}); - setIsTeamModalVisible(false); - } - } catch (error) { - console.error("Error creating the team:", error); - NotificationsManager.fromBackend("Error creating the team: " + error); - } - }; - - return ( - -
- <> - - - - - Organization{" "} - - Organizations can have multiple teams. Learn more about{" "} - e.stopPropagation()} - > - user management hierarchy - - - } - > - - - - } - name="organization_id" - initialValue={currentOrg ? currentOrg.organization_id : null} - className="mt-8" - > - { - form.setFieldValue("organization_id", value); - setCurrentOrgForCreateTeam(organizations?.find((org) => org.organization_id === value) || null); - }} - filterOption={(input, option) => { - if (!option) return false; - const optionValue = option.children?.toString() || ""; - return optionValue.toLowerCase().includes(input.toLowerCase()); - }} - optionFilterProp="children" - > - {organizations?.map((org) => ( - - {org.organization_alias}{" "} - ({org.organization_id}) - - ))} - - - - Models{" "} - - - - - } - name="models" - > - - - All Proxy Models - - {modelsToPick.map((model) => ( - - {getModelDisplayName(model)} - - ))} - - - - - - Team Member Settings - - - - Optional defaults applied when members join this team. All fields can be overridden per member. - - prev.models !== cur.models} - > - {({ getFieldValue }) => { - const teamModels: string[] = getFieldValue("models") || []; - const opts = teamModels.length > 0 ? teamModels : modelsToPick; - return ( - - Default Model Access{" "} - - - - - } - name="default_team_member_models" - > - - {opts.map((m) => ( - - {getModelDisplayName(m)} - - ))} - - - ); - }} - - (value ? Number(value) : undefined)} - tooltip="Default spend budget for each member in this team." - > - - - - - - - - - - - - - - - - - - - - daily - weekly - monthly - - - - - - - - - - { - if (!mcpAccessGroupsLoaded) { - fetchMcpAccessGroups(); - setMcpAccessGroupsLoaded(true); - } - }} - > - - Additional Settings - - - - { - e.target.value = e.target.value.trim(); - }} - /> - - - - - { - if (!value) { - return Promise.resolve(); - } - try { - JSON.parse(value); - return Promise.resolve(); - } catch (error) { - return Promise.reject(new Error("Please enter valid JSON")); - } - }, - }, - ]} - > - - - - Guardrails{" "} - - e.stopPropagation()} - > - - - - - } - name="guardrails" - className="mt-8" - help="Select existing guardrails or enter new ones" - > - ({ - value: name, - label: name, - }))} - /> - - - Disable Global Guardrails{" "} - - - - - } - name="disable_global_guardrails" - className="mt-4" - valuePropName="checked" - help="Bypass global guardrails for this team" - > - - - - Policies{" "} - - e.stopPropagation()} - > - - - - - } - name="policies" - className="mt-8" - help="Select existing policies or enter new ones" - > - ({ - value: name, - label: name, - }))} - /> - - - Allowed Vector Stores{" "} - - - - - } - name="allowed_vector_store_ids" - className="mt-8" - help="Select vector stores this team can access. Leave empty for access to all vector stores" - > - form.setFieldValue("allowed_vector_store_ids", values)} - value={form.getFieldValue("allowed_vector_store_ids")} - accessToken={accessToken || ""} - placeholder="Select vector stores (optional)" - /> - - - - - - - MCP Settings - - - - Allowed MCP Servers{" "} - - - - - } - name="allowed_mcp_servers_and_groups" - className="mt-4" - help="Select MCP servers or access groups this team can access" - > - form.setFieldValue("allowed_mcp_servers_and_groups", val)} - value={form.getFieldValue("allowed_mcp_servers_and_groups")} - accessToken={accessToken || ""} - placeholder="Select MCP servers or access groups (optional)" - /> - - - {/* Hidden field to register mcp_tool_permissions with the form */} - - - - prevValues.allowed_mcp_servers_and_groups !== currentValues.allowed_mcp_servers_and_groups || - prevValues.mcp_tool_permissions !== currentValues.mcp_tool_permissions - } - > - {() => ( -
- form.setFieldsValue({ mcp_tool_permissions: toolPerms })} - /> -
- )} - - - - - - - Agent Settings - - - - Allowed Agents{" "} - - - - - } - name="allowed_agents_and_groups" - className="mt-4" - help="Select agents or access groups this team can access" - > - form.setFieldValue("allowed_agents_and_groups", val)} - value={form.getFieldValue("allowed_agents_and_groups")} - accessToken={accessToken || ""} - placeholder="Select agents or access groups (optional)" - /> - - - - - - - Search Tool Settings - - - - Allowed Search Tools{" "} - - - - - } - name="object_permission_search_tools" - className="mt-4" - help="Restrict which configured search tools keys on this team may call." - > - form.setFieldValue("object_permission_search_tools", vals)} - value={form.getFieldValue("object_permission_search_tools")} - accessToken={accessToken || ""} - placeholder="Select search tools (optional, empty = all allowed)" - /> - - - - - - - Logging Settings - - -
- -
- - - - - - Model Aliases - - -
- - Create custom aliases for models that can be used by team members in API calls. This allows you to - create shortcuts for specific models. - - -
- - - -
- Create Team -
-
- - ); -}; - -export default CreateTeamModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/DeleteTeamModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/DeleteTeamModal.test.tsx deleted file mode 100644 index 1e4907dcca4..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/DeleteTeamModal.test.tsx +++ /dev/null @@ -1,171 +0,0 @@ -import { render, screen } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import React from "react"; -import { describe, expect, it, vi } from "vitest"; -import { Team } from "@/components/key_team_helpers/key_list"; -import DeleteTeamModal from "./DeleteTeamModal"; - -const makeTeam = (overrides: Partial = {}): Team => ({ - team_id: "team-1", - team_alias: "Engineering", - models: [], - max_budget: null, - budget_duration: null, - tpm_limit: null, - rpm_limit: null, - organization_id: "org-1", - created_at: "2024-01-01T00:00:00Z", - keys: [], - members_with_roles: [], - spend: 0, - ...overrides, -}); - -const renderModal = (props: Partial[0]> = {}) => { - const defaults = { - teams: [makeTeam()], - teamToDelete: "team-1", - onCancel: vi.fn(), - onConfirm: vi.fn(), - }; - return render(); -}; - -describe("DeleteTeamModal", () => { - it("should render the title, team name label, and confirmation input", () => { - renderModal(); - - expect(screen.getByText("Delete Team")).toBeInTheDocument(); - expect(screen.getByText("Engineering")).toBeInTheDocument(); - expect(screen.getByPlaceholderText("Enter team name exactly")).toBeInTheDocument(); - }); - - it("should render Cancel and Force Delete buttons", () => { - renderModal(); - - expect(screen.getByRole("button", { name: /^cancel$/i })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /force delete/i })).toBeInTheDocument(); - }); - - it("should not show the warning banner when the team has no keys", () => { - renderModal({ teams: [makeTeam({ keys: [] })] }); - - expect(screen.queryByText(/Warning/i)).not.toBeInTheDocument(); - }); - - it("should show a warning with singular 'key' when the team has exactly 1 key", () => { - const team = makeTeam({ keys: [{ token: "tok-1" } as any] }); - renderModal({ teams: [team] }); - - expect(screen.getByText(/This team has 1 associated key\./)).toBeInTheDocument(); - }); - - it("should show a warning with plural 'keys' when the team has multiple keys", () => { - const team = makeTeam({ - keys: [{ token: "tok-1" } as any, { token: "tok-2" } as any, { token: "tok-3" } as any], - }); - renderModal({ teams: [team] }); - - expect(screen.getByText(/This team has 3 associated keys\./)).toBeInTheDocument(); - }); - - it("should note that associated keys will also be deleted in the warning", () => { - const team = makeTeam({ keys: [{ token: "tok-1" } as any] }); - renderModal({ teams: [team] }); - - expect(screen.getByText(/Deleting the team will also delete all associated keys/)).toBeInTheDocument(); - }); - - it("should disable Force Delete when the input is empty", () => { - renderModal(); - - expect(screen.getByRole("button", { name: /force delete/i })).toBeDisabled(); - }); - - it("should keep Force Delete disabled when the input does not exactly match the team name", async () => { - const user = userEvent.setup(); - renderModal(); - - await user.type(screen.getByPlaceholderText("Enter team name exactly"), "engineer"); - - expect(screen.getByRole("button", { name: /force delete/i })).toBeDisabled(); - }); - - it("should enable Force Delete only after typing the exact team name (case-sensitive)", async () => { - const user = userEvent.setup(); - renderModal(); - - const input = screen.getByPlaceholderText("Enter team name exactly"); - - await user.type(input, "Engineering"); - - expect(screen.getByRole("button", { name: /force delete/i })).toBeEnabled(); - }); - - it("should call onConfirm when Force Delete is clicked with a valid input", async () => { - const user = userEvent.setup(); - const onConfirm = vi.fn(); - renderModal({ onConfirm }); - - await user.type(screen.getByPlaceholderText("Enter team name exactly"), "Engineering"); - await user.click(screen.getByRole("button", { name: /force delete/i })); - - expect(onConfirm).toHaveBeenCalledTimes(1); - }); - - it("should not call onConfirm when Force Delete is clicked with an invalid input", async () => { - const user = userEvent.setup(); - const onConfirm = vi.fn(); - renderModal({ onConfirm }); - - // Button is disabled so click has no effect - await user.click(screen.getByRole("button", { name: /force delete/i })); - - expect(onConfirm).not.toHaveBeenCalled(); - }); - - it("should call onCancel when the Cancel button is clicked", async () => { - const user = userEvent.setup(); - const onCancel = vi.fn(); - renderModal({ onCancel }); - - await user.click(screen.getByRole("button", { name: /^cancel$/i })); - - expect(onCancel).toHaveBeenCalledTimes(1); - }); - - it("should call onCancel when the Close button is clicked", async () => { - const user = userEvent.setup(); - const onCancel = vi.fn(); - renderModal({ onCancel }); - - await user.click(screen.getByRole("button", { name: /^close$/i })); - - expect(onCancel).toHaveBeenCalledTimes(1); - }); - - it("should reset the confirmation input when Cancel is clicked", async () => { - const user = userEvent.setup(); - renderModal(); - - const input = screen.getByPlaceholderText("Enter team name exactly"); - await user.type(input, "Engineering"); - expect(input).toHaveValue("Engineering"); - - await user.click(screen.getByRole("button", { name: /^cancel$/i })); - - expect(input).toHaveValue(""); - }); - - it("should reset the confirmation input when the Close button is clicked", async () => { - const user = userEvent.setup(); - renderModal(); - - const input = screen.getByPlaceholderText("Enter team name exactly"); - await user.type(input, "Engineering"); - - await user.click(screen.getByRole("button", { name: /^close$/i })); - - expect(input).toHaveValue(""); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/DeleteTeamModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/DeleteTeamModal.tsx deleted file mode 100644 index 28d80faacdc..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/DeleteTeamModal.tsx +++ /dev/null @@ -1,96 +0,0 @@ -import { AlertTriangleIcon, XIcon } from "lucide-react"; -import React, { useState } from "react"; -import { Team } from "@/components/key_team_helpers/key_list"; - -interface DeleteTeamModalProps { - teams: Team[] | null; - teamToDelete: string | null; - onCancel: () => void; - onConfirm: () => void; -} - -const DeleteTeamModal = ({ teams, teamToDelete, onCancel, onConfirm }: DeleteTeamModalProps) => { - const [deleteConfirmInput, setDeleteConfirmInput] = useState(""); - - const team = teams?.find((t) => t.team_id === teamToDelete); - const teamName = team?.team_alias || ""; - const keyCount = team?.keys?.length || 0; - const isValid = deleteConfirmInput === teamName; - - return ( -
-
-
-
-

Delete Team

- -
-
- {keyCount > 0 && ( -
-
- -
-
-

- Warning: This team has {keyCount} associated key{keyCount > 1 ? "s" : ""}. -

-

- Deleting the team will also delete all associated keys. This action is irreversible. -

-
-
- )} -

- Are you sure you want to force delete this team and all its keys? -

-
- - setDeleteConfirmInput(e.target.value)} - placeholder="Enter team name exactly" - className="w-full px-4 py-3 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500 text-base" - autoFocus - /> -
-
-
-
- - -
-
-
- ); -}; - -export default DeleteTeamModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/hooks/useFetchTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/teams/hooks/useFetchTeams.ts deleted file mode 100644 index c02787896f9..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/hooks/useFetchTeams.ts +++ /dev/null @@ -1,30 +0,0 @@ -import { useCallback, useEffect, useState } from "react"; -import { fetchTeams } from "@/components/common_components/fetch_teams"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { Organization, Team } from "@/components/networking"; - -interface useFetchTeamsProps { - currentOrg: Organization | null; - setTeams: (teams: Team[] | null) => void; -} - -const useFetchTeams = ({ currentOrg, setTeams }: useFetchTeamsProps) => { - const [lastRefreshed, setLastRefreshed] = useState(""); - const { accessToken, userId, userRole } = useAuthorized(); - - const onRefreshClick = useCallback(() => { - const currentDate = new Date(); - setLastRefreshed(currentDate.toLocaleString()); - }, []); - - useEffect(() => { - if (accessToken) { - fetchTeams(accessToken, userId, userRole, currentOrg, setTeams).then(); - } - onRefreshClick(); - }, [accessToken, currentOrg, lastRefreshed, onRefreshClick, setTeams, userId, userRole]); - - return { lastRefreshed, setLastRefreshed, onRefreshClick }; -}; - -export default useFetchTeams; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx deleted file mode 100644 index 041c50dd32a..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/page.tsx +++ /dev/null @@ -1,31 +0,0 @@ -"use client"; - -import TeamsView from "@/app/(dashboard)/teams/TeamsView"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import useTeams from "@/app/(dashboard)/hooks/useTeams"; -import { useEffect, useState } from "react"; -import { Organization } from "@/components/networking"; -import { fetchOrganizations } from "@/components/organizations"; - -const TeamsPage = () => { - const { accessToken, userId, userRole } = useAuthorized(); - const { teams, setTeams } = useTeams(); - const [organizations, setOrganizations] = useState([]); - - useEffect(() => { - fetchOrganizations(accessToken, setOrganizations).then(() => {}); - }, [accessToken]); - - return ( - - ); -}; - -export default TeamsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/test-key/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/test-key/page.tsx deleted file mode 100644 index 0f984686c40..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/test-key/page.tsx +++ /dev/null @@ -1,45 +0,0 @@ -"use client"; - -import ChatUI from "@/components/playground/chat_ui/ChatUI"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useState, useEffect } from "react"; -import { fetchProxySettings } from "@/utils/proxyUtils"; - -interface ProxySettings { - PROXY_BASE_URL?: string; - LITELLM_UI_API_DOC_BASE_URL?: string | null; -} - -const TestKeyPage = () => { - const { token, accessToken, userRole, userId, disabledPersonalKeyCreation } = useAuthorized(); - const [proxySettings, setProxySettings] = useState(undefined); - - useEffect(() => { - const initializeProxySettings = async () => { - if (accessToken) { - const settings = await fetchProxySettings(accessToken); - if (settings) { - setProxySettings({ - PROXY_BASE_URL: settings.PROXY_BASE_URL || undefined, - LITELLM_UI_API_DOC_BASE_URL: settings.LITELLM_UI_API_DOC_BASE_URL, - }); - } - } - }; - - initializeProxySettings(); - }, [accessToken]); - - return ( - - ); -}; - -export default TestKeyPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tools/mcp-servers/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tools/mcp-servers/page.tsx deleted file mode 100644 index 9b94de6c9f2..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/tools/mcp-servers/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import { MCPServers } from "@/components/mcp_tools"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const MCPServersPage = () => { - const { accessToken, userRole, userId } = useAuthorized(); - - return ; -}; - -export default MCPServersPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tools/vector-stores/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tools/vector-stores/page.tsx deleted file mode 100644 index 8516a0faa1a..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/tools/vector-stores/page.tsx +++ /dev/null @@ -1,12 +0,0 @@ -"use client"; - -import VectorStoreManagement from "@/components/vector_store_management"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -const VectorStoresPage = () => { - const { accessToken, userId, userRole } = useAuthorized(); - - return ; -}; - -export default VectorStoresPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/page.tsx deleted file mode 100644 index 477c1163ce7..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/page.tsx +++ /dev/null @@ -1,14 +0,0 @@ -"use client"; - -import UsagePageView from "@/components/UsagePage/components/UsagePageView"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import useTeams from "@/app/(dashboard)/hooks/useTeams"; - -const UsagePage = () => { - const { accessToken, userRole, userId, premiumUser } = useAuthorized(); - const { teams } = useTeams(); - - return ; -}; - -export default UsagePage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx deleted file mode 100644 index 9874dd48865..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/page.tsx +++ /dev/null @@ -1,53 +0,0 @@ -"use client"; - -import ViewUserDashboard from "@/components/view_users"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import useTeams from "@/app/(dashboard)/hooks/useTeams"; -import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { isProxyAdminRole } from "@/utils/roles"; -import { useState, useMemo } from "react"; -import { Organization } from "@/components/networking"; - -const UsersPage = () => { - const { accessToken, userRole, userId, token } = useAuthorized(); - const [keys, setKeys] = useState([]); - - const { teams } = useTeams(); - const { data: organizations, isLoading: isOrgsLoading } = useOrganizations(); - - // Three states: - // - undefined: org data still loading (non-proxy-admin) — query should wait - // - null: proxy admin or no org filtering needed — query runs unfiltered - // - Array<{organization_id, organization_alias}>: org admin orgs — query runs filtered - const orgAdminOrgIds = useMemo((): Array<{organization_id: string, organization_alias: string}> | null | undefined => { - if (!userId || !userRole) return null; - // Proxy admins see all users — no org filtering - if (isProxyAdminRole(userRole)) return null; - - // Still loading org data — signal "not ready yet" - if (isOrgsLoading || !organizations) return undefined; - - const adminOrgs = organizations - .filter((org: Organization) => - org.members?.some((member) => member.user_id === userId && member.user_role === "org_admin") - ) - .map((org: Organization) => ({ organization_id: org.organization_id, organization_alias: org.organization_alias })); - - return adminOrgs.length > 0 ? adminOrgs : null; - }, [userId, organizations, userRole, isOrgsLoading]); - - return ( - - ); -}; - -export default UsersPage; diff --git a/ui/litellm-dashboard/src/app/layout.tsx b/ui/litellm-dashboard/src/app/layout.tsx index a4ed17cde39..f79b7eb7028 100644 --- a/ui/litellm-dashboard/src/app/layout.tsx +++ b/ui/litellm-dashboard/src/app/layout.tsx @@ -3,6 +3,7 @@ import { Inter } from "next/font/google"; import "./globals.css"; import AntdGlobalProvider from "@/contexts/AntdGlobalProvider"; +import { AuthProvider } from "@/contexts/AuthContext"; import ReactQueryProvider from "@/contexts/ReactQueryProvider"; const inter = Inter({ subsets: ["latin"] }); @@ -22,7 +23,9 @@ export default function RootLayout({ - {children} + + {children} + diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 06bf3b68d05..ce12967c911 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -1,5 +1,6 @@ "use client"; +import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; import PlaygroundPage from "@/app/(dashboard)/playground/page"; @@ -19,7 +20,7 @@ import { Team } from "@/components/key_team_helpers/key_list"; import { MCPServers } from "@/components/mcp_tools"; import ModelHubTable from "@/components/AIHub/ModelHubTable"; import Navbar from "@/components/navbar"; -import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName, getInProductNudgesCall } from "@/components/networking"; +import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button"; @@ -44,24 +45,14 @@ import WorkflowRuns from "@/components/workflow_runs"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; -import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; -import { isJwtExpired } from "@/utils/jwtUtils"; +import { useAuth } from "@/contexts/AuthContext"; import { buildLoginUrlWithReturn, consumeReturnUrl, isValidReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; -import { formatUserRole, isAdminRole } from "@/utils/roles"; +import { isAdminRole } from "@/utils/roles"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { jwtDecode } from "jwt-decode"; import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; import { ConfigProvider, theme } from "antd"; -function deleteCookie(name: string, path = "/") { - // Best-effort client-side clear (works for non-HttpOnly cookies without Domain) - document.cookie = `${name}=; Max-Age=0; Path=${path}`; - if (name === "token") { - clearTokenCookies(); - } -} - interface ProxySettings { PROXY_BASE_URL: string; PROXY_LOGOUT_URL: string; @@ -73,16 +64,21 @@ interface ProxySettings { * When a user visits ?page=, they are redirected to /ui/. * Add entries here as pages are migrated from the if/else chain to path-based routes. */ -const LEGACY_REDIRECTS: Record = { - api_ref: "api-reference", - "api-reference": "api-reference", -}; +const LEGACY_REDIRECTS: Record = {}; function CreateKeyPageContent() { - const [userRole, setUserRole] = useState(""); - const [premiumUser, setPremiumUser] = useState(false); - const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false); - const [userEmail, setUserEmail] = useState(null); + const { + authLoading, + token, + userID, + userRole, + userEmail, + accessToken, + premiumUser, + setUserRole, + setUserEmail, + } = useAuth(); + const [teams, setTeams] = useState(null); const [keys, setKeys] = useState([]); const [organizations, setOrganizations] = useState([]); @@ -92,14 +88,10 @@ function CreateKeyPageContent() { PROXY_LOGOUT_URL: "", }); - const [showSSOBanner, setShowSSOBanner] = useState(true); const router = useRouter(); const searchParams = useSearchParams()!; const [modelData, setModelData] = useState({ data: [] }); - const [token, setToken] = useState(null); const [createClicked, setCreateClicked] = useState(false); - const [authLoading, setAuthLoading] = useState(true); - const [userID, setUserID] = useState(null); // Survey state - always show by default const [showSurveyPrompt, setShowSurveyPrompt] = useState(true); @@ -187,7 +179,6 @@ function CreateKeyPageContent() { setPage(newPage); }; - const [accessToken, setAccessToken] = useState(null); const [sidebarCollapsed, setSidebarCollapsed] = useState(false); // Track if we've already attempted a return URL redirect to prevent race conditions @@ -203,38 +194,6 @@ function CreateKeyPageContent() { }; const redirectToLogin = authLoading === false && token === null && invitation_id === null; - useEffect(() => { - let cancelled = false; - - (async () => { - try { - await getUiConfig(); // ensures proxyBaseUrl etc. are ready - } catch { - // proceed regardless; we still need to decide auth state - } - - if (cancelled) return; - - const raw = getCookie("token"); - const valid = raw && !isJwtExpired(raw) ? raw : null; - - // If token exists but is invalid/expired, clear it so downstream code - // doesn't keep trying to use it and cause redirect spasms. - if (raw && !valid) { - deleteCookie("token", "/"); - } - - if (!cancelled) { - setToken(valid); - setAuthLoading(false); - } - })(); - - return () => { - cancelled = true; - }; - }, []); - useEffect(() => { if (redirectToLogin) { // Store the current URL so we can redirect back after login @@ -293,62 +252,6 @@ function CreateKeyPageContent() { } }, [token]); - useEffect(() => { - if (!token) { - return; - } - - // Defensive: re-check expiry in case cookie changed after mount - if (isJwtExpired(token)) { - deleteCookie("token", "/"); - setToken(null); - return; - } - - let decoded: any = null; - try { - decoded = jwtDecode(token); - } catch { - // Malformed token → treat as unauthenticated - deleteCookie("token", "/"); - setToken(null); - return; - } - - if (decoded) { - // set accessToken - setAccessToken(decoded.key); - - setDisabledPersonalKeyCreation(decoded.disabled_non_admin_personal_key_creation); - - // check if userRole is defined - if (decoded.user_role) { - const formattedUserRole = formatUserRole(decoded.user_role); - setUserRole(formattedUserRole); - } - - if (decoded.user_email) { - setUserEmail(decoded.user_email); - } - - if (decoded.login_method) { - setShowSSOBanner(decoded.login_method == "username_password" ? true : false); - } - - if (decoded.premium_user) { - setPremiumUser(decoded.premium_user); - } - - if (decoded.auth_header_name) { - setGlobalLitellmHeaderName(decoded.auth_header_name); - } - - if (decoded.user_id) { - setUserID(decoded.user_id); - } - } - }, [token]); - useEffect(() => { if (accessToken && userID && userRole) { fetchUserModels(userID, userRole, accessToken, setUserModels); @@ -473,18 +376,12 @@ function CreateKeyPageContent() { ) : (
@@ -553,6 +450,8 @@ function CreateKeyPageContent() { + ) : page == "api_ref" || page == "api-reference" ? ( + ) : page == "logging-and-alerts" ? ( ) : page == "budgets" ? ( diff --git a/ui/litellm-dashboard/src/components/Navbar/BlogDropdown/BlogDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/BlogDropdown/BlogDropdown.tsx index ddb2a33cdaa..657e5a39bab 100644 --- a/ui/litellm-dashboard/src/components/Navbar/BlogDropdown/BlogDropdown.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/BlogDropdown/BlogDropdown.tsx @@ -1,6 +1,7 @@ import { useDisableBlogPosts } from "@/app/(dashboard)/hooks/useDisableBlogPosts"; import { useBlogPosts, type BlogPost } from "@/app/(dashboard)/hooks/blogPosts/useBlogPosts"; -import { LoadingOutlined } from "@ant-design/icons"; +import { NAV_PRODUCT_LINK_CLASS } from "@/components/Navbar/navProductLinkClass"; +import { DownOutlined, LoadingOutlined } from "@ant-design/icons"; import { Button, Dropdown, Space, Typography } from "antd"; import type { MenuProps } from "antd"; import React from "react"; @@ -74,9 +75,13 @@ export const BlogDropdown: React.FC = () => { ]; } + // Blog opens a post list; Docs is a single outbound link — navbar adds a layout-only chevron there for alignment. return ( - + ); }; diff --git a/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.test.tsx b/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.test.tsx index 6994def858b..4f07f0e2daa 100644 --- a/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.test.tsx @@ -29,14 +29,14 @@ describe("CommunityEngagementButtons", () => { expect(joinSlackLink).toHaveAttribute("rel", "noopener noreferrer"); }); - it("should render Star us on GitHub button with correct link", () => { + it("should render GitHub link with correct href", () => { renderWithProviders(); - const starOnGithubLink = screen.getByRole("link", { name: /star us on github/i }); - expect(starOnGithubLink).toBeInTheDocument(); - expect(starOnGithubLink).toHaveAttribute("href", "https://github.com/BerriAI/litellm"); - expect(starOnGithubLink).toHaveAttribute("target", "_blank"); - expect(starOnGithubLink).toHaveAttribute("rel", "noopener noreferrer"); + const githubLink = screen.getByRole("link", { name: /litellm on github/i }); + expect(githubLink).toBeInTheDocument(); + expect(githubLink).toHaveAttribute("href", "https://github.com/BerriAI/litellm"); + expect(githubLink).toHaveAttribute("target", "_blank"); + expect(githubLink).toHaveAttribute("rel", "noopener noreferrer"); }); it("should not render buttons when prompts are disabled", () => { @@ -45,6 +45,6 @@ describe("CommunityEngagementButtons", () => { renderWithProviders(); expect(screen.queryByRole("link", { name: /join slack/i })).not.toBeInTheDocument(); - expect(screen.queryByRole("link", { name: /star us on github/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("link", { name: /litellm on github/i })).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx b/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx index 649bcc0b589..f6a43196a32 100644 --- a/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx @@ -1,36 +1,45 @@ import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; import { GithubOutlined, SlackOutlined } from "@ant-design/icons"; -import { Button } from "antd"; +import { Tooltip } from "antd"; import React from "react"; +const iconBtnClass = + "inline-flex h-9 w-9 shrink-0 items-center justify-center rounded-md border-0 bg-transparent text-gray-500 transition-colors hover:bg-gray-100 hover:text-gray-700 cursor-pointer"; + export const CommunityEngagementButtons: React.FC = () => { const disableShowPrompts = useDisableShowPrompts(); - // Hide buttons if prompts are disabled if (disableShowPrompts) { return null; } return ( - <> - - - +
+ + + + + + + + + + +
); }; diff --git a/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx new file mode 100644 index 00000000000..4ead085d977 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx @@ -0,0 +1,69 @@ +import { renderWithProviders, screen } from "../../../../tests/test-utils"; +import { NotificationsBell, AGENT_PLATFORM_URL } from "./NotificationsBell"; +import React from "react"; +import userEvent from "@testing-library/user-event"; + +describe("NotificationsBell", () => { + beforeEach(() => { + localStorage.clear(); + }); + + it("should open notifications with Agent Platform details and GitHub link", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /^notifications$/i })); + expect(screen.getByText(/LiteLLM Agent Platform/i)).toBeInTheDocument(); + const githubBtn = screen.getByRole("link", { name: /^GitHub$/i }); + expect(githubBtn).toHaveAttribute("href", AGENT_PLATFORM_URL); + expect(githubBtn).toHaveAttribute("target", "_blank"); + expect(githubBtn).toHaveAttribute("rel", "noopener noreferrer"); + }); + + it("should offer mark as read when announcement is unread", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /^notifications$/i })); + expect(screen.getByRole("button", { name: /^mark as read$/i })).toBeInTheDocument(); + }); + + it("should hide mark as read and persist after marking read", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /^notifications$/i })); + await user.click(screen.getByRole("button", { name: /^mark as read$/i })); + expect(localStorage.getItem("litellmHideAgentPlatformBanner")).toBe("true"); + await user.click(screen.getByRole("button", { name: /^notifications$/i })); + expect(screen.queryByRole("button", { name: /^mark as read$/i })).not.toBeInTheDocument(); + }); + + it("should not show mark as read when previously dismissed", async () => { + localStorage.setItem("litellmHideAgentPlatformBanner", "true"); + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: /^notifications$/i })); + expect(screen.queryByRole("button", { name: /^mark as read$/i })).not.toBeInTheDocument(); + }); + + it("should sync sibling instances when one is dismissed", async () => { + const user = userEvent.setup(); + renderWithProviders( + <> +
+ +
+
+ +
+ , + ); + + // Both bells start unread → both render the "Mark as read" affordance once opened. + const [bellA, bellB] = screen.getAllByRole("button", { name: /^notifications$/i }); + await user.click(bellA); + await user.click(screen.getByRole("button", { name: /^mark as read$/i })); + + // Dismissing in bell A must also clear bell B without a remount. + await user.click(bellB); + expect(screen.queryByRole("button", { name: /^mark as read$/i })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx new file mode 100644 index 00000000000..a3b7678e1e7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx @@ -0,0 +1,59 @@ +"use client"; + +import { + HIDE_AGENT_PLATFORM_BANNER_KEY, + useHideAgentPlatformBanner, +} from "@/app/(dashboard)/hooks/useHideAgentPlatformBanner"; +import { emitLocalStorageChange, setLocalStorageItem } from "@/utils/localStorageUtils"; +import { BellOutlined } from "@ant-design/icons"; +import { Badge, Button, Popover, Typography } from "antd"; +import React, { useState } from "react"; + +export const AGENT_PLATFORM_URL = "https://github.com/BerriAI/litellm-agent-platform"; + +export const NotificationsBell: React.FC = () => { + const hidden = useHideAgentPlatformBanner(); + const hasUnread = !hidden; + const [open, setOpen] = useState(false); + + const markDismissed = () => { + setLocalStorageItem(HIDE_AGENT_PLATFORM_BANNER_KEY, "true"); + emitLocalStorageChange(HIDE_AGENT_PLATFORM_BANNER_KEY); + setOpen(false); + }; + + const content = ( +
+ + LiteLLM Agent Platform + + + Open-source agent infra — sandboxes, durable sessions, and workers on AWS Fargate. + +
+ + {hasUnread ? ( + + ) : null} +
+
+ ); + + return ( + + + + ); +}; diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx index de853303c15..31ddae31798 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx @@ -37,6 +37,8 @@ vi.mock("@/utils/localStorageUtils", () => ({ describe("UserDropdown", () => { const mockOnLogout = vi.fn(); + const getAccountTrigger = () => screen.getByRole("button", { name: /account menu/i }); + beforeEach(() => { vi.clearAllMocks(); mockUseAuthorizedImpl = () => ({ @@ -55,22 +57,23 @@ describe("UserDropdown", () => { it("should render", () => { renderWithProviders(); - expect(screen.getByRole("button")).toBeInTheDocument(); + expect(getAccountTrigger()).toBeInTheDocument(); }); - it("should display user button with User text", () => { + it("should surface initials and account menu affordance", () => { renderWithProviders(); - expect(screen.getByText("User")).toBeInTheDocument(); + expect(getAccountTrigger()).toBeInTheDocument(); + expect(screen.getByText("TE")).toBeInTheDocument(); }); it("should show user email when dropdown is opened", async () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); }); @@ -78,7 +81,7 @@ describe("UserDropdown", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { expect(screen.getByText("test-user-id")).toBeInTheDocument(); @@ -89,10 +92,10 @@ describe("UserDropdown", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("Admin")).toBeInTheDocument(); + expect(screen.getAllByText("Admin").length).toBeGreaterThan(0); }); }); @@ -100,7 +103,7 @@ describe("UserDropdown", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { expect(screen.getByText("Standard")).toBeInTheDocument(); @@ -118,7 +121,7 @@ describe("UserDropdown", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { expect(screen.getByText("Premium")).toBeInTheDocument(); @@ -129,10 +132,10 @@ describe("UserDropdown", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); await user.click(screen.getByText("Logout")); @@ -144,10 +147,10 @@ describe("UserDropdown", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); const toggle = screen.getByLabelText("Toggle hide new feature indicators"); @@ -169,10 +172,10 @@ describe("UserDropdown", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); const toggle = screen.getByLabelText("Toggle hide new feature indicators"); @@ -189,10 +192,10 @@ describe("UserDropdown", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); const toggle = screen.getByLabelText("Toggle hide all prompts"); @@ -215,10 +218,10 @@ describe("UserDropdown", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); const toggle = screen.getByLabelText("Toggle hide all prompts"); @@ -231,6 +234,17 @@ describe("UserDropdown", () => { expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowPrompts"); }); + it("should show Account in the trigger when user id is the default placeholder", () => { + mockUseAuthorizedImpl = () => ({ + userId: "default_user_id", + userEmail: null as any, + userRole: "Admin", + premiumUser: false, + }); + renderWithProviders(); + expect(screen.getByText("Account")).toBeInTheDocument(); + }); + it("should display dash when user email is not available", async () => { const user = userEvent.setup(); mockUseAuthorizedImpl = () => ({ @@ -242,7 +256,7 @@ describe("UserDropdown", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { expect(screen.getByText("-")).toBeInTheDocument(); @@ -260,7 +274,7 @@ describe("UserDropdown", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { const dashElements = screen.getAllByText("-"); @@ -277,10 +291,10 @@ describe("UserDropdown", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(getAccountTrigger()); await waitFor(() => { - expect(screen.getByText("test@example.com")).toBeInTheDocument(); + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); }); const toggle = screen.getByLabelText("Toggle hide new feature indicators"); diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx index 6490cd32fa7..64a2f1260ba 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx @@ -9,6 +9,7 @@ import { removeLocalStorageItem, setLocalStorageItem, } from "@/utils/localStorageUtils"; +import { navAccountDisplayName } from "@/components/Navbar/navDisplayName"; import { CrownOutlined, DownOutlined, @@ -23,6 +24,39 @@ import React, { useEffect, useState } from "react"; const { Text } = Typography; +function hueFromString(seed: string): number { + let h = 0; + for (let i = 0; i < seed.length; i += 1) { + h = seed.charCodeAt(i) + ((h << 5) - h); + } + return Math.abs(h) % 360; +} + +function initialsFromIdentity(email: string | null, userId: string | null): string { + const local = email?.split("@")[0]?.trim(); + if (local) { + const parts = local + .replace(/[^a-zA-Z0-9]+/g, " ") + .trim() + .split(/\s+/) + .filter(Boolean); + if (parts.length >= 2) { + return `${parts[0]!.charAt(0)}${parts[1]!.charAt(0)}`.toUpperCase(); + } + if (parts.length === 1) { + const p = parts[0]!; + return p.length >= 2 ? p.slice(0, 2).toUpperCase() : `${p.charAt(0)}`.toUpperCase(); + } + } + if (userId && userId.length >= 2) { + return userId.slice(0, 2).toUpperCase(); + } + if (userId && userId.length === 1) { + return `${userId.toUpperCase()}•`; + } + return "?"; +} + interface UserDropdownProps { onLogout: () => void; } @@ -61,19 +95,12 @@ const UserDropdown: React.FC = ({ onLogout }) => { {userEmail || "-"} {premiumUser ? ( - } - color="gold" - > + } color="gold"> Premium ) : ( - } - > - Standard - + }>Standard )} @@ -83,12 +110,7 @@ const UserDropdown: React.FC = ({ onLogout }) => { User ID - + {userId || "-"} @@ -189,13 +211,17 @@ const UserDropdown: React.FC = ({ onLogout }) => { ); + const seed = userEmail || userId || "user"; + const initials = initialsFromIdentity(userEmail, userId); + const hue = hueFromString(seed); + const displayName = navAccountDisplayName(userEmail, userId); + return ( ( -
+
{renderUserInfoSection()} {React.cloneElement(menu as React.ReactElement, { @@ -204,12 +230,23 @@ const UserDropdown: React.FC = ({ onLogout }) => {
)} > - ); diff --git a/ui/litellm-dashboard/src/components/Navbar/navDisplayName.test.ts b/ui/litellm-dashboard/src/components/Navbar/navDisplayName.test.ts new file mode 100644 index 00000000000..96e768c5d31 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/navDisplayName.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from "vitest"; +import { navAccountDisplayName } from "./navDisplayName"; + +describe("navAccountDisplayName", () => { + it("should prefer email when present", () => { + expect(navAccountDisplayName("x@y.com", "ignored")).toBe("x@y.com"); + }); + + it("should map default_user_id placeholder to Account", () => { + expect(navAccountDisplayName(null, "default_user_id")).toBe("Account"); + expect(navAccountDisplayName(null, "DEFAULT_USER_ID")).toBe("Account"); + }); + + it("should show a sensible token when user id is non-placeholder", () => { + expect(navAccountDisplayName(null, "user-uuid-123")).toBe("user-uuid-123"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Navbar/navDisplayName.ts b/ui/litellm-dashboard/src/components/Navbar/navDisplayName.ts new file mode 100644 index 00000000000..d6f51dc37a8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/navDisplayName.ts @@ -0,0 +1,15 @@ +/** Primary label for the navbar account control — avoids raw placeholder JWT/user IDs in the UI. */ +export function navAccountDisplayName(userEmail: string | null, userId: string | null): string { + const email = userEmail?.trim(); + if (email) { + return email; + } + const id = userId?.trim(); + if (!id) { + return "Account"; + } + if (/^default[_\s-]?user[_\s-]?id$/i.test(id)) { + return "Account"; + } + return id; +} diff --git a/ui/litellm-dashboard/src/components/Navbar/navProductLinkClass.ts b/ui/litellm-dashboard/src/components/Navbar/navProductLinkClass.ts new file mode 100644 index 00000000000..ca4b2e5d1f3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/navProductLinkClass.ts @@ -0,0 +1,3 @@ +/** Shared styling for Docs / Blog in the top nav (product navigation zone). */ +export const NAV_PRODUCT_LINK_CLASS = + "inline-flex h-9 shrink-0 items-center justify-center gap-1 rounded-md px-2 text-sm font-medium leading-none text-gray-800 transition-colors hover:bg-gray-100 hover:text-gray-950"; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index f2fe1fe96ec..46373ab24db 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -49,12 +49,9 @@ const { Sider } = Layout; /** * Pages migrated to path-based routing under (dashboard)/. * Key = legacy page id, Value = route segment. - * Keep in sync with MIGRATED_PAGES in (dashboard)/layout.tsx and - * LEGACY_REDIRECTS in app/page.tsx. + * Keep in sync with MIGRATED_PAGES in (dashboard)/layout.tsx. */ -const MIGRATED_PAGES: Record = { - "api-reference": "api-reference", -}; +const MIGRATED_PAGES: Record = {}; /** Build an absolute href for a migrated page, respecting base URL + serverRootPath. */ function migratedHref(routeSegment: string): string { @@ -295,8 +292,8 @@ const menuGroups: MenuGroup[] = [ groupLabel: "DEVELOPER TOOLS", items: [ { - key: "api-reference", - page: "api-reference", + key: "api_ref", + page: "api_ref", label: "API Reference", icon: , }, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 5ed4c0468b8..768083be6e9 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -256,14 +256,26 @@ export default function ModelInfoView({ tags: values.tags, }; - if (form.isFieldTouched("input_cost") && values.input_cost !== undefined && values.input_cost !== null) { - updatedLitellmParams.input_cost_per_token = Number(values.input_cost) / 1_000_000; + if (form.isFieldTouched("input_cost")) { + if (values.input_cost !== undefined && values.input_cost !== null && values.input_cost !== "") { + updatedLitellmParams.input_cost_per_token = Number(values.input_cost) / 1_000_000; + } else { + // Explicit null signals the backend to remove the pricing override. + updatedLitellmParams.input_cost_per_token = null; + } } - if (form.isFieldTouched("output_cost") && values.output_cost !== undefined && values.output_cost !== null) { - updatedLitellmParams.output_cost_per_token = Number(values.output_cost) / 1_000_000; + if (form.isFieldTouched("output_cost")) { + if (values.output_cost !== undefined && values.output_cost !== null && values.output_cost !== "") { + updatedLitellmParams.output_cost_per_token = Number(values.output_cost) / 1_000_000; + } else { + updatedLitellmParams.output_cost_per_token = null; + } } - // Cache Read Cost: explicit value if provided, else fall back to input cost (when input cost touched). + // Cache Read Cost: + // - explicit value provided → use it + // - field touched but empty → explicit null (signals backend to remove override) + // - only input_cost touched → fall back to input_cost (guarded against null) if (form.isFieldTouched("cache_read_cost") || form.isFieldTouched("input_cost")) { if ( values.cache_read_cost !== undefined && @@ -271,14 +283,19 @@ export default function ModelInfoView({ values.cache_read_cost !== "" ) { updatedLitellmParams.cache_read_input_token_cost = Number(values.cache_read_cost) / 1_000_000; - } else if (updatedLitellmParams.input_cost_per_token !== undefined) { + } else if (form.isFieldTouched("cache_read_cost")) { + updatedLitellmParams.cache_read_input_token_cost = null; + } else if ( + updatedLitellmParams.input_cost_per_token !== undefined && + updatedLitellmParams.input_cost_per_token !== null + ) { updatedLitellmParams.cache_read_input_token_cost = updatedLitellmParams.input_cost_per_token; } } - // Cache Write Cost: explicit value if provided, else clear the override - // so the backend falls back to the model-level default. Sending 0 here - // would persist a zero rate even when the user intended to unset it. + // Cache Write Cost: explicit value if provided, else explicit null so the + // backend removes the override and falls back to the model-level default. + // Sending 0 here would persist a zero rate even when the user intended to unset it. if (form.isFieldTouched("cache_write_cost")) { if ( values.cache_write_cost !== undefined && @@ -287,7 +304,7 @@ export default function ModelInfoView({ ) { updatedLitellmParams.cache_creation_input_token_cost = Number(values.cache_write_cost) / 1_000_000; } else { - delete updatedLitellmParams.cache_creation_input_token_cost; + updatedLitellmParams.cache_creation_input_token_cost = null; } } diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index 2e122164969..274e81db527 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -30,6 +30,7 @@ const mockUserDropdownData = vi.hoisted(() => ({ vi.mock("./Navbar/UserDropdown/UserDropdown", async (importOriginal) => { const React = await import("react"); const { useState } = React; + const { Button } = await import("antd"); const localStorageUtils = await import("@/utils/localStorageUtils"); return { default: function MockUserDropdown({ onLogout }: { onLogout: () => void }) { @@ -37,9 +38,9 @@ vi.mock("./Navbar/UserDropdown/UserDropdown", async (importOriginal) => { const [open, setOpen] = useState(false); return (
- + {open && (
{userId} @@ -136,30 +137,25 @@ Object.defineProperty(window, "location", { describe("Navbar", () => { const defaultProps = { - userID: "test-user", - userEmail: "test@example.com", - userRole: "Admin", - premiumUser: false, proxySettings: {}, setProxySettings: vi.fn(), accessToken: "test-token", isPublicPage: false, - isDarkMode: false, - toggleDarkMode: vi.fn(), }; it("should render without crashing", () => { renderWithProviders(); + expect(screen.getByRole("button", { name: /^notifications$/i })).toBeInTheDocument(); expect(screen.getByText("Docs")).toBeInTheDocument(); - expect(screen.getByText("User")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /open account menu/i })).toBeInTheDocument(); }); it("should display user information in dropdown", async () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /open account menu/i })); await waitFor(() => { expect(screen.getByText("test-user")).toBeInTheDocument(); @@ -198,7 +194,7 @@ describe("Navbar", () => { }); renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /open account menu/i })); await waitFor(() => { expect(screen.getByText("Premium")).toBeInTheDocument(); @@ -247,11 +243,12 @@ describe("Navbar", () => { mockUseThemeImpl = () => ({ logoUrl: null }); }); - it("should hide user dropdown on public pages", () => { + it("should hide user dropdown and notifications on public pages", () => { const publicPageProps = { ...defaultProps, isPublicPage: true }; renderWithProviders(); - expect(screen.queryByText("User")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /open account menu/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /^notifications$/i })).not.toBeInTheDocument(); }); it("should handle hide new features toggle", async () => { @@ -265,7 +262,7 @@ describe("Navbar", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /open account menu/i })); await waitFor(() => { expect(screen.getByText("test-user")).toBeInTheDocument(); @@ -290,7 +287,7 @@ describe("Navbar", () => { renderWithProviders(); - await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /open account menu/i })); await waitFor(() => { expect(screen.getByText("test-user")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 055038b6d88..e5a1490788c 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,47 +1,39 @@ import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; import { useDisableBouncingIcon } from "@/app/(dashboard)/hooks/useDisableBouncingIcon"; +import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; +import { useWorker } from "@/hooks/useWorker"; import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; import { clearStoredReturnUrl } from "@/utils/returnUrlUtils"; import { fetchProxySettings } from "@/utils/proxyUtils"; -import { MenuFoldOutlined, MenuUnfoldOutlined, MoonOutlined, SunOutlined } from "@ant-design/icons"; -import { Button, Switch, Tag } from "antd"; +import { DownOutlined, MenuFoldOutlined, MenuUnfoldOutlined } from "@ant-design/icons"; +import { Tag } from "antd"; import Link from "next/link"; import React, { useEffect, useState } from "react"; import { BlogDropdown } from "./Navbar/BlogDropdown/BlogDropdown"; import { CommunityEngagementButtons } from "./Navbar/CommunityEngagementButtons/CommunityEngagementButtons"; +import { NAV_PRODUCT_LINK_CLASS } from "./Navbar/navProductLinkClass"; +import { NotificationsBell } from "./Navbar/NotificationsBell/NotificationsBell"; import UserDropdown from "./Navbar/UserDropdown/UserDropdown"; import WorkerDropdown from "./Navbar/WorkerDropdown/WorkerDropdown"; interface NavbarProps { - userID: string | null; - userEmail: string | null; - userRole: string | null; - premiumUser: boolean; proxySettings: any; setProxySettings: React.Dispatch>; accessToken: string | null; isPublicPage: boolean; sidebarCollapsed?: boolean; onToggleSidebar?: () => void; - isDarkMode: boolean; - toggleDarkMode: () => void; } const Navbar: React.FC = ({ - userID, - userEmail, - userRole, - premiumUser, proxySettings, setProxySettings, accessToken, isPublicPage = false, sidebarCollapsed = false, onToggleSidebar, - isDarkMode, - toggleDarkMode, }) => { const baseUrl = getProxyBaseUrl(); const [logoutUrl, setLogoutUrl] = useState(""); @@ -49,8 +41,10 @@ const Navbar: React.FC = ({ const { data: healthData } = useHealthReadinessDetails(accessToken); const version = healthData?.litellm_version; const disableBouncingIcon = useDisableBouncingIcon(); + const hideCommunityLinks = useDisableShowPrompts(); + const { isControlPlane, selectedWorker } = useWorker(); + const showWorkerSwitch = isControlPlane && selectedWorker !== null; - // Simple logo URL: use custom logo if available, otherwise default const imageUrl = logoUrl || `${baseUrl}/get_image`; useEffect(() => { @@ -87,14 +81,14 @@ const Navbar: React.FC = ({ }; return ( -