diff --git a/.github/workflows/osv-scan.yml b/.github/workflows/osv-scan.yml index 003f4ab1d8a..5016957bd6f 100644 --- a/.github/workflows/osv-scan.yml +++ b/.github/workflows/osv-scan.yml @@ -6,6 +6,8 @@ on: - main - "litellm_**" - "rc/**" + - "stable/**" + - "lts/**" schedule: - cron: "23 6 * * *" workflow_dispatch: diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0a490bb8f45..94152aeb7a0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -5,6 +5,9 @@ on: branches: - main - "litellm_**" + - "rc/**" + - "stable/**" + - "lts/**" push: branches: - main @@ -280,6 +283,7 @@ jobs: tests/unit/llms/e2b tests/unit/llms/edenai tests/unit/llms/elevenlabs + tests/unit/llms/empiriolabs tests/unit/llms/exa_ai tests/unit/llms/fal_ai tests/unit/llms/fastcrw @@ -973,10 +977,10 @@ jobs: dist: loadscope steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }} + timeout-minutes: 3 with: persist-credentials: false - fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }} + fetch-depth: 1 - name: Detect relevant changes id: changes @@ -1009,13 +1013,11 @@ jobs: with: artifact: ${{ needs.rust-bridge.outputs.artifact }} - - name: Install dependencies and check schema + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' - timeout-minutes: 8 + timeout-minutes: 5 env: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} - UNIT_SHARD: ${{ matrix.shard }} - SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }} run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then @@ -1027,9 +1029,21 @@ jobs: uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" fi uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - if [ "$UNIT_SHARD" = "db-and-spend" ]; then - uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF" - fi + + - name: Check schema.prisma against the merge base + if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request' + timeout-minutes: 3 + env: + GH_TOKEN: ${{ github.token }} + PR_NUMBER: ${{ github.event.pull_request.number }} + EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref) + merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') + test -n "$merge_base" + git fetch --no-tags --depth=1 origin "$merge_base" + uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base" - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' diff --git a/README.md b/README.md index c4aeee80fb4..31a93bf8c67 100644 --- a/README.md +++ b/README.md @@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ | | | | | | | | | [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ | | | ✅ | ✅ | | | | +| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ | | ✅ | | | | | | | | | [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ | | | | | | | | | [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) | | | | | | | | | | | | [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ | | ✅ | | | | | | diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index cfe4ca98145..69e36a0c076 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -200,7 +200,6 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi unpriced: Final = _without_read_price(billed.response) cost: Final = logging_obj.response_cost_calculator(result=unpriced) end_time: Final = datetime.now() - logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time) GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( async_coroutine=logging_obj.async_success_handler( result=unpriced, start_time=logging_obj.start_time, end_time=end_time diff --git a/helm/litellm-helm/Chart.yaml b/helm/litellm-helm/Chart.yaml index a3cb388ffc6..059889cb667 100644 --- a/helm/litellm-helm/Chart.yaml +++ b/helm/litellm-helm/Chart.yaml @@ -18,7 +18,7 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 1.1.3 +version: 1.1.4 # This is the version number of the application being deployed. This version number should be # incremented each time you make changes to the application. Versions are not expected to diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index 694cca05b52..bc6e29632f1 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -56,8 +56,8 @@ If `db.useStackgresOperator` is used (not yet implemented): | `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` | | `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` | -| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` | -| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` | +| `billingMetrics.tokenUrl` | Cognito OAuth2 client-credentials token endpoint. | `https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token` | +| `billingMetrics.secretName` | Name of an existing Secret with `client-id` and `client-secret` keys. | `litellm-billing-metrics` | | `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` | #### Example `proxy_config` ConfigMap from values (default): @@ -109,10 +109,10 @@ type: Opaque #### Enterprise billable-request metering -Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on: +Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector with Cognito bearer authentication. Create a Secret with `client-id` and `client-secret` keys, then enable the chart values: ``` -kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key +kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=... ``` ``` @@ -120,7 +120,8 @@ billingMetrics: enabled: true ``` -Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time. +The endpoint and token URL must use HTTPS. The chart requires both URLs and +the Secret name when billing metrics are enabled ### Database Settings diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 929fc415163..230d2acaeef 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -50,53 +50,27 @@ app.kubernetes.io/name: {{ include "litellm.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} {{- end }} -{{/* -Enterprise billable-request metering. The client certificate identifies the -deployment to LiteLLM's collector, so it is mounted read-only from an existing -Secret rather than passed through the environment. -*/}} -{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} -{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} - {{- define "litellm.billingMetricsEnv" -}} - name: LITELLM_BILLING_METRICS_ENDPOINT value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} -{{- if .Values.billingMetrics.caSecretName }} -- name: LITELLM_BILLING_METRICS_CA_CERT - value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} -{{- end }} +- name: LITELLM_BILLING_METRICS_TOKEN_URL + value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-id +- name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-secret {{- with .Values.billingMetrics.exportIntervalMs }} - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: {{ . | quote }} {{- end }} {{- end -}} -{{- define "litellm.billingMetricsVolumes" -}} -- name: billing-metrics-mtls - secret: - secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - secret: - secretName: {{ .Values.billingMetrics.caSecretName }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumeMounts" -}} -- name: billing-metrics-mtls - mountPath: {{ include "litellm.billingMetrics.certDir" . }} - readOnly: true -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - mountPath: {{ include "litellm.billingMetrics.caDir" . }} - readOnly: true -{{- end }} -{{- end -}} - {{/* Create the name of the service account to use */}} diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index 37e7eaec9e9..a37193f0efa 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -148,9 +148,6 @@ spec: - name: npm mountPath: /.npm {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- if include "litellm.collector.socketDir" . }} - name: collector-socket mountPath: {{ include "litellm.collector.socketDir" . }} @@ -234,9 +231,6 @@ spec: items: - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }} path: "config.yaml" - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} {{- if include "litellm.collector.socketDir" . }} - name: collector-socket emptyDir: diff --git a/helm/litellm-helm/tests/billing_metrics_tests.yaml b/helm/litellm-helm/tests/billing_metrics_tests.yaml index 71803df378c..52185c8dd9d 100644 --- a/helm/litellm-helm/tests/billing_metrics_tests.yaml +++ b/helm/litellm-helm/tests/billing_metrics_tests.yaml @@ -1,31 +1,21 @@ -suite: test billingMetrics wiring on the proxy deployment +suite: test bearer-only billingMetrics wiring templates: - deployment.yaml - configmap-litellm.yaml - - migrations-job.yaml tests: - - it: is off by default, adding no env, volume, or mount + - it: is disabled by default template: deployment.yaml asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - notContains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - notContains: path: spec.template.spec.containers[0].env content: name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics - - it: renders the endpoint and the mounted cert paths when enabled + - it: renders required URLs and Cognito secret references when enabled template: deployment.yaml set: billingMetrics: @@ -39,119 +29,34 @@ tests: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: /etc/litellm/billing-mtls/tls.crt + name: LITELLM_BILLING_METRICS_TOKEN_URL + value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: /etc/litellm/billing-mtls/tls.key - - # The conventional Secret name is the default, so enabling the block is enough. - - it: mounts the default cert secret read-only alongside the config volume - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - - it: honours a secretName override - template: deployment.yaml - set: - billingMetrics: - enabled: true - secretName: my-billing-mtls - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: my-billing-mtls - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - - it: honours an endpoint override - template: deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: https://collector.internal:4318 - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-id - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://collector.internal:4318 - - # The production collector presents a public web-PKI certificate, so the CA - # override must stay absent unless a private collector is configured. - - it: omits the CA env, volume, and mount when no caSecretName is set - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca + name: billing-metrics - notContains: path: spec.template.spec.containers[0].volumeMounts content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt + name: billing-metrics - - it: mounts the CA secret when caSecretName is set - template: deployment.yaml - set: - billingMetrics: - enabled: true - caSecretName: billing-ca - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true - - - it: passes the export interval through only when set + - it: renders the export interval only when set template: deployment.yaml set: billingMetrics: @@ -164,129 +69,7 @@ tests: name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: "5000" - - it: omits the export interval when unset - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS - value: "60000" - - # Kubernetes resolves duplicate env names last-wins, so the chart-owned billing - # entries must render after .Values.envVars or a user could silently redirect - # the metering export. The three billing entries are the last ones emitted here - # (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case). - - it: renders the billing endpoint after envVars so it cannot be shadowed - template: deployment.yaml - set: - migrationJob: - enabled: false - billingMetrics: - enabled: true - envVars: - LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://shadowed.example - - equal: - path: spec.template.spec.containers[0].env[-3] - value: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - equal: - path: spec.template.spec.containers[0].env[-2].name - value: LITELLM_BILLING_METRICS_CLIENT_CERT - - equal: - path: spec.template.spec.containers[0].env[-1].name - value: LITELLM_BILLING_METRICS_CLIENT_KEY - - - it: keeps user-supplied volumes and mounts alongside the billing secret - template: deployment.yaml - set: - billingMetrics: - enabled: true - volumes: - - name: custom-callbacks - configMap: - name: my-callbacks - volumeMounts: - - name: custom-callbacks - mountPath: /app/callbacks - asserts: - - contains: - path: spec.template.spec.volumes - content: - name: custom-callbacks - configMap: - name: my-callbacks - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: custom-callbacks - mountPath: /app/callbacks - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - - it: still mounts the proxy config when enabled - template: deployment.yaml - set: - billingMetrics: - enabled: true - asserts: - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: litellm-config - mountPath: /etc/litellm/config.yaml - subPath: config.yaml - - # Only the proxy serves billable traffic. The migrations Job must never mount - # the client certificate, and it renders its own env and volumes, so nothing - # stops a future edit from wiring the billing include into it by mistake. - - it: does not touch the migrations job when enabled - template: migrations-job.yaml - set: - billingMetrics: - enabled: true - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - notExists: - path: spec.template.spec.containers[0].volumeMounts - - notExists: - path: spec.template.spec.volumes - - - it: fails loudly when enabled with an emptied secretName - template: deployment.yaml - set: - billingMetrics: - enabled: true - secretName: "" - asserts: - - failedTemplate: - errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) - - - it: fails loudly when enabled without an endpoint + - it: requires the endpoint when enabled template: deployment.yaml set: billingMetrics: @@ -295,3 +78,23 @@ tests: asserts: - failedTemplate: errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true + + - it: requires the token URL when enabled + template: deployment.yaml + set: + billingMetrics: + enabled: true + tokenUrl: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true + + - it: requires the secret name when enabled + template: deployment.yaml + set: + billingMetrics: + enabled: true + secretName: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true diff --git a/helm/litellm-helm/tests/collector_tests.yaml b/helm/litellm-helm/tests/collector_tests.yaml index 0340b1161b7..1365290e16c 100644 --- a/helm/litellm-helm/tests/collector_tests.yaml +++ b/helm/litellm-helm/tests/collector_tests.yaml @@ -191,7 +191,7 @@ tests: metricsServer.port: 9090 billingMetrics.enabled: true billingMetrics.endpoint: https://metering.example.com - billingMetrics.secretName: billing-mtls + billingMetrics.secretName: billing-metrics asserts: - contains: path: spec.template.spec.containers[0].env @@ -214,9 +214,9 @@ tests: name: LITELLM_BILLING_METRICS_ENDPOINT any: true - notContains: - path: spec.template.spec.containers[1].volumeMounts + path: spec.template.spec.containers[1].env content: - name: billing-metrics-mtls + name: LITELLM_BILLING_METRICS_CLIENT_SECRET any: true - it: should give the sidecar the same scratch mounts as the proxy on a read-only root diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index f51f8bdd173..9afce050261 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -154,19 +154,13 @@ masterkeySecretName: "" # if set, use this secret key for the master key; otherwise, use the default key masterkeySecretKey: "" -# Optional: enterprise billable-request metering. When enabled, the proxy counts -# successful requests to inference, MCP, and A2A endpoints and pushes them to -# LiteLLM's collector over mutual TLS. Requires an enterprise license. -# The client certificate identifies the deployment, so it is mounted read-only -# from an existing Secret and never passed through the environment. +# Optional: enterprise billable-request metering. Requires an enterprise license. billingMetrics: enabled: false - endpoint: https://telemetry.litellm.ai # collector to push the counter to - secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key - # Only for private or test collectors whose server certificate is not on the - # public web PKI. The production collector needs no CA override. - caSecretName: "" # existing Secret holding ca.crt - exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + endpoint: https://telemetry.litellm.ai + tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token + secretName: litellm-billing-metrics + exportIntervalMs: "" proxyConfigMap: # when true, creates a new configmap diff --git a/helm/litellm/Chart.yaml b/helm/litellm/Chart.yaml index e67f5790c7e..341172da925 100644 --- a/helm/litellm/Chart.yaml +++ b/helm/litellm/Chart.yaml @@ -2,7 +2,7 @@ apiVersion: v2 name: litellm description: LiteLLM componentized — gateway, UI backend, and UI as separate services type: application -version: 0.1.0 +version: 0.1.1 appVersion: "0.1.0" annotations: org.opencontainers.image.source: "https://github.com/BerriAI/litellm" diff --git a/helm/litellm/README.md b/helm/litellm/README.md index 97b94754402..d66bb3284bf 100644 --- a/helm/litellm/README.md +++ b/helm/litellm/README.md @@ -22,6 +22,23 @@ backend: gatewayUrl: https://gateway.example.com ``` +## Enterprise billing metrics + +Billing metrics use Cognito bearer authentication. Create the Secret with +`client-id` and `client-secret` keys, then enable billing metrics: + +```bash +kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=... +``` + +```yaml +billingMetrics: + enabled: true +``` + +The chart requires an HTTPS endpoint, an HTTPS token URL, and the Secret name +when billing metrics are enabled + ## Routes to another Service For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 18e938564e9..b704782f5d1 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -58,56 +58,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }} helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }} {{- end -}} -{{/* -Enterprise billable-request metering. Wired into gateway and backend, not the -migrations job. The gateway serves nearly all billable traffic, but the backend -keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a -SpendLogs row, so metering only the gateway would silently drop that traffic. -The client certificate identifies the deployment to LiteLLM's collector, so it is -mounted read-only from an existing Secret rather than passed through the -environment. -*/}} -{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} -{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} - {{- define "litellm.billingMetricsEnv" -}} - name: LITELLM_BILLING_METRICS_ENDPOINT value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} -- name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} -{{- if .Values.billingMetrics.caSecretName }} -- name: LITELLM_BILLING_METRICS_CA_CERT - value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} -{{- end }} +- name: LITELLM_BILLING_METRICS_TOKEN_URL + value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-id +- name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }} + key: client-secret {{- with .Values.billingMetrics.exportIntervalMs }} - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: {{ . | quote }} {{- end }} {{- end -}} -{{- define "litellm.billingMetricsVolumes" -}} -- name: billing-metrics-mtls - secret: - secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - secret: - secretName: {{ .Values.billingMetrics.caSecretName }} -{{- end }} -{{- end -}} - -{{- define "litellm.billingMetricsVolumeMounts" -}} -- name: billing-metrics-mtls - mountPath: {{ include "litellm.billingMetrics.certDir" . }} - readOnly: true -{{- if .Values.billingMetrics.caSecretName }} -- name: billing-metrics-mtls-ca - mountPath: {{ include "litellm.billingMetrics.caDir" . }} - readOnly: true -{{- end }} -{{- end -}} {{/* Per-component selector labels — used in both Service selectors and Deployment matchLabels. diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 7360261dab1..7db14ad87a5 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -79,16 +79,13 @@ spec: {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} {{- include "litellm.envFrom" .Values.backend | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.backend.volumeMounts }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- with .Values.backend.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -114,16 +111,13 @@ spec: {{- with .Values.backend.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.backend.volumes }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: name: {{ include "litellm.gateway.fullname" . }}-config {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} {{- with .Values.backend.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index b0e7f0dcebb..987109e361e 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -86,7 +86,7 @@ spec: {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} {{- end }} {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config @@ -101,9 +101,6 @@ spec: - name: collector-socket mountPath: {{ include "litellm.gateway.collectorSocketDir" . }} {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} - {{- end }} {{- with .Values.gateway.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -206,7 +203,7 @@ spec: {{- with .Values.gateway.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} + {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config @@ -222,9 +219,6 @@ spec: emptyDir: sizeLimit: 1Mi {{- end }} - {{- if .Values.billingMetrics.enabled }} - {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} - {{- end }} {{- with .Values.gateway.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm/tests/billing_metrics_tests.yaml b/helm/litellm/tests/billing_metrics_tests.yaml index ceba0bd1430..d17b8192fce 100644 --- a/helm/litellm/tests/billing_metrics_tests.yaml +++ b/helm/litellm/tests/billing_metrics_tests.yaml @@ -1,39 +1,30 @@ -suite: test billingMetrics wiring on gateway and backend +suite: test bearer-only billingMetrics wiring templates: - gateway/deployment.yaml - gateway/configmap.yaml - backend/deployment.yaml + - backend/configmap.yaml - migrations-job.yaml values: - ./values/required.yaml tests: - - it: is off by default, adding no env, volume, or mount + - it: is disabled by default template: gateway/deployment.yaml asserts: - - notContains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - equal: - path: spec.template.spec.containers[0].volumeMounts - value: - - name: gateway-config - mountPath: /app/config/config.yaml - subPath: config.yaml - notContains: path: spec.template.spec.containers[0].env content: name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics - - it: renders the endpoint and the mounted cert paths when enabled + - it: wires gateway URLs and bearer credentials from Secret keys template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls asserts: - contains: path: spec.template.spec.containers[0].env @@ -43,87 +34,68 @@ tests: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_CERT - value: /etc/litellm/billing-mtls/tls.crt + name: LITELLM_BILLING_METRICS_TOKEN_URL + value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CLIENT_KEY - value: /etc/litellm/billing-mtls/tls.key - - - it: mounts the cert secret read-only alongside the config volume - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_ID + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-id - contains: - path: spec.template.spec.volumes + path: spec.template.spec.containers[0].env content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - # The production collector presents a public web-PKI certificate, so the CA - # override must stay absent unless a private collector is configured. - - it: omits the CA env, volume, and mount when no caSecretName is set - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca + name: billing-metrics - notContains: - path: spec.template.spec.containers[0].env + path: spec.template.spec.containers[0].volumeMounts content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt + name: billing-metrics - - it: mounts the CA secret when caSecretName is set - template: gateway/deployment.yaml + - it: wires billing metrics on the backend without volumes + template: backend/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls - caSecretName: billing-ca asserts: - contains: path: spec.template.spec.containers[0].env content: - name: LITELLM_BILLING_METRICS_CA_CERT - value: /etc/litellm/billing-mtls-ca/ca.crt - - contains: + name: LITELLM_BILLING_METRICS_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: litellm-billing-metrics + key: client-secret + - notContains: path: spec.template.spec.volumes content: - name: billing-metrics-mtls-ca - secret: - secretName: billing-ca - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls-ca - mountPath: /etc/litellm/billing-mtls-ca - readOnly: true + name: billing-metrics - - it: passes the export interval through only when set + - it: does not add billing metrics to migrations + template: migrations-job.yaml + set: + billingMetrics: + enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + + - it: renders an optional export interval template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls exportIntervalMs: 5000 asserts: - contains: @@ -132,102 +104,27 @@ tests: name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS value: "5000" - - it: keeps user-supplied gateway volumes alongside the billing secret + - it: requires the endpoint when enabled template: gateway/deployment.yaml set: billingMetrics: enabled: true - secretName: billing-mtls - gateway.volumes: - - name: custom-callbacks - configMap: - name: my-callbacks - gateway.volumeMounts: - - name: custom-callbacks - mountPath: /app/callbacks + endpoint: "" asserts: - - contains: - path: spec.template.spec.volumes - content: - name: custom-callbacks - configMap: - name: my-callbacks - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls + - failedTemplate: + errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true - # The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp), - # which writes a SpendLogs row, so it must meter too or that traffic is lost. - - it: meters the backend as well, since it serves the MCP transport - template: backend/deployment.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - contains: - path: spec.template.spec.containers[0].volumeMounts - content: - name: billing-metrics-mtls - mountPath: /etc/litellm/billing-mtls - readOnly: true - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: billing-mtls - - - it: leaves the backend alone when metering is off - template: backend/deployment.yaml - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - # The migrations job runs prisma and serves no traffic; it must never receive - # the client key. - - it: never mounts the billing cert on the migrations job - template: migrations-job.yaml - set: - billingMetrics: - enabled: true - secretName: billing-mtls - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: LITELLM_BILLING_METRICS_ENDPOINT - value: https://telemetry.litellm.ai - - isNull: - path: spec.template.spec.volumes - - # The conventional Secret name is the default, so enabling metering needs no - # secretName at all; the guard below only fires on an explicitly blanked one. - - it: uses the conventional secret name by default + - it: requires the token URL when enabled template: gateway/deployment.yaml set: billingMetrics: enabled: true + tokenUrl: "" asserts: - - contains: - path: spec.template.spec.volumes - content: - name: billing-metrics-mtls - secret: - secretName: litellm-billing-metrics-mtls + - failedTemplate: + errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true - - it: fails loudly when the secretName is explicitly blanked + - it: requires the secret name when enabled template: gateway/deployment.yaml set: billingMetrics: @@ -235,15 +132,4 @@ tests: secretName: "" asserts: - failedTemplate: - errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) - - - it: fails loudly when enabled without an endpoint - template: gateway/deployment.yaml - set: - billingMetrics: - enabled: true - endpoint: "" - secretName: billing-mtls - asserts: - - failedTemplate: - errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true diff --git a/helm/litellm/tests/migration_job_tests.yaml b/helm/litellm/tests/migration_job_tests.yaml index 2ebb1b44926..b9e52ba7123 100644 --- a/helm/litellm/tests/migration_job_tests.yaml +++ b/helm/litellm/tests/migration_job_tests.yaml @@ -91,7 +91,7 @@ tests: app.kubernetes.io/name: litellm app.kubernetes.io/instance: RELEASE-NAME app.kubernetes.io/managed-by: Helm - helm.sh/chart: litellm-0.1.0 + helm.sh/chart: litellm-0.1.1 app.kubernetes.io/component: migrations - it: renders pod-level and container-level securityContext in their own scopes diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 92a48153c5a..cb37656cb77 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -180,23 +180,14 @@ masterKey: secretKey: master-key # Optional: enterprise billable-request metering. When enabled, the gateway and -# backend count successful requests to inference, MCP, and A2A endpoints and push -# them to LiteLLM's collector over mutual TLS. Both components serve billable -# routes: the backend keeps the named-server MCP transport. Requires an -# enterprise license. The client certificate identifies the deployment, so it is -# mounted read-only from an existing Secret and never passed through the env. +# backend count successful requests to inference, MCP, and A2A endpoints. Requires +# an enterprise license. billingMetrics: enabled: false - endpoint: https://telemetry.litellm.ai # collector to push the counter to - # An existing Secret holding the client certificate under tls.crt and its key - # under tls.key, usually created from the onboarding artifact. The default is - # the conventional name, so the common path is to create that Secret and set - # enabled: true. Override only if yours is named differently. - secretName: litellm-billing-metrics-mtls - # Only for private or test collectors whose server certificate is not on the - # public web PKI. The production collector needs no CA override. - caSecretName: "" # existing Secret holding ca.crt - exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + endpoint: https://telemetry.litellm.ai + tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token + secretName: litellm-billing-metrics + exportIntervalMs: "" # External Postgres connection. database: diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql new file mode 100644 index 00000000000..b00cf892b68 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql @@ -0,0 +1,19 @@ +SET LOCAL lock_timeout = '10s'; +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" ( + "delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text, + "request_id" TEXT NOT NULL, + "state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending', + "attempts" INTEGER NOT NULL DEFAULT 0, + "next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "applied_at" TIMESTAMP(3), + "last_error" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "accounting_payload" JSONB NOT NULL, + "is_batch" BOOLEAN NOT NULL +); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx" + ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx" + ON "LiteLLM_SpendLogDelivery" ("state", "applied_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx" + ON "LiteLLM_SpendLogDelivery" ("request_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c007b047087..3aaa74d4fc2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery { @@index([target, state, next_attempt_at]) } +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) +} + model LiteLLM_SubscriptionAccountTable { subscription_account_id String @id @default(uuid()) custom_llm_provider String diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b93a8827428..ba1131d7a92 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3389,15 +3389,19 @@ version = "0.1.0" dependencies = [ "aws-config", "aws-credential-types", + "aws-runtime", "aws-sdk-sts", "aws-sigv4", "aws-smithy-runtime-api", + "aws-smithy-types", "aws-types", "bytes", "futures-util", "litellm-auth-types", "litellm-error", "litellm-http", + "litellm-secret-value", + "litellm-secrets", "moka", "reqwest", "rstest", @@ -3406,6 +3410,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", "veil", ] @@ -3418,6 +3423,7 @@ dependencies = [ "litellm-auth-types", "litellm-error", "litellm-secret-value", + "litellm-secrets", "moka", "rstest", "serde_json", @@ -3431,6 +3437,7 @@ dependencies = [ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ + "axum", "base64 0.22.1", "gcp_auth", "google-cloud-auth", @@ -3440,10 +3447,12 @@ dependencies = [ "litellm-secret-value", "moka", "rstest", + "rustls 0.23.42", "serde", "serde_json", "sha2 0.10.9", "tempfile", + "time", "tokio", "veil", ] diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index e4578777a08..c8ed7bb0f14 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -11,14 +11,16 @@ litellm-error.workspace = true litellm-http.workspace = true reqwest.workspace = true -moka = { workspace = true, features = ["sync"] } +moka.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } sha2.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true -aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] } +aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] } aws-sigv4 = "1.5.1" @@ -26,6 +28,9 @@ aws-types = "1.4.0" aws-smithy-runtime-api = "1.13.0" [dev-dependencies] +litellm-secrets.workspace = true +aws-smithy-types = "1.6.4" +aws-runtime = { version = "1.9.4", default-features = false } rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } bytes.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 81aa4b068bc..e6e32754b4f 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,13 +1,16 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; -use moka::sync::Cache; -use sha2::{Digest, Sha256}; +use aws_credential_types::provider::{SharedCredentialsProvider, error::CredentialsError}; +use litellm_secret_value::SecretValue; +use moka::{Expiry, future::Cache}; +use std::{future::Future, pin::Pin, sync::Arc}; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - SignableBody, SignableRequest, SigningParams, SigningSettings, sign, + SignableBody, SignableRequest, SignatureLocation, SigningInstructions, SigningParams, + SigningSettings, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -23,72 +26,106 @@ use super::constants::{ }; use crate::error::{CredentialFlow, Error, SigningStage}; -const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); -const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); +const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(600); +const EXPIRY_BUFFER: Duration = Duration::from_secs(60); #[derive(Clone)] pub struct AwsAuthService { - static_credentials: Cache, - ambient_credentials: Cache, + providers: Cache, + credentials: Cache, + loader: Arc, } impl Default for AwsAuthService { fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl AwsAuthService { + fn new(loader: Arc) -> Self { Self { - static_credentials: Cache::builder() + providers: Cache::builder().max_capacity(200).build(), + credentials: Cache::builder() .max_capacity(200) - .time_to_live(STATIC_CREDENTIALS_TTL) - .build(), - ambient_credentials: Cache::builder() - .max_capacity(200) - .time_to_live(AMBIENT_CREDENTIALS_TTL) + .expire_after(CredentialExpiry) .build(), + loader, } } } -fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { - match flow { - AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL), - AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL), - AwsAuthFlow::WebIdentity { .. } - | AwsAuthFlow::AssumeRole { .. } - | AwsAuthFlow::Profile { .. } - | AwsAuthFlow::SessionToken { .. } => None, +struct CredentialExpiry; + +impl Expiry for CredentialExpiry { + fn expire_after_create( + &self, + _: &AwsAuthConfig, + credentials: &Credentials, + _: std::time::Instant, + ) -> Option { + Some(credentials.expiry().map_or(CREDENTIAL_CACHE_TTL, |expiry| { + expiry + .duration_since(SystemTime::now()) + .unwrap_or_default() + .saturating_sub(EXPIRY_BUFFER) + .min(CREDENTIAL_CACHE_TTL) + })) } } -#[derive(Clone, Debug, Default, PartialEq, Eq)] +type ProviderFuture<'a> = + Pin> + Send + 'a>>; + +trait ProviderLoader: Send + Sync { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_>; +} + +struct NativeProviderLoader; + +impl ProviderLoader for NativeProviderLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + Box::pin(build_provider(config, flow)) + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct AwsAuthConfig { - pub access_key_id: Option, - pub secret_access_key: Option, - pub session_token: Option, + pub access_key_id: Option, + pub secret_access_key: Option, + pub session_token: Option, pub region_name: Option, pub session_name: Option, pub profile_name: Option, pub role_name: Option, - pub web_identity_token: Option, + pub web_identity_token: Option, pub sts_endpoint: Option, - pub external_id: Option, + pub external_id: Option, } impl AwsAuthConfig { fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + Sync)) -> Self { Self { - access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), + access_key_id: self + .access_key_id + .or_else(|| env_lookup(AWS_ACCESS_KEY_ID).map(SecretValue::new)), secret_access_key: self .secret_access_key - .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)), - session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)), + .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY).map(SecretValue::new)), + session_token: self + .session_token + .or_else(|| env_lookup(AWS_SESSION_TOKEN).map(SecretValue::new)), region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)), session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)), profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)), role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)), web_identity_token: self .web_identity_token - .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)), + .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN).map(SecretValue::new)), sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)), - external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)), + external_id: self + .external_id + .or_else(|| env_lookup(AWS_EXTERNAL_ID).map(SecretValue::new)), } } } @@ -96,52 +133,31 @@ impl AwsAuthConfig { #[derive(Clone, Debug, PartialEq, Eq)] pub enum AwsAuthFlow { WebIdentity { - token: String, + token: SecretValue, role: String, session_name: String, }, AssumeRole { role: String, session_name: Option, + web_identity_role: Option, }, Profile { name: String, }, SessionToken { - access_key_id: String, - secret_access_key: String, - session_token: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, + session_token: SecretValue, }, StaticKeys { - access_key_id: String, - secret_access_key: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, region_name: String, }, DefaultChain, } -fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { - let mut hasher = Sha256::new(); - hasher.update(format!("{config:?}:{flow:?}")); - format!("{:x}", hasher.finalize()) -} - -impl AwsAuthService { - fn get_cached_credentials(&self, key: &str) -> Option { - self.static_credentials - .get(key) - .or_else(|| self.ambient_credentials.get(key)) - } - - fn set_cached_credentials(&self, key: String, credentials: Credentials, ttl: Duration) { - if ttl == STATIC_CREDENTIALS_TTL { - self.static_credentials.insert(key, credentials); - } else { - self.ambient_credentials.insert(key, credentials); - } - } -} - fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { let mut parts = arn.splitn(6, ':'); let ("arn", partition, _, _, account, resource) = ( @@ -170,6 +186,9 @@ pub fn classify_auth( config: AwsAuthConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> AwsAuthFlow { + let web_identity_role = env_lookup(AWS_WEB_IDENTITY_TOKEN_FILE) + .filter(|token_file| !token_file.is_empty()) + .and_then(|_| env_lookup(AWS_ROLE_ARN)); let config = config.with_environment(env_lookup); if let (Some(token), Some(role), Some(session_name)) = ( config.web_identity_token.clone(), @@ -186,6 +205,7 @@ pub fn classify_auth( return AwsAuthFlow::AssumeRole { role, session_name: config.session_name.clone(), + web_identity_role, }; } if let Some(name) = config.profile_name { @@ -216,6 +236,30 @@ pub fn classify_auth( AwsAuthFlow::DefaultChain } +async fn profile_provider( + name: &str, + region: Option, + config: aws_config::provider_config::ProviderConfig, + builder: aws_config::profile::credentials::Builder, +) -> Result { + let region = match region { + Some(region) => Some(aws_types::region::Region::new(region)), + None => { + aws_config::default_provider::region::DefaultRegionChain::builder() + .profile_name(name) + .build() + .region() + .await + } + }; + Ok(SharedCredentialsProvider::new( + builder + .configure(&config.with_region(region)) + .profile_name(name) + .build(), + )) +} + impl AwsAuthService { pub async fn resolve_credentials( &self, @@ -224,183 +268,240 @@ impl AwsAuthService { ) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); - match flow { - AwsAuthFlow::SessionToken { - access_key_id, - secret_access_key, - session_token, - } => Ok(Credentials::new( - access_key_id, - secret_access_key, - Some(session_token), - None, - "litellm-static-session", - )), - AwsAuthFlow::StaticKeys { - access_key_id, - secret_access_key, - region_name, - } => { - let flow = AwsAuthFlow::StaticKeys { - access_key_id: access_key_id.clone(), - secret_access_key: secret_access_key.clone(), - region_name, - }; - let key = cache_key(&resolved, &flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let credentials = Credentials::new( - access_key_id, - secret_access_key, - None, - None, - "litellm-static", - ); - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), - ); - Ok(credentials) - } - AwsAuthFlow::Profile { name } => { - let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() - .profile_name(name) - .build(); + self.fresh_credentials(&resolved, &flow).await + } + + /// The credentials botocore's default chain finds in `region`, without LiteLLM's Bedrock settings + /// (`AWS_ROLE_NAME`, `AWS_PROFILE_NAME`, `AWS_SESSION_NAME`) that `resolve_credentials` also reads. + pub async fn default_chain_credentials(&self, region: &str) -> Result { + let config = AwsAuthConfig { + region_name: Some(region.to_owned()), + ..Default::default() + }; + self.fresh_credentials(&config, &AwsAuthFlow::DefaultChain) + .await + } + + async fn fresh_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + let cached = self.cached_credentials(resolved, flow).await?; + if fresh(&cached) { + return Ok(cached); + } + self.credentials.invalidate(resolved).await; + let refetched = self.cached_credentials(resolved, flow).await?; + if fresh(&refetched) { + return Ok(refetched); + } + self.credentials.invalidate(resolved).await; + Err(Error::ExpiredCredentials) + } + + async fn cached_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + self.credentials + .try_get_with(resolved.clone(), async { + let provider = self + .providers + .try_get_with( + resolved.clone(), + self.loader.load(resolved.clone(), flow.clone()), + ) + .await + .map_err(|error| (*error).clone())?; provider .provide_credentials() .await - .map_err(Error::credentials(CredentialFlow::Profile)) - } - AwsAuthFlow::AssumeRole { role, session_name } => { - if is_already_running_as_role(&role, &resolved).await? { - let ambient_flow = AwsAuthFlow::DefaultChain; - let key = cache_key(&resolved, &ambient_flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - return Ok(credentials); - } - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name.clone() { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint.clone() { - loader = loader.endpoint_url(endpoint); - } - if let (Some(access_key_id), Some(secret_access_key)) = - (resolved.access_key_id, resolved.secret_access_key) - { - loader = loader.credentials_provider(Credentials::new( - access_key_id, - secret_access_key, - resolved.session_token, - None, - "litellm-role-source", - )); - } - let sdk_config = loader.load().await; - let builder = aws_config::sts::AssumeRoleProvider::builder(role); - let builder = match session_name { - Some(name) => builder.session_name(name), - None => builder.session_name(default_session_name()), - }; - let builder = match resolved.external_id { - Some(id) => builder.external_id(id), - None => builder, - }; - let provider = builder.configure(&sdk_config).build().await; - provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::AssumeRole)) - } - AwsAuthFlow::WebIdentity { - token, - role, - session_name, - } => { - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint { - loader = loader.endpoint_url(endpoint); - } - let sdk_config = loader.load().await; - let client = aws_sdk_sts::Client::new(&sdk_config); - let response = client - .assume_role_with_web_identity() - .role_arn(role) - .role_session_name(session_name) - .web_identity_token(token) - .send() - .await - .map_err(Error::credentials(CredentialFlow::WebIdentity))?; - let credentials = response - .credentials() - .ok_or(Error::MissingWebIdentityCredentials)?; - let expiration = - SystemTime::try_from(*credentials.expiration()).map_err(|source| { - Error::WebIdentityExpiration(litellm_error::Source::new(source)) - })?; - Ok(Credentials::new( - credentials.access_key_id(), - credentials.secret_access_key(), - Some(credentials.session_token().to_string()), - Some(expiration), - "litellm-web-identity", - )) - } - AwsAuthFlow::DefaultChain => { - let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&AwsAuthFlow::DefaultChain) - .unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - Ok(credentials) + .map_err(Error::credentials(flow.credential_flow())) + }) + .await + .map_err(|error| (*error).clone()) + } +} + +fn fresh(credentials: &Credentials) -> bool { + credentials + .expiry() + .is_none_or(|expiry| expiry > SystemTime::now() + EXPIRY_BUFFER) +} + +impl AwsAuthFlow { + fn credential_flow(&self) -> CredentialFlow { + match self { + Self::Profile { .. } => CredentialFlow::Profile, + Self::AssumeRole { .. } => CredentialFlow::AssumeRole, + Self::WebIdentity { .. } => CredentialFlow::WebIdentity, + Self::DefaultChain | Self::StaticKeys { .. } | Self::SessionToken { .. } => { + CredentialFlow::DefaultChain } } } } -async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result { +async fn build_provider( + resolved: AwsAuthConfig, + flow: AwsAuthFlow, +) -> Result { + match flow { + AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + Some(session_token.expose().to_owned()), + None, + "litellm-static-session", + ))), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + .. + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + None, + None, + "litellm-static", + ))), + AwsAuthFlow::Profile { name } => { + profile_provider( + &name, + resolved.region_name, + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder(), + ) + .await + } + AwsAuthFlow::DefaultChain => { + let builder = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder(); + let builder = match resolved.region_name { + Some(region) => builder.region(aws_types::region::Region::new(region)), + None => builder, + }; + Ok(SharedCredentialsProvider::new(builder.build().await)) + } + AwsAuthFlow::AssumeRole { + role, + session_name, + web_identity_role, + } => { + if is_already_running_as_role(&role, web_identity_role.as_deref(), &resolved).await? { + return Ok(SharedCredentialsProvider::new( + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await, + )); + } + let sdk_config = sdk_config(resolved.clone()).await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role) + .session_name(session_name.unwrap_or_else(default_session_name)); + let builder = match resolved.external_id { + Some(id) => builder.external_id(id.expose()), + None => builder, + }; + Ok(SharedCredentialsProvider::new( + builder.configure(&sdk_config).build().await, + )) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => Ok(SharedCredentialsProvider::new(WebIdentityProvider { + client: aws_sdk_sts::Client::new(&sdk_config(resolved).await), + token, + role, + session_name, + })), + } +} + +async fn sdk_config(config: AwsAuthConfig) -> aws_types::SdkConfig { + let loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + let loader = match config.region_name { + Some(region) => loader.region(aws_types::region::Region::new(region)), + None => loader, + }; + let loader = match config.sts_endpoint { + Some(endpoint) => loader.endpoint_url(endpoint), + None => loader, + }; + let loader = match (config.access_key_id, config.secret_access_key) { + (Some(key), Some(secret)) => loader.credentials_provider(Credentials::new( + key.expose(), + secret.expose(), + config.session_token.map(|token| token.expose().to_owned()), + None, + "litellm-role-source", + )), + _ => loader, + }; + loader.load().await +} + +#[derive(Debug)] +struct WebIdentityProvider { + client: aws_sdk_sts::Client, + token: SecretValue, + role: String, + session_name: String, +} + +impl ProvideCredentials for WebIdentityProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async { + let response = self + .client + .assume_role_with_web_identity() + .role_arn(&self.role) + .role_session_name(&self.session_name) + .web_identity_token(self.token.expose()) + .send() + .await + .map_err(CredentialsError::provider_error)?; + let credentials = response.credentials().ok_or_else(|| { + CredentialsError::provider_error(Error::MissingWebIdentityCredentials) + })?; + let expiration = SystemTime::try_from(*credentials.expiration()) + .map_err(CredentialsError::provider_error)?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_owned()), + Some(expiration), + "litellm-web-identity", + )) + }) + } +} + +async fn is_already_running_as_role( + role: &str, + web_identity_role: Option<&str>, + config: &AwsAuthConfig, +) -> Result { + if config.access_key_id.is_some() || config.secret_access_key.is_some() { + return Ok(false); + } if role_identity(role).is_none() { return Ok(false); } - if let (Ok(current_role), Ok(token_file)) = ( - std::env::var(AWS_ROLE_ARN), - std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) && !token_file.is_empty() - { - return Ok(same_role_arns(role, ¤t_role)); + if let Some(current_role) = web_identity_role { + return Ok(same_role_arns(role, current_role)); } let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); @@ -473,16 +574,6 @@ pub fn sign_headers( headers: &HeaderMap, scope: SigningScope<'_>, ) -> Result { - let identity: Identity = scope.credentials.clone().into(); - let params = v4::SigningParams::builder() - .identity(&identity) - .region(scope.region) - .name(scope.service) - .time(scope.time) - .settings(SigningSettings::default()) - .build() - .map(SigningParams::from) - .map_err(Error::signing(SigningStage::Parameters))?; let header_refs = headers .iter() .map(|(name, value)| { @@ -494,16 +585,14 @@ pub fn sign_headers( }) }) .collect::, _>>()?; - let request = SignableRequest::new( + let instructions = sign_request( method, url, header_refs.into_iter(), - SignableBody::Bytes(body), - ) - .map_err(Error::signing(SigningStage::Request))?; - let (instructions, _) = sign(request, ¶ms) - .map_err(Error::signing(SigningStage::Signature))? - .into_parts(); + body, + SigningSettings::default(), + scope, + )?; instructions .headers() .map(|(name, value)| { @@ -516,6 +605,60 @@ pub fn sign_headers( .collect() } +pub fn presign_url( + url: &str, + expires_in: Duration, + scope: SigningScope<'_>, +) -> Result { + let mut signed = url::Url::parse(url).map_err(Error::signing(SigningStage::Request))?; + let mut settings = SigningSettings::default(); + settings.signature_location = SignatureLocation::QueryParams; + settings.expires_in = Some(expires_in); + let instructions = sign_request( + "GET", + signed.as_str(), + std::iter::empty(), + &[], + settings, + scope, + )?; + let (signature, auth): (Vec<_>, Vec<_>) = instructions + .params() + .iter() + .map(|(name, value)| (*name, value.as_ref())) + .partition(|(name, _)| *name == "X-Amz-Signature"); + signed + .query_pairs_mut() + .extend_pairs(auth.into_iter().chain(signature)); + Ok(signed.into()) +} + +fn sign_request<'a>( + method: &'a str, + url: &'a str, + headers: impl Iterator, + body: &'a [u8], + settings: SigningSettings, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(scope.region) + .name(scope.service) + .time(scope.time) + .settings(settings) + .build() + .map(SigningParams::from) + .map_err(Error::signing(SigningStage::Parameters))?; + let request = SignableRequest::new(method, url, headers, SignableBody::Bytes(body)) + .map_err(Error::signing(SigningStage::Request))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(Error::signing(SigningStage::Signature))? + .into_parts(); + Ok(instructions) +} + /// Model-id and region parsing shared by every Bedrock route. pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { let mut stripped = model; @@ -586,16 +729,26 @@ pub fn aws_auth_config( env_lookup: &dyn Fn(&str) -> Option, ) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: params.resolve(&AwsParams::ACCESS_KEY_ID, env_lookup), - secret_access_key: params.resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup), - session_token: params.resolve(&AwsParams::SESSION_TOKEN, env_lookup), + access_key_id: params + .resolve(&AwsParams::ACCESS_KEY_ID, env_lookup) + .map(SecretValue::new), + secret_access_key: params + .resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup) + .map(SecretValue::new), + session_token: params + .resolve(&AwsParams::SESSION_TOKEN, env_lookup) + .map(SecretValue::new), region_name: params.resolve(&AwsParams::REGION, env_lookup), session_name: params.resolve(&AwsParams::SESSION_NAME, env_lookup), profile_name: params.resolve(&AwsParams::PROFILE_NAME, env_lookup), role_name: params.resolve(&AwsParams::ROLE_NAME, env_lookup), - web_identity_token: params.resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup), + web_identity_token: params + .resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup) + .map(SecretValue::new), sts_endpoint: params.resolve(&AwsParams::STS_ENDPOINT, env_lookup), - external_id: params.resolve(&AwsParams::EXTERNAL_ID, env_lookup), + external_id: params + .resolve(&AwsParams::EXTERNAL_ID, env_lookup) + .map(SecretValue::new), } } @@ -660,10 +813,133 @@ mod tests { use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; use reqwest::header::CONTENT_TYPE; + #[cfg(unix)] + #[rstest::fixture] + fn process_profile() -> aws_runtime::env_config::file::EnvConfigFiles { + let credentials = serde_json::json!({ + "Version": 1, + "AccessKeyId": "process-access", + "SecretAccessKey": "process-secret", + "SessionToken": "process-session", + }); + aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + format!( + "[profile smoke]\ncredential_process = /usr/bin/printf '%s' '{credentials}'\n" + ), + ) + .build() + } + + #[cfg(unix)] + #[rstest::rstest] + #[tokio::test] + async fn profile_process_credentials_can_be_loaded( + process_profile: aws_runtime::env_config::file::EnvConfigFiles, + ) { + let provider = profile_provider( + "smoke", + Some("us-west-2".into()), + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_files(process_profile), + ) + .await + .unwrap(); + let credentials = provider.provide_credentials().await.unwrap(); + + assert_eq!(credentials.access_key_id(), "process-access"); + assert_eq!(credentials.secret_access_key(), "process-secret"); + assert_eq!(credentials.session_token(), Some("process-session")); + } + + #[derive(Debug)] + struct StsConnector; + + impl aws_smithy_runtime_api::client::http::HttpConnector for StsConnector { + fn call( + &self, + request: aws_smithy_runtime_api::client::orchestrator::HttpRequest, + ) -> aws_smithy_runtime_api::client::http::HttpConnectorFuture { + assert!(request.uri().contains("sts.us-west-2.")); + aws_smithy_runtime_api::client::http::HttpConnectorFuture::new(async { + Ok( + aws_smithy_runtime_api::client::orchestrator::HttpResponse::new( + 200_u16.try_into().unwrap(), + aws_smithy_types::body::SdkBody::from( + r#" + + role-accessrole-secret + role-session2100-01-01T00:00:00Z + arn:aws:sts::123456789012:assumed-role/smoke/session + smoke:session + smoke"#, + ), + ), + ) + }) + } + } + + #[rstest::rstest] + #[tokio::test] + async fn role_profile_uses_region_when_acquiring_sts_credentials() { + use aws_smithy_runtime_api::client::http::{SharedHttpConnector, http_client_fn}; + let files = aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + "[profile smoke]\nrole_arn = arn:aws:iam::123456789012:role/smoke\nsource_profile = source\n\n[profile source]\naws_access_key_id = source-access\naws_secret_access_key = source-secret\n", + ).build(); + let config = aws_config::provider_config::ProviderConfig::without_region() + .with_http_client(http_client_fn(|_, _| { + SharedHttpConnector::new(StsConnector) + })); + let credentials = profile_provider( + "smoke", + Some("us-west-2".into()), + config, + aws_config::profile::ProfileFileCredentialsProvider::builder().profile_files(files), + ) + .await + .unwrap(); + let credentials = credentials.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "role-access"); + assert_eq!(credentials.secret_access_key(), "role-secret"); + assert_eq!(credentials.session_token(), Some("role-session")); + assert!(credentials.expiry().is_some()); + } + fn no_env(_: &str) -> Option { None } + #[rstest::rstest] + #[case::access_key(Some("source-access"), None)] + #[case::secret_key(None, Some("source-secret"))] + #[case::source_pair(Some("source-access"), Some("source-secret"))] + #[tokio::test] + async fn explicit_source_keys_cannot_reuse_an_ambient_role( + #[case] access_key: Option<&str>, + #[case] secret_key: Option<&str>, + ) { + let role = "arn:aws:iam::123456789012:role/target"; + let config = AwsAuthConfig { + access_key_id: access_key.map(SecretValue::new), + secret_access_key: secret_key.map(SecretValue::new), + ..Default::default() + }; + assert!( + !is_already_running_as_role(role, Some(role), &config) + .await + .unwrap() + ); + } + #[test] fn secret_names_cover_environment_reads() { let seen = std::sync::Arc::new(std::sync::Mutex::new( @@ -706,7 +982,10 @@ mod tests { }; assert_eq!( - aws_auth_config(¶ms, &env).secret_access_key.as_deref(), + aws_auth_config(¶ms, &env) + .secret_access_key + .as_ref() + .map(SecretValue::expose), expected ); assert_eq!( @@ -831,6 +1110,53 @@ mod tests { ); } + #[rstest::rstest] + #[case::running_with_web_identity("/var/run/token", Some("arn:aws:iam::1:role/current"))] + #[case::empty_token_file("", None)] + #[tokio::test] + async fn assume_role_reads_the_ambient_role_from_the_secret_snapshot( + #[case] token_file: &str, + #[case] expected: Option<&str>, + ) { + use litellm_secrets::source::SecretSource; + + struct WebIdentitySecrets<'a>(&'a str); + + impl SecretSource for WebIdentitySecrets<'_> { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> futures_util::future::BoxFuture< + 'a, + Result, litellm_secrets::Error>, + > { + Box::pin(async move { + Ok(match name { + AWS_WEB_IDENTITY_TOKEN_FILE => Some(SecretValue::new(self.0)), + AWS_ROLE_ARN => Some(SecretValue::new("arn:aws:iam::1:role/current")), + _ => None, + }) + }) + } + } + + let snapshot = WebIdentitySecrets(token_file) + .resolve(AwsParams::secret_names()) + .await + .unwrap(); + let config = AwsAuthConfig { + role_name: Some("arn:aws:iam::1:role/target".into()), + ..Default::default() + }; + let AwsAuthFlow::AssumeRole { + web_identity_role, .. + } = classify_auth(config, &|name| snapshot.get(name)) + else { + panic!("a role name selects AssumeRole"); + }; + assert_eq!(web_identity_role.as_deref(), expected); + } + #[tokio::test] async fn static_credentials_do_not_use_network() { let credentials = AwsAuthService::default() @@ -849,110 +1175,231 @@ mod tests { assert_eq!(credentials.session_token(), None); } - #[test] - fn cache_policy_matches_python_flows() { - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::StaticKeys { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - region_name: "us-east-1".into(), - }), - Some(STATIC_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::DefaultChain), - Some(AMBIENT_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::SessionToken { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - session_token: "token".into(), - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::Profile { - name: "profile".into() - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::AssumeRole { - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: None, - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::WebIdentity { - token: "token".into(), - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: "session".into(), - }), - None - ); + struct CountingLoader { + loads: Arc, + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + } + + impl ProviderLoader for CountingLoader { + fn load(&self, config: AwsAuthConfig, _: AwsAuthFlow) -> ProviderFuture<'_> { + let loads = self.loads.clone(); + let provider = CountingProvider { + calls: self.calls.clone(), + short_lifetime: self.short_lifetime, + short_calls: self.short_calls, + key: config.access_key_id.unwrap(), + }; + Box::pin(async move { + loads.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(SharedCredentialsProvider::new(provider)) + }) + } + } + + #[derive(Debug)] + struct CountingProvider { + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + key: SecretValue, + } + + impl ProvideCredentials for CountingProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async move { + tokio::task::yield_now().await; + let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let lifetime = if call < self.short_calls { + self.short_lifetime + } else { + Duration::from_secs(3600) + }; + Ok(Credentials::new( + self.key.expose(), + "credential-value", + Some(format!("session-{call}")), + Some(SystemTime::now() + lifetime), + "counting", + )) + }) + } + } + + #[derive(Default)] + struct RecordingLoader { + loads: std::sync::Mutex>, + } + + impl ProviderLoader for RecordingLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + self.loads.lock().unwrap().push((config, flow)); + Box::pin(async { + Ok(SharedCredentialsProvider::new(Credentials::new( + "ak", + "sk", + None, + None, + "recording", + ))) + }) + } } #[rstest::rstest] - fn cache_round_trip_preserves_credentials() { - let auth = AwsAuthService::default(); - let key = format!("cache-test-{}", std::process::id()); - let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); - auth.set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); + #[tokio::test] + async fn default_chain_credentials_do_not_share_a_configured_role_in_the_same_region() { + let loader = Arc::new(RecordingLoader::default()); + let auth = AwsAuthService::new(loader.clone()); + let bedrock_env = |name: &str| match name { + AWS_ROLE_NAME => Some("arn:aws:iam::123456789012:role/bedrock".to_string()), + AWS_REGION_NAME => Some("us-east-1".to_string()), + _ => None, + }; + auth.resolve_credentials(AwsAuthConfig::default(), &bedrock_env) + .await + .unwrap(); + + auth.default_chain_credentials("us-east-1").await.unwrap(); + auth.default_chain_credentials("us-east-1").await.unwrap(); + + let loads = loader.loads.lock().unwrap(); + assert_eq!(loads.len(), 2); assert_eq!( - auth.get_cached_credentials(&key) - .map(|value| value.access_key_id().to_string()), - Some("cache-ak".to_string()) + loads[1], + ( + AwsAuthConfig { + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + AwsAuthFlow::DefaultChain + ) ); } #[rstest::rstest] #[tokio::test] - async fn cloned_services_reuse_credentials_but_independent_services_do_not() { - let auth = AwsAuthService::default(); + async fn near_expiry_credentials_are_refreshed_without_rebuilding_the_provider() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(30), + short_calls: 1, + })); let config = AwsAuthConfig { - access_key_id: Some("configured-key".into()), - secret_access_key: Some("configured-secret".into()), - region_name: Some("us-east-1".into()), - ..AwsAuthConfig::default() + access_key_id: Some("identity".into()), + ..Default::default() }; - let flow = classify_auth(config.clone(), &no_env); - let cached = Credentials::new("cached-key", "cached-secret", None, None, "test"); - auth.set_cached_credentials( - cache_key(&config, &flow), - cached.clone(), - STATIC_CREDENTIALS_TTL, - ); - - let reused = auth + let first = auth + .resolve_credentials(config.clone(), &no_env) + .await + .unwrap(); + let second = auth .clone() .resolve_credentials(config.clone(), &no_env) .await .unwrap(); - let independent = AwsAuthService::default() - .resolve_credentials(config.clone(), &no_env) - .await - .unwrap(); - let different = AwsAuthConfig { - access_key_id: Some("different-key".into()), - ..config.clone() - }; - let other_identity = auth - .resolve_credentials(different.clone(), &no_env) - .await - .unwrap(); + assert_eq!(first.session_token(), Some("session-1")); + assert_eq!(first.session_token(), second.session_token()); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } - assert_eq!(reused.access_key_id(), cached.access_key_id()); - assert_eq!(reused.secret_access_key(), cached.secret_access_key()); + #[rstest::rstest] + #[case::recovered_by_one_refetch(1, Some("session-1"))] + #[case::still_expired_after_refetch(2, None)] + #[tokio::test] + async fn expired_credentials_are_refetched_once( + #[case] short_calls: usize, + #[case] session: Option<&str>, + ) { + use std::sync::atomic::{AtomicUsize, Ordering}; + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: Arc::new(AtomicUsize::new(0)), + calls: calls.clone(), + short_lifetime: Duration::ZERO, + short_calls, + })); + let config = AwsAuthConfig { + access_key_id: Some("identity".into()), + ..Default::default() + }; + let result = auth.resolve_credentials(config, &no_env).await; assert_eq!( - Some(independent.access_key_id()), - config.access_key_id.as_deref() + result.as_ref().ok().and_then(Credentials::session_token), + session ); - assert_eq!( - Some(other_identity.access_key_id()), - different.access_key_id.as_deref() + assert!(session.is_some() || matches!(result, Err(Error::ExpiredCredentials))); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test] + async fn concurrent_calls_share_acquisition_and_distinct_identities_stay_isolated() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(3600), + short_calls: 0, + })); + let first = AwsAuthConfig { + access_key_id: Some("first".into()), + ..Default::default() + }; + let second = AwsAuthConfig { + access_key_id: Some("second".into()), + ..Default::default() + }; + let (a, b, c) = tokio::join!( + auth.resolve_credentials(first.clone(), &no_env), + auth.resolve_credentials(first, &no_env), + auth.resolve_credentials(second, &no_env) ); + let a = a.unwrap(); + let b = b.unwrap(); + let c = c.unwrap(); + assert_eq!(a.session_token(), b.session_token()); + assert_ne!(a.access_key_id(), c.access_key_id()); + assert_eq!(loads.load(Ordering::SeqCst), 2); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + fn auth_config_debug_never_discloses_credentials() { + let config = AwsAuthConfig { + access_key_id: Some("credential-value".into()), + secret_access_key: Some("credential-value".into()), + session_token: Some("credential-value".into()), + web_identity_token: Some("credential-value".into()), + external_id: Some("credential-value".into()), + ..Default::default() + }; + let output = format!("{config:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); + } + + #[rstest::rstest] + #[case::session(AwsAuthFlow::SessionToken { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), session_token: "credential-value".into() })] + #[case::static_keys(AwsAuthFlow::StaticKeys { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), region_name: "region".into() })] + #[case::web_identity(AwsAuthFlow::WebIdentity { token: "credential-value".into(), role: "role".into(), session_name: "session".into() })] + fn auth_flow_debug_never_discloses_credentials(#[case] flow: AwsAuthFlow) { + let output = format!("{flow:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); } #[test] @@ -1142,8 +1589,8 @@ mod tests { let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { - access_key_id: Some(access_key_id), - secret_access_key: Some(secret_access_key), + access_key_id: Some(access_key_id.into()), + secret_access_key: Some(secret_access_key.into()), region_name: Some("us-west-2".to_string()), ..Default::default() }, @@ -1191,4 +1638,44 @@ mod tests { failures.join("; ") ); } + + #[rstest::rstest] + #[case::provisioned( + None, + "https://cache-1/?Action=connect&User=iam-user", + "cache-1/?Action=connect&User=iam-user&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Signature=dfa58492a6ae9ba065088724170a52437329d03f3e9bc9a4c90a80cc8a18e2f6" + )] + #[case::serverless_with_a_session_token( + Some("session-token"), + "https://cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache", + "cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Security-Token=session-token&X-Amz-Signature=e48bdbb3bf9d5e8f1d325bf33d5260e333df1f6b8868e8edd1e4c3cc01b4c3f9" + )] + fn presigned_urls_match_botocore_sigv4_query_auth( + #[case] session_token: Option<&str>, + #[case] url: &str, + #[case] expected: &str, + ) { + // Expected values from botocore 1.43.117 SigV4QueryAuth with get_current_datetime patched to + // 2026-10-29T12:00:00Z (generated 2026-10-29), the signer Python's ElastiCache IAM provider uses + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + session_token.map(str::to_owned), + None, + "test", + ); + let signed = presign_url( + url, + Duration::from_secs(900), + SigningScope { + region: "us-east-1", + service: "elasticache", + credentials: &credentials, + time: UNIX_EPOCH + Duration::from_secs(1_793_275_200), + }, + ) + .unwrap(); + + assert_eq!(signed.strip_prefix("https://"), Some(expected)); + } } diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 789fde156bd..0e4662eda9e 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0]; pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0]; pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0]; pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0]; -pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; -pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV; +pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV; pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0]; pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0]; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index 5f88fc8a171..5445de1e6c3 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -2,6 +2,8 @@ use litellm_error::Source; #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { + #[error("AWS credential provider returned expired credentials")] + ExpiredCredentials, #[error("could not load AWS {flow} credentials")] Credentials { flow: CredentialFlow, diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index dff917d7c13..2d9ad9ba123 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer { ) .map_err(|error| match error { Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error), - Error::Credentials { .. } + Error::ExpiredCredentials + | Error::Credentials { .. } | Error::MissingWebIdentityCredentials | Error::WebIdentityExpiration(_) | Error::Signing { .. } => litellm_http::Error::signature(error), diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 7dcb0e513a7..36c10a74c80 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -8,17 +8,19 @@ repository.workspace = true [dependencies] litellm-secret-value.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true moka.workspace = true serde_json.workspace = true sha2.workspace = true strum.workspace = true +tokio.workspace = true url.workspace = true azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] -litellm-error.workspace = true +litellm-secrets.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index 9ed505e4160..6319406ce9c 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -3,5 +3,5 @@ mod native; mod resolve; mod types; -pub use resolve::{AzureAuthService, SECRET_NAMES}; +pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES}; pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index f40745ab7a2..03021428ea7 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest { pub(crate) struct ValidatedAzureRequest { request: NativeAzureRequest, credential_source: InputSource, + probe_imds: bool, } impl ValidatedAzureRequest { @@ -68,9 +69,18 @@ impl ValidatedAzureRequest { Ok(Self { request, credential_source, + probe_imds: false, }) } + pub(crate) fn with_imds_probe(self, enabled: bool) -> Self { + Self { + probe_imds: enabled + && matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }), + ..self + } + } + pub(crate) fn credential_source(&self) -> InputSource { self.credential_source } @@ -87,9 +97,34 @@ impl ValidatedAzureRequest { } } +pub(crate) enum CredentialAttempt { + Resolved(ResolvedCredential), + Unavailable(Error), + Failed(Error), +} + +impl CredentialAttempt { + pub(crate) fn into_result(self) -> Result { + match self { + Self::Resolved(credential) => Ok(credential), + Self::Unavailable(error) | Self::Failed(error) => Err(error), + } + } +} + pub(crate) struct NativeAzureTokenAcquirer { cache: AzureCredentialProviderCache, transport: Option, + imds_available: tokio::sync::Mutex, +} + +enum ImdsAvailability { + Unknown, + Available, + Unavailable { + retry_at: tokio::time::Instant, + error: Error, + }, } impl Default for NativeAzureTokenAcquirer { @@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: None, + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } @@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: Some(transport), + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } - pub(crate) async fn acquire( + pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt { + let discover_imds = request.probe_imds; + if discover_imds && let Err(error) = self.discover_imds().await { + return CredentialAttempt::Unavailable(error); + } + match self.acquire_token(request).await { + Ok(credential) => CredentialAttempt::Resolved(credential), + Err(error) => { + let missing_identity = discover_imds && unassigned_identity(&error); + if missing_identity { + CredentialAttempt::Unavailable(error) + } else { + CredentialAttempt::Failed(error) + } + } + } + } + + async fn discover_imds(&self) -> Result<(), Error> { + let mut availability = self.imds_available.lock().await; + match &*availability { + ImdsAvailability::Available => return Ok(()), + ImdsAvailability::Unavailable { retry_at, error } + if tokio::time::Instant::now() < *retry_at => + { + return Err(error.clone()); + } + ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {} + } + let result = probe_imds(self.transport.clone().unwrap_or_default()).await; + *availability = match &result { + Ok(()) => ImdsAvailability::Available, + Err(error) => ImdsAvailability::Unavailable { + retry_at: tokio::time::Instant::now() + Duration::from_secs(30), + error: error.clone(), + }, + }; + result + } + + async fn acquire_token( &self, request: ValidatedAzureRequest, ) -> Result { @@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer { } } +fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool { + std::iter::successors(Some(error), |cause| cause.source()).any(|cause| { + cause + .downcast_ref::() + .or_else(|| { + cause + .downcast_ref::()? + .downcast_ref::() + }) + .is_some_and(|error| { + error.http_status() == Some(azure_core::http::StatusCode::BadRequest) + }) + }) +} + +async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> { + use azure_core::http::{Context, Method, Request, Url}; + + let endpoint = + Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01") + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error)) + })?; + let mut request = Request::new(endpoint, Method::Get); + let response = tokio::time::timeout( + Duration::from_secs(1), + transport.send(&Context::new(), &mut request), + ) + .await + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })? + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })?; + // Headerless IMDS discovery expects 400, checked 2026-11-01 + // https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging + if response.status() != azure_core::http::StatusCode::BadRequest { + return Err(Error::CredentialAcquisition( + ErrorDetail::UnexpectedStatus { + operation: "Azure IMDS availability probe", + status: u16::from(response.status()), + }, + )); + } + Ok(()) +} + impl NativeAzureRequest { fn scope(&self) -> &str { match self { @@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion { #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; + use std::time::Duration; use azure_core::http::headers::Headers; use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport}; @@ -569,6 +695,296 @@ mod tests { } } + #[derive(Debug)] + struct AvailableImdsClient { + probe_status: StatusCode, + probe_delay: Duration, + recover_after: Option, + response_status: Option, + probes: std::sync::atomic::AtomicUsize, + } + + impl Default for AvailableImdsClient { + fn default() -> Self { + Self { + probe_status: StatusCode::BadRequest, + probe_delay: Duration::ZERO, + recover_after: None, + response_status: None, + probes: std::sync::atomic::AtomicUsize::new(0), + } + } + } + + impl HttpClient for AvailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + request: &'life1 Request, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if request + .headers() + .get_optional_str(&azure_core::http::headers::HeaderName::from_static( + "metadata", + )) + .is_none() + { + let probe = self + .probes + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if self.recover_after.is_some_and(|failures| probe >= failures) { + return Ok(AsyncRawResponse::from_bytes( + StatusCode::BadRequest, + Headers::new(), + "metadata header required", + )); + } + tokio::time::sleep(self.probe_delay).await; + return Ok(AsyncRawResponse::from_bytes( + self.probe_status, + Headers::new(), + "metadata header required", + )); + } + if let Some(status) = self.response_status { + return Ok(AsyncRawResponse::from_bytes( + status, + Headers::new(), + "identity endpoint failure", + )); + } + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + let expires = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + Headers::new(), + serde_json::json!({ + "access_token": "managed-token", "expires_on": expires.to_string(), + "token_type": "Bearer", "resource": "https://service.test", + }) + .to_string(), + )) + }) + } + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn available_imds_retains_full_sdk_token_acquisition() { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: None, + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + let credential = acquirer.acquire(request).await.into_result().unwrap(); + assert_eq!(credential.secret().expose(), "managed-token"); + } + + #[rstest::rstest] + #[case::another_metadata_service(StatusCode::NotFound)] + #[case::token_required(StatusCode::Unauthorized)] + #[case::proxy_rejection(StatusCode::Forbidden)] + #[case::unavailable(StatusCode::ServiceUnavailable)] + #[tokio::test(start_paused = true)] + async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition( + #[case] status: StatusCode, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + assert!(matches!( + acquirer.acquire(request).await, + super::CredentialAttempt::Unavailable(_) + )); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn an_imds_timeout_is_not_repeated_on_the_next_request() { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request)) + .await + .unwrap(); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::fixture] + fn imds_request() -> ValidatedAzureRequest { + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true) + } + + #[rstest::rstest] + #[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))] + #[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)] + #[tokio::test(start_paused = true)] + async fn failed_imds_discovery_recovers_without_restarting_the_acquirer( + imds_request: ValidatedAzureRequest, + #[case] status: StatusCode, + #[case] delay: Duration, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + probe_delay: delay, + recover_after: Some(1), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert_eq!( + first.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!( + second.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(Duration::from_secs(60)).await; + assert_eq!( + acquirer + .acquire(imds_request) + .await + .into_result() + .unwrap() + .secret() + .expose(), + "managed-token", + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn a_persistent_imds_outage_reuses_each_failed_probe( + imds_request: ValidatedAzureRequest, + ) { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + assert!(matches!( + acquirer.acquire(imds_request.clone()).await, + super::CredentialAttempt::Unavailable(_), + )); + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)] + #[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)] + #[case::default_identity_forbidden(true, StatusCode::Forbidden, false)] + #[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)] + #[tokio::test(start_paused = true)] + async fn only_unavailable_default_identities_allow_fallback( + #[case] discovery: bool, + #[case] status: StatusCode, + #[case] unavailable: bool, + ) { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: Some(status), + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(discovery); + let outcome = acquirer.acquire(request).await; + let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_)); + let error = outcome.into_result().unwrap_err(); + assert_eq!( + is_unavailable, + unavailable, + "{}", + litellm_error::Report::new(&error) + ); + } + + #[rstest::rstest] #[tokio::test] async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() { let transport = Arc::new(RecordingTokenClient::default()); @@ -582,8 +998,12 @@ mod tests { "https://login.test", ); - let first = acquirer.acquire(request.clone()).await.unwrap(); - let second = acquirer.acquire(request).await.unwrap(); + let first = acquirer + .acquire(request.clone()) + .await + .into_result() + .unwrap(); + let second = acquirer.acquire(request).await.into_result().unwrap(); assert_eq!(first.secret().expose(), "native-token"); assert_eq!(second.secret().expose(), "native-token"); @@ -599,6 +1019,7 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] async fn credential_provider_cache_isolates_every_client_secret_identity_field() { let transport = Arc::new(RecordingTokenClient::default()); @@ -645,10 +1066,10 @@ mod tests { ), ]; - acquirer.acquire(base.clone()).await.unwrap(); - acquirer.acquire(base).await.unwrap(); + acquirer.acquire(base.clone()).await.into_result().unwrap(); + acquirer.acquire(base).await.into_result().unwrap(); for request in variants.into_iter().skip(1) { - acquirer.acquire(request).await.unwrap(); + acquirer.acquire(request).await.into_result().unwrap(); } assert_eq!(transport.requests.lock().unwrap().len(), 6); diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index a92630af26a..9ff4eda540a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,3 +1,4 @@ +use crate::native::CredentialAttempt; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, TokenProviderHandle, @@ -12,25 +13,31 @@ use std::sync::Arc; use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; -const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; -const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID"; -const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID"; -const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET"; -const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE"; -const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; -const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; -const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; +#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)] +enum AzureSecretName { + #[strum(serialize = "AZURE_AD_TOKEN")] + AdToken, + #[strum(serialize = "AZURE_TENANT_ID")] + TenantId, + #[strum(serialize = "AZURE_CLIENT_ID")] + ClientId, + #[strum(serialize = "AZURE_CLIENT_SECRET")] + ClientSecret, + #[strum(serialize = "AZURE_SCOPE")] + Scope, + #[strum(serialize = "AZURE_AUTHORITY_HOST")] + AuthorityHost, + #[strum(serialize = "AZURE_CREDENTIAL")] + Credential, + #[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")] + FederatedTokenFile, + #[strum(serialize = "IDENTITY_ENDPOINT")] + IdentityEndpoint, + #[strum(serialize = "MSI_ENDPOINT")] + MsiEndpoint, +} -pub const SECRET_NAMES: &[&str] = &[ - AZURE_AD_TOKEN_ENV, - AZURE_TENANT_ID_ENV, - AZURE_CLIENT_ID_ENV, - AZURE_CLIENT_SECRET_ENV, - AZURE_SCOPE_ENV, - AZURE_AUTHORITY_HOST_ENV, - AZURE_CREDENTIAL_ENV, - AZURE_FEDERATED_TOKEN_FILE_ENV, -]; +pub const SECRET_NAMES: &[&str] = ::VARIANTS; #[derive(Clone, Debug)] pub(crate) enum AzureCredentialPlan { @@ -57,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>>; + ) -> Pin + Send + '_>>; } impl AzureTokenAcquirer for NativeAzureTokenAcquirer { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>> { + ) -> Pin + Send + '_>> { Box::pin(NativeAzureTokenAcquirer::acquire(self, request)) } } @@ -83,12 +90,36 @@ impl AzureAuthService { Self { native } } + pub async fn identity_token( + &self, + identity: &AzureIdentity, + scope: &str, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let plan = identity_plan(identity, scope, env_lookup)?; + let inputs = AzureAuthInputs::default_credential_for_scope(scope); + Ok(self + .run_plan(plan, &inputs, env_lookup) + .await? + .map(Sourced::into_value)) + } + pub async fn get_azure_ad_token( &self, inputs: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result>, Error> { - match select_auth_plan(inputs, env_lookup)? { + self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup) + .await + } + + async fn run_plan( + &self, + plan: AzureCredentialPlan, + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result>, Error> { + match plan { AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; @@ -123,6 +154,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Sourced::new(credential, source)) .map(Some) } @@ -131,6 +163,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Some(Sourced::new(credential, source))) } AzureCredentialPlan::Chain(requests) => { @@ -138,8 +171,16 @@ impl AzureAuthService { for request in requests { let source = request.credential_source(); match self.native.acquire(request).await { - Ok(credential) => return Ok(Some(Sourced::new(credential, source))), - Err(error) => failures.push(error), + CredentialAttempt::Resolved(credential) => { + return Ok(Some(Sourced::new(credential, source))); + } + CredentialAttempt::Unavailable(error) => failures.push(error), + CredentialAttempt::Failed(error) => { + return Err(Error::CredentialChain { + last: Box::new(error), + earlier: failures, + }); + } } } match failures.pop() { @@ -155,39 +196,161 @@ impl AzureAuthService { } } +/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings +/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads. +#[derive(Clone, Debug)] +pub enum AzureIdentity { + ClientSecret { + tenant_id: String, + client_id: String, + client_secret: SecretValue, + }, + ManagedIdentity { + client_id: String, + }, + Default, +} + +fn identity_plan( + identity: &AzureIdentity, + scope: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment); + let environment = |name: &str| { + env_lookup(name) + .filter(|value| !value.trim().is_empty()) + .map(|value| Sourced::new(value, InputSource::Environment)) + }; + let scope = deployment(scope); + let authority = environment(AzureSecretName::AuthorityHost.into()); + match identity { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + NativeAzureRequest::ClientSecret { + tenant_id: deployment(tenant_id), + client_id: deployment(client_id), + client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment), + scope, + authority, + }, + )?)), + AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native( + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: Some(deployment(client_id)), + scope, + selection_source: InputSource::Deployment, + })?, + )), + AzureIdentity::Default => { + let tenant_id = environment(AzureSecretName::TenantId.into()); + let client_id = environment(AzureSecretName::ClientId.into()); + let from_environment = match ( + tenant_id.clone(), + client_id.clone(), + environment(AzureSecretName::ClientSecret.into()), + ) { + (Some(tenant_id), Some(client_id), Some(client_secret)) => { + Some(NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret: client_secret.map(SecretValue::new), + scope: scope.clone(), + authority: authority.clone(), + }) + } + _ => None, + }; + let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none() + && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none(); + let AzureCredentialPlan::Chain(chain) = select_native_plan( + Some(Sourced::new( + AzureCredentialType::DefaultAzureCredential, + InputSource::Deployment, + )), + tenant_id, + client_id, + environment(AzureSecretName::FederatedTokenFile.into()), + scope, + authority, + InputSource::Environment, + )? + else { + unreachable!("DefaultAzureCredential always plans a chain") + }; + let environment_request = from_environment + .map(ValidatedAzureRequest::new) + .transpose()?; + Ok(AzureCredentialPlan::Chain( + environment_request + .into_iter() + .chain(chain) + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )) + } + } +} + pub(crate) fn select_auth_plan( inputs: &AzureAuthInputs, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup); - let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup); - let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup); - let client_secret = - configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup); - let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup) - .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); - let authority = configured_string( - &inputs.azure_authority_host, - AZURE_AUTHORITY_HOST_ENV, + let token = configured_secret( + &inputs.azure_ad_token, + AzureSecretName::AdToken.into(), env_lookup, ); - let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup) - .map(|value| { - value - .value() - .parse::() - .map(|selector| Sourced::new(selector, value.source())) - }) - .transpose() - .map_err(|error| { - Error::InvalidConfiguration(ErrorDetail::failed( - "Azure credential selector parsing", - error, - )) - })?; + let tenant_id = configured_string( + &inputs.tenant_id, + AzureSecretName::TenantId.into(), + env_lookup, + ); + let client_id = configured_string( + &inputs.client_id, + AzureSecretName::ClientId.into(), + env_lookup, + ); + let client_secret = configured_secret( + &inputs.client_secret, + AzureSecretName::ClientSecret.into(), + env_lookup, + ); + let scope = configured_string( + &inputs.azure_scope, + AzureSecretName::Scope.into(), + env_lookup, + ) + .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); + let authority = configured_string( + &inputs.azure_authority_host, + AzureSecretName::AuthorityHost.into(), + env_lookup, + ); + let selector = configured_string( + &inputs.azure_credential, + AzureSecretName::Credential.into(), + env_lookup, + ) + .map(|value| { + value + .value() + .parse::() + .map(|selector| Sourced::new(selector, value.source())) + }) + .transpose() + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed( + "Azure credential selector parsing", + error, + )) + })?; let federated_token_file = configured_string( &inputs.federated_token_file, - AZURE_FEDERATED_TOKEN_FILE_ENV, + AzureSecretName::FederatedTokenFile.into(), env_lookup, ); @@ -237,7 +400,12 @@ pub(crate) fn select_auth_plan( return Ok(AzureCredentialPlan::Missing); } - select_native_plan( + let probe_imds = selector + .as_ref() + .is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential) + && env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none() + && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none(); + match select_native_plan( selector, tenant_id, client_id, @@ -245,7 +413,15 @@ pub(crate) fn select_auth_plan( scope, authority, inputs.enable_azure_ad_token_refresh.source(), - ) + )? { + AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain( + requests + .into_iter() + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )), + plan => Ok(plan), + } } fn select_native_plan( @@ -509,7 +685,10 @@ mod tests { #[derive(Debug)] struct FileResolver; + use crate::native::CredentialAttempt; + struct ChainAcquirer { + terminal_on: Option<&'static str>, requests: Mutex>, succeed_on: Option<&'static str>, } @@ -518,20 +697,25 @@ mod tests { fn acquire( &self, request: ValidatedAzureRequest, - ) -> std::pin::Pin> + Send + '_>> - { + ) -> std::pin::Pin + Send + '_>> { let kind = request.kind(); self.requests.lock().unwrap().push(kind); Box::pin(async move { if self.succeed_on == Some(kind) { - Ok(ResolvedCredential::AccessToken { + CredentialAttempt::Resolved(ResolvedCredential::AccessToken { token: SecretValue::new("chain-token"), expires_on: None, }) } else { - Err(Error::CredentialAcquisition( - litellm_auth_types::ErrorDetail::Unresolved { subject: kind }, - )) + let error = + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved { + subject: kind, + }); + if self.terminal_on == Some(kind) { + CredentialAttempt::Failed(error) + } else { + CredentialAttempt::Unavailable(error) + } } }) } @@ -690,9 +874,210 @@ mod tests { assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); } + #[derive(Debug)] + struct UnavailableImdsClient; + + impl azure_core::http::HttpClient for UnavailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + _request: &'life1 azure_core::http::Request, + ) -> std::pin::Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(std::future::pending()) + } + } + + struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer); + + impl AzureTokenAcquirer for OfflineChainAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> std::pin::Pin + Send + '_>> { + Box::pin(async move { + if request.kind() == "developer-tools" { + return CredentialAttempt::Resolved(ResolvedCredential::AccessToken { + token: SecretValue::new("cli-token"), + expires_on: None, + }); + } + self.0.acquire(request).await + }) + } + } + + #[rstest::rstest] + #[case::default_off_azure("DefaultAzureCredential", None, true)] + #[case::explicit_managed_identity("ManagedIdentityCredential", None, false)] + #[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)] + #[tokio::test(start_paused = true)] + async fn only_default_imds_discovery_is_bounded( + #[case] selector: &str, + #[case] identity_endpoint: Option<&str>, + #[case] expect_cli: bool, + ) { + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let inputs = AzureAuthInputs { + azure_credential: crate::types::ConfigValue::Value(Sourced::new( + selector.into(), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let environment = |name: &str| { + (name == "IDENTITY_ENDPOINT") + .then_some(identity_endpoint) + .flatten() + .map(str::to_string) + }; + let outcome = tokio::time::timeout( + std::time::Duration::from_secs(2), + service.get_azure_ad_token(&inputs, &environment), + ) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().value().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] + #[case::app_service(Some("IDENTITY_ENDPOINT"), false)] + #[case::legacy_app_service(Some("MSI_ENDPOINT"), false)] + #[case::off_azure(None, true)] + #[tokio::test(start_paused = true)] + async fn default_chain_uses_endpoints_from_the_secret_snapshot( + #[case] endpoint: Option<&str>, + #[case] expect_cli: bool, + #[values(false, true)] consumer_identity: bool, + ) { + use litellm_secrets::source::SecretSource; + + struct EndpointSecrets<'a>(Option<&'a str>); + + impl SecretSource for EndpointSecrets<'_> { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> std::pin::Pin< + Box< + dyn Future, litellm_secrets::Error>> + + Send + + 'a, + >, + > { + Box::pin(async move { + Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test"))) + }) + } + } + + let snapshot = EndpointSecrets(endpoint) + .resolve(SECRET_NAMES) + .await + .unwrap(); + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let lookup = |name: &str| snapshot.get(name); + let scope = "https://service.test/.default"; + let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async { + if consumer_identity { + service + .identity_token(&super::AzureIdentity::Default, scope, &lookup) + .await + } else { + service + .get_azure_ad_token( + &AzureAuthInputs::default_credential_for_scope(scope), + &lookup, + ) + .await + .map(|token| token.map(Sourced::into_value)) + } + }) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] + #[case::litellm_provider_settings_are_ignored( + &[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")], + &["managed-identity", "developer-tools"] + )] + #[case::environment_client_secret_comes_first( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")], + &["client-secret", "managed-identity", "developer-tools"] + )] + #[case::aks_workload_identity( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")], + &["workload-identity", "managed-identity", "developer-tools"] + )] + #[tokio::test] + async fn the_default_identity_follows_azure_identity_defaults( + #[case] environment: &[(&str, &str)], + #[case] expected: &[&str], + ) { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let credential = service + .identity_token( + &super::AzureIdentity::Default, + "https://redis.azure.com/.default", + &lookup, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!(credential.secret().expose(), "chain-token"); + assert_eq!(*acquirer.requests.lock().unwrap(), expected); + } + + #[rstest::rstest] #[tokio::test] async fn default_chain_uses_declared_order_and_stops_after_success() { let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, requests: Mutex::new(Vec::new()), succeed_on: Some("developer-tools"), }); @@ -715,9 +1100,13 @@ mod tests { ); } + #[rstest::rstest] + #[case::all_unavailable(false)] + #[case::terminal_failure_after_unavailable(true)] #[tokio::test] - async fn chain_reports_each_acquisition_failure() { + async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) { let acquirer = Arc::new(ChainAcquirer { + terminal_on: terminal.then_some("developer-tools"), requests: Mutex::new(Vec::new()), succeed_on: None, }); @@ -744,6 +1133,27 @@ mod tests { ); } + #[rstest::rstest] + #[tokio::test] + async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: Some("managed-identity"), + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let error = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap_err(); + assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]); + assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty())); + } + #[derive(Debug)] struct CallerToken(&'static str); diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index ad60f0db817..90aab8d281c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -google-sdk = ["dep:google-cloud-auth", "dep:http"] +google-sdk = ["dep:http"] [dependencies] litellm-secret-value.workspace = true @@ -21,10 +21,14 @@ tokio.workspace = true veil.workspace = true gcp_auth = "0.12.7" -google-cloud-auth = { workspace = true, optional = true } +google-cloud-auth.workspace = true http = { workspace = true, optional = true } [dev-dependencies] +time = { workspace = true, features = ["formatting"] } +axum = { workspace = true, features = ["json"] } litellm-error.workspace = true rstest.workspace = true +rustls.workspace = true tempfile.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 650d5831796..5976ef82344 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -1,20 +1,11 @@ -use std::{future::Future, pin::Pin, sync::Arc}; +use std::sync::Arc; -use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential}; -use moka::future::Cache; -use serde_json::Value; -use sha2::{Digest, Sha256}; +use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential}; use crate::{ - config::{ - CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, - is_base64_encoded_json, names_a_credentials_file, non_empty_env, - }, - constants::{ - CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV, - VERTEXAI_API_KEY_ENV, - }, + config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env}, + constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV}, + identity::{GcpAuth, GcpProviderLoader, GcpTokenSource}, }; pub struct VertexEnvironment { @@ -22,28 +13,15 @@ pub struct VertexEnvironment { pub project_id: String, } -struct VertexAccessToken { - token: String, - project_id: String, -} - -#[derive(Clone)] +#[derive(Clone, Default)] pub struct VertexAuth { - providers: Cache>, - loader: Arc, -} - -impl Default for VertexAuth { - fn default() -> Self { - Self::new(Arc::new(GcpProviderLoader)) - } + identity: GcpAuth, } impl VertexAuth { - pub fn new(loader: Arc) -> Self { + pub fn new(loader: Arc) -> Self { Self { - providers: Cache::builder().max_capacity(64).build(), - loader, + identity: GcpAuth::new(loader), } } @@ -52,7 +30,10 @@ impl VertexAuth { config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { - self.load_provider(config, env_lookup).await?.token().await + self.identity + .access_token(credential_source(config, env_lookup)) + .await + .map(|token| token.expose().to_owned()) } pub async fn validate_environment( @@ -77,10 +58,18 @@ impl VertexAuth { (true, _) => None, (false, Some(token)) => Some(token), (false, None) => { - let access = self.get_access_token(config, env_lookup).await?; + let provider = self.load_provider(config, env_lookup).await?; + let (token, project_id) = match project_id { + Some(project_id) => (provider.token().await?, project_id), + None => tokio::try_join!(provider.token(), provider.project_id())?, + }; return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), + headers: apply_credential( + headers, + token.expose(), + CredentialPlacement::Bearer, + )?, + project_id, }); } }; @@ -104,248 +93,171 @@ impl VertexAuth { }) } - async fn get_access_token( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - let provider = self.load_provider(config, env_lookup).await?; - let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; - Ok(VertexAccessToken { token, project_id }) - } - async fn load_provider( &self, config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, Error> { - let source = credential_source(config, env_lookup); - let key = source.cache_key(); - self.providers - .try_get_with(key, self.loader.load(source)) + ) -> Result, Error> { + self.identity + .load_provider(credential_source(config, env_lookup)) .await - .map_err(|error| (*error).clone()) } } -pub trait VertexTokenSource: Send + Sync { - fn project_id(&self) -> VertexAuthFuture<'_, String>; - fn token(&self) -> VertexAuthFuture<'_, String>; -} - -pub trait VertexProviderLoader: Send + Sync { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; -} - -pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; - -struct GcpTokenSource(Arc); - -impl VertexTokenSource for GcpTokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .project_id() - .await - .map(|project| project.to_string()) - .map_err(auth_acquisition_error) - }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .token(&[CLOUD_PLATFORM_SCOPE]) - .await - .map(|token| token.as_str().to_string()) - .map_err(auth_acquisition_error) - }) - } -} - -struct GcpProviderLoader; - -impl VertexProviderLoader for GcpProviderLoader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - Box::pin(async move { - let provider: Arc = match source { - CredentialSource::Inline(configured) => Arc::new( - CustomServiceAccount::from_json(validate_request_credentials( - configured.expose(), - )?) - .map_err(auth_acquisition_error)?, - ), - CredentialSource::Trusted(configured) => { - Arc::new(load_trusted_service_account(configured.expose())?) - } - CredentialSource::ApplicationCredentials(path) => { - Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) - } - CredentialSource::Adc => { - gcp_auth::provider().await.map_err(auth_acquisition_error)? - } - }; - Ok(Arc::new(GcpTokenSource(provider)) as Arc) - }) - } -} - -fn validate_request_credentials(configured: &str) -> Result<&str, Error> { - let token_uri = serde_json::from_str::(configured) - .ok() - .and_then(|credentials| { - credentials - .get("token_uri") - .and_then(Value::as_str) - .map(str::to_string) - }); - if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { - return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { - field: "request-controlled vertex_credentials token_uri".into(), - expected: "the canonical Google OAuth token endpoint", - })); - } - Ok(configured) -} - -impl CredentialSource { - fn cache_key(&self) -> CredentialCacheKey { - match self { - Self::Inline(configured) => { - CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) - } - Self::Trusted(configured) => { - CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) - } - Self::ApplicationCredentials(path) => { - CredentialCacheKey::ApplicationCredentials(path.clone()) - } - Self::Adc => CredentialCacheKey::Adc, - } - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -enum CredentialCacheKey { - Inline([u8; 32]), - Trusted([u8; 32]), - ApplicationCredentials(String), - Adc, -} - -fn load_trusted_service_account(configured: &str) -> Result { - if is_base64_encoded_json(configured) { - return Err(Error::InvalidConfiguration( - ErrorDetail::Base64EncodedJson { - field: "vertex_credentials", - }, - )); - } - if names_a_credentials_file(configured) { - CustomServiceAccount::from_file(configured) - } else { - CustomServiceAccount::from_json(configured) - } - .map_err(auth_acquisition_error) -} - -fn auth_acquisition_error(error: gcp_auth::Error) -> Error { - Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( - "Vertex AI credentials", - error, - )) -} - #[cfg(test)] mod tests { + use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use crate::{ + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, + VertexConfig, get_vertex_ai_location, secret_names, + }; use litellm_secret_value::SecretValue; + use rstest::rstest; + use serde_json::{Value, json}; - use super::*; - - #[rstest::rstest] - #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] - #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] - #[case::missing_endpoint("{}", false)] - fn request_credentials_require_canonical_token_endpoint( - #[case] credentials: &str, - #[case] accepted: bool, - ) { - assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + struct FakeProvider { + project: Option<&'static str>, + calls: Arc, } - #[rstest::rstest] - #[case::missing_file( - "/nonexistent/vertex-key.json", - "failed to open application credentials file" - )] - #[case::malformed_json( - " {\"type\": \"service_account\",", - "failed to deserialize ApplicationCredentials" - )] - #[case::quoted_json( - "'{\"type\": \"service_account\"}'", - "failed to deserialize ApplicationCredentials" - )] - #[tokio::test] - async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( - #[case] configured: &str, - #[case] expected: &str, - ) { - let Err(error) = GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - else { - panic!("{configured:?} should not load"); - }; - let report = litellm_error::Report::new(&error).to_string(); - assert!(report.contains(expected), "{report}"); - } + impl GcpTokenSource for FakeProvider { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { + self.project.map(str::to_owned).ok_or( + litellm_auth_types::Error::CredentialAcquisition( + litellm_auth_types::ErrorDetail::Unresolved { + subject: "credential project", + }, + ), + ) + }) + } - #[rstest::rstest] - #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] - #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] - #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] - #[tokio::test] - async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { - assert!( - trusted_load_error(configured) - .await - .contains("is base64-encoded JSON, not the JSON itself") - ); - } - - async fn trusted_load_error(configured: &str) -> String { - match GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - { - Ok(_) => panic!("{configured:?} should not load"), - Err(error) => litellm_error::Report::new(&error).to_string(), + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(SecretValue::new("adc-token")) }) } } - #[rstest::rstest] - #[tokio::test] - async fn a_relative_key_file_named_like_json_is_read_as_a_file() { - let key = tempfile::Builder::new() - .prefix("{vertex") - .suffix(".json") - .tempfile_in(".") - .unwrap(); - std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); - let relative_name = key.path().file_name().unwrap().to_str().unwrap(); - let absolute_path = std::fs::canonicalize(key.path()).unwrap(); - assert_eq!( - trusted_load_error(relative_name).await, - trusted_load_error(absolute_path.to_str().unwrap()).await - ); + struct FakeLoader { + loads: Arc, + provider: Arc, } - #[rstest::rstest] - fn inline_and_trusted_credentials_never_share_a_cache_entry() { - assert_ne!( - CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), - CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() - ); + impl GcpProviderLoader for FakeLoader { + fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } + } + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { + calls, + project: Some("adc-project"), + }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) + } + + #[rstest] + #[tokio::test] + async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + let allowed = secret_names() + .iter() + .map(|name| (*name).to_owned()) + .collect::>(); + assert!(seen.lock().unwrap().is_subset(&allowed)); + } + + #[rstest] + #[tokio::test] + async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); + } + + #[rstest] + #[tokio::test] + async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + let config = VertexConfig::default(); + let environment = |_: &str| None; + let (first, second) = tokio::try_join!( + auth.validate_environment(Vec::new(), None, &config, &environment), + auth.validate_environment(Vec::new(), None, &config, &environment), + ) + .unwrap(); + assert_eq!(first.project_id, "adc-project"); + assert_eq!(first.headers[0].1, "Bearer adc-token"); + assert_eq!(first.project_id, second.project_id); + assert_eq!(first.headers, second.headers); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } + + #[rstest] + #[tokio::test] + async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() { + let auth = VertexAuth::new(Arc::new(FakeLoader { + loads: Arc::new(AtomicUsize::new(0)), + provider: Arc::new(FakeProvider { + calls: Arc::new(AtomicUsize::new(0)), + project: None, + }), + })); + let configured = VertexConfig::new(None, Some("configured-project".into()), None); + let environment = auth + .validate_environment(Vec::new(), None, &configured, &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "configured-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); } } diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 133e1c26903..048e40b0a2f 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -146,6 +146,20 @@ pub enum CredentialSource { Adc, } +impl CredentialSource { + pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option) -> Self { + VertexParams::CREDENTIALS + .resolve(&|_| None, env_lookup) + .map(|configured| Self::Trusted(SecretValue::new(configured))) + .unwrap_or_else(|| Self::application_default(env_lookup)) + } + + pub fn application_default(env_lookup: &dyn Fn(&str) -> Option) -> Self { + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map_or(Self::Adc, Self::ApplicationCredentials) + } +} + pub(crate) fn credential_source( config: &VertexConfig, env_lookup: &dyn Fn(&str) -> Option, @@ -158,12 +172,7 @@ pub(crate) fn credential_source( } }; } - if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) { - return CredentialSource::Trusted(SecretValue::new(configured)); - } - non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) - .map(CredentialSource::ApplicationCredentials) - .unwrap_or(CredentialSource::Adc) + CredentialSource::from_vertex_environment(env_lookup) } fn optional_credentials( @@ -319,6 +328,24 @@ mod tests { )); } + #[rstest::rstest] + fn application_default_ignores_vertex_credentials() { + let env = |name: &str| match name { + VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()), + GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()), + _ => None, + }; + assert!( + matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + CredentialSource::application_default( + &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into()) + ), + CredentialSource::Adc + )); + } + #[rstest::rstest] fn params_become_a_deployment_sourced_config() { let params = VertexParams { diff --git a/litellm-rust/crates/auth-gcp/src/identity.rs b/litellm-rust/crates/auth-gcp/src/identity.rs new file mode 100644 index 00000000000..5144221b218 --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/identity.rs @@ -0,0 +1,445 @@ +use crate::{ + config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file}, + constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT}, +}; +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use litellm_auth_types::{Error, ErrorDetail}; +use litellm_secret_value::SecretValue; +use moka::future::Cache; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use std::{future::Future, pin::Pin, sync::Arc}; + +#[derive(Clone)] +pub struct GcpAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for GcpAuth { + fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl GcpAuth { + pub fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + pub async fn access_token(&self, source: CredentialSource) -> Result { + self.load_provider(source).await?.token().await + } + + pub(crate) async fn load_provider( + &self, + source: CredentialSource, + ) -> Result, Error> { + self.providers + .try_get_with(source.cache_key(), self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +pub trait GcpTokenSource: Send + Sync { + fn project_id(&self) -> GcpAuthFuture<'_, String>; + fn token(&self) -> GcpAuthFuture<'_, SecretValue>; +} + +pub trait GcpProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc>; +} + +pub type GcpAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct ServiceAccountTokenSource(Arc); + +impl GcpTokenSource for ServiceAccountTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| SecretValue::new(token.as_str())) + .map_err(auth_acquisition_error) + }) + } +} + +struct SdkTokenSource { + credentials: google_cloud_auth::credentials::AccessTokenCredentials, + project_id: Option, +} + +impl GcpTokenSource for SdkTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.project_id + .clone() + .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project", + })) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.credentials + .access_token() + .await + .map(|token| SecretValue::new(token.token)) + .map_err(auth_acquisition_error) + }) + } +} + +async fn load_application_credentials(path: &str) -> Result, Error> { + let contents = tokio::fs::read_to_string(path) + .await + .map_err(auth_acquisition_error)?; + load_credentials_json(&contents) +} + +fn load_credentials_json(contents: &str) -> Result, Error> { + let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?; + let kind = configured.get("type").and_then(Value::as_str); + let project_id = configured + .get("project_id") + .and_then(Value::as_str) + .map(str::to_string); + let credentials = match kind { + Some("authorized_user") => { + google_cloud_auth::credentials::user_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("impersonated_service_account") => { + google_cloud_auth::credentials::impersonated::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("external_account") => { + google_cloud_auth::credentials::external_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("service_account") | None => { + let provider = + CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?; + return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider)))); + } + Some(_) => { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "Google application credentials type".into(), + expected: "a supported Google credential type", + })); + } + } + .map_err(auth_acquisition_error)?; + Ok(Arc::new(SdkTokenSource { + credentials, + project_id, + })) +} + +struct NativeProviderLoader; + +impl GcpProviderLoader for NativeProviderLoader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + return load_trusted_credentials(configured.expose()).await; + } + CredentialSource::ApplicationCredentials(path) => { + return load_application_credentials(&path).await; + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, Error> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "request-controlled vertex_credentials token_uri".into(), + expected: "the canonical Google OAuth token endpoint", + })); + } + Ok(configured) +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +async fn load_trusted_credentials(configured: &str) -> Result, Error> { + if is_base64_encoded_json(configured) { + return Err(Error::InvalidConfiguration( + ErrorDetail::Base64EncodedJson { + field: "vertex_credentials", + }, + )); + } + if names_a_credentials_file(configured) { + load_application_credentials(configured).await + } else { + load_credentials_json(configured) + } +} + +fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error { + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( + "Google credentials", + error, + )) +} + +#[cfg(test)] +mod tests { + use litellm_secret_value::SecretValue; + + use super::*; + + #[rstest::rstest] + #[case::authorized_user("authorized_user", 2)] + #[case::external_account("external_account", 2)] + #[case::impersonated_service_account("impersonated_service_account", 4)] + #[tokio::test] + async fn explicit_application_credentials_refresh_expired_tokens( + #[case] kind: &str, + #[case] expected_requests: usize, + #[values(false, true)] inline: bool, + ) { + use axum::{Json, Router, routing::post}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let _ = rustls::crypto::ring::default_provider().install_default(); + let requests = Arc::new(AtomicUsize::new(0)); + let recorded = requests.clone(); + let impersonations = requests.clone(); + let app = Router::new().route("/token", post(move || { + let recorded = recorded.clone(); + async move { + Json(serde_json::json!({ + "access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expires_in": 3600, + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + })) + } + })); + let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || { + let recorded = impersonations.clone(); + async move { + let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap(); + Json(serde_json::json!({ + "accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expireTime": expiration, + })) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let file = tempfile::NamedTempFile::new().unwrap(); + let subject_token = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(subject_token.path(), "subject-token").unwrap(); + let configured = if kind == "external_account" { + serde_json::json!({ + "type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"), + "credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project" + }) + } else { + let source = serde_json::json!({ + "type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret", + "refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token") + }); + if kind == "impersonated_service_account" { + serde_json::json!({ + "type": kind, "source_credentials": source, "quota_project_id": "quota-project", + "service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken") + }) + } else { + source + } + }; + std::fs::write(file.path(), configured.to_string()).unwrap(); + let source = if inline { + CredentialSource::Trusted(SecretValue::new(configured.to_string())) + } else { + CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into()) + }; + let provider = NativeProviderLoader.load(source).await.unwrap(); + + assert!(matches!( + provider.project_id().await, + Err(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project" + })) + )); + let original = provider.token().await.unwrap(); + tokio::time::pause(); + tokio::time::advance(std::time::Duration::from_secs(3601)).await; + tokio::time::resume(); + let refreshed = provider.token().await.unwrap(); + assert_ne!(original, refreshed); + assert_eq!(requests.load(Ordering::SeqCst), expected_requests); + server.abort(); + } + + #[rstest::rstest] + #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] + #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] + #[case::missing_endpoint("{}", false)] + fn request_credentials_require_canonical_token_endpoint( + #[case] credentials: &str, + #[case] accepted: bool, + ) { + assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + } + + #[rstest::rstest] + #[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")] + #[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")] + #[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")] + #[tokio::test] + async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( + #[case] configured: &str, + #[case] expected: &str, + ) { + let Err(error) = NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + else { + panic!("{configured:?} should not load"); + }; + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); + } + + #[rstest::rstest] + #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] + #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] + #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] + #[tokio::test] + async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { + assert!( + trusted_load_error(configured) + .await + .contains("is base64-encoded JSON, not the JSON itself") + ); + } + + async fn trusted_load_error(configured: &str) -> String { + match NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + { + Ok(_) => panic!("{configured:?} should not load"), + Err(error) => litellm_error::Report::new(&error).to_string(), + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_relative_key_file_named_like_json_is_read_as_a_file() { + let key = tempfile::Builder::new() + .prefix("{vertex") + .suffix(".json") + .tempfile_in(".") + .unwrap(); + std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); + let relative_name = key.path().file_name().unwrap().to_str().unwrap(); + let absolute_path = std::fs::canonicalize(key.path()).unwrap(); + assert_eq!( + trusted_load_error(relative_name).await, + trusted_load_error(absolute_path.to_str().unwrap()).await + ); + } + + #[rstest::rstest] + #[tokio::test] + async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + file.path(), + r#"{"type":"unsupported","private_key":"credential-value"}"#, + ) + .unwrap(); + let result = GcpAuth::default() + .access_token(CredentialSource::ApplicationCredentials( + file.path().to_str().unwrap().into(), + )) + .await; + let error = result.unwrap_err(); + assert!(matches!(error, Error::InvalidConfiguration(_))); + assert!( + !litellm_error::Report::new(&error) + .to_string() + .contains("credential-value") + ); + } + + #[rstest::rstest] + fn inline_and_trusted_credentials_never_share_a_cache_entry() { + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } +} diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 2b389d45926..03c660e6eb8 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,16 +1,16 @@ mod auth; mod config; pub mod constants; +mod identity; #[cfg(feature = "google-sdk")] mod sdk; -pub use auth::{ - VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource, -}; +pub use auth::{VertexAuth, VertexEnvironment}; pub use config::{ CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, }; pub use constants::secret_names; +pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource}; #[cfg(feature = "google-sdk")] pub use sdk::GoogleCredentials; diff --git a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs b/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs deleted file mode 100644 index 5991749ab8a..00000000000 --- a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs +++ /dev/null @@ -1,120 +0,0 @@ -use std::{ - collections::{BTreeMap, BTreeSet}, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, -}; - -use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader, - VertexTokenSource, get_vertex_ai_location, secret_names, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -struct FakeProvider { - calls: Arc, -} - -impl VertexTokenSource for FakeProvider { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-project".into()) }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-token".into()) }) - } -} - -struct FakeLoader { - loads: Arc, - provider: Arc, -} - -impl VertexProviderLoader for FakeLoader { - fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - let loads = self.loads.clone(); - let provider = self.provider.clone(); - Box::pin(async move { - loads.fetch_add(1, Ordering::SeqCst); - Ok(provider) - }) - } -} - -fn config(value: Value) -> VertexConfig { - VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) - .unwrap() -} - -fn auth(calls: Arc, loads: Arc) -> VertexAuth { - let provider: Arc = Arc::new(FakeProvider { calls }); - VertexAuth::new(Arc::new(FakeLoader { loads, provider })) -} - -#[rstest] -#[tokio::test] -async fn secret_names_cover_environment_reads() { - let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); - let recorded = seen.clone(); - let env = |name: &str| { - recorded.lock().unwrap().insert(name.to_string()); - None - }; - let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); - auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) - .await - .unwrap(); - get_vertex_ai_location(&VertexConfig::default(), &env); - assert!( - seen.lock() - .unwrap() - .iter() - .all(|name| secret_names().contains(&name.as_str())) - ); -} - -#[rstest] -#[tokio::test] -async fn explicit_token_and_header_do_not_acquire_adc() { - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); - let configured = config(json!({"vertex_project":"project-1"})); - let explicit = auth - .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) - .await - .unwrap(); - assert_eq!(explicit.headers[0].1, "Bearer access-token"); - let existing = auth - .validate_environment( - vec![("authorization".into(), "Bearer existing".into())], - None, - &configured, - &|_| None, - ) - .await - .unwrap(); - assert_eq!(existing.headers[0].1, "Bearer existing"); - assert_eq!(loads.load(Ordering::SeqCst), 0); -} - -#[rstest] -#[tokio::test] -async fn provider_is_reused_across_authentication_calls() { - let calls = Arc::new(AtomicUsize::new(0)); - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(calls.clone(), loads.clone()); - for _ in 0..2 { - let environment = auth - .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) - .await - .unwrap(); - assert_eq!(environment.project_id, "adc-project"); - assert_eq!(environment.headers[0].1, "Bearer adc-token"); - } - assert_eq!(loads.load(Ordering::SeqCst), 1); - assert_eq!(calls.load(Ordering::SeqCst), 4); -} diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index a8ba7b64e03..a2e5629bf3f 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -92,6 +92,11 @@ impl Error { #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { + #[error("{operation} returned HTTP {status}")] + UnexpectedStatus { + operation: &'static str, + status: u16, + }, #[error("{field} must be {expected}")] InvalidType { field: String, diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs index 8075fc15e9f..de21fbf94c9 100644 --- a/litellm-rust/crates/auth-types/src/params/aws.rs +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -53,6 +53,9 @@ pub struct AwsParams { } impl AwsParams { + pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN"; + pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; + pub const ACCESS_KEY_ID: ParamSpec = spec( "access key id", ParamKind::Credential, @@ -149,6 +152,10 @@ impl AwsParams { AwsParams::SPECS .iter() .flat_map(|spec| spec.env.iter().copied()) + .chain([ + AwsParams::ROLE_ARN_ENV, + AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV, + ]) .fold(Vec::new(), |mut names, name| { if !names.contains(&name) { names.push(name); diff --git a/litellm-rust/crates/auth-types/tests/aws_params.rs b/litellm-rust/crates/auth-types/tests/aws_params.rs index ed9a4f78ec2..b438984c199 100644 --- a/litellm-rust/crates/auth-types/tests/aws_params.rs +++ b/litellm-rust/crates/auth-types/tests/aws_params.rs @@ -3,21 +3,13 @@ use rstest::rstest; use serde_json::{Map, Value}; #[rstest] -fn fields_name_every_param_once_and_in_declaration_order() { +fn every_named_param_roundtrips_without_loss() { let filled: Map = AwsParams::fields() .map(|name| (name.to_string(), Value::from(format!("value-of-{name}")))) .collect(); let typed = AwsParams::from_optional_params(&filled); let serialized = serde_json::to_value(&typed).unwrap(); - assert_eq!( - serialized - .as_object() - .unwrap() - .keys() - .map(String::as_str) - .collect::>(), - AwsParams::fields().collect::>() - ); + assert_eq!(filled.len(), AwsParams::fields().count()); assert_eq!(serialized, Value::Object(filled.clone())); assert_eq!( serde_json::from_value::(Value::Object(filled)).unwrap(), diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index a6b08445bc8..b8c6720c234 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -8,7 +8,7 @@ use litellm_cache::{ use litellm_http::Client; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; -use crate::{GcpTokenSource, TokenSource}; +use crate::{GcsTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; @@ -62,7 +62,7 @@ pub struct GcsCache { impl GcsCache { pub fn new(config: GcsConfig, client: Client, codec: S) -> Self { - let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone())); Self::with_token_source(config, client, codec, token) } diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs index cbb61cf0685..d40bca081a4 100644 --- a/litellm-rust/crates/cache-gcs/src/lib.rs +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -2,4 +2,4 @@ mod cache; mod token; pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; -pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; +pub use token::{GcsTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 852d8bbf035..4f639e7b865 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,7 +1,6 @@ use std::{future::Future, pin::Pin}; -use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, Sourced}; +use litellm_auth_gcp::{CredentialSource, GcpAuth}; use litellm_secret_value::SecretValue; use litellm_cache::Error; @@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static { fn bearer_token(&self) -> Pin> + Send + '_>>; } -pub struct GcpTokenSource { - auth: VertexAuth, - config: VertexConfig, +pub struct GcsTokenSource { + auth: GcpAuth, + credentials: Option, } -impl GcpTokenSource { +impl GcsTokenSource { pub fn new(path_service_account: Option) -> Self { - let credentials = path_service_account - .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); Self { - auth: VertexAuth::default(), - config: VertexConfig::new(credentials, None, None), + auth: GcpAuth::default(), + credentials: path_service_account.filter(|path| !path.trim().is_empty()), } } } -impl TokenSource for GcpTokenSource { +impl TokenSource for GcsTokenSource { fn bearer_token(&self) -> Pin> + Send + '_>> { Box::pin(async move { self.auth - .access_token(&self.config, &|name| std::env::var(name).ok()) + .access_token(match &self.credentials { + Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())), + None => CredentialSource::application_default(&|name| std::env::var(name).ok()), + }) .await + .map(|token| token.expose().to_owned()) .map_err(Error::unavailable) }) } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index 659390a5659..02707c24881 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials { self.config.secret_access_key.clone(), ) { return Ok(AwsCredentials::new( - access_key_id, - secret_access_key, - self.config.session_token.clone(), + access_key_id.expose(), + secret_access_key.expose(), + self.config + .session_token + .as_ref() + .map(|value| value.expose().to_owned()), None, "litellm-s3-cache", )); diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/cache-s3/tests/auth.rs index 59c5aa2e091..5aa251ac749 100644 --- a/litellm-rust/crates/cache-s3/tests/auth.rs +++ b/litellm-rust/crates/cache-s3/tests/auth.rs @@ -5,9 +5,9 @@ use rstest::rstest; fn explicit(session_token: Option<&str>) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), - session_token: session_token.map(str::to_string), + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + session_token: session_token.map(Into::into), region_name: Some("us-east-1".to_string()), ..Default::default() } diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index 7378dde21bf..311d2402e72 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -23,8 +23,8 @@ pub fn config(endpoint: &str) -> S3CacheConfig { url: endpoint.to_string(), }), auth: AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), region_name: Some("us-east-1".to_string()), ..Default::default() }, diff --git a/litellm-rust/crates/inference-ocr/tests/resources.rs b/litellm-rust/crates/inference-ocr/tests/resources.rs index d618647e89b..3132a837eaf 100644 --- a/litellm-rust/crates/inference-ocr/tests/resources.rs +++ b/litellm-rust/crates/inference-ocr/tests/resources.rs @@ -7,7 +7,7 @@ use std::sync::{ use litellm_auth::AuthServices; use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource, + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, }; use litellm_http::{HttpSettings, Resolution}; use litellm_inference::resources::CoreResources; @@ -20,28 +20,28 @@ use support::{ReceivedRequest, json_response, upstream}; struct TokenSource(String); -impl VertexTokenSource for TokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { +impl GcpTokenSource for TokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { Box::pin(async { Ok(self.0.clone()) }) } - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async { Ok(self.0.clone()) }) + fn token(&self) -> GcpAuthFuture<'_, litellm_secret_value::SecretValue> { + Box::pin(async { Ok(litellm_secret_value::SecretValue::new(self.0.clone())) }) } } #[derive(Default)] struct Loader(AtomicUsize); -impl VertexProviderLoader for Loader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { +impl GcpProviderLoader for Loader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { Box::pin(async move { self.0.fetch_add(1, Ordering::SeqCst); let identity = match source { CredentialSource::Trusted(secret) => secret.expose().to_string(), other => panic!("unexpected credential source: {other:?}"), }; - Ok(Arc::new(TokenSource(identity)) as Arc) + Ok(Arc::new(TokenSource(identity)) as Arc) }) } } diff --git a/litellm-rust/crates/secret-value/src/secret_value.rs b/litellm-rust/crates/secret-value/src/secret_value.rs index 27e4227080d..33c559b80ec 100644 --- a/litellm-rust/crates/secret-value/src/secret_value.rs +++ b/litellm-rust/crates/secret-value/src/secret_value.rs @@ -16,6 +16,18 @@ impl SecretValue { } } +impl From for SecretValue { + fn from(value: String) -> Self { + Self::new(value) + } +} + +impl From<&str> for SecretValue { + fn from(value: &str) -> Self { + Self::new(value) + } +} + impl PartialEq for SecretValue { fn eq(&self, other: &Self) -> bool { subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into() diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs index 66f4f51bda9..6bec8ad96a1 100644 --- a/litellm-rust/crates/secrets-aws/src/auth.rs +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -39,30 +39,15 @@ impl Credentials { Self { auth, config: AwsAuthConfig { - access_key_id: context - .access_key_id - .as_ref() - .map(|value| value.expose().to_owned()), - secret_access_key: context - .secret_access_key - .as_ref() - .map(|value| value.expose().to_owned()), - session_token: context - .session_token - .as_ref() - .map(|value| value.expose().to_owned()), + access_key_id: context.access_key_id.as_ref().cloned(), + secret_access_key: context.secret_access_key.as_ref().cloned(), + session_token: context.session_token.as_ref().cloned(), region_name: region(settings, environment.as_ref()).ok(), role_name: settings.aws_role_name.clone(), session_name: settings.aws_session_name.clone(), - external_id: settings - .aws_external_id - .as_ref() - .map(|v| v.expose().to_owned()), + external_id: settings.aws_external_id.as_ref().cloned(), profile_name: settings.aws_profile_name.clone(), - web_identity_token: settings - .aws_web_identity_token - .as_ref() - .map(|v| v.expose().to_owned()), + web_identity_token: settings.aws_web_identity_token.as_ref().cloned(), sts_endpoint: settings.aws_sts_endpoint.clone(), }, environment, diff --git a/litellm/__init__.py b/litellm/__init__.py index c1f4bb6a3df..8e5dd68a186 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1796,6 +1796,9 @@ if TYPE_CHECKING: from .llms.bespoke.decisions.transformation import ( BespokeDecisionsConfig as BespokeDecisionsConfig, ) + from .llms.empiriolabs.decisions.transformation import ( + EmpirioLabsDecisionsConfig as EmpirioLabsDecisionsConfig, + ) from .llms.eu_jev.decisions.transformation import ( EuJevDecisionsConfig as EuJevDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 65074093dbb..bf903e23bc0 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = ( "OllamaDecisionsConfig", "LayaDecisionsConfig", "BespokeDecisionsConfig", + "EmpirioLabsDecisionsConfig", "EuJevDecisionsConfig", "OpenAIDecisionsConfig", "AzureDecisionsConfig", @@ -759,6 +760,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "OllamaDecisionsConfig": (".llms.ollama.decisions.transformation", "OllamaDecisionsConfig"), "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"), "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"), + "EmpirioLabsDecisionsConfig": (".llms.empiriolabs.decisions.transformation", "EmpirioLabsDecisionsConfig"), "EuJevDecisionsConfig": (".llms.eu_jev.decisions.transformation", "EuJevDecisionsConfig"), "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"), "AzureDecisionsConfig": (".llms.azure.decisions.transformation", "AzureDecisionsConfig"), diff --git a/litellm/accounting/__init__.py b/litellm/accounting/__init__.py index d4e56f61279..e69de29bb2d 100644 --- a/litellm/accounting/__init__.py +++ b/litellm/accounting/__init__.py @@ -1,62 +0,0 @@ -"""Public accounting contracts. Importing these values does not enable accounting.""" - -from litellm.accounting.measurements import measurements_for -from litellm.accounting.protocols import BudgetEnforcer, CostAggregator, CostCompiler, UsageMeter -from litellm.accounting.types import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - DeliveryState, - DeliveryTarget, - Dimension, - DimensionFilter, - DimensionName, - EventKind, - Granularity, - InboundUsageEvent, - Measurement, - Metric, - MetricName, - Multiplier, - NamedCost, - OutboundUsageEvent, - Outcome, - PricedQuantity, - PricingSnapshot, - Rate, - Unpriced, - validate_aggregate_rows, -) - -__all__ = [ - "AdmissionDecision", - "AggregateQuery", - "AggregateRow", - "BudgetEnforcer", - "CostAggregator", - "CostCompiler", - "CostEvent", - "DeliveryState", - "DeliveryTarget", - "Dimension", - "DimensionFilter", - "DimensionName", - "EventKind", - "Granularity", - "InboundUsageEvent", - "Measurement", - "Metric", - "MetricName", - "Multiplier", - "NamedCost", - "OutboundUsageEvent", - "Outcome", - "PricedQuantity", - "PricingSnapshot", - "Rate", - "Unpriced", - "UsageMeter", - "measurements_for", - "validate_aggregate_rows", -] diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py new file mode 100644 index 00000000000..1e5f23e2997 --- /dev/null +++ b/litellm/accounting/application.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING, Final + +from litellm.accounting.delivery_projection import entity_transactions, project_deliveries +from litellm.accounting.projection import Contribution +from litellm.proxy.db.budget_window_spend_writer import commit_window_spend_updates +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction +from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert +from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction + +if TYPE_CHECKING: + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +async def apply_deliveries( + transaction: LedgerTransaction, + deliveries: tuple[DeliveryProjection, ...], + *, + prisma_client: "PrismaClient", + writer: "DBSpendUpdateWriter", + proxy_logging_obj: "ProxyLogging", +) -> tuple[Contribution, ...]: + from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row + + contributions: Final = project_deliveries(deliveries) + await writer.apply_accounting_entities( + prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions) + ) + for table in sorted(DAILY_SPEND_TABLES.values(), key=lambda item: item.name): + rows = tuple( + {**dict(item.key), **dict(item.values), **dict(item.increments)} + for item in contributions + if item.table == table.name + ) + merged = merge_by_conflict_key(table, rows) + for offset in range(0, len(merged), 100): + sql, params = build_bulk_upsert(table, merged[offset : offset + 100]) + async with db_span("upsert_daily_spend", table.name): + await transaction.execute_raw(sql, *params) + async with transaction.batch_() as batcher: + for item in contributions: + if item.table == "LiteLLM_DailyModelUsage": + queue_model_usage_upsert(batcher, dict(item.key), dict(item.increments)) + for item in contributions: + if item.table == "LiteLLM_DailyGuardrailUsageUnits": + key = dict(item.key) + counters = dict(item.increments) + await upsert_usage_unit_row( + prisma_client, + UsageUnitKey(key["guardrail_id"], key["date"], key["team_id"], key["api_key"], key["usage_unit"]), + UsageUnitIncrement(int(counters["units"]), float(counters["cost"]), int(counters["untracked_units"])), + outer_transaction=transaction, + ) + windows: Final = tuple(_window(item) for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + await commit_window_spend_updates(prisma_client, windows, outer_transaction=transaction) + return contributions + + +def _window(item: Contribution) -> WindowSpendTransaction: + key: Final = dict(item.key) + values: Final = dict(item.values) + return WindowSpendTransaction( + entity_type=key["entity_type"], + entity_id=key["entity_id"], + window_duration=key["window_duration"], + window_start=values["window_start"] or "", + spend=float(dict(item.increments)["spend"]), + started_at=values["started_at"], + ) diff --git a/litellm/accounting/delivery_projection.py b/litellm/accounting/delivery_projection.py new file mode 100644 index 00000000000..da63542d5b8 --- /dev/null +++ b/litellm/accounting/delivery_projection.py @@ -0,0 +1,78 @@ +from collections.abc import Iterable +from itertools import chain, groupby +from typing import TYPE_CHECKING, Final +from urllib.parse import quote + +from litellm.accounting.projection import Contribution, project +from litellm.repositories.accounting_delivery_repository import DeliveryProjection + + +def _destination(contribution: Contribution) -> tuple[str, tuple[tuple[str, str], ...], str]: + return ( + contribution.table, + contribution.key, + (dict(contribution.values).get("window_start") or "") + if contribution.table == "LiteLLM_BudgetWindowSpend" + else "", + ) + + +def _sum(contributions: tuple[Contribution, ...]) -> Contribution: + first: Final = contributions[0] + increments: Final = tuple( + (column, sum(dict(item.increments).get(column, 0) for item in contributions)) + for column in dict.fromkeys( + column for column, _ in chain.from_iterable(item.increments for item in contributions) + ) + ) + starts: Final = tuple( + value + for column, value in chain.from_iterable(item.values for item in contributions) + if column == "started_at" and value + ) + values: Final = tuple( + (column, min(starts) if column == "started_at" and starts else value) for column, value in first.values + ) + return Contribution(first.table, first.key, increments, values) + + +def project_deliveries(deliveries: Iterable[DeliveryProjection]) -> tuple[Contribution, ...]: + ordered: Final = sorted(chain.from_iterable(project(delivery.row) for delivery in deliveries), key=_destination) + return tuple(_sum(tuple(group)) for _, group in groupby(ordered, key=_destination)) + + +def entity_transactions(contributions: tuple[Contribution, ...]) -> "DBSpendUpdateTransactions": + return { + "user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_UserTable")), + "end_user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_EndUserTable")), + "key_list_transactions": dict(_entity_totals(contributions, "LiteLLM_VerificationToken")), + "team_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamTable")), + "team_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamMembership")), + "org_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationTable")), + "org_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationMembership")), + "project_list_transactions": dict(_entity_totals(contributions, "LiteLLM_ProjectTable")), + "tag_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TagTable")), + "agent_list_transactions": dict(_entity_totals(contributions, "LiteLLM_AgentsTable")), + "model_access_group_list_transactions": dict( + _entity_totals(contributions, "LiteLLM_ModelAccessGroupBudgetTable") + ), + } + + +def _entity_totals(contributions: tuple[Contribution, ...], table: str) -> tuple[tuple[str, float], ...]: + return tuple( + (_entity_key(item), float(dict(item.increments)["spend"])) for item in contributions if item.table == table + ) + + +def _entity_key(item: Contribution) -> str: + fields: Final = dict(item.key) + if item.table == "LiteLLM_TeamMembership": + return f"team_id::{fields['team_id']}::user_id::{fields['user_id']}" + if item.table == "LiteLLM_OrganizationMembership": + return f"org_id::{quote(fields['organization_id'], safe='')}::user_id::{quote(fields['user_id'], safe='')}" + return item.key[0][1] + + +if TYPE_CHECKING: + from litellm.proxy._types import DBSpendUpdateTransactions diff --git a/litellm/accounting/measurements.py b/litellm/accounting/measurements.py deleted file mode 100644 index 83011a280da..00000000000 --- a/litellm/accounting/measurements.py +++ /dev/null @@ -1,13 +0,0 @@ -from typing import Final - -from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent - - -def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]: - """Return the complete contribution set for one event, ready for one apply call.""" - metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics - occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at - return tuple( - Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric) - for metric in metrics or () - ) diff --git a/litellm/accounting/metrics.py b/litellm/accounting/metrics.py new file mode 100644 index 00000000000..ade715d5d14 --- /dev/null +++ b/litellm/accounting/metrics.py @@ -0,0 +1,30 @@ +from typing import Final + +from prometheus_client import Gauge, Histogram + +from litellm.repositories.accounting_delivery_repository import DeliveryMetrics + +_PENDING: Final = Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax") +_PARKED: Final = Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax") +_OLDEST: Final = Gauge( + "litellm_accounting_oldest_pending_seconds", + "Age of the oldest pending SpendLogs delivery", + multiprocess_mode="livemax", +) +_BATCH_TIME: Final = Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time") +_BATCH_COUNT: Final = Gauge( + "litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax" +) +_TABLE_BYTES: Final = Gauge( + "litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax" +) + + +def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None: + _PENDING.set(metrics.pending_count) + _PARKED.set(metrics.parked_count) + _OLDEST.set(max(0.0, metrics.oldest_pending_seconds)) + _BATCH_TIME.observe(seconds) + + _BATCH_COUNT.set(metrics.batch_count) + _TABLE_BYTES.set(metrics.table_bytes) diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py new file mode 100644 index 00000000000..dd52b49f101 --- /dev/null +++ b/litellm/accounting/projection.py @@ -0,0 +1,326 @@ +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.constants import ( + BILLED_RETRY_ATTEMPT_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MODEL_INSIGHTS_DEFAULT_TASK, + MODEL_INSIGHTS_TASK_TAG_PREFIX, +) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCostByUnitEntry +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, daily_spend_entity_ids +from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks +from litellm.proxy.db.model_usage_rollup import date_from_start_time +from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING +from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens +from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +_FIELDS: Final = TypeAdapter(dict[str, object]) +_STRINGS: Final = TypeAdapter(tuple[str, ...]) +_VALUES: Final = TypeAdapter(tuple[object, ...]) +_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...]) +_TASK_TYPES: Final = frozenset(load_model_insight_tasks()) +_MAX_INDEXED_TEXT_BYTES: Final = 2000 +_MAX_COUNTER: Final = 2**40 + + +@dataclass(frozen=True, slots=True) +class Contribution: + table: str + key: tuple[tuple[str, str], ...] + increments: tuple[tuple[str, int | float], ...] + values: tuple[tuple[str, str | None], ...] = () + + +def _fields(value: object) -> Mapping[str, object]: + if isinstance(value, str): + return _FIELDS.validate_json(value) + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else MappingProxyType({}) + + +def _text(value: object) -> str: + return "" if value is None else str(value) + + +def _number(value: object) -> int | float: + return value if isinstance(value, (int, float)) else 0 + + +def _tags(row: Mapping[str, object]) -> tuple[str, ...]: + value: Final = row.get("request_tags") + if value is None: + return () + parsed: Final = _VALUES.validate_json(value) if isinstance(value, str) else _VALUES.validate_python(value) + return tuple(tag for tag in parsed if isinstance(tag, str)) + + +def _date(row: Mapping[str, object]) -> str | None: + start: Final = row.get("startTime") + return date_from_start_time(start) if isinstance(start, (datetime, str)) else None + + +def _spend(table: str, key: tuple[tuple[str, str], ...], spend: float) -> Contribution: + return Contribution(table, key, (("spend", spend),)) + + +def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + spend: Final = _number(row.get("spend")) + entities: Final = _fields(metadata.get("accounting_entities")) + identities: Final = ( + ("LiteLLM_VerificationToken", "token", entities.get("api_key")), + ("LiteLLM_UserTable", "user_id", entities.get("user")), + ("LiteLLM_UserTable", "user_id", entities.get("proxy_budget_name")), + ("LiteLLM_TeamTable", "team_id", entities.get("team_id")), + ("LiteLLM_OrganizationTable", "organization_id", entities.get("organization_id")), + ("LiteLLM_EndUserTable", "user_id", entities.get("end_user")), + ("LiteLLM_AgentsTable", "agent_id", entities.get("agent_id")), + ("LiteLLM_ProjectTable", "project_id", entities.get("project_id")), + ) + groups: Final = _STRINGS.validate_python(metadata.get("accounting_model_access_groups", ())) + return ( + *( + _spend(table, ((column, identity),), spend) + for table, column, identity in identities + if isinstance(identity, str) + ), + *(_spend("LiteLLM_TagTable", (("tag_name", tag),), spend) for tag in _tags(row) if tag), + *(_spend("LiteLLM_ModelAccessGroupBudgetTable", (("access_group_name", group),), spend) for group in groups), + *_membership_contributions(entities, spend), + ) + + +def _membership_contributions(row: Mapping[str, object], spend: float) -> tuple[Contribution, ...]: + user: Final = row.get("user") + if not isinstance(user, str): + return () + owners: Final = ( + ("LiteLLM_TeamMembership", "team_id", row.get("team_id")), + ("LiteLLM_OrganizationMembership", "organization_id", row.get("organization_id")), + ) + return tuple( + _spend(table, ((column, identity), ("user_id", user)), spend) + for table, column, identity in owners + if isinstance(identity, str) + ) + + +def _daily_counters(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[tuple[str, int | float], ...]: + usage: Final = _fields(metadata.get("usage_object")) + savings: Final = _fields(metadata.get("accounting_savings")) + internal: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) + caller: Final = not internal and metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is not True + success: Final = metadata.get("status") != "failure" + duration: Final = row.get("request_duration_ms") + timed: Final = not internal and success and isinstance(duration, int) and duration >= 0 + return ( + ("spend", _number(row.get("spend"))), + ("prompt_tokens", _number(row.get("prompt_tokens"))), + ("completion_tokens", _number(row.get("completion_tokens"))), + ("api_requests", int(caller)), + ("successful_requests", int(caller and success)), + ("failed_requests", int(caller and not success)), + ("cache_read_input_tokens", extract_cache_read_tokens(usage)), + ("cache_creation_input_tokens", extract_cache_creation_tokens(usage)), + ("compression_saved_tokens", extract_compression_saved_tokens(metadata)), + ("compression_savings_spend", _number(savings.get("compression"))), + ("prompt_caching_savings_spend", _number(savings.get("prompt_caching"))), + ("gateway_injected_caching_savings_spend", _number(savings.get("gateway_injected_caching"))), + ("autorouter_savings_spend", 0.0 if internal else _number(savings.get("autorouter"))), + ("total_response_time_ms", _number(duration) if timed else 0), + ("timed_requests", int(timed)), + ) + + +def _daily_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _date(row) + if date is None or "api_key" not in row: + return () + if "mcp_namespaced_tool_name" not in row and ( + "model" not in row or "custom_llm_provider" not in row or "model_group" not in row + ): + return () + counters: Final = _daily_counters(row, metadata) + endpoint: Final = ROUTE_ENDPOINT_MAPPING.get(_text(row.get("call_type"))) + shared_key: Final = ( + ("date", date), + ("api_key", _text(row.get("api_key"))), + ("model", _text(row.get("model"))), + ("custom_llm_provider", _text(row.get("custom_llm_provider"))), + ("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))), + ("endpoint", _text(endpoint)), + ) + # The legacy daily organization writer uses the charged org argument, not + # SpendLogs' normalized empty-string identity for a request with no org. + daily_row: Final = {**row, "organization_id": _fields(metadata.get("accounting_entities")).get("organization_id")} + return tuple(_daily_rows(daily_row, shared_key, counters)) + + +def _daily_rows( + row: Mapping[str, object], + shared_key: tuple[tuple[str, str], ...], + counters: tuple[tuple[str, int | float], ...], +) -> Iterator[Contribution]: + for entity, table in DAILY_SPEND_TABLES.items(): + for identity in daily_spend_entity_ids(row, entity): + yield Contribution( + table.name, + ((table.entity_id_column, _text(identity)), *shared_key), + counters, + ( + ("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))), + *((("request_id", _text(row.get("request_id"))),) if table.carries_request_id else ()), + ), + ) + + +def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + windows: Final = _ROWS.validate_python(metadata.get("accounting_windows", ())) + return tuple( + Contribution( + "LiteLLM_BudgetWindowSpend", + tuple((column, _text(window[column])) for column in ("entity_type", "entity_id", "window_duration")), + (("spend", _number(window.get("spend", row.get("spend")))),), + (("window_start", _text(window["window_start"])), ("started_at", _text(row.get("startTime")))), + ) + for window in windows + ) + + +def project(row: Mapping[str, object]) -> tuple[Contribution, ...]: + metadata: Final = _fields(row.get("metadata")) + contributions: Final = ( + *_entity_contributions(row, metadata), + *_daily_contributions(row, metadata), + *_window_contributions(row, metadata), + *_model_contributions(row, metadata), + *_guardrail_contributions(row, metadata), + ) + return tuple(_capped(contribution) for contribution in contributions if _indexable(contribution)) + + +def _indexable(contribution: Contribution) -> bool: + texts: Final = (value for _, value in (*contribution.key, *contribution.values) if value is not None) + return all(len(text.encode("utf-8", "surrogatepass")) <= _MAX_INDEXED_TEXT_BYTES for text in texts) + + +def _capped(contribution: Contribution) -> Contribution: + """Counters are summed into bigint columns in the charge's transaction, so one capped far below that range + can never overflow a day's total.""" + return replace( + contribution, + increments=tuple( + (name, max(-_MAX_COUNTER, min(_MAX_COUNTER, value)) if isinstance(value, int) else value) + for name, value in contribution.increments + ), + ) + + +def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _date(row) + origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) + if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): + return () + model: Final = _text(row.get("model")) or "unknown" + caller: Final = not origin + success: Final = row.get("status") == "success" + return ( + Contribution( + "LiteLLM_DailyModelUsage", + ( + ("date", date), + ("model_group", _text(row.get("model_group")) or model), + ("model", model), + ("custom_llm_provider", _text(row.get("custom_llm_provider")) or "unknown"), + ("task_type", _model_task(row)), + ), + ( + ("spend", _number(row.get("spend"))), + ("prompt_tokens", _number(row.get("prompt_tokens"))), + ("completion_tokens", _number(row.get("completion_tokens"))), + ("request_count", int(caller)), + ("successful_requests", int(caller and success)), + ("failed_requests", int(caller and not success)), + ), + ), + ) + + +def _utc_date(row: Mapping[str, object]) -> str | None: + start: Final = row.get("startTime") + try: + parsed: Final = datetime.fromisoformat(start.replace("Z", "+00:00")) if isinstance(start, str) else start + except ValueError: + return None + if not isinstance(parsed, datetime): + return None + return ( + (parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed) + .astimezone(timezone.utc) + .date() + .isoformat() + ) + + +def _guardrail_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _utc_date(row) + info: Final = metadata.get("guardrail_information") or metadata.get("standard_logging_guardrail_information") + if not row.get("request_id") or date is None or not isinstance(info, list): + return () + return tuple(chain.from_iterable(_guardrail_entry(row, date, entry) for entry in _ROWS.validate_python(info))) + + +def _guardrail_entry(row: Mapping[str, object], date: str, entry: Mapping[str, object]) -> tuple[Contribution, ...]: + identity: Final = _text(entry.get("guardrail_id") or entry.get("guardrail_name")) + usage: Final = entry.get("guardrail_usage") + if not identity or not isinstance(usage, Mapping): + return () + costs: Final = _guardrail_unit_costs(entry) + return tuple( + _guardrail_unit(row, (identity, date, unit), units, None if costs is None else costs.get(unit)) + for unit, units in _fields(usage).items() + if isinstance(units, int) and not isinstance(units, bool) and units > 0 + ) + + +def _guardrail_unit( + row: Mapping[str, object], identity: tuple[str, str, str], units: int, cost: float | None +) -> Contribution: + guardrail_id, date, unit = identity + return Contribution( + "LiteLLM_DailyGuardrailUsageUnits", + ( + ("guardrail_id", guardrail_id), + ("date", date), + ("team_id", _text(row.get("team_id"))), + ("api_key", _text(row.get("api_key"))), + ("usage_unit", unit), + ), + (("units", units), ("cost", cost or 0.0), ("untracked_units", units if cost is None else 0)), + ) + + +def _model_task(row: Mapping[str, object]) -> str: + return next( + ( + task + for tag in _tags(row) + if tag.startswith(MODEL_INSIGHTS_TASK_TAG_PREFIX) + and (task := tag.removeprefix(MODEL_INSIGHTS_TASK_TAG_PREFIX)) in _TASK_TYPES + ), + MODEL_INSIGHTS_DEFAULT_TASK, + ) + + +def _guardrail_unit_costs(entry: Mapping[str, object]) -> Mapping[str, float | None] | None: + try: + cost: Final = GuardrailCostByUnitEntry.model_validate(entry) + except ValidationError: + return None + return None if cost.guardrail_cost_in_spend is False else cost.guardrail_cost_by_unit diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py deleted file mode 100644 index 0208e37fbdf..00000000000 --- a/litellm/accounting/protocols.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Contracts only. Implementations own persistence, pricing, projections and live budgets. - -The event row's ``payload`` is a private envelope versioned by the row's ``schema_version``. -Usage envelopes retain the public event and request-scoped pricing context: quantities with units, -model/provider, requested and billed tier/region when known, and per-request/deployment overrides. -These inputs, including the effective cost-map rates, must survive request teardown. -Cost envelopes retain the public CostEvent, including its required finalized pricing snapshot. -T5/T6 own envelope serialization; this public API does not introduce a second persisted schema. -""" - -from typing import Protocol - -from litellm.accounting.types import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - Dimension, - InboundUsageEvent, - Measurement, - Metric, - OutboundUsageEvent, - Unpriced, -) - - -class UsageMeter(Protocol): - async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: - """Acknowledge commit, no-op identical IDs and reject conflicting content. - - Persist the event and its request-scoped pricing context before acknowledging. - HTTP acknowledgement is a separate policy. Unknown usage remains None. - """ - ... - - -class CostCompiler(Protocol): - async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: - """Reuse the committed charge or finalize one charge and its delivery intent. - - Resolve captured request-scoped context and rates by event identity. - Persist the resolved inputs in the required snapshot. Replay - reuses the committed result and never consumes free allowance again. - Missing pricing is Unpriced, never an inferred zero. - """ - ... - - -class CostAggregator(Protocol): - async def apply(self, measurements: tuple[Measurement, ...]) -> None: - """Apply one source event's complete contribution set per target/version. - - Each target commits all contributions and its receipt atomically. A receipt - for the first metric must not suppress the rest. Usage need not wait for cost. - """ - ... - - async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]: - """Return disjoint [start, end) rows per group using existing daily projections. - - Filters combine with AND, values within a filter with IN. DAY emits daily - buckets; None emits whole-interval totals. Reject unsupported metric/scope - combinations. Different groups may overlap, including non-additive tags. - """ - ... - - -class BudgetEnforcer(Protocol): - async def reserve( - self, reservation_id: str, dimensions: tuple[Dimension, ...], estimated: tuple[Metric, ...] - ) -> AdmissionDecision: - """Apply existing admission policies and estimates; duplicate IDs add no hold.""" - ... - - async def settle(self, reservation_id: str | None, cost_event: CostEvent) -> None: - """Settle the committed charge with atomic deduplication and hold mutation. - - Preserve scope/window attribution and retry unfinished settlement after - restart. SQL projection receipts do not provide live-counter atomicity. - """ - ... - - async def release(self, reservation_id: str) -> None: - """Idempotently release unused exposure at the request boundary, preserving charges.""" - ... diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py new file mode 100644 index 00000000000..d5652f79ed0 --- /dev/null +++ b/litellm/accounting/snapshot.py @@ -0,0 +1,161 @@ +from collections.abc import Callable, Mapping, Sequence +from datetime import datetime +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens +from litellm.proxy.spend_tracking.savings import SavingsSpend, compute_savings_spend, marks_gateway_injection + +if TYPE_CHECKING: + from litellm.router import Router + +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def _mapping(value: object) -> Mapping[str, object] | None: + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else None + + +_METADATA_FIELDS: Final = frozenset( + { + "usage_object", + "status", + "guardrail_information", + "standard_logging_guardrail_information", + "compression_savings", + "internal_call_origin", + "billed_retry_attempt", + } +) + + +_ACCOUNTING_FIELDS: Final = frozenset( + { + "request_id", + "call_type", + "api_key", + "spend", + "prompt_tokens", + "completion_tokens", + "startTime", + "model", + "model_group", + "custom_llm_provider", + "user", + "team_id", + "organization_id", + "end_user", + "agent_id", + "request_tags", + "status", + "request_duration_ms", + "mcp_namespaced_tool_name", + } +) + + +def build_accounting_payload( + payload: SpendLogsPayload, + *, + entities: Mapping[str, str | None], + model_access_groups: Sequence[str], + windows: Sequence[Mapping[str, object]], + router: "Callable[[], Router | None]", +) -> Mapping[str, object]: + metadata: Final = _FIELDS.validate_json(payload["metadata"]) + savings: Final = _savings_or_zero(payload, metadata, router) + accounting_metadata: Final = { + **{key: _accounting_metadata_field(key, value) for key, value in metadata.items() if key in _METADATA_FIELDS}, + "accounting_entities": entities, + "accounting_savings": savings._asdict(), + "accounting_model_access_groups": model_access_groups, + "accounting_windows": windows, + } + return _FIELDS.validate_python( + safe_json_structure( + { + **{ + key: value.isoformat() if isinstance(value, datetime) else value + for key, value in payload.items() + if key in _ACCOUNTING_FIELDS + }, + "metadata": accounting_metadata, + }, + non_finite_as_null=True, + ) + ) + + +def _savings_or_zero( + payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]" +) -> SavingsSpend: + """Savings are analytics computed from caller-supplied counts, so one that cannot be computed is stored as zero + rather than losing the charge it rides with.""" + try: + return _savings(payload, metadata, router) + except (ArithmeticError, ValueError, TypeError) as error: + verbose_proxy_logger.warning( + "Recording zero savings for request_id=%s, savings could not be computed: %s", payload["request_id"], error + ) + return SavingsSpend(compression=0.0, prompt_caching=0.0) + + +def _savings( + payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]" +) -> SavingsSpend: + return compute_savings_spend( + model=payload["model"], + custom_llm_provider=payload["custom_llm_provider"], + compression_saved_tokens=extract_compression_saved_tokens(metadata), + gateway_injected_cache=marks_gateway_injection(metadata, payload["model_id"]), + routing_decision=_mapping(metadata.get("routing_decision")), + model_id=payload["model_id"], + llm_router=router, + usage_object=_mapping(metadata.get("usage_object")), + cost_breakdown=_mapping(metadata.get("cost_breakdown")), + recorded_autorouter_savings=metadata.get("autorouter_savings"), + recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")), + billed_at=payload["endTime"], + ) + + +_GUARDRAIL_ENTRIES: Final = TypeAdapter(tuple[dict[str, object], ...]) +_GUARDRAIL_METADATA_FIELDS: Final = frozenset({"guardrail_information", "standard_logging_guardrail_information"}) +_GUARDRAIL_ACCOUNTING_FIELDS: Final = frozenset( + { + "guardrail_id", + "guardrail_name", + "guardrail_provider", + "guardrail_usage", + "guardrail_cost_by_unit", + "guardrail_cost_in_spend", + } +) +_GUARDRAIL_RESPONSE: Final = TypeAdapter(dict[str, object]) + + +def _accounting_guardrail_entries(value: object) -> object: + try: + entries: Final = _GUARDRAIL_ENTRIES.validate_python((value,) if isinstance(value, dict) else value) + except ValidationError: + return None + return [_accounting_guardrail_entry(entry) for entry in entries] + + +def _accounting_guardrail_entry(entry: Mapping[str, object]) -> Mapping[str, object]: + kept: Final = {key: item for key, item in entry.items() if key in _GUARDRAIL_ACCOUNTING_FIELDS} + response: Final = entry.get("guardrail_response") + if not isinstance(response, dict) or "tokens_saved" not in response: + return kept + return { + **kept, + "guardrail_response": {"tokens_saved": _GUARDRAIL_RESPONSE.validate_python(response)["tokens_saved"]}, + } + + +def _accounting_metadata_field(key: str, value: object) -> object: + return _accounting_guardrail_entries(value) if key in _GUARDRAIL_METADATA_FIELDS else value diff --git a/litellm/accounting/types.py b/litellm/accounting/types.py deleted file mode 100644 index 2ee6f60f756..00000000000 --- a/litellm/accounting/types.py +++ /dev/null @@ -1,428 +0,0 @@ -"""Accounting values use UTC instants, decimal strings on the wire and USD charges.""" - -from collections.abc import Mapping -from datetime import datetime, timezone -from decimal import Decimal -from enum import Enum -from typing import Annotated, Final, Literal, TypeAlias - -from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator -from typing_extensions import Self - -from litellm.types.utils import CallTypes - -_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True) -Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")] -Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)] - - -def _utc(value: datetime) -> datetime: - return value.astimezone(timezone.utc) - - -Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)] - - -class MetricName(str, Enum): - """Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens.""" - - REQUESTS = "requests" - SUCCESSFUL_REQUESTS = "successful_requests" - FAILED_REQUESTS = "failed_requests" - PROMPT_TOKENS = "prompt_tokens" - COMPLETION_TOKENS = "completion_tokens" - REASONING_TOKENS = "reasoning_tokens" - TOTAL_TOKENS = "total_tokens" - CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens" - CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens" - CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr" - SPEND = "spend" - - -class DimensionName(str, Enum): - KEY = "key" - USER = "user" - TEAM = "team" - ORGANIZATION = "organization" - END_USER = "end_user" - MODEL = "model" - MODEL_GROUP = "model_group" - DEPLOYMENT = "deployment" - PROVIDER = "provider" - TOOL = "tool" - AGENT = "agent" - PROJECT = "project" - MODEL_ACCESS_GROUP = "model_access_group" - API_BASE = "api_base" - SERVICE_TIER = "service_tier" - TAG = "tag" - - -class Outcome(str, Enum): - SUCCEEDED = "succeeded" - FAILED = "failed" - CANCELLED = "cancelled" - - -class Granularity(str, Enum): - DAY = "day" - - -class EventKind(str, Enum): - INBOUND_USAGE = "inbound_usage" - OUTBOUND_USAGE = "outbound_usage" - COST = "cost" - - -class DeliveryState(str, Enum): - PENDING = "pending" - APPLIED = "applied" - PARKED = "parked" - - -class DeliveryTarget(str, Enum): - """One member per existing table the projector writes; each is applied and receipted independently.""" - - KEY_SPEND = "LiteLLM_VerificationToken" - USER_SPEND = "LiteLLM_UserTable" - TEAM_SPEND = "LiteLLM_TeamTable" - TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership" - ORGANIZATION_SPEND = "LiteLLM_OrganizationTable" - ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership" - END_USER_SPEND = "LiteLLM_EndUserTable" - AGENT_SPEND = "LiteLLM_AgentsTable" - TAG_SPEND = "LiteLLM_TagTable" - PROJECT_SPEND = "LiteLLM_ProjectTable" - MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable" - BUDGET_SPEND = "LiteLLM_BudgetTable" - BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend" - DAILY_USER_SPEND = "LiteLLM_DailyUserSpend" - DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend" - DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend" - DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend" - DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend" - DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend" - DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage" - DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend" - AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend" - - -_SUBSET_PAIRS: Final = ( - (MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS), - (MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS), -) - - -def _check_subsets(values: Mapping[str, Decimal]) -> None: - for part, whole in _SUBSET_PAIRS: - if part.value in values and whole.value in values and values[part.value] > values[whole.value]: - raise ValueError(f"{part.value} cannot exceed {whole.value}") - - -class Metric(BaseModel): - model_config = _CONFIG - - name: MetricName - value: Quantity - unit: Literal["count", "USD"] = "count" - - @model_validator(mode="after") - def validate_unit(self) -> Self: - if (self.name is MetricName.SPEND) != (self.unit == "USD"): - raise ValueError("SPEND requires USD; usage metrics require count") - if self.unit == "count" and self.value != self.value.to_integral_value(): - raise ValueError("Count metrics require integral values") - return self - - -class Dimension(BaseModel): - model_config = _CONFIG - - name: DimensionName - value: Identifier - - -_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP}) - - -def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]: - if len(frozenset(values)) != len(values): - raise ValueError("Duplicate dimension values") - single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS) - if len(frozenset(single_names)) != len(single_names): - raise ValueError("Only tag and model access group dimensions may have multiple values") - return values - - -Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)] -_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS)) - - -def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]: - if len(frozenset(value.name for value in values)) != len(values): - raise ValueError("Duplicate metric names") - return values - - -Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)] - - -def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None: - if end < start or (end == start and not allow_empty): - raise ValueError("End must follow start") - - -class InboundUsageEvent(BaseModel): - model_config = _CONFIG - - event_id: Identifier - request_id: Identifier - call_type: CallTypes - started_at: Instant - ended_at: Instant - outcome: Outcome - dimensions: Dimensions - metrics: Metrics - - @model_validator(mode="after") - def validate_caller_activity(self) -> Self: - _interval(self.started_at, self.ended_at, allow_empty=True) - if any(metric.name not in _CALLER_METRICS for metric in self.metrics): - raise ValueError("Inbound events own only caller request metrics") - if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics): - raise ValueError("Inbound events count exactly one caller request") - for metric in self.metrics: - if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED): - raise ValueError("Successful request count must match outcome") - if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED): - raise ValueError("Failed request count must match outcome") - return self - - -class OutboundUsageEvent(BaseModel): - model_config = _CONFIG - - event_id: Identifier - attempt_id: Identifier - request_id: Identifier | None - retry_of_attempt_id: Identifier | None - call_type: CallTypes - started_at: Instant - ended_at: Instant - outcome: Outcome - dimensions: Dimensions - metrics: Metrics | None - - @model_validator(mode="after") - def validate_attempt_usage(self) -> Self: - _interval(self.started_at, self.ended_at, allow_empty=True) - if self.retry_of_attempt_id == self.attempt_id: - raise ValueError("An attempt cannot retry itself") - if self.metrics is not None and any( - metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics - ): - raise ValueError("Outbound events own only attempt usage metrics") - if self.metrics is not None: - _check_subsets({metric.name.value: metric.value for metric in self.metrics}) - return self - - -PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"] - - -class Rate(BaseModel): - model_config = _CONFIG - - measure: Identifier - unit: PricingUnit - usd_per_unit: Quantity - - -class PricedQuantity(BaseModel): - """A pricing input in its billing unit. Only seconds may be fractional.""" - - model_config = _CONFIG - - measure: Identifier - unit: PricingUnit - value: Quantity - - @model_validator(mode="after") - def validate_discrete_quantity(self) -> Self: - if self.unit != "second" and self.value != self.value.to_integral_value(): - raise ValueError("Discrete priced quantities require integral values") - return self - - -class Multiplier(BaseModel): - """A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5.""" - - model_config = _CONFIG - - name: Identifier - value: Identifier - factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)] - - -class NamedCost(BaseModel): - """A free-form extra charge by name, e.g. azure_model_router_flat_cost.""" - - model_config = _CONFIG - - name: Identifier - amount: Quantity - - -class PricingSnapshot(BaseModel): - """Retained inputs and result of authoritative compilation, reproducible after cost-map rates change. - - Request-scoped inputs, overrides and effective rates come from the usage row's versioned payload. - Retries reuse the committed charge and snapshot. - Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge. - Component sums and float tolerances belong to reconciliation, not this structural contract. - """ - - model_config = _CONFIG - - schema_version: Literal[1] = 1 - calculator_version: Identifier - model: Identifier - provider: Identifier | None - pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"] - priced_model_key: Identifier - multipliers: tuple[Multiplier, ...] = () - rates: tuple[Rate, ...] = () - usage: tuple[PricedQuantity, ...] = () - input_cost: Quantity | None = None - output_cost: Quantity | None = None - cache_read_cost: Quantity | None = None - cache_creation_cost: Quantity | None = None - reasoning_cost: Quantity | None = None - tool_usage_cost: Quantity | None = None - guardrail_cost: Quantity | None = None - additional_costs: tuple[NamedCost, ...] = () - original_cost: Quantity | None = None - discount_percent: Quantity | None = None - discount_amount: Quantity | None = None - margin_percent: Quantity | None = None - margin_fixed_amount: Quantity | None = None - margin_total_amount: Quantity | None = None - upstream_reported_cost: Quantity | None = None - total_cost: Quantity - - @model_validator(mode="after") - def validate_snapshot(self) -> Self: - rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates) - quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage) - if len(rate_keys) != len(self.rates): - raise ValueError("Duplicate rate measures and units") - if len(quantity_keys) != len(self.usage): - raise ValueError("Duplicate priced quantity measures and units") - if not rate_keys.issubset(quantity_keys): - raise ValueError("Every rate requires one priced quantity with the same measure and unit") - if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs): - raise ValueError("Duplicate additional cost names") - if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers): - raise ValueError("Duplicate multiplier names") - _check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"}) - return self - - -class CostEvent(BaseModel): - model_config = _CONFIG - - event_id: Identifier - usage_event_id: Identifier - occurred_at: Instant - dimensions: Dimensions - charges: Metrics - snapshot: PricingSnapshot - - @model_validator(mode="after") - def validate_charge(self) -> Self: - if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND: - raise ValueError("A cost event requires one authoritative USD spend metric") - if self.event_id == self.usage_event_id: - raise ValueError("Cost and usage events require distinct identities") - if self.snapshot.total_cost != self.charges[0].value: - raise ValueError("Snapshot total must equal the authoritative spend") - return self - - -class Unpriced(BaseModel): - model_config = _CONFIG - - usage_event_id: Identifier - reason: Identifier - - -class Measurement(BaseModel): - model_config = _CONFIG - - source_event_id: Identifier - occurred_at: Instant - dimensions: Dimensions - metric: Metric - - -class DimensionFilter(BaseModel): - model_config = _CONFIG - - name: DimensionName - values: Annotated[tuple[Identifier, ...], Field(min_length=1)] - - -class AggregateQuery(BaseModel): - model_config = _CONFIG - - metric: MetricName - filters: tuple[DimensionFilter, ...] - group_by: tuple[DimensionName, ...] - start: Instant - end: Instant - granularity: Granularity | None - - @model_validator(mode="after") - def validate_query(self) -> Self: - _interval(self.start, self.end) - if len(frozenset(item.name for item in self.filters)) != len(self.filters): - raise ValueError("Use one IN filter per dimension; filters combine with AND") - if len(frozenset(self.group_by)) != len(self.group_by): - raise ValueError("Duplicate group-by dimensions") - if any( - (instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0) - for instant in (self.start, self.end) - ): - raise ValueError("Initial reporting queries require UTC day boundaries") - return self - - -class AggregateRow(BaseModel): - model_config = _CONFIG - - start: Instant - end: Instant - dimensions: Dimensions - value: Quantity - - @model_validator(mode="after") - def validate_bounds(self) -> Self: - _interval(self.start, self.end) - return self - - -def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]: - """Validate [start, end) buckets per group; overlapping tag groups are not additive.""" - groups: Final = frozenset(frozenset(row.dimensions) for row in rows) - for group in groups: - ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start) - if any(left.end > right.start for left, right in zip(ordered, ordered[1:])): - raise ValueError("Aggregate buckets overlap within a dimension group") - return rows - - -class AdmissionDecision(BaseModel): - model_config = _CONFIG - - allowed: bool - reason: Identifier | None = None diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py new file mode 100644 index 00000000000..df71f2da31c --- /dev/null +++ b/litellm/accounting/worker.py @@ -0,0 +1,102 @@ +import os +import time +from typing import ( + TYPE_CHECKING, + Final, + Literal, + Protocol, + cast, # noqa: TID251 # generated Prisma transaction signatures cannot express the borrowed transaction protocol +) + +from litellm._logging import verbose_proxy_logger +from litellm.accounting.application import apply_deliveries +from litellm.accounting.metrics import record_deliveries +from litellm.accounting.projection import Contribution +from litellm.repositories.accounting_delivery_repository import ( + AccountingDeliveryRepository, + DeliveryProjection, + LedgerDatabase, + LedgerTransaction, +) +from litellm.types.integrations.slack_alerting import AlertType + +if TYPE_CHECKING: + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class _AccountingAlerts(Protocol): + async def alerting_handler( + self, message: str, level: Literal["Low", "Medium", "High"], alert_type: AlertType + ) -> None: ... + + +_JOB: Final = "accounting_delivery" + + +async def flush_accounting_deliveries( + writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging" +) -> None: + acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False) + if acquired is False: + return + applied: tuple[Contribution, ...] = () # rebind-ok: keep committed contributions for cache invalidation + + async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None: + nonlocal applied + applied = await apply_deliveries( # rebind-ok: invalidate caches only after the outer transaction commits + transaction, deliveries, prisma_client=prisma_client, writer=writer, proxy_logging_obj=proxy_logging_obj + ) + + try: + database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction + LedgerDatabase, prisma_client + ) + repository: Final = AccountingDeliveryRepository(database) + started: Final = time.monotonic() + batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000")) + + async def run_batch(*, retry_only: bool) -> bool: + result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only) + if result.applied: + await writer.invalidate_accounting_entities(applied, proxy_logging_obj) + if result.parked: + verbose_proxy_logger.error( + "Accounting delivery parked %d requests after isolated failures", len(result.parked) + ) + await _alert( + proxy_logging_obj, + message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.", + ) + return bool(result.applied or result.failed) + + await run_batch(retry_only=False) + retry_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True): + break + await repository.cleanup_applied( + retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")), + batch_size=batch_size * 2, + ) + metrics: Final = await repository.metrics() + record_deliveries(metrics, time.monotonic() - started) + if metrics.batch_count >= int( + os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000") + ) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")): + await _alert( + proxy_logging_obj, + message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.", + ) + if metrics.oldest_pending_seconds > 300: + await _alert( + proxy_logging_obj, + message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.", + ) + finally: + if acquired: + await writer.pod_lock_manager.release_lock(cronjob_id=_JOB) + + +async def _alert(alerts: _AccountingAlerts, message: str) -> None: + await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a0c0a90be5f..d272a235758 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,6 +10,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils import batch_utils as core_batch_utils from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo @@ -38,8 +39,8 @@ class BatchCostUsageResult: completion_cost: float = 0.0 -_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) _DECODED_BATCH_LINE: Final = TypeAdapter(object) +batch_cost_is_final: Final = core_batch_utils.batch_cost_is_final def _uses_native_vertex_output( @@ -54,26 +55,6 @@ def _uses_native_vertex_output( return first_row is not None and is_native_vertex_batch_output_row(first_row) -_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) - - -def batch_cost_is_final(batch: Batch) -> bool: - """Whether this retrieve of the batch is the one to account its cost from. - - A batch still in flight has nothing to price, and a "completed" batch can report - no output_file_id for a moment before the output populates; pricing either records - $0 under the batch's single spend row and pins it there. Final means a completed - batch whose output file has arrived or whose counts prove no line succeeded, or - any other terminal status (failed, cancelled, expired). - """ - if batch.status not in _TERMINAL_BATCH_STATUSES: - return False - if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: - return True - request_counts: Final = batch.request_counts - return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 - - async def calculate_batch_cost_and_usage( file_content_dictionary: Sequence[dict[str, object]], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 8dd4d87ae13..8e6341805db 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1180,11 +1180,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): include_namespace=True, replay_openai_history=provider == "openai", ) - # OpenAI's Responses API rejects an empty input. For a system-only - # request, carry the system message as a system-role input item instead - # of instructions, mirroring how non-string system content is already - # handled in convert_chat_completion_messages_to_responses_api. - is_system_only_request: Final = not converted_input_items and converted_instructions is not None + # Keep the system text in input instead of instructions when the Responses API needs it there. + # It rejects an empty input, and json_object requires the word "json" in input, ignoring instructions. + typed_optional_params: Final = cast("Mapping[str, object]", optional_params) # cast-ok: untyped dict + system_goes_in_input: Final = converted_instructions is not None and ( + not converted_input_items or _payload_type(typed_optional_params.get("response_format")) == "json_object" + ) target_input_items: Final = ( _without_audio_input_parts(converted_input_items) if _drops_audio_input(model, litellm_params) @@ -1196,12 +1197,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "type": "message", "role": "system", "content": [{"type": "input_text", "text": converted_instructions}], - } + }, + *target_input_items, ] - if is_system_only_request + if system_goes_in_input else target_input_items ) - instructions: Final = None if is_system_only_request else converted_instructions + instructions: Final = None if system_goes_in_input else converted_instructions # Build responses API request using the reverse transformation logic responses_api_request: Final = ResponsesAPIOptionalRequestParams() diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9154a6fdc30..3910a554905 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -616,11 +616,17 @@ def cost_per_token( ) if per_second is not None: return per_second + if custom_llm_provider == "anthropic": + return anthropic_cost_per_token( + model=model, usage=usage_block, service_tier=service_tier, model_info=requested_model_info_override + ) return generic_cost_per_token( model=model, usage=usage_block, custom_llm_provider=custom_llm_provider or "", service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, model_info=requested_model_info_override, ) @@ -3323,15 +3329,35 @@ def batch_cost_calculator( return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift +def _standard_rate(model_info: ModelInfo, key: str, prompt_tokens: int) -> float | None: + return get_cost_per_unit(model_info, get_modality_token_cost_key(model_info, key, prompt_tokens), None) + + +def _cached_audio_batch_fallback( + model_info: ModelInfo, batch_cache_read_rate: float | None, cache_read_rate: float | None, prompt_tokens: int +) -> float: + if batch_cache_read_rate is not None: + return batch_cache_read_rate + standard_audio_rate: Final = _standard_rate(model_info, "cache_read_input_audio_token_cost", prompt_tokens) + if standard_audio_rate is not None: + return standard_audio_rate + return 0.0 if cache_read_rate is None else cache_read_rate + + def _batch_prompt_cost( usage: Usage, model_info: ModelInfo, input_rate: float, - cache_read_rate: float | None, + batch_cache_read_rate: float | None, cache_creation_rate: float | None, cache_creation_rates_by_ttl: Mapping[str, float], ) -> float: details: Final = parse_prompt_tokens_details(usage) + cache_read_rate: Final = ( + _standard_rate(model_info, "cache_read_input_token_cost", usage.prompt_tokens) + if batch_cache_read_rate is None + else batch_cache_read_rate + ) cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0 cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0 bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl) @@ -3352,7 +3378,11 @@ def _batch_prompt_cost( + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate) + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0) + cached_audio_tokens - * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0) + * _batch_rate( + model_info, + "cache_read_input_audio_token_cost_batches", + _cached_audio_batch_fallback(model_info, batch_cache_read_rate, cache_read_rate, usage.prompt_tokens), + ) + ( calculate_cache_writing_cost( written_tokens, diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index ece85b09432..c4f5c1efbcd 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -340,6 +340,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_AccessGroupTable", "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent", + "LiteLLM_SpendLogDelivery", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginUpload", "LiteLLM_MemoryTable", diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1c9a443b546..fc1a601f9da 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -54,6 +54,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import is_registered_cust from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, ) +from litellm.llms.azure.common_utils import caused_by_entra_rejection from litellm.llms.github_copilot.per_user_auth import github_copilot_user_session_from from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.models.end_user import LiteLLM_EndUserTable @@ -448,7 +449,7 @@ def _is_deployment_failure(exception: openai.APIError, logged_call: Mapping[str, return not ( isinstance(exception, CallerCredentialAuthenticationError) or bool(logged_call.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)) - or logged_call.get("api_call_start_time") is None + or (logged_call.get("api_call_start_time") is None and not caused_by_entra_rejection(exception)) or _used_a_caller_credential(logged_call) ) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 5b94d4b4e5b..7c7b7ea5c4b 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -3,21 +3,32 @@ import json from collections.abc import Callable, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar +from dataclasses import dataclass from os import PathLike -from typing import Final, TypeVar +from typing import Final, Literal, TypeVar from openai.types import Batch from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import finite_cost +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG, canonical_video_id, is_video_status_call_type _DECODED_JSONL_ROW: Final = TypeAdapter(object) _DEPLOYMENT_FIELD: Final = TypeAdapter(Mapping[str, object]) +_PAYLOAD_FIELDS: Final = TypeAdapter(Mapping[str, object]) _ENDPOINT_PATHS: Final = TypeAdapter(tuple[str, ...]) +_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) +_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) BATCH_ENDPOINT: Final = "/v1/batch" +POLLED_BATCH_ID_KWARG: Final = "polled_batch_id" +POLLED_BATCH_COST_IS_FINAL_KWARG: Final = "polled_batch_cost_is_final" +BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" _AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value}) _batch_create_routing: Final[ContextVar[bool]] = ContextVar("batch_create_routing", default=False) @@ -289,12 +300,151 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES -def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str: - return f"{spend_key}:batch:{batch_id}" +@dataclass(frozen=True, slots=True) +class ChargeOnceJob: + kind: Literal["batch", "video"] + job_id: str + repeat_poll_cost: float = 0.0 -def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None: - """A finished batch reports its whole cost on every poll, so its id is charged once per counter.""" - if response_cost <= 0 or not is_batch_retrieve_call_type(call_type): +def job_charged_once_marker_key(spend_key: str, job: ChargeOnceJob) -> str: + return f"{spend_key}:{job.kind}:{job.job_id}" + + +def _guardrail_cost_of_poll(standard_logging_payload: object) -> float: + guardrail_information: Final = ( + _PAYLOAD_FIELDS.validate_python(standard_logging_payload).get("guardrail_information") + if isinstance(standard_logging_payload, Mapping) + else None + ) + return guardrail_information_cost(guardrail_information) + + +def charged_poll_cost(kwargs: Mapping[str, object]) -> object: + """The poll's whole cost as the cost callback charges it, guardrails included, which the logging kwargs' own + response_cost leaves out.""" + standard_logging_payload: Final = kwargs.get("standard_logging_object") + logged_cost: Final = ( + _PAYLOAD_FIELDS.validate_python(standard_logging_payload).get("response_cost") + if isinstance(standard_logging_payload, Mapping) + else None + ) + return kwargs.get("response_cost") if logged_cost is None else logged_cost + + +def job_cost_of_poll(response_cost: object, standard_logging_payload: object) -> float: + """What the polled job itself cost, since a guardrail on the poll charges every poll, running or finished.""" + return (finite_cost(response_cost) or 0.0) - _guardrail_cost_of_poll(standard_logging_payload) + + +def repeat_poll_cost(call_type: object, response_cost: object, standard_logging_payload: object) -> float: + """What a poll still costs once an earlier poll charged its job: a video poll's guardrails, which every video + poll has always paid, and nothing for a batch poll, whose repeats have always been free.""" + if not is_video_status_call_type(call_type): + return 0.0 + return max(0.0, min(finite_cost(response_cost) or 0.0, _guardrail_cost_of_poll(standard_logging_payload))) + + +def batch_cost_is_final(batch: Batch) -> bool: + """Whether this retrieve of the batch is the one to account its cost from. + + A batch still in flight has nothing to price, and a "completed" batch can report + no output_file_id for a moment before the output populates; pricing either records + $0 under the batch's single spend row and pins it there. Final means a completed + batch whose output file has arrived or whose counts prove no line succeeded, or + any other terminal status (failed, cancelled, expired). + """ + if batch.status not in _TERMINAL_BATCH_STATUSES: + return False + if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: + return True + request_counts: Final = batch.request_counts + return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 + + +@dataclass(frozen=True, slots=True) +class PolledBatch: + id: str + cost_is_final: bool + + +def _polled_batch(kwargs: Mapping[str, object], response_obj: object) -> PolledBatch | None: + """The retrieved batch, read from the logging kwargs when message redaction replaced the logged response.""" + if not is_batch_retrieve_call_type(kwargs.get("call_type")): return None - return response_obj.id if isinstance(response_obj, Batch) else None + if isinstance(response_obj, Batch): + return PolledBatch(id=response_obj.id, cost_is_final=batch_cost_is_final(response_obj)) + stamped_id: Final = kwargs.get(POLLED_BATCH_ID_KWARG) + stamped_cost_is_final: Final = kwargs.get(POLLED_BATCH_COST_IS_FINAL_KWARG) + if not isinstance(stamped_id, str) or not stamped_id or not isinstance(stamped_cost_is_final, bool): + return None + return PolledBatch(id=stamped_id, cost_is_final=stamped_cost_is_final) + + +def defers_batch_cost(kwargs: Mapping[str, object], response_obj: object) -> bool: + """A batch retrieve before the batch's cost is final charges nothing, guardrails included, since the retrieve + that settles the batch charges its whole cost once.""" + polled: Final = _polled_batch(kwargs, response_obj) + return polled is not None and not polled.cost_is_final + + +def pays_only_its_guardrail(call_type: object, kwargs: Mapping[str, object], response_cost: object) -> bool: + """A video poll whose whole cost is its guardrails', charged like any other request.""" + if not is_video_status_call_type(call_type): + return False + return ( + job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) + <= 0 + < (finite_cost(response_cost) or 0.0) + ) + + +def _requested_video_id(kwargs: Mapping[str, object]) -> str | None: + litellm_params: Final = kwargs.get("litellm_params") + requested: Final = ( + _PAYLOAD_FIELDS.validate_python(litellm_params).get("video_id") if isinstance(litellm_params, Mapping) else None + ) + return requested if isinstance(requested, str) and requested else None + + +def charged_video_id( + call_type: object, kwargs: Mapping[str, object], fallback_video_id: str | None, response_cost: object +) -> str | None: + """The id a finished video's own cost is charged once under. + + That is the id the client polled, which names the deployment that made the video, since a + custom provider's own id can repeat across deployments. The status comes from the logging + kwargs because message redaction replaces the logged video object. + """ + if not is_video_status_call_type(call_type) or kwargs.get(POLLED_VIDEO_STATUS_KWARG) != "completed": + return None + if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: + return None + video_id: Final = _requested_video_id(kwargs) or fallback_video_id + return None if video_id is None else canonical_video_id(video_id) + + +def _charged_batch(kwargs: Mapping[str, object], response_obj: object, response_cost: float) -> ChargeOnceJob | None: + polled: Final = _polled_batch(kwargs, response_obj) + if polled is None or not polled.cost_is_final or response_cost <= 0: + return None + return ChargeOnceJob(kind="batch", job_id=polled.id) + + +def job_to_charge_once( + kwargs: Mapping[str, object], response_obj: object, response_cost: float +) -> ChargeOnceJob | None: + """A finished batch or video reports its whole cost on every poll, so it is charged once per counter.""" + call_type: Final = kwargs.get("call_type") + if not is_video_status_call_type(call_type): + return _charged_batch(kwargs, response_obj, response_cost) + video_id: Final = charged_video_id( + call_type, kwargs, response_obj.id if isinstance(response_obj, VideoObject) else None, response_cost + ) + if video_id is None: + return None + return ChargeOnceJob( + kind="video", + job_id=video_id, + repeat_poll_cost=repeat_poll_cost(call_type, response_cost, kwargs.get("standard_logging_object")), + ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f8080531961..ce57a9bd4a5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ from datetime import datetime as dt_object from enum import Enum from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeVar, cast from httpx import Response, Timeout from openai import AsyncStream as OpenAIAsyncStream @@ -38,7 +38,7 @@ from litellm._logging import ( verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import batch_cost_is_final, handle_completed_batch +from litellm.batches.batch_utils import handle_completed_batch from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches @@ -79,6 +79,11 @@ from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.litellm_core_utils.agentic_followup_failure import ServedRound +from litellm.litellm_core_utils.batch_utils import ( + POLLED_BATCH_COST_IS_FINAL_KWARG, + POLLED_BATCH_ID_KWARG, + batch_cost_is_final, +) from litellm.litellm_core_utils.classifier_logging import ( classifier_audit_fields, classifier_input_snapshot, @@ -220,6 +225,7 @@ from litellm.types.utils import ( batch_chat_usage, ) from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG from litellm.utils import ( get_base_model_from_metadata, print_verbose, @@ -283,8 +289,6 @@ if TYPE_CHECKING: from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -662,54 +666,7 @@ class UsageObservation(Enum): ESTIMATED = "estimated" -@dataclass(frozen=True, slots=True) -class UsageCaptureObservation: - usage_json: str | None - provenance: UsageObservation - response_cost: float | None = None - - @classmethod - def from_usage( - cls, usage: Usage | None, provenance: UsageObservation, response_cost: float | None = None - ) -> "UsageCaptureObservation | UsageCaptureError": - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError - - try: - serialized: Final = ( - None if usage is None or provenance is UsageObservation.UNKNOWN else usage.model_dump_json() - ) - except Exception as error: # noqa: BLE001 # observation failures must not interrupt legacy retry handling - return UsageCaptureError(str(error)) - else: - return cls(serialized, provenance, response_cost) - - class Logging(LiteLLMLoggingBaseClass): - usage_capture_callback: ClassVar[ - Callable[ - ["Logging", object, datetime.datetime, datetime.datetime, bool], "CapturedUsage | UsageCaptureError | None" - ] - | None - ] = None - - def capture_usage_before_logging( - self, - response: object, - start_time: datetime.datetime | None = None, - end_time: datetime.datetime | None = None, - *, - freeze_only: bool = False, - ) -> None: - capture: Final = type(self).usage_capture_callback - if capture is not None and (not self.stream or self._is_assembled_stream_success(response)): - capture( - self, - response, - start_time or self.start_time, - end_time or datetime.datetime.now(datetime.timezone.utc), - freeze_only, - ) - global \ supabaseClient, \ promptLayerLogger, \ @@ -949,7 +906,6 @@ class Logging(LiteLLMLoggingBaseClass): self._stash_owed_on_success: bool = False self._stash_billed_by_continuation: bool = False self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = () - self._failed_attempt_ids_owed_on_success: tuple[str, ...] = () self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = () self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None @@ -2582,11 +2538,6 @@ class Logging(LiteLLMLoggingBaseClass): ) return - self.capture_usage_before_logging( - TypeAdapter(object).validate_python(result), - TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(start_time), - TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(end_time), - ) await self.async_success_handler( result, start_time=start_time, @@ -2698,20 +2649,12 @@ class Logging(LiteLLMLoggingBaseClass): self.finalized_cost = finalized self.model_call_details["response_cost"] = finalized.amount - def record_usage_capture(self, captured: "CapturedUsage") -> None: - self.model_call_details["_accounting_capture"] = captured - - def record_usage_pricing(self, pricing: object) -> None: - self.model_call_details["_accounting_pricing_snapshot"] = pricing - def record_partial_usage_for_failure( self, usage: Usage, response_cost: float, *, observation: UsageObservation = UsageObservation.KNOWN ) -> None: """Stash what an interrupted stream already consumed so the failure log bills it instead of zero.""" self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost - if type(self).usage_capture_callback is not None: - self.model_call_details["_accounting_usage_observed"] = observation def record_failure_after_earlier_rounds( self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float @@ -2795,12 +2738,6 @@ class Logging(LiteLLMLoggingBaseClass): owed_on_success: Final = self._stash_owed_on_success self._stash_billed_by_continuation = False self._stash_owed_on_success = False - captured_attempt_ids: Final = ( - self.capture_failed_attempt_usage(usage) - if isinstance(usage, Usage) and type(self).usage_capture_callback is not None - else () - ) - self.model_call_details.pop("_accounting_capture", None) self.model_call_details.pop("_accounting_pricing_snapshot", None) self.model_call_details.pop("_accounting_usage_observed", None) self.model_call_details.pop("combined_usage_object", None) @@ -2816,60 +2753,10 @@ class Logging(LiteLLMLoggingBaseClass): if not owed_on_success or billed_by_continuation: return self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt) - self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids) def mark_stash_billed_by_continuation(self) -> None: self._stash_billed_by_continuation = True - def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]: - from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage - - capture: Final = type(self).usage_capture_callback - if capture is None: - return () - attempt: Final = self.attempt_scoped_copy() - attempt.prepared_cost = None - attempt.finalized_cost = None - identity: Final = ( - f"{self.litellm_call_id}-attempt-{len(self.model_call_details.get('failed_attempts_usage', ()))}" - ) - params: Final = _LITELLM_PARAMS.validate_python(self.model_call_details.get("litellm_params") or {}) - metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(params) - observed: Final[object] = self.model_call_details.get("_accounting_usage_observed") - cost: Final[object] = self.model_call_details.get("response_cost") - attempt.model_call_details = { - **attempt.model_call_details, - "litellm_call_id": identity, - "parent_request_id": self.litellm_call_id, - "_accounting_capture": None, - "_accounting_prior_captures": (), - "_accounting_usage_snapshot": UsageCaptureObservation.from_usage( - usage, - observed if isinstance(observed, UsageObservation) else UsageObservation.KNOWN, - float(cost) if isinstance(cost, (int, float)) else None, - ), - "litellm_params": { - **params, - "metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, - "litellm_metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, - }, - } - captured: Final = capture( - attempt, - {"id": identity, "usage": usage.model_dump(), "status": "failed"}, - self.start_time, - datetime.datetime.now(datetime.timezone.utc), - False, - ) - if not isinstance(captured, CapturedUsage): - return () - self.model_call_details["_accounting_prior_captures"] = ( - *self.model_call_details.get("_accounting_prior_captures", ()), - captured, - ) - return tuple(write.event.event_id for write in captured.writes) - def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" self.model_call_details["continued_partial_usage"] = usage @@ -2909,8 +2796,6 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details.pop("failed_attempts_usage", None) self.failed_attempts_billed_on_success = failed_attempts self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts - if self._failed_attempt_ids_owed_on_success: - self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)] ) @@ -3073,45 +2958,6 @@ class Logging(LiteLLMLoggingBaseClass): return self._handle_a2a_response_logging(result=result) return result - def usage_capture_response(self, result: object) -> object: - if self.stream and not self._is_assembled_stream_success(result): - return None - copied: Final = result.model_copy(deep=True) if isinstance(result, BaseModel) else result - normalized: Final = self.normalize_logging_result(self.normalize_endpoint_logging_result(copied)) - if self.model_call_details.get("call_type") == CallTypes.arealtime.value and isinstance(result, list): - frames: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(result) - measured: Final = any( - isinstance(_LITELLM_PARAMS.validate_python(frame.get("response") or {}).get("usage"), Mapping) - for frame in frames - if frame.get("type") == "response.done" - ) - if not measured: - metered: Final = self.model_call_details.get(REALTIME_METERED_USAGE_KEY) - return { - "id": self.litellm_call_id, - "usage": {"metered_usage": TypeAdapter(object).dump_python(metered, mode="json")} - if metered is not None - else None, - } - assembled: Final = self._get_assembled_streaming_response( - normalized, self.start_time, self.start_time, True, [] - ) - return normalized if assembled is None else assembled - - def usage_capture_pricing(self, result: object) -> Mapping[str, JsonValue]: - from litellm.utils import count_characters - - return TypeAdapter(dict[str, JsonValue]).validate_python( - { - "custom_pricing": self.custom_pricing_for(result), - "custom_cost_per_token": self.custom_cost_per_token, - "custom_cost_per_second": self.custom_cost_per_second, - "prompt_characters": count_characters(self._prompt_for_cost_calculation()) - if self.model_call_details.get("call_type") in (CallTypes.speech.value, CallTypes.aspeech.value) - else None, - } - ) - def normalize_logging_result(self, result: object) -> object: """ Some endpoints return a different type of result than what is expected by the logging system. @@ -3468,6 +3314,11 @@ class Logging(LiteLLMLoggingBaseClass): result = self.normalize_endpoint_logging_result(result) if self.stream is not True: result = self._as_served_response(result) # rebind-ok: bridged chat calls log the served response + if isinstance(result, VideoObject): + self.model_call_details[POLLED_VIDEO_STATUS_KWARG] = result.status + if isinstance(result, LiteLLMBatch): + self.model_call_details[POLLED_BATCH_ID_KWARG] = result.id + self.model_call_details[POLLED_BATCH_COST_IS_FINAL_KWARG] = batch_cost_is_final(result) logging_result: Final = self.normalize_logging_result(result=result) if cache_hit is None and _served_from_cache_at_no_charge(logging_result): @@ -3826,7 +3677,6 @@ class Logging(LiteLLMLoggingBaseClass): self._as_served_chat_response(complete_streaming_response) ) self.model_call_details["complete_streaming_response"] = logged_streaming_response - self.capture_usage_before_logging(logged_streaming_response, start_time, end_time) ## STANDARDIZED LOGGING PAYLOAD if not for_async_request: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( @@ -4352,8 +4202,6 @@ class Logging(LiteLLMLoggingBaseClass): cost_for_built_in_tools_cost_usd_dollar=0.0, ) - self.capture_usage_before_logging(result, start_time, end_time) - self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( kwargs=self.model_call_details, messages=self.model_call_details.get("messages") @@ -7295,6 +7143,7 @@ class StandardLoggingPayloadSetup: init_response_obj: object, api_base: str | None = None, router_model_id: str | None = None, + custom_model_info: ModelInfo | None = None, ) -> StandardLoggingModelInformation: model_cost_name: Final = select_model_name_for_cost_calc( model=base_model if custom_pricing else None, @@ -7313,6 +7162,8 @@ class StandardLoggingPayloadSetup: if router_model_id is not None else model_cost_name ) + if model_cost_name is not None and custom_model_info is not None: + return StandardLoggingModelInformation(model_map_key=model_cost_name, model_map_value=custom_model_info) if model_cost_name is None: model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None) else: @@ -8063,6 +7914,7 @@ def get_standard_logging_object_payload( init_response_obj=init_response_obj, api_base=litellm_params.get("api_base"), router_model_id=logging_obj.get_router_model_id(), + custom_model_info=logging_obj.custom_model_info if custom_pricing else None, ) error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 54213f40ac0..e3aecda1deb 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -224,7 +224,20 @@ class StandardBuiltInToolCostTracking: web_search_options=standard_built_in_tools_params.get("web_search_options", None), model_info=model_info, ) - return per_call_cost * StandardBuiltInToolCostTracking._count_web_search_calls(response_object) + return per_call_cost * StandardBuiltInToolCostTracking._billable_web_search_calls( + resolved_usage, response_object + ) + + @staticmethod + def _billable_web_search_calls(usage: Usage | None, response_object: object) -> int: + """ + Providers without their own calculator but whose usage carries server_tool_use.web_search_requests + (OpenAI-compatible providers that report a per-tool count) are billed that many calls on the per-call path. + """ + reported: Final = None if usage is None else get_web_search_requests_from_usage(usage) + if isinstance(reported, int) and not isinstance(reported, bool) and reported > 0: + return reported + return StandardBuiltInToolCostTracking._count_web_search_calls(response_object) @staticmethod def _count_web_search_calls(response_object: object) -> int: diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8e6e71a31d3..b6af181980e 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -836,8 +836,6 @@ class RealTimeStreaming: if self.session_tools or self.tool_calls: self.logging_obj.model_call_details["realtime_tools"] = self.session_tools self.logging_obj.model_call_details["realtime_tool_calls"] = self.tool_calls - if LiteLLMLogging.usage_capture_callback is not None: - self.logging_obj.capture_usage_before_logging(self.messages) # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 9da0e35bdda..ac5c0720eff 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -272,15 +272,31 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ return { **fields, **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), + **({"reasoning": REDACTED_BY_LITELLM} if fields.get("reasoning") is not None else {}), + **({"reasoning_text": REDACTED_BY_LITELLM} if fields.get("reasoning_text") is not None else {}), **{ key: None for key in ( "thinking_blocks", "reasoningContentBlocks", + "reasoningContent", + "annotations", "citations", + "citation", + "citationsContent", "compaction_blocks", "compaction_start", "compaction_delta", + "reasoning_details", + "web_search_calls", + "web_search_results", + "tool_results", + "code_interpreter_results", + "search_results", + "rerank_results", + "executed_tools", + "custom_outputs", + "tool_search_items", "token_ids", ) if key in fields diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py index 5ee96ea9133..10b6f07bc4d 100644 --- a/litellm/litellm_core_utils/tool_fields.py +++ b/litellm/litellm_core_utils/tool_fields.py @@ -10,7 +10,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params ROUTES_WITH_TOOLS: Final = frozenset( { @@ -58,6 +58,19 @@ _TOOL_ENTRY: Final = TypeAdapter(_ToolEntry) _WEB_SEARCH_OPTIONS: Final = TypeAdapter(dict[str, object] | None, config=ConfigDict(strict=True)) +def declared_function_name(tool: object) -> str | None: + fields: Final = as_str_mapping(tool) + if fields is None: + return None + function: Final = as_str_mapping(fields.get("function")) + name: Final = (fields if function is None else function).get("name") + return name if isinstance(name, str) and name else None + + +def qualified_member_name(namespace: str, name: str) -> str: + return f"{namespace}__{name}" + + @dataclass(frozen=True, slots=True) class DropPolicy: drop_params: bool diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0db6a41b503..eae8ac04e8e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -102,6 +102,14 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" +def is_https_url_with_host(url: str) -> bool: + try: + parsed_url: Final = urlsplit(url) + return parsed_url.scheme == "https" and bool(parsed_url.hostname) + except ValueError: + return False + + def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7e3fda06b1a..6436cd832c7 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -781,12 +781,17 @@ class AnthropicMessagesHandler(BaseTranslation): for event in AnthropicMessagesHandler._parse_sse_data_line(line) ) - def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: + def _translate_to_openai( + self, + data: dict, # mutable-ok: the request body the adapter reads + *, + keep_cache_control: bool = False, + ) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" ( chat_completion_compatible_request, _tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + ) = LiteLLMAnthropicMessagesAdapter(keep_cache_control=keep_cache_control).translate_anthropic_to_openai( anthropic_message_request=cast(AnthropicMessagesRequest, scannable_tool_result_request(data)), preserve_midturn_system=True, ) @@ -813,9 +818,13 @@ class AnthropicMessagesHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + *, + request_data: dict[str, object] | None = None, # mutable-ok: the request dict guardrails receive + keep_cache_control: bool = False, ) -> Mapping[str, object]: """ - Process input messages by applying guardrails to text content. + Process input messages by applying guardrails to text content. ``request_data`` is what the guardrail sees as + the request when ``data`` is only its body, e.g. a provider pass-through route that nests the body """ messages: Final = data.get("messages") if messages is None: @@ -827,8 +836,17 @@ class AnthropicMessagesHandler(BaseTranslation): # The top-level prompt is translated on its own below so it can be hoisted in front of # any mid-turn system entries and scanned first, aligned with that structured position. - translation_source: Final = {key: value for key, value in data.items() if key != "system"} - chat_completion_compatible_request: Final = self._translate_to_openai(translation_source) + body_model: Final = data.get("model") + request_model: Final = (request_data or {}).get("model") + model: Final[str] = ( + body_model + if isinstance(body_model, str) and body_model + else (request_model if isinstance(request_model, str) else "") + ) + translation_source: Final = {**{key: value for key, value in data.items() if key != "system"}, "model": model} + chat_completion_compatible_request: Final = self._translate_to_openai( + translation_source, keep_cache_control=keep_cache_control + ) full_structured_messages: Final = cast( list[AllMessageValues], @@ -837,7 +855,9 @@ class AnthropicMessagesHandler(BaseTranslation): has_midturn_system_message: Final = any( str(message.get("role") or "").lower() == "system" for message in full_structured_messages ) - hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data) + hoisted_system_message: Final = ( + None if skip_system else self._hoisted_top_level_system_message(data, model, keep_cache_control) + ) if hoisted_system_message is not None: full_structured_messages.insert(0, hoisted_system_message) # skip_system already excluded the trusted top-level prompt (it is simply not hoisted); @@ -899,7 +919,6 @@ class AnthropicMessagesHandler(BaseTranslation): if structured_messages: inputs["structured_messages"] = structured_messages # Include model information if available - model: Final = data.get("model") if model: inputs["model"] = model original_tool_names: Final = _original_tool_names( @@ -907,7 +926,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data=data, + request_data=data if request_data is None else request_data, input_type="request", logging_obj=litellm_logging_obj, ) @@ -955,6 +974,7 @@ class AnthropicMessagesHandler(BaseTranslation): hoisted_system_message=hoisted_system_message, preserve_system_messages=has_midturn_system_message, original_tool_names=original_tool_names.used, + model=model, ) else: if guardrailed_texts and len(guardrailed_texts) != len(scanned): @@ -977,17 +997,20 @@ class AnthropicMessagesHandler(BaseTranslation): return data - def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None: + def _hoisted_top_level_system_message( + self, data: Mapping[str, object], model: str, keep_cache_control: bool + ) -> AllMessageValues | None: """Return the system message produced by translating the top-level prompt.""" system: Final = data.get("system") if not system: return None probe: Final = self._translate_to_openai( { - "model": data.get("model") or "", + "model": model, "messages": [], "system": system, - } + }, + keep_cache_control=keep_cache_control, ) hoisted: Final = probe.get("messages") or [] return hoisted[0] if hoisted else None @@ -1121,6 +1144,7 @@ class AnthropicMessagesHandler(BaseTranslation): preserve_system_messages: bool = False, *, original_tool_names: Mapping[str, str], + model: str, ) -> None: """Write a guardrail's structured-message rewrite back without losing corrections.""" from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -1129,7 +1153,6 @@ class AnthropicMessagesHandler(BaseTranslation): ) _is_system: Final = AnthropicMessagesHandler._is_system - model: Final = str(data.get("model") or "") converted: Final[list] = [] # mutable-ok: API message payload def _convert_run(run: list) -> None: # mutable-ok: API message payload diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 4b1b5a8a11a..afcd94c5b2c 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -380,8 +380,8 @@ _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.va class LiteLLMAnthropicMessagesAdapter: - def __init__(self): - pass + def __init__(self, *, keep_cache_control: bool = False) -> None: + self.keep_cache_control: Final = keep_cache_control ### FOR [BETA] `/v1/messages` endpoint support @@ -437,7 +437,7 @@ class LiteLLMAnthropicMessagesAdapter: cache_control: Final = ( source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) ) - if cache_control and model and self.target_consumes_cache_control(model): + if cache_control and (self.keep_cache_control or (model and self.target_consumes_cache_control(model))): # TypedDict objects support dict operations at runtime # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) if isinstance(target, dict): diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6bb6d584385..16c8bf6f9f1 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,15 +1,16 @@ import asyncio import hashlib +import itertools import json import os import re -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from contextlib import closing from types import MappingProxyType -from typing import Final, Literal, NamedTuple, cast +from typing import Final, Literal, NamedTuple, NoReturn, cast import httpx -from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from openai import APIError, AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict @@ -20,7 +21,7 @@ from litellm.constants import DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler, MaskedHTTPStatusError from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -52,6 +53,13 @@ class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object): _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams) +_AZURE_AD_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter[Callable[[], str | Awaitable[str]] | None]( + Callable[[], str | Awaitable[str]] | None +) + + +def _callable_azure_ad_token_provider(provider: object) -> Callable[[], str | Awaitable[str]] | None: + return _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python(provider) if callable(provider) else None class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object): @@ -81,6 +89,160 @@ class AzureOpenAIError(BaseLLMException): ) +class EntraRejectionError(AzureOpenAIError): + pass + + +_GUID: Final = re.compile(r"[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}") +_AADSTS_REASON: Final = re.compile(r"\bAADSTS(\d+)\b.*?(?:\.(?=[\s'\"\\]|$)|$)", re.MULTILINE) +_ENTRA_CREDENTIAL_REJECTION_CODES: Final = frozenset( + { + 50012, + 50013, + 50027, + 50034, + 50053, + 50055, + 50057, + 50059, + 50076, + 50079, + 50126, + 50128, + 53003, + 65001, + 70001, + 70002, + 70011, + 70021, + 70025, + 70052, + 90002, + 90094, + 500011, + 500014, + 700016, + 700023, + 700024, + 700025, + 700027, + 700211, + 700213, + 700223, + 700238, + 7000112, + 7000215, + 7000218, + 7000222, + } +) + + +def _entra_rejection_reason(message: str) -> str | None: + reason: Final = _AADSTS_REASON.search(message) + if reason is None or int(reason.group(1)) not in _ENTRA_CREDENTIAL_REJECTION_CODES: + return None + return _GUID.sub("", reason.group(0)) + + +def _entra_rejection(reason: str) -> EntraRejectionError: + return EntraRejectionError(status_code=401, message=f"Failed to get Azure AD token: {reason}") + + +def entra_rejection_as_auth_error(error: Exception) -> EntraRejectionError | None: + try: + from azure.core.exceptions import ClientAuthenticationError + except ImportError: + return None + if not isinstance(error, ClientAuthenticationError): + return None + status_code: Final = error.status_code + if status_code is not None and (status_code == 429 or status_code >= 500): + return None + reason: Final = _entra_rejection_reason(str(error)) + return None if reason is None else _entra_rejection(reason) + + +def _raise_entra_rejection_as_auth_error(error: Exception) -> NoReturn: + auth_error: Final = entra_rejection_as_auth_error(error) + if auth_error is None: + raise error + raise auth_error from error + + +def entra_token(provider: Callable[[], str]) -> str: + try: + return provider() + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + + +def entra_rejection_types() -> tuple[type[Exception], ...]: + try: + from azure.core.exceptions import ClientAuthenticationError + except ImportError: + return () + return (ClientAuthenticationError,) + + +_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16 + + +def _exception_chain(error: BaseException) -> Iterator[BaseException]: + link: BaseException | None = error # rebind-ok: advances one link of the cause chain per iteration + for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH): + if link is None: + return + yield link + link = link.__cause__ or link.__context__ + + +_REQUEST_FAILURE_TYPES: Final = (APIError, httpx.HTTPError) + + +def _ends_the_failed_attempt(link: BaseException, failed_attempt: BaseException) -> bool: + return not isinstance(link, EntraRejectionError) and ( + isinstance(link, _REQUEST_FAILURE_TYPES) or type(link) is type(failed_attempt) + ) + + +def _links_within_the_failed_attempt(error: BaseException) -> Iterator[BaseException]: + yield error + yield from itertools.takewhile( + lambda link: not _ends_the_failed_attempt(link, error), itertools.islice(_exception_chain(error), 1, None) + ) + + +def entra_rejection_in(error: BaseException) -> EntraRejectionError | None: + return next( + (link for link in _links_within_the_failed_attempt(error) if isinstance(link, EntraRejectionError)), None + ) + + +def caused_by_entra_rejection(error: BaseException) -> bool: + return entra_rejection_in(error) is not None + + +async def _awaited_token(token: Awaitable[str]) -> str: + try: + return await token + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + + +def _entra_rejections_as_auth_errors( + provider: Callable[[], str | Awaitable[str]], +) -> Callable[[], str | Awaitable[str]]: + def call_provider() -> str | Awaitable[str]: + try: + token: Final = provider() + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + return token if isinstance(token, str) else _awaited_token(token) + + return call_provider + + def process_azure_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "x-ratelimit-limit-requests" in headers: @@ -91,7 +253,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"] if "x-ratelimit-remaining-tokens" in headers: openai_headers["x-ratelimit-remaining-tokens"] = headers["x-ratelimit-remaining-tokens"] - llm_response_headers: Final = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()} + llm_response_headers: Final = {f"llm_provider-{k}": v for k, v in headers.items()} return {**llm_response_headers, **openai_headers} @@ -257,22 +419,34 @@ def _cached_oidc_token(cache_key: str) -> str | None: return cached if isinstance(cached, str) else None +def _oidc_exchange_rejection_reason(error: MaskedHTTPStatusError) -> str | None: + if error.response.status_code not in (400, 401): + return None + return _entra_rejection_reason(error.text or error.response.text) + + def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str: cached_access_token: Final = _cached_oidc_token(cache_key) if cached_access_token is not None: return cached_access_token with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: - req_token: Final = client.post( - token_url, - data={ - "client_id": client_id, - "grant_type": "client_credentials", - "scope": scope, - "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", - "client_assertion": oidc_token, - }, - ) + try: + req_token: Final = client.post( + token_url, + data={ + "client_id": client_id, + "grant_type": "client_credentials", + "scope": scope, + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": oidc_token, + }, + ) + except MaskedHTTPStatusError as error: + reason: Final = _oidc_exchange_rejection_reason(error) + if reason is None: + raise + raise _entra_rejection(reason) from error if req_token.status_code != 200: raise AzureOpenAIError( @@ -483,6 +657,9 @@ def get_azure_ad_token( raise except Exception as e: verbose_logger.error("Error calling Azure AD token provider: %s", e) + auth_error: Final = entra_rejection_as_auth_error(e) + if auth_error is not None: + raise auth_error from e raise RuntimeError(f"Failed to get Azure AD token: {e}") from e return azure_ad_token @@ -696,7 +873,7 @@ class BaseAzureLLM(BaseOpenAILLM): api_version: str | None, is_async: bool, ) -> dict: - azure_ad_token_provider = litellm_params.get("azure_ad_token_provider") + azure_ad_token_provider = _callable_azure_ad_token_provider(litellm_params.get("azure_ad_token_provider")) # If we have api_key, then we have higher priority azure_ad_token = litellm_params.get("azure_ad_token") @@ -775,7 +952,7 @@ class BaseAzureLLM(BaseOpenAILLM): azure_client_params["timeout"] = timeout if azure_ad_token_provider is not None: - azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider + azure_client_params["azure_ad_token_provider"] = _entra_rejections_as_auth_errors(azure_ad_token_provider) # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client diff --git a/litellm/llms/azure/decisions/transformation.py b/litellm/llms/azure/decisions/transformation.py index 3fb40e39702..bdbd784185f 100644 --- a/litellm/llms/azure/decisions/transformation.py +++ b/litellm/llms/azure/decisions/transformation.py @@ -8,6 +8,7 @@ import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, cached_azure_ad_token_refresh_provider, + entra_token, ) from litellm.llms.openai.decisions.transformation import OpenAIDecisionsConfig from litellm.secret_managers.main import get_secret_str @@ -61,7 +62,7 @@ class AzureDecisionsConfig(OpenAIDecisionsConfig): scope: Final = ( params.azure_scope or get_secret_str("AZURE_SCOPE") or "https://cognitiveservices.azure.com/.default" ) - azure_headers["Authorization"] = f"Bearer {cached_azure_ad_token_refresh_provider(scope)()}" + azure_headers["Authorization"] = f"Bearer {entra_token(cached_azure_ad_token_refresh_provider(scope))}" return azure_headers async def avalidate_environment( diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 6c954d00e64..f95dc1e4388 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -3,11 +3,10 @@ from collections.abc import Callable, Mapping from copy import deepcopy from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import urlsplit import httpx -from openai.types.responses import ResponseReasoningItem from pydantic import TypeAdapter from litellm._logging import verbose_logger @@ -16,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import prompt_cach from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.llms.base_llm.responses.transformation import SUMMARY_ONLY_REASONING_FIELD_VALUES, ReasoningFieldValues from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import * @@ -72,6 +72,21 @@ def _with_responses_api_version(api_base: str | None, litellm_params: Mapping[st return MappingProxyType({**litellm_params, "api_version": _resolve_responses_api_version(api_base, litellm_params)}) +_AZURE_REJECTED_INPUT_FIELDS: Final = MappingProxyType( + {"message": frozenset({"status"}), "reasoning": frozenset({"status"})} +) +_AZURE_REJECTED_INPUT_FIELDS_UNDER_DROP_PARAMS: Final = MappingProxyType( + {**_AZURE_REJECTED_INPUT_FIELDS, "reasoning": frozenset({"status", "content"})} +) + + +def _without_fields_azure_rejects(item: object, rejected_fields: Mapping[str, frozenset[str]]) -> object: + if not is_str_keyed_mapping(item): + return item + rejected: Final = rejected_fields.get(str(item.get("type")), frozenset()) + return {key: value for key, value in item.items() if key not in rejected} if rejected else item + + class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Parameters not supported by Azure Responses API AZURE_UNSUPPORTED_PARAMS = ["context_management"] @@ -122,68 +137,17 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def get_stripped_model_name(self, model: str) -> str: return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "") - def _handle_reasoning_item(self, item: dict[str, Any]) -> Mapping[str, object]: - """ - Handle reasoning items to filter out the status field. - Issue: https://github.com/BerriAI/litellm/issues/13484 - - Azure OpenAI API does not accept 'status' field in reasoning input items. - """ - if item.get("type") == "reasoning": - try: - # Ensure required fields are present for ResponseReasoningItem - item_data: Final = dict(item) - if "summary" not in item_data: - item_data["summary"] = ( - item_data.get("reasoning_content", "")[:100] + "..." - if len(item_data.get("reasoning_content", "")) > 100 - else item_data.get("reasoning_content", "") - ) - - # Create ResponseReasoningItem object from the item data - reasoning_item: Final = ResponseReasoningItem.model_validate(item_data) - - # Convert back to dict with exclude_none=True to exclude None fields - dict_reasoning_item: Final = reasoning_item.model_dump(exclude_none=True) - dict_reasoning_item.pop("status", None) - - return dict_reasoning_item - except Exception as e: - verbose_logger.debug("Failed to create ResponseReasoningItem, falling back to manual filtering: %s", e) - # Fallback: manually filter out known None fields - filtered_item: Final = { - k: v - for k, v in item.items() - if v is not None or k not in {"status", "content", "encrypted_content"} - } - return filtered_item - return item - def _validate_input_param( - self, input: str | ResponseInputParam, *, drops_set_created_by: bool = False + self, input: str | ResponseInputParam, *, drop_params: bool = False ) -> str | ResponseInputParam: - """ - Override parent method to also filter out 'status' field from message items. - Azure OpenAI API does not accept 'status' field in input messages. - """ - from typing import cast - - # First call parent's validation - validated_input: Final = super()._validate_input_param(input, drops_set_created_by=drops_set_created_by) - - # Then filter out status from message items - if isinstance(validated_input, list): - filtered_input: Final[list[object]] = [] - for item in validated_input: - if isinstance(item, dict) and item.get("type") == "message": - # Filter out status field from message items - filtered_item = {k: v for k, v in item.items() if k != "status"} - filtered_input.append(filtered_item) - else: - filtered_input.append(item) - return cast(ResponseInputParam, filtered_input) - - return validated_input + validated_input: Final = super()._validate_input_param(input, drop_params=drop_params) + if not isinstance(validated_input, list): + return validated_input + rejected_fields: Final = ( + _AZURE_REJECTED_INPUT_FIELDS_UNDER_DROP_PARAMS if drop_params else _AZURE_REJECTED_INPUT_FIELDS + ) + sanitized_items: Final = [_without_fields_azure_rejects(item, rejected_fields) for item in validated_input] + return cast("ResponseInputParam", sanitized_items) # cast-ok: items keep their shape minus rejected fields def transform_responses_api_request( self, diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index e255a667606..6d0241d4d3b 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -36,7 +36,7 @@ import httpx from pydantic import ConfigDict, ValidationError from litellm.litellm_core_utils.url_utils import ensure_api_route -from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider +from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider, entra_rejection_as_auth_error from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.responses import ( ResponsesSearchAnnotation, @@ -270,6 +270,9 @@ class BingGroundingSearchConfig(BaseSearchConfig): try: return minter() except Exception as e: + rejection: Final = entra_rejection_as_auth_error(e) + if rejection is not None: + raise rejection from e raise ValueError( f"Grounding with Bing Search: no credential available. Pass api_key, set {TOKEN_ENV} " f"to an Entra bearer token, or configure azure-identity (AZURE_CLIENT_ID / " diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 97ee44bd3bb..3609871fdf8 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -16,6 +16,7 @@ from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.azure.common_utils import EntraRejectionError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, ) @@ -135,6 +136,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): except AnthropicError: # Re-raise Anthropic exceptions as-is raise + except EntraRejectionError as rejection: + raise AnthropicError(status_code=rejection.status_code, message=rejection.message) from rejection except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code verbose_logger.error("HTTP error in CountTokens handler: %s", e) diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 147bf69e91e..1635b9b5679 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,5 +1,7 @@ import asyncio -from collections.abc import Mapping +import os +from collections.abc import Callable, Mapping +from functools import partial from types import MappingProxyType from typing import Final, Literal from urllib.parse import urlparse @@ -52,6 +54,28 @@ AZURE_ENTRA_CREDENTIAL_PARAM_KEYS: Final = frozenset({"azure_ad_token", "client_ AZURE_ENTRA_LITELLM_PARAM_KEYS: Final = AZURE_ENTRA_CREDENTIAL_PARAM_KEYS | frozenset( {"tenant_id", "client_id", "azure_username", "azure_scope"} ) +AZURE_DEFAULT_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +_AZURE_ENTRA_ENVIRONMENT_VARIABLES: Final = MappingProxyType( + { + "tenant_id": "AZURE_TENANT_ID", + "client_id": "AZURE_CLIENT_ID", + "client_secret": "AZURE_CLIENT_SECRET", + "azure_username": "AZURE_USERNAME", + "azure_password": "AZURE_PASSWORD", + "azure_ad_token": "AZURE_AD_TOKEN", + } +) + + +def azure_entra_environment_defaults() -> Mapping[str, object]: + return MappingProxyType( + { + **{key: os.getenv(variable) for key, variable in _AZURE_ENTRA_ENVIRONMENT_VARIABLES.items()}, + "azure_scope": os.getenv("AZURE_SCOPE", AZURE_DEFAULT_SCOPE), + } + ) + + AZURE_ENTRA_CREDENTIAL_HELP: Final = ( "Set `tenant_id` + `client_id` + `client_secret`, `azure_ad_token` (an `oidc/` token also needs " "`tenant_id` + `client_id`), or `client_id` + `azure_username` + `azure_password` in the agent's `litellm_params`" @@ -70,6 +94,17 @@ def _resolve_config_secret(value: object) -> str | None: return get_secret_str(value) if value.startswith("os.environ/") else value +def _agent_entra_token(provider: Callable[[], str], model: str) -> str: + from litellm.llms.azure.common_utils import EntraRejectionError, entra_token + + try: + return entra_token(provider) + except EntraRejectionError as rejection: + raise litellm.AuthenticationError( + message=rejection.message, llm_provider="azure_ai", model=model + ) from rejection + + def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str: """Mints the Entra bearer from the agent's own litellm_params, never from process-wide AZURE_* env vars.""" from litellm.llms.azure.common_utils import ( @@ -88,18 +123,32 @@ def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str: azure_username: Final = resolved["azure_username"] azure_password: Final = resolved["azure_password"] azure_ad_token: Final = resolved["azure_ad_token"] + model: Final = str(litellm_params.get("model") or "") if tenant_id and client_id and client_secret: - return get_azure_ad_token_from_entra_id( - tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope - )() + return _agent_entra_token( + get_azure_ad_token_from_entra_id( + tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope + ), + model, + ) if client_id and azure_username and azure_password: - return get_azure_ad_token_from_username_password( - client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope - )() + return _agent_entra_token( + get_azure_ad_token_from_username_password( + client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope + ), + model, + ) federated: Final = azure_ad_token is not None and azure_ad_token.startswith("oidc/") if azure_ad_token and federated and tenant_id and client_id: - return get_azure_ad_token_from_oidc( - azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope + return _agent_entra_token( + partial( + get_azure_ad_token_from_oidc, + azure_ad_token=azure_ad_token, + azure_client_id=client_id, + azure_tenant_id=tenant_id, + scope=scope, + ), + model, ) if azure_ad_token and not federated: return azure_ad_token diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 3a65a8928bf..3514a5e6e5c 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -14,6 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -60,6 +61,8 @@ if TYPE_CHECKING: _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream" +_INVOKE_ACTIONS: Final = frozenset({"invoke", _INVOKE_STREAM_ACTION}) +_BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN: Final = frozenset({"document", "search_result"}) _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream" _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream" @@ -73,6 +76,83 @@ def _is_invoke_stream_endpoint(endpoint: str) -> bool: return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION +def _anthropic_invoke_body( + data: Mapping[str, object], +) -> dict[str, object] | None: # mutable-ok: the body guardrails write back into + """The body of an InvokeModel call in Anthropic Messages format, which Bedrock requires to carry anthropic_version""" + endpoint: Final = data.get("endpoint") + body: Final = data.get("data") + if not isinstance(endpoint, str) or endpoint.rstrip("/").split("/")[-1] not in _INVOKE_ACTIONS: + return None + if not isinstance(body, dict): + return None + invoke_body: Final = _BODY_BY_REFERENCE.validate_python(body) + if not isinstance(invoke_body.get("anthropic_version"), str) or not isinstance(invoke_body.get("messages"), list): + return None + messages: Final = _ARRAY_BY_REFERENCE.validate_python(invoke_body["messages"]) + if not all(_is_well_formed_message(message) for message in messages): + return None + blocks: Final = _anthropic_message_blocks(messages) + if any(block.get("type") in _BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN for block in blocks): + return None + return invoke_body + + +def _body_level_request_data( + passthrough_request: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: the request dict guardrails receive + """What `/v1/messages` hands a guardrail as `request_data`: the Anthropic body at the top level, with the proxy's own keys (model, metadata, headers) on top""" + body: Final = as_str_mapping(passthrough_request.get("data")) or {} + proxy_keys: Final = {key: value for key, value in passthrough_request.items() if key != "data"} + return {**body, **proxy_keys} + + +def _is_well_formed_message(message: object) -> bool: + """Whether the Messages translation can read `message`, so a malformed body still reaches Bedrock to be rejected there""" + fields: Final = as_str_mapping(message) + if fields is None or not isinstance(fields.get("role"), str): + return False + content: Final = fields.get("content") + if isinstance(content, str): + return True + return isinstance(content, list) and all( + as_str_mapping(block) is not None for block in _ARRAY_BY_REFERENCE.validate_python(content) + ) + + +def _anthropic_message_blocks(messages: Sequence[object]) -> tuple[Mapping[str, object], ...]: + fields: Final = (as_str_mapping(message) for message in messages) + return tuple(chain.from_iterable(_anthropic_content_blocks(field.get("content")) for field in fields if field)) + + +def _anthropic_content_blocks(content: object) -> tuple[Mapping[str, object], ...]: + blocks: Final[list[Mapping[str, object]]] = [] # mutable-ok: collected while walking the content + pending: Final[list[object]] = [content] # mutable-ok: content lists still to walk, nested tool results included + while pending: + current = pending.pop() + if not isinstance(current, list): + continue + for item in _ARRAY_BY_REFERENCE.validate_python(current): + block = as_str_mapping(item) + if block is None: + continue + blocks.append(block) + if block.get("type") == "tool_result": + pending.append(block.get("content")) + return tuple(blocks) + + +def _anthropic_invoke_reply( + response: object, + request_data: object, +) -> dict[str, object] | None: # mutable-ok: the reply guardrails write back into + if not isinstance(response, dict) or not isinstance(request_data, dict): + return None + if _anthropic_invoke_body(_OBJECT_BY_REFERENCE.validate_python(request_data)) is None: + return None + return _BODY_BY_REFERENCE.validate_python(response) + + def _generic_passthrough_handler() -> "PassThroughEndpointHandler": """ Fallback for non-Converse Bedrock routes (e.g. invoke). The generic @@ -1285,6 +1365,17 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): body: Final = data.get("data") if not _is_converse_endpoint(endpoint): + passthrough_request: Final = _BODY_BY_REFERENCE.validate_python(data) + anthropic_body: Final = _anthropic_invoke_body(passthrough_request) + if anthropic_body is not None: + await AnthropicMessagesHandler().process_input_messages( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + data=anthropic_body, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + request_data=_body_level_request_data(passthrough_request), + keep_cache_control=True, + ) + return passthrough_request return await _generic_passthrough_handler().process_input_messages( data=data, guardrail_to_apply=guardrail_to_apply, @@ -1410,6 +1501,15 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): and not _is_converse_endpoint(endpoint) and not _is_invoke_stream_endpoint(endpoint) ): + anthropic_reply: Final = _anthropic_invoke_reply(response, request_data) # pyright: ignore[reportUnknownArgumentType] # the base signature types request_data as a bare dict + if anthropic_reply is not None: + return await AnthropicMessagesHandler().process_output_response( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + response=anthropic_reply, # pyright: ignore[reportArgumentType] # the Anthropic reply dict, written back in place + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=_body_level_request_data(_OBJECT_BY_REFERENCE.validate_python(request_data or {})), + ) return await _generic_passthrough_handler().process_output_response( response=response, guardrail_to_apply=guardrail_to_apply, diff --git a/litellm/llms/empiriolabs/decisions/transformation.py b/litellm/llms/empiriolabs/decisions/transformation.py new file mode 100644 index 00000000000..ce3464310c4 --- /dev/null +++ b/litellm/llms/empiriolabs/decisions/transformation.py @@ -0,0 +1,10 @@ +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class EmpirioLabsDecisionsConfig(BaseDecisionsConfig): + path = "/v1/decisions" + api_key_env = ("EMPIRIOLABS_API_KEY",) + api_base_env = ("EMPIRIOLABS_API_BASE",) + + def get_default_api_base(self) -> str | None: + return "https://api.empiriolabs.ai" diff --git a/litellm/llms/github_copilot/github_host.py b/litellm/llms/github_copilot/github_host.py index 4fc921d490d..cfe122ebac4 100644 --- a/litellm/llms/github_copilot/github_host.py +++ b/litellm/llms/github_copilot/github_host.py @@ -21,8 +21,10 @@ class GithubHostEndpoints: device_code_url: str access_token_url: str api_key_url: str + copilot_user_url: str user_url: str client_id: str + exchanges_copilot_token: bool copilot_api_hosts: frozenset[str] copilot_api_host_suffix: str | None default_copilot_api_base: str @@ -49,14 +51,17 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: if normalized is None: raise ValueError("host must be github.com or a single-subdomain GHE.com hostname") client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) + exchanges_copilot_token: Final = client_id == DEFAULT_GITHUB_CLIENT_ID if normalized == GITHUB_COPILOT_DEFAULT_HOST: return GithubHostEndpoints( host=normalized, device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL), access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL), api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL), + copilot_user_url="https://api.github.com/copilot_internal/user", user_url="https://api.github.com/user", client_id=client_id, + exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({"githubcopilot.com"}), copilot_api_host_suffix=".githubcopilot.com", default_copilot_api_base="https://api.githubcopilot.com", @@ -70,8 +75,10 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: device_code_url=f"https://{normalized}/login/device/code", access_token_url=f"https://{normalized}/login/oauth/access_token", api_key_url=f"https://{api_host}/copilot_internal/v2/token", + copilot_user_url=f"https://{api_host}/copilot_internal/user", user_url=f"https://{api_host}/user", client_id=client_id, + exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({copilot_api_host}), copilot_api_host_suffix=None, default_copilot_api_base=f"https://{copilot_api_host}", diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index 75a10f2e671..e210414ce8f 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -46,6 +46,7 @@ _LLM_PROVIDER: Final = "github_copilot" _DEVICE_FLOW_SCOPE: Final = "read:user" _SESSION_CACHE_MAX_SIZE: Final = 1000 _SYNC_LOCK_COUNT: Final = 64 +_ACCOUNT_LOOKUP_TTL_SECONDS: Final = 1800 class _SessionCache(Protocol): @@ -125,12 +126,20 @@ class _CopilotEndpointsPayload(TypedDict): api: ReadOnly[NotRequired[object]] +@with_config(ConfigDict(extra="allow", strict=True)) +class _CopilotAccountPayload(TypedDict): + access_type_sku: ReadOnly[NotRequired[object]] + chat_enabled: ReadOnly[NotRequired[object]] + endpoints: ReadOnly[NotRequired[object]] + + @with_config(ConfigDict(extra="allow", strict=True)) class _GithubUserPayload(TypedDict): login: ReadOnly[str] _COPILOT_TOKEN: Final = TypeAdapter(_CopilotTokenPayload) +_COPILOT_ACCOUNT: Final = TypeAdapter(_CopilotAccountPayload) _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload) _GITHUB_USER: Final = TypeAdapter(_GithubUserPayload) @@ -273,6 +282,8 @@ def _session_from_response( llm_provider=_LLM_PROVIDER, model="", ) + if not endpoints.exchanges_copilot_token: + return _session_from_account(response, github_token, credential_name, cache_key, endpoints) try: payload: Final = _COPILOT_TOKEN.validate_python(response.json()) except Exception as e: @@ -282,13 +293,7 @@ def _session_from_response( model="", ) from e - payload_endpoints: Final[object] = payload.get("endpoints") - endpoints_api: Final[object] = ( - _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") - if isinstance(payload_endpoints, Mapping) - else None - ) - api_base: Final = validated_copilot_api_base(endpoints_api, endpoints) + api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base) expires_at: Final = _expires_at_seconds(payload.get("expires_at")) if expires_at > 0: @@ -296,6 +301,41 @@ def _session_from_response( return session +def _session_from_account( + response: httpx.Response, + github_token: str, + credential_name: str, + cache_key: str, + endpoints: GithubHostEndpoints, +) -> GithubCopilotUserSession: + try: + payload: Final = _COPILOT_ACCOUNT.validate_python(response.json()) + except Exception as e: + raise ServiceUnavailableError( + message="GitHub Copilot account lookup returned an invalid response", + llm_provider=_LLM_PROVIDER, + model="", + ) from e + if payload.get("chat_enabled") is False or payload.get("access_type_sku") == "no_access": + _SESSION_CACHE.delete_cache(cache_key) + raise _reject_connection(credential_name) + + api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) + session: Final = GithubCopilotUserSession(token=github_token, api_base=api_base) + _cache_session(cache_key, github_token, session, time.time() + _ACCOUNT_LOOKUP_TTL_SECONDS) + return session + + +def _endpoints_api(payload_endpoints: object) -> object: + if not isinstance(payload_endpoints, Mapping): + return None + return _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") + + +def _session_lookup_url(endpoints: GithubHostEndpoints) -> str: + return endpoints.api_key_url if endpoints.exchanges_copilot_token else endpoints.copilot_user_url + + def _to_int(raw: object) -> int: if isinstance(raw, bool): return 0 @@ -337,7 +377,7 @@ def _do_exchange( _SyncGetClient, litellm.module_level_client ) response: Final = client.get( - endpoints.api_key_url, + _session_lookup_url(endpoints), headers=github_api_headers(github_token), ) except Exception as e: @@ -367,7 +407,7 @@ async def _do_aexchange( client: Final = _new_async_github_client() try: response: Final = await client.get( - endpoints.api_key_url, + _session_lookup_url(endpoints), headers=github_api_headers(github_token), ) except Exception as e: diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index cafeece0a4e..d2ef33b89dd 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -13,24 +13,26 @@ The handler is intentionally a thin orchestration layer: 2. Build ``RecognitionConfig`` / ``StreamingRecognitionConfig`` protobufs from the structured dict produced by :class:`NvidiaRivaAudioTranscriptionConfig`. -3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS) - vs self-hosted (any host:port, optional TLS) modes. -4. Stream the audio through Riva's ``streaming_response_generator`` and - aggregate ``is_final`` results into a single transcript, or, when - ``riva_offline`` is set, send the whole buffer through the unary - ``offline_recognize`` RPC for NIMs that only ship offline engines. +3. Open a gRPC channel honoring NVCF (TLS by default) vs self-hosted (any + host:port, optional TLS) modes, and attach the ``authorization`` and NVCF + ``function-id`` metadata to each call exactly once. +4. Stream the audio through Riva's ``StreamingRecognize`` RPC and aggregate + ``is_final`` results into a single transcript, or, when ``riva_offline`` + is set, send the whole buffer through the unary ``Recognize`` RPC for NIMs + that only ship offline engines. 5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed on ``_hidden_params`` so cost calculation works. -``riva-client`` is imported lazily so ``litellm`` core remains usable -without the optional STT extras installed. +``grpcio`` and ``protobuf`` are imported lazily so ``litellm`` core remains +usable without the optional STT extras installed. """ +from __future__ import annotations + import asyncio -import inspect -from collections.abc import Callable, Iterable, Mapping, Sequence -from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol +from collections.abc import Iterable, Sequence +from contextlib import closing +from typing import TYPE_CHECKING, Final import httpx @@ -64,89 +66,19 @@ if TYPE_CHECKING: Logging as LiteLLMLoggingObj, ) -# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches -# NVIDIA's recommended chunk size for streaming ASR — small enough for -# responsive endpointing, large enough to keep per-RPC overhead low. -_DEFAULT_CHUNK_SAMPLES: Final = 1600 -_DEFAULT_CHUNK_BYTES: Final = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/sample - - _RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`." -class _RivaAuth(Protocol): - """Opaque ``riva.client.Auth`` handle.""" - - -class _RecognizeFuture(Protocol): - def result(self, timeout: float | None = None) -> object: ... - - def done(self) -> bool: ... - - def cancel(self) -> bool: ... - - -class _AsrService(Protocol): - @property - def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ... - - def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ... - - -class _EndpointingConfig(Protocol): - """Opaque ``EndpointingConfig`` protobuf message.""" - - -class _EndpointingConfigField(Protocol): - CopyFrom: Callable[[_EndpointingConfig], None] - - -class _RecognitionConfig(Protocol): - @property - def endpointing_config(self) -> _EndpointingConfigField: ... - - -class _StreamingRecognitionConfig(Protocol): - """Opaque ``StreamingRecognitionConfig`` protobuf message.""" - - -class _AudioEncoding(Protocol): - @property - def LINEAR_PCM(self) -> object: ... - - -class _RivaClientModule(Protocol): - """The ``riva.client`` entry points this handler calls.""" - - @property - def Auth(self) -> Callable[..., _RivaAuth]: ... - - @property - def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ... - - -class _RivaAsrModule(Protocol): - """The protobuf constructors this handler calls, from whichever module exposes them.""" - - @property - def AudioEncoding(self) -> _AudioEncoding: ... - - @property - def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ... - - @property - def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ... - - @property - def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ... - - -def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]: - return riva_module.Auth - - -def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: - return riva_asr_module.AudioEncoding +def _use_ssl(optional_params: dict) -> bool: + """ + NVCF endpoints are TLS-only, so ``nvcf_function_id`` defaults ``use_ssl`` + to True. Self-hosted Riva defaults to plaintext but honors an explicit + ``use_ssl``, since TLS-terminating ingresses in front of it are common. + """ + use_ssl_override: Final = optional_params.get("use_ssl") + if use_ssl_override is not None: + return bool(use_ssl_override) + return bool(optional_params.get("nvcf_function_id")) def _parse_riva_offline(value: object, drop_params: bool) -> bool: @@ -192,7 +124,7 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, atranscription: bool = False, @@ -237,13 +169,11 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, provider_config: NvidiaRivaAudioTranscriptionConfig | None = None, ) -> TranscriptionResponse: - # ``riva-client`` exposes a sync streaming generator, so we offload - # the blocking call to a worker thread to keep the event loop free. return await asyncio.to_thread( self._run_sync, model=model, @@ -267,7 +197,7 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, provider_config: NvidiaRivaAudioTranscriptionConfig, @@ -317,18 +247,12 @@ class NvidiaRivaAudioTranscription: ) or litellm.drop_params is True riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params) - riva_module, riva_asr_module = _import_riva() - auth_obj: Final = self._construct_auth( - riva_module=riva_module, - api_base=api_base, - api_key=api_key, - optional_params=optional_params, - ) + try: + from litellm.llms.nvidia_riva.audio_transcription import riva_grpc + except ImportError as e: + raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - recognition_config: Final = self._build_recognition_config_proto( - riva_asr_module=riva_asr_module, - recognition_config_dict=recognition_config_dict, - ) + recognition_config: Final = riva_grpc.recognition_config(recognition_config_dict) logging_obj.pre_call( input=None, api_key=api_key, @@ -344,26 +268,18 @@ class NvidiaRivaAudioTranscription: }, ) + metadata: Final = riva_grpc.call_metadata(api_key, optional_params.get("nvcf_function_id")) try: - asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj) - responses: Final = ( - self._offline_responses( - asr_service=asr_service, - pcm_bytes=resampled.pcm_bytes, - recognition_config=recognition_config, - timeout=None if timeout is None else _deadline_seconds(timeout, drop_params), + deadline: Final = None if timeout is None else _deadline_seconds(timeout, drop_params) + with closing(riva_grpc.open_channel(api_base, _use_ssl(optional_params))) as channel: + responses: Final = ( + riva_grpc.recognize(channel, recognition_config, resampled.pcm_bytes, metadata, deadline) + if riva_offline + else riva_grpc.streaming_recognize( + channel, recognition_config, resampled.pcm_bytes, metadata, deadline + ) ) - if riva_offline - else self._streaming_responses( - asr_service=asr_service, - riva_asr_module=riva_asr_module, - pcm_bytes=resampled.pcm_bytes, - recognition_config=recognition_config, - timeout=timeout, - drop_params=drop_params, - ) - ) - final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) + final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) except NvidiaRivaException: raise except Exception as e: @@ -400,142 +316,6 @@ class NvidiaRivaAudioTranscription: return final_response - def _construct_auth( - self, - riva_module: _RivaClientModule, - api_base: str, - api_key: str | None, - optional_params: dict, - ) -> _RivaAuth: - """ - Build a ``riva.client.Auth`` object. - - - When ``nvcf_function_id`` is provided we attach the NVCF - ``function-id`` and bearer ``authorization`` metadata, and default - ``use_ssl`` to True (NVCF endpoints are TLS-only). - - Otherwise (self-hosted) we default ``use_ssl`` to False but still - honor an explicit override — self-hosted Riva behind an ingress - with TLS termination is a real deployment topology. - """ - nvcf_function_id: Final = optional_params.get("nvcf_function_id") - use_ssl_override: Final = optional_params.get("use_ssl") - use_ssl: Final = bool(use_ssl_override) if use_ssl_override is not None else bool(nvcf_function_id) - - metadata: Final[list[tuple[str, str]]] = [] - if nvcf_function_id: - metadata.append(("function-id", str(nvcf_function_id))) - if api_key: - metadata.append(("authorization", f"Bearer {api_key}")) - - try: - return _auth_factory(riva_module)(uri=api_base, use_ssl=use_ssl, metadata_args=metadata) - except TypeError: - # Older riva-client signatures used positional-only args. - return _auth_factory(riva_module)(None, use_ssl, api_base, metadata) - - def _build_recognition_config_proto( - self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any] - ) -> _RecognitionConfig: - encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper() - encoding_enum: Final[object] = getattr( - _audio_encoding(riva_asr_module), - encoding_name, - _audio_encoding(riva_asr_module).LINEAR_PCM, - ) - - config: Final[_RecognitionConfig] = riva_asr_module.RecognitionConfig( - encoding=encoding_enum, - sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]), - language_code=recognition_config_dict["language_code"], - audio_channel_count=int(recognition_config_dict["audio_channel_count"]), - enable_automatic_punctuation=bool(recognition_config_dict.get("enable_automatic_punctuation", True)), - enable_word_time_offsets=bool(recognition_config_dict.get("enable_word_time_offsets", False)), - max_alternatives=int(recognition_config_dict.get("max_alternatives", 1)), - model=recognition_config_dict.get("model", "") or "", - verbatim_transcripts=bool(recognition_config_dict.get("verbatim_transcripts", False)), - profanity_filter=bool(recognition_config_dict.get("profanity_filter", False)), - ) - - endpointing: Final = recognition_config_dict.get("endpointing_config") - if isinstance(endpointing, dict) and endpointing: - try: - ep: Final[_EndpointingConfig] = riva_asr_module.EndpointingConfig(**endpointing) - config.endpointing_config.CopyFrom(ep) - except Exception: - # If the user supplied an unknown EndpointingConfig field - # (older Riva server), fall back to Riva's defaults rather - # than failing the whole request. - pass - - return config - - def _streaming_responses( - self, - asr_service: _AsrService, - riva_asr_module: _RivaAsrModule, - pcm_bytes: bytes, - recognition_config: _RecognitionConfig, - timeout: float | str | httpx.Timeout | None, - drop_params: bool, - ) -> Iterable[object]: - streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig( - config=recognition_config, interim_results=False - ) - return asr_service.streaming_response_generator( - audio_chunks=self._iter_audio_chunks(pcm_bytes), - streaming_config=streaming_config, - **self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params), - ) - - def _offline_responses( - self, - asr_service: _AsrService, - pcm_bytes: bytes, - recognition_config: _RecognitionConfig, - timeout: float | None, - ) -> Iterable[object]: - pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True) - try: - return (pending.result(timeout=timeout),) - except Exception as e: - if pending.done(): - raise - pending.cancel() - raise NvidiaRivaException( - status_code=408, - message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds", - ) from e - - @classmethod - def _timeout_kwargs( - cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool - ) -> Mapping[str, float]: - # Forward the deadline so the call cannot block forever if the - # server stalls. Older riva-client versions do not accept a - # ``timeout`` kwarg, so pass it only when supported. - if timeout is None or not cls._supports_timeout_kwarg(rpc): - return MappingProxyType({}) - return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)}) - - @staticmethod - def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool: - try: - sig: Final = inspect.signature(callable_obj) - except (TypeError, ValueError): - return False - params: Final = sig.parameters - if "timeout" in params: - return True - return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) - - @staticmethod - def _iter_audio_chunks(pcm_bytes: bytes): - for offset in range(0, len(pcm_bytes), _DEFAULT_CHUNK_BYTES): - chunk = pcm_bytes[offset : offset + _DEFAULT_CHUNK_BYTES] - if not chunk: - continue - yield chunk - @staticmethod def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> Sequence[RivaFinalResult]: """ @@ -570,27 +350,3 @@ class NvidiaRivaAudioTranscription: } ) return final_results - - -def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]: - """ - Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``. - - We try the SDK first (preferred) and fall back to importing the proto - module separately when the SDK packaging changes between versions. - """ - try: - import riva.client as riva_client - except ImportError as e: - raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - - riva_asr_module = riva_client - if not hasattr(riva_asr_module, "RecognitionConfig"): - try: - from riva.client.proto import riva_asr_pb2 - - riva_asr_module = riva_asr_pb2 - except ImportError as e: - raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - - return riva_client, riva_asr_module diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py b/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py new file mode 100644 index 00000000000..13af5820f07 --- /dev/null +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py @@ -0,0 +1,199 @@ +""" +The part of NVIDIA Riva's speech recognition gRPC schema that LiteLLM sends and reads. + +Names and numbers match ``riva/proto/riva_asr.proto`` and ``riva/proto/riva_audio.proto`` in +https://github.com/nvidia-riva/common. Fields LiteLLM never sets or reads are left out, and the +wire format skips them. The messages live in their own descriptor pool so they never collide +with nvidia-riva-client's copies when both are imported in one process. +""" + +from dataclasses import dataclass +from enum import Enum, IntEnum +from typing import Final + +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from google.protobuf.message import Message + +_PACKAGE: Final = "nvidia.riva.asr" +_SERVICE: Final = f"/{_PACKAGE}.RivaSpeechRecognition" +RECOGNIZE_METHOD: Final = f"{_SERVICE}/Recognize" +STREAMING_RECOGNIZE_METHOD: Final = f"{_SERVICE}/StreamingRecognize" + + +class _Kind(Enum): + FLOAT = "TYPE_FLOAT" + INT32 = "TYPE_INT32" + BOOL = "TYPE_BOOL" + STRING = "TYPE_STRING" + MESSAGE = "TYPE_MESSAGE" + BYTES = "TYPE_BYTES" + ENUM = "TYPE_ENUM" + + +class AudioEncoding(IntEnum): + ENCODING_UNSPECIFIED = 0 + LINEAR_PCM = 1 + + +LINEAR_PCM: Final = AudioEncoding.LINEAR_PCM.value + + +@dataclass(frozen=True, slots=True) +class _Field: + name: str + number: int + kind: _Kind + type_name: str = "" + repeated: bool = False + optional: bool = False + oneof: str = "" + + +_MESSAGES: Final[tuple[tuple[str, tuple[_Field, ...]], ...]] = ( + ( + "RecognitionConfig", + ( + _Field("encoding", 1, _Kind.ENUM, "AudioEncoding"), + _Field("sample_rate_hertz", 2, _Kind.INT32), + _Field("language_code", 3, _Kind.STRING), + _Field("max_alternatives", 4, _Kind.INT32), + _Field("profanity_filter", 5, _Kind.BOOL), + _Field("audio_channel_count", 7, _Kind.INT32), + _Field("enable_word_time_offsets", 8, _Kind.BOOL), + _Field("enable_automatic_punctuation", 11, _Kind.BOOL), + _Field("model", 13, _Kind.STRING), + _Field("verbatim_transcripts", 14, _Kind.BOOL), + _Field("endpointing_config", 25, _Kind.MESSAGE, "EndpointingConfig", optional=True), + ), + ), + ( + "EndpointingConfig", + ( + _Field("start_history", 1, _Kind.INT32, optional=True), + _Field("start_threshold", 2, _Kind.FLOAT, optional=True), + _Field("stop_history", 3, _Kind.INT32, optional=True), + _Field("stop_threshold", 4, _Kind.FLOAT, optional=True), + _Field("stop_history_eou", 5, _Kind.INT32, optional=True), + _Field("stop_threshold_eou", 6, _Kind.FLOAT, optional=True), + ), + ), + ( + "StreamingRecognitionConfig", + ( + _Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"), + _Field("interim_results", 2, _Kind.BOOL), + ), + ), + ( + "StreamingRecognizeRequest", + ( + _Field("streaming_config", 1, _Kind.MESSAGE, "StreamingRecognitionConfig", oneof="streaming_request"), + _Field("audio_content", 2, _Kind.BYTES, oneof="streaming_request"), + ), + ), + ( + "WordInfo", + ( + _Field("start_time", 1, _Kind.INT32), + _Field("end_time", 2, _Kind.INT32), + _Field("word", 3, _Kind.STRING), + ), + ), + ( + "SpeechRecognitionAlternative", + ( + _Field("transcript", 1, _Kind.STRING), + _Field("words", 3, _Kind.MESSAGE, "WordInfo", repeated=True), + ), + ), + ( + "StreamingRecognitionResult", + ( + _Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True), + _Field("is_final", 2, _Kind.BOOL), + ), + ), + ( + "StreamingRecognizeResponse", + (_Field("results", 1, _Kind.MESSAGE, "StreamingRecognitionResult", repeated=True),), + ), + ( + "RecognizeRequest", + ( + _Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"), + _Field("audio", 2, _Kind.BYTES), + ), + ), + ( + "SpeechRecognitionResult", + (_Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True),), + ), + ( + "RecognizeResponse", + (_Field("results", 1, _Kind.MESSAGE, "SpeechRecognitionResult", repeated=True),), + ), +) + + +def _oneof_names(fields: tuple[_Field, ...]) -> tuple[str, ...]: + declared: Final = tuple(dict.fromkeys(field.oneof for field in fields if field.oneof)) + synthetic: Final = tuple(f"_{field.name}" for field in fields if field.optional) + return declared + synthetic + + +def _field_proto(field: _Field, oneofs: tuple[str, ...]) -> descriptor_pb2.FieldDescriptorProto: + oneof: Final = f"_{field.name}" if field.optional else field.oneof + return descriptor_pb2.FieldDescriptorProto( + name=field.name, + number=field.number, + type=descriptor_pb2.FieldDescriptorProto.Type.Value(field.kind.value), + label=( + descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED + if field.repeated + else descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL + ), + type_name=f".{_PACKAGE}.{field.type_name}" if field.type_name else None, + oneof_index=oneofs.index(oneof) if oneof else None, + proto3_optional=field.optional or None, + ) + + +def _message_proto(name: str, fields: tuple[_Field, ...]) -> descriptor_pb2.DescriptorProto: + oneofs: Final = _oneof_names(fields) + return descriptor_pb2.DescriptorProto( + name=name, + field=[_field_proto(field, oneofs) for field in fields], + oneof_decl=[descriptor_pb2.OneofDescriptorProto(name=oneof) for oneof in oneofs], + ) + + +_FILE: Final = descriptor_pb2.FileDescriptorProto( + name="litellm/nvidia_riva/riva_asr.proto", + package=_PACKAGE, + syntax="proto3", + enum_type=[ + descriptor_pb2.EnumDescriptorProto( + name="AudioEncoding", + value=[ + descriptor_pb2.EnumValueDescriptorProto(name=encoding.name, number=encoding.value) + for encoding in AudioEncoding + ], + ) + ], + message_type=[_message_proto(name, fields) for name, fields in _MESSAGES], +) +_POOL: Final = descriptor_pool.DescriptorPool() +_POOL.AddSerializedFile(_FILE.SerializeToString()) + + +def _message_class(name: str) -> type[Message]: + return message_factory.GetMessageClass(_POOL.FindMessageTypeByName(f"{_PACKAGE}.{name}")) + + +RecognitionConfig: Final = _message_class("RecognitionConfig") +EndpointingConfig: Final = _message_class("EndpointingConfig") +StreamingRecognitionConfig: Final = _message_class("StreamingRecognitionConfig") +StreamingRecognizeRequest: Final = _message_class("StreamingRecognizeRequest") +StreamingRecognizeResponse: Final = _message_class("StreamingRecognizeResponse") +RecognizeRequest: Final = _message_class("RecognizeRequest") +RecognizeResponse: Final = _message_class("RecognizeResponse") diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py new file mode 100644 index 00000000000..8367853673f --- /dev/null +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py @@ -0,0 +1,113 @@ +""" +gRPC calls to NVIDIA Riva's speech recognition service. + +The ``authorization`` and NVCF ``function-id`` metadata ride on each call, never on the channel, so every +request carries them exactly once over both plaintext and TLS. +""" + +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, Final, TypeAlias + +import grpc +from google.protobuf.message import Message +from pydantic import TypeAdapter + +from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema as schema +from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException + +# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches +# NVIDIA's recommended chunk size for streaming ASR, small enough for +# responsive endpointing, large enough to keep per-RPC overhead low. +_CHUNK_SAMPLES: Final = 1600 +_CHUNK_BYTES: Final = _CHUNK_SAMPLES * 2 + +CallMetadata: TypeAlias = tuple[tuple[str, str], ...] + +_ENDPOINTING_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def call_metadata(api_key: str | None, nvcf_function_id: object) -> CallMetadata: + function_id: Final = (("function-id", str(nvcf_function_id)),) if nvcf_function_id else () + authorization: Final = (("authorization", f"Bearer {api_key}"),) if api_key else () + return function_id + authorization + + +def open_channel(api_base: str, use_ssl: bool) -> grpc.Channel: + if use_ssl: + return grpc.secure_channel(api_base, grpc.ssl_channel_credentials()) + return grpc.insecure_channel(api_base) + + +def recognition_config(config: Mapping[str, Any]) -> Message: + endpointing: Final = _endpointing_config(config.get("endpointing_config")) + return schema.RecognitionConfig( + encoding=schema.LINEAR_PCM, + sample_rate_hertz=int(config["sample_rate_hertz"]), + language_code=config["language_code"], + audio_channel_count=int(config["audio_channel_count"]), + enable_automatic_punctuation=bool(config.get("enable_automatic_punctuation", True)), + enable_word_time_offsets=bool(config.get("enable_word_time_offsets", False)), + max_alternatives=int(config.get("max_alternatives", 1)), + model=config.get("model", "") or "", + verbatim_transcripts=bool(config.get("verbatim_transcripts", False)), + profanity_filter=bool(config.get("profanity_filter", False)), + **({} if endpointing is None else {"endpointing_config": endpointing}), + ) + + +def _endpointing_config(endpointing: object) -> Message | None: + try: + fields: Final = _ENDPOINTING_FIELDS.validate_python(endpointing) + return schema.EndpointingConfig(**fields) if fields else None + except (ValueError, TypeError): + # A missing, unknown or mistyped field falls back to Riva's default + # endpointing rather than failing the whole request. + return None + + +def _streaming_requests(config: Message, pcm_bytes: bytes) -> Iterator[Message]: + yield schema.StreamingRecognizeRequest( + streaming_config=schema.StreamingRecognitionConfig(config=config, interim_results=False) + ) + for offset in range(0, len(pcm_bytes), _CHUNK_BYTES): + yield schema.StreamingRecognizeRequest(audio_content=pcm_bytes[offset : offset + _CHUNK_BYTES]) + + +def streaming_recognize( + channel: grpc.Channel, + config: Message, + pcm_bytes: bytes, + metadata: CallMetadata, + timeout: float | None, +) -> Iterator[Message]: + rpc: Final = channel.stream_stream( + schema.STREAMING_RECOGNIZE_METHOD, + request_serializer=schema.StreamingRecognizeRequest.SerializeToString, + response_deserializer=schema.StreamingRecognizeResponse.FromString, + ) + return rpc(_streaming_requests(config, pcm_bytes), metadata=metadata, timeout=timeout) + + +def recognize( + channel: grpc.Channel, + config: Message, + pcm_bytes: bytes, + metadata: CallMetadata, + timeout: float | None, +) -> Sequence[Message]: + rpc: Final = channel.unary_unary( + schema.RECOGNIZE_METHOD, + request_serializer=schema.RecognizeRequest.SerializeToString, + response_deserializer=schema.RecognizeResponse.FromString, + ) + pending: Final = rpc.future(schema.RecognizeRequest(config=config, audio=pcm_bytes), metadata=metadata) + try: + return (pending.result(timeout=timeout),) + except Exception as e: + if pending.done(): + raise + pending.cancel() + raise NvidiaRivaException( + status_code=408, + message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds", + ) from e diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 6bbf3598bd8..3cb798b7318 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -669,7 +669,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): prompt_cache_options: object = None, ) -> tuple[str | ResponseInputParam, Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None]: validated_input: Final = self._validate_input_param( - self._drop_bridge_minted_reasoning_items(input), drops_set_created_by=drops_params(litellm_params) + self._drop_bridge_minted_reasoning_items(input), drop_params=drops_params(litellm_params) ) reads_breakpoints: Final = self._reads_prompt_cache_breakpoints( model=model, litellm_params=litellm_params, prompt_cache_options=prompt_cache_options @@ -951,7 +951,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return tools if all(new is old for new, old in zip(sanitized, tools, strict=True)) else sanitized def _validate_input_param( - self, input: str | ResponseInputParam, *, drops_set_created_by: bool = False + self, input: str | ResponseInputParam, *, drop_params: bool = False ) -> str | ResponseInputParam: """ Ensure all input fields if pydantic are converted to dict @@ -965,9 +965,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): # if it's pydantic, convert to dict if isinstance(item, BaseModel): validated_input.append( - replayable_responses_item( - item.model_dump(exclude_none=True), drops_set_created_by=drops_set_created_by - ) + replayable_responses_item(item.model_dump(exclude_none=True), drops_set_created_by=drop_params) ) elif isinstance(item, dict): validated_input.append( @@ -975,7 +973,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): self._handle_reasoning_item(cast(dict[str, object], item)) if item.get("type") == "reasoning" else item, - drops_set_created_by=drops_set_created_by, + drops_set_created_by=drop_params, ) ) else: diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index d7aafe9f06e..66b41967ca7 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -2,9 +2,11 @@ Dynamic configuration class generator for JSON-based providers. """ -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from typing import TYPE_CHECKING, Final, Literal, overload +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, @@ -20,6 +22,26 @@ from .json_loader import SimpleProviderConfig if TYPE_CHECKING: from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig +_PAID_SEARCH_TOOLS: Final = frozenset({"web_search", "web_search_image", "image_search"}) +_PAYLOAD_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_TOOL_USAGE_ADAPTER: Final[TypeAdapter[Mapping[str, int]]] = TypeAdapter(Mapping[str, int]) + + +def web_search_requests_from_tool_usage(payload: object) -> int | None: + """ + Billable search calls in a parsed OpenAI-compatible chat response whose usage carries a + tool_usage {tool_name: count} map, as EmpirioLabs reports built-in tool invocations. + Free tools (code_interpreter, web_extractor) are not counted. + """ + try: + body: Final = _PAYLOAD_ADAPTER.validate_python(payload) + usage: Final = _PAYLOAD_ADAPTER.validate_python(body.get("usage")) + tool_usage: Final = _TOOL_USAGE_ADAPTER.validate_python(usage.get("tool_usage"), strict=True) + except ValidationError: + return None + calls: Final = sum(count for tool, count in tool_usage.items() if tool in _PAID_SEARCH_TOOLS and count > 0) + return calls if calls > 0 else None + def create_config_class(provider: SimpleProviderConfig): """Generate config class dynamically from JSON configuration""" @@ -179,6 +201,25 @@ def create_config_class(provider: SimpleProviderConfig): return optional_params + def transform_parsed_response_dict( + self, + parsed_response: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + """ + Carry the provider's usage.tool_usage search count as usage.server_tool_use.web_search_requests, + the field the built-in tool cost path bills per call. + """ + repaired: Final = _PAYLOAD_ADAPTER.validate_python( + super().transform_parsed_response_dict(parsed_response) # pyright: ignore[reportUnknownMemberType] # base_class is picked at runtime + ) + if not provider.special_handling.get("web_search_requests_from_tool_usage"): + return {**repaired} + web_search_requests: Final = web_search_requests_from_tool_usage(repaired) + if web_search_requests is None: + return {**repaired} + usage: Final = _PAYLOAD_ADAPTER.validate_python(repaired["usage"]) + return {**repaired, "usage": {**usage, "server_tool_use": {"web_search_requests": web_search_requests}}} + @property def custom_llm_provider(self) -> str | None: return provider.slug diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index c89ede6e77e..f937e4d2f75 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -178,7 +178,10 @@ "param_mappings": { "max_completion_tokens": "max_tokens" }, - "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] + "special_handling": { + "web_search_requests_from_tool_usage": true + }, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/embeddings"] }, "meta": { "base_url": "https://api.meta.ai/v1", diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index afd27b0f696..e5e3744dd80 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -89,7 +89,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): """Handle preset/ model prefix: send as {"preset": name} instead of {"model": name}.""" input = self._ensure_message_type(input) if model.startswith("preset/"): - input = self._validate_input_param(input, drops_set_created_by=drops_params(litellm_params)) + input = self._validate_input_param(input, drop_params=drops_params(litellm_params)) data: Final[dict] = { "preset": model[len("preset/") :], "input": input, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 32d7b39a20e..80244cfc3eb 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -5,9 +5,11 @@ import json import re import threading import time +from collections import Counter from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import lru_cache, partial +from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args @@ -34,7 +36,11 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.factory import encode_tool_call_id_with_signature from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning -from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp +from litellm.litellm_core_utils.tool_fields import ( + declared_function_name, + qualified_member_name, + tool_params_without_hosted_mcp, +) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -544,6 +550,138 @@ def mapped_gemini_speech_config( } +def _as_chat_tool(member: JsonValue) -> JsonValue: + if not isinstance(member, dict) or "function" in member or "name" not in member: + return member + return { + "type": member.get("type", "function"), + "function": {key: value for key, value in member.items() if key != "type"}, + } + + +def _namespace_members_or_tool(tool: JsonValue) -> tuple[JsonValue, ...]: + if not isinstance(tool, dict) or tool.get("type") != "namespace": + return (tool,) + members: Final = tool.get("tools") + return _members_named_once(tuple(map(_as_chat_tool, members))) if isinstance(members, list) else () + + +def _members_named_once( + members: tuple[JsonValue, ...], repeatable: frozenset[str] = frozenset() +) -> tuple[JsonValue, ...]: + names: Final = tuple(map(declared_function_name, members)) + first_index: Final = {name: index for index, name in reversed(tuple(enumerate(names)))} + return tuple( + member + for index, (member, name) in enumerate(zip(members, names, strict=True)) + if name is None or name in repeatable or first_index[name] == index + ) + + +def _is_namespace(tool: JsonValue) -> bool: + return isinstance(tool, dict) and tool.get("type") == "namespace" + + +def _top_level_name(tool: JsonValue) -> str | None: + name: Final = tool.get("name") if isinstance(tool, dict) else None + return name if isinstance(name, str) and name else None + + +def _listed(tools: JsonValue) -> tuple[JsonValue, ...]: + return tuple(tools) if isinstance(tools, list) else () + + +def _declarable_names(tool: JsonValue) -> tuple[str | None, ...]: + members: Final = _namespace_members_or_tool(tool) if _is_namespace(tool) else () + return (declared_function_name(tool), *map(declared_function_name, members)) + + +def _declarable_names_in(tools: tuple[JsonValue, ...]) -> tuple[str | None, ...]: + return tuple(chain.from_iterable(map(_declarable_names, tools))) + + +def _inlines_members(tool: JsonValue, declarable: Counter[str | None]) -> bool: + members: Final = _namespace_members_or_tool(tool) + return _is_namespace(tool) and all( + name is None or declarable[name] == 1 for name in map(declared_function_name, members) + ) + + +def _qualified_member(namespace: str, member: JsonValue) -> JsonValue: + function: Final = member.get("function") if isinstance(member, dict) else None + name: Final = declared_function_name(member) + if not isinstance(member, dict) or not isinstance(function, dict) or name is None: + return member + return {**member, "function": {**function, "name": qualified_member_name(namespace, name)}} + + +def _as_declared( + tool: JsonValue, declarable: Counter[str | None], top_level_names: frozenset[str] +) -> tuple[JsonValue, ...]: + if _inlines_members(tool, declarable): + return _namespace_members_or_tool(tool) + namespace: Final = _top_level_name(tool) + if not _is_namespace(tool) or namespace is None: + return (tool,) + qualified: Final = (_qualified_member(namespace, member) for member in _namespace_members_or_tool(tool)) + return tuple(member for member in qualified if declared_function_name(member) not in top_level_names) + + +def _top_level_names_in(tools: tuple[JsonValue, ...]) -> frozenset[str]: + top_level_tools: Final = (tool for tool in tools if not _is_namespace(tool)) + return frozenset(name for name in map(declared_function_name, top_level_tools) if name is not None) + + +def _with_namespace_members_inlined(tools: JsonValue, outside_declarations: JsonValue = None) -> JsonValue: + if not isinstance(tools, list) or not any(map(_is_namespace, tools)): + return tools + candidates: Final = (*tools, *_listed(outside_declarations)) + declarable: Final = Counter(_declarable_names_in(candidates)) + top_level_names: Final = _top_level_names_in(candidates) + declared: Final = chain.from_iterable(_as_declared(tool, declarable, top_level_names) for tool in tools) + return list(_members_named_once(tuple(declared), top_level_names)) + + +def _names_each_tool_declares(candidates: tuple[JsonValue, ...]) -> tuple[tuple[str, ...], ...]: + declarable: Final = Counter(_declarable_names_in(candidates)) + top_level_names: Final = _top_level_names_in(candidates) + declared: Final = tuple( + tuple(filter(None, map(declared_function_name, _as_declared(tool, declarable, top_level_names)))) + for tool in candidates + ) + owned_names: Final = chain.from_iterable(zip(repeat(index), names) for index, names in enumerate(declared)) + first_owner: Final = {name: index for index, name in reversed(tuple(owned_names))} + return tuple(tuple(name for name in names if first_owner[name] == index) for index, names in enumerate(declared)) + + +def _qualified_names_of( + member_name: str, tool: JsonValue, declarable: Counter[str | None], declared: tuple[str, ...] +) -> tuple[str, ...]: + namespace: Final = _top_level_name(tool) + if not _is_namespace(tool) or namespace is None or _inlines_members(tool, declarable): + return () + qualified: Final = qualified_member_name(namespace, member_name) + return (qualified,) if qualified in declared else () + + +def _forced_function_names(name: str, tools: JsonValue) -> tuple[str, ...]: + candidates: Final = _listed(tools) + if any(not _is_namespace(tool) and declared_function_name(tool) == name for tool in candidates): + return (name,) + declarable: Final = Counter(_declarable_names_in(candidates)) + declared_by_tool: Final = tuple(zip(candidates, _names_each_tool_declares(candidates), strict=True)) + named: Final = next( + (declared for tool, declared in declared_by_tool if _is_namespace(tool) and _top_level_name(tool) == name), + None, + ) + if named is not None: + return named + qualified: Final = chain.from_iterable( + _qualified_names_of(name, tool, declarable, declared) for tool, declared in declared_by_tool + ) + return tuple(qualified) or (name,) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -815,7 +953,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return supported_params - def map_tool_choice_values(self, model: str, tool_choice: str | dict) -> ToolConfig | None: + def map_tool_choice_values( + self, model: str, tool_choice: str | dict, tools: JsonValue = None, drop_params: bool = False + ) -> ToolConfig | None: if tool_choice == "none": return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="NONE")) elif tool_choice == "required": @@ -825,7 +965,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif isinstance(tool_choice, dict): # only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html name: Final = tool_choice.get("function", {}).get("name", "") - return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=[name])) + forced: Final = _forced_function_names(name, tools) + if forced: + return ToolConfig( + functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=list(forced)) + ) + if drops_unsupported_params(drop_params): + return None + raise litellm.utils.UnsupportedParamsError( + message=f"VertexAI can't force tool_choice '{name}', since none of that namespace's members can be sent to Gemini as a function. Name one of its functions instead, or set `litellm.drop_params = True` to drop tool_choice.", + status_code=400, + ) else: raise litellm.utils.UnsupportedParamsError( message=f"VertexAI doesn't support tool_choice={tool_choice}. Supported tool_choice values=['auto', 'required', json object]. To drop it from the call, set `litellm.drop_params = True.", @@ -1068,7 +1218,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext - def _map_function(self, value: list[dict], optional_params: dict, send_integer_enums: bool = False) -> list[Tools]: + def _map_function( + self, + value: list[dict], + optional_params: dict, + send_integer_enums: bool = False, + outside_declarations: JsonValue = None, + ) -> list[Tools]: """ Map OpenAI-style tools/functions to Vertex AI format. @@ -1099,7 +1255,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # remove 'strict' from tools remove_strict_from_schema(schema_value) - for tool in cast(list[dict], schema_value): # cast-ok: a deep copy keeps the caller's tool shape + inlined_tools: Final = _with_namespace_members_inlined(schema_value, outside_declarations) + for tool in cast(list[dict], inlined_tools): # cast-ok: a deep copy keeps the caller's tool shape openai_function_object: ChatCompletionToolParamFunctionChunk | None = None if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk(**tool["function"]) @@ -1734,7 +1891,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif (param == "tools" or param == "functions") and isinstance(value, list) and value: # Pass optional_params so _map_function can add toolConfig if needed mapped_tools = with_integer_enums_that_fit( - partial(self._map_function, value, optional_params), + partial( + self._map_function, + value, + optional_params, + outside_declarations=request_params.get("functions" if param == "tools" else "tools"), + ), allow_integer_enums and self._is_gemini_3_or_newer(model), ) optional_params = self._add_tools_to_optional_params(optional_params, mapped_tools) @@ -1742,6 +1904,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tool_choice_value = self.map_tool_choice_values( model=model, tool_choice=value, + tools=[*_listed(non_default_params.get("tools")), *_listed(non_default_params.get("functions"))], + drop_params=drop_params, ) if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7419f48f3fb..4d576443b30 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9003,11 +9003,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9039,7 +9035,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -9055,11 +9055,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9091,7 +9087,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -9107,12 +9107,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9139,7 +9135,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-luna-2026-09-22": { "cache_creation_input_token_cost": 1.25e-07, @@ -9155,12 +9159,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9187,7 +9187,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, @@ -9239,7 +9247,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6-sol-2026-09-22": { "cache_creation_input_token_cost": 2.5e-06, @@ -9291,7 +9307,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.25e-06, @@ -9354,11 +9378,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9386,7 +9406,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-sol-2026-09-29": { "cache_creation_input_token_cost": 2.5e-06, @@ -9402,11 +9430,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9434,7 +9458,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/us/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, @@ -15566,7 +15598,7 @@ "output_vector_size": 1536 }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 3e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "mode": "embedding", @@ -15683,8 +15715,10 @@ "input_cost_per_image": 7.2e-05, "input_cost_per_video_per_second": 0.00084, "input_cost_per_audio_per_second": 0.000168, + "input_cost_per_token_batches": 8.1e-08, "output_cost_per_token": 0.0, "output_vector_size": 3072, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_embedding_image_input": true, "supports_image_input": true, "supports_video_input": true, @@ -15693,11 +15727,14 @@ "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.8e-08, "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -15708,11 +15745,14 @@ "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 1.05e-08, "input_cost_per_token": 4.2e-08, + "input_cost_per_token_batches": 2.1e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.68e-07, + "output_cost_per_token_batches": 8.4e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -15720,12 +15760,21 @@ }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, "input_cost_per_token": 9.6e-07, + "input_cost_per_token_batches": 4.8e-07, + "input_cost_per_token_flex": 4.8e-07, + "input_cost_per_token_priority": 1.68e-06, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.84e-06, + "output_cost_per_token_batches": 1.92e-06, + "output_cost_per_token_flex": 1.92e-06, + "output_cost_per_token_priority": 6.72e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -22047,6 +22096,7 @@ "output_cost_per_token": 2.8e-07, "output_dbu_cost_per_token": 4e-06, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -22060,8 +22110,12 @@ "supports_vision": false }, "databricks/system.ai.deepseek-v4-flash-0731": { + "deprecation_date": "2026-11-05", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "litellm_provider": "databricks", "mode": "chat", + "source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { @@ -22081,6 +22135,7 @@ "output_cost_per_token": 3.95997e-06, "output_dbu_cost_per_token": 5.6571e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -23158,6 +23213,7 @@ "output_cost_per_token": 4.04999e-06, "output_dbu_cost_per_token": 5.7857e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-glm-5-3", "supported_modalities": [ "text", "image" @@ -36252,6 +36308,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_image_token_batches": 1.6e-05, "output_cost_per_token_batches": 5e-06, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ @@ -36275,6 +36332,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_image_token_batches": 1.6e-05, "output_cost_per_token_batches": 5e-06, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ @@ -36294,6 +36352,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", @@ -36329,6 +36388,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -36364,6 +36424,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -38126,7 +38187,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38165,7 +38226,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38204,7 +38265,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "responses", "output_cost_per_token": 7.5e-05, @@ -39718,6 +39779,7 @@ "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 8e-06, + "output_cost_per_image_token_batches": 4e-06, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", @@ -55961,6 +56023,7 @@ "openai/sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://platform.openai.com/docs/api-reference/videos", @@ -59823,6 +59886,7 @@ "sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://developers.openai.com/api/docs/pricing", @@ -59848,6 +59912,7 @@ "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 3.2e-05, + "output_cost_per_image_token_batches": 1.6e-05, "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "source": "https://developers.openai.com/api/docs/pricing", @@ -64661,7 +64726,8 @@ { "name": "claude-family-baseline", "pattern": "claude-[a-z]+-\\d+(?:[-.]\\d+)?(?:-\\d{8})?", - "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number. On mapped Claude entries from Anthropic, Azure Foundry, Bedrock and Vertex it only fills the keys an entry leaves out, so the entry's own values always win.", "model_info": { "mode": "chat", "max_input_tokens": 200000, @@ -69076,12 +69142,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69094,12 +69166,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69110,12 +69188,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69124,12 +69208,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69141,12 +69231,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69160,12 +69256,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69371,12 +69473,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69389,12 +69497,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69405,12 +69519,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69419,12 +69539,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69436,12 +69562,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69455,12 +69587,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69664,16 +69802,32 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-06, - "input_cost_per_token_above_272k_tokens": 5.4e-06, - "cache_creation_input_token_cost": 3.375e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-06, - "cache_read_input_token_cost": 2.7e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, - "output_cost_per_token": 1.62e-05, - "output_cost_per_token_above_272k_tokens": 2.43e-05, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-05, + "cache_creation_input_token_cost_flex": 1.65e-06, + "cache_creation_input_token_cost_priority": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-06, + "cache_read_input_token_cost_flex": 1.32e-07, + "cache_read_input_token_cost_priority": 5.28e-07, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-05, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 5.28e-06, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-05, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-05, + "output_cost_per_token_flex": 7.92e-06, + "output_cost_per_token_priority": 3.168e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { @@ -69701,16 +69855,32 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-07, - "input_cost_per_token_above_272k_tokens": 5.4e-07, - "cache_creation_input_token_cost": 3.375e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-07, - "cache_read_input_token_cost": 2.7e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, - "output_cost_per_token": 1.62e-06, - "output_cost_per_token_above_272k_tokens": 2.43e-06, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-06, + "cache_creation_input_token_cost_flex": 1.65e-07, + "cache_creation_input_token_cost_priority": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-07, + "cache_read_input_token_cost_flex": 1.32e-08, + "cache_read_input_token_cost_priority": 5.28e-08, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-07, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-06, + "input_cost_per_token_flex": 1.32e-07, + "input_cost_per_token_priority": 5.28e-07, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-06, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-06, + "output_cost_per_token_flex": 7.92e-07, + "output_cost_per_token_priority": 3.168e-06, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { @@ -69737,9 +69907,9 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 3.375e-06, - "cache_read_input_token_cost": 3.375e-07, - "output_cost_per_token": 2.025e-05, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", "supports_web_search": true @@ -69759,20 +69929,36 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "input_cost_per_token": 1.5e-06, "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 2.4e-07 + "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_batches": 1.2e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.625e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.25e-06 }, "bedrock_mantle/us-gov-west-1/xai.grok-4.6": { "use_openai_responses_path": true, "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69786,11 +69972,17 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": { "input_cost_per_token": 4.8e-08, + "input_cost_per_token_flex": 2.4e-08, + "input_cost_per_token_priority": 8.4e-08, "output_cost_per_token": 9.6e-08, "litellm_provider": "bedrock_mantle", "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 4.8e-08, + "output_cost_per_token_flex": 4.8e-08, + "output_cost_per_token_priority": 1.68e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69804,11 +69996,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": { "input_cost_per_token": 1.56e-07, + "input_cost_per_token_batches": 7.8e-08, + "input_cost_per_token_flex": 7.8e-08, + "input_cost_per_token_priority": 2.73e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69822,11 +70021,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": { "input_cost_per_token": 1.68e-07, + "input_cost_per_token_batches": 8.4e-08, + "input_cost_per_token_flex": 8.4e-08, + "input_cost_per_token_priority": 2.94e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69841,12 +70047,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69860,12 +70072,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69981,10 +70199,17 @@ "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69999,12 +70224,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70018,12 +70249,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -81729,7 +81966,8 @@ "input_cost_per_token_priority": 5.25e-07, "output_cost_per_token_priority": 2.1e-06, "cache_read_input_token_cost_priority": 5.25e-08, - "cache_creation_input_token_cost_priority": 5.25e-07 + "cache_creation_input_token_cost_priority": 5.25e-07, + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017" }, "databricks/databricks-deepseek-v4-1-flash-1017": { "cache_creation_input_token_cost": 3.0002e-07, @@ -82181,6 +82419,8 @@ "databricks/databricks-kimi-k2-7": { "cache_creation_input_token_cost": 9.4997e-07, "cache_read_input_token_cost": 1.8998e-07, + "deprecation_date": "2026-10-30", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "input_cost_per_token": 9.4997e-07, "input_dbu_cost_per_token": 1.3571e-05, "litellm_provider": "databricks", @@ -82193,6 +82433,7 @@ "output_cost_per_token": 4.00001e-06, "output_dbu_cost_per_token": 5.7143e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "successor_model": "databricks/databricks-kimi-k3", "supported_modalities": [ "text", "image" @@ -82594,6 +82835,1603 @@ "supports_function_calling": true, "source": "https://tokenfactory.nebius.com/api/public/models_info" }, + "empiriolabs/aplomb-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 0.0, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1/decisions" + ], + "supports_vision": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/aplomb-1" + }, + "empiriolabs/deepseek-v4-1-flash-1017": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017" + }, + "empiriolabs/deepseek-v4-flash-0731": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 4.24e-07, + "output_cost_per_token": 1.272e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + }, + "empiriolabs/deepseek-v4-flash-0731:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 4.4e-07, + "output_cost_per_token": 1.32e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + }, + "empiriolabs/deepseek-v4-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 3.3e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/deepseek-v4-pro-0813": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 1.32e-06, + "output_cost_per_token": 3.96e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813" + }, + "empiriolabs/deepseek-v4-pro:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/deepseek-v4-pro:variant2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 3.301e-06, + "cache_read_input_token_cost": 1.38e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/deepseek-v4-pro:variant3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/fugu-max": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 2.5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.007, + "search_context_size_medium": 0.007, + "search_context_size_high": 0.007 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-max" + }, + "empiriolabs/fugu-ultra-v1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-06, + "output_cost_per_token": 4.5e-05, + "input_cost_per_token_above_272k_tokens": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 6.75e-05, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 1.5e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-ultra-v1-0" + }, + "empiriolabs/fugu-ultra-v1-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 1e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost": 5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-ultra-v1-1" + }, + "empiriolabs/fugu-ultra-v2-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 1e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost": 5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-ultra-v2-0" + }, + "empiriolabs/gemma-4-26b-a4b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 2.9e-07, + "cache_read_input_token_cost": 2.5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/gemma-4-26b-a4b" + }, + "empiriolabs/glm-4-6v-flashx": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 3.6e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-4-6v-flashx" + }, + "empiriolabs/glm-4-7-flashx": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 202752, + "max_output_tokens": 131072, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 6.3e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-4-7-flashx" + }, + "empiriolabs/glm-5-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 202745, + "input_cost_per_token": 8.25e-07, + "output_cost_per_token": 3.301e-06, + "input_cost_per_token_above_32k_tokens": 1.1e-06, + "output_cost_per_token_above_32k_tokens": 3.851e-06, + "cache_read_input_token_cost_above_32k_tokens": 2.2e-07, + "cache_read_input_token_cost": 1.65e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-1" + }, + "empiriolabs/glm-5-2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-2" + }, + "empiriolabs/glm-5-2:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 3.851e-06, + "cache_read_input_token_cost": 2.75e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-2" + }, + "empiriolabs/glm-5-2:variant2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-2" + }, + "empiriolabs/glm-5-3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-3" + }, + "empiriolabs/glm-5-3-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 2.5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-3-flash" + }, + "empiriolabs/kimi-k2-6": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 16000, + "input_cost_per_token": 8.939e-07, + "output_cost_per_token": 3.7131e-06, + "cache_read_input_token_cost": 1.788e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-6" + }, + "empiriolabs/kimi-k2-7-code": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 9.5e-07, + "output_cost_per_token": 4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + }, + "empiriolabs/kimi-k2-7-code-highspeed": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1.9e-06, + "output_cost_per_token": 8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed" + }, + "empiriolabs/kimi-k2-7-code:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "input_cost_per_token": 8.939e-07, + "output_cost_per_token": 3.7131e-06, + "cache_read_input_token_cost": 1.788e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + }, + "empiriolabs/kimi-k3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k3" + }, + "empiriolabs/mimo-v2-6-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/mimo-v2-6-flash" + }, + "empiriolabs/mimo-v2-6-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 4.35e-07, + "output_cost_per_token": 8.7e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro" + }, + "empiriolabs/mimo-v2-6-pro-ultraspeed": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 4.35e-06, + "output_cost_per_token": 8.7e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed" + }, + "empiriolabs/minimax-m2-7": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 3e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m2-7" + }, + "empiriolabs/minimax-m2-7-highspeed": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 3e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m2-7-highspeed" + }, + "empiriolabs/minimax-m3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 524288, + "input_cost_per_token": 2.25e-07, + "output_cost_per_token": 9e-07, + "input_cost_per_token_above_512k_tokens": 4.5e-07, + "output_cost_per_token_above_512k_tokens": 1.8e-06, + "cache_read_input_token_cost_above_512k_tokens": 9e-08, + "cache_read_input_token_cost": 4.5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m3" + }, + "empiriolabs/minimax-m3-1-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 524288, + "input_cost_per_token": 1.125e-07, + "output_cost_per_token": 4.5e-07, + "cache_read_input_token_cost": 1.125e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m3-1-flash" + }, + "empiriolabs/minimax-m3:priority": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048458, + "max_output_tokens": 524288, + "input_cost_per_token": 3.375e-07, + "output_cost_per_token": 1.35e-06, + "input_cost_per_token_above_512k_tokens": 6.75e-07, + "output_cost_per_token_above_512k_tokens": 2.7e-06, + "cache_read_input_token_cost_above_512k_tokens": 1.35e-07, + "cache_read_input_token_cost": 6.75e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m3" + }, + "empiriolabs/mistral-small-4": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.084, + "search_context_size_medium": 0.084, + "search_context_size_high": 0.14 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/mistral-small-4" + }, + "empiriolabs/muse-glimmer-30b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, + "cache_read_input_token_cost": 5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/muse-glimmer-30b" + }, + "empiriolabs/muse-spark-1-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 4.25e-06, + "cache_read_input_token_cost": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/muse-spark-1-1" + }, + "empiriolabs/muse-spark-1-2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 4.25e-06, + "cache_read_input_token_cost": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/muse-spark-1-2" + }, + "empiriolabs/muse-spark-1-3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 4.25e-06, + "cache_read_input_token_cost": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/muse-spark-1-3" + }, + "empiriolabs/nova-lite-1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 317379, + "max_output_tokens": 10000, + "input_cost_per_token": 6.9e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 3.86e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_prompt_caching": true, + "source": "https://empiriolabs.ai/models/nova-lite-1-0" + }, + "empiriolabs/nova-lite-2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 3.16e-06, + "cache_read_input_token_cost": 2.128e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "source": "https://empiriolabs.ai/models/nova-lite-2" + }, + "empiriolabs/nova-micro-1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 133003, + "max_output_tokens": 10000, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, + "cache_read_input_token_cost": 2.24e-08, + "supports_system_messages": true, + "supports_prompt_caching": true, + "source": "https://empiriolabs.ai/models/nova-micro-1-0" + }, + "empiriolabs/nova-pro-1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 317319, + "max_output_tokens": 10000, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 9.6e-06, + "supports_system_messages": true, + "supports_vision": true, + "source": "https://empiriolabs.ai/models/nova-pro-1-0" + }, + "empiriolabs/perplexity-advanced-deep-research": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_output_tokens": 16000, + "input_cost_per_token": 1.2e-05, + "output_cost_per_token": 6e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012, + "search_context_size_high": 0.012 + }, + "source": "https://empiriolabs.ai/models/perplexity-advanced-deep-research" + }, + "empiriolabs/perplexity-deep-research": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919327, + "input_cost_per_token": 4.8e-06, + "output_cost_per_token": 1.9e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-deep-research" + }, + "empiriolabs/perplexity-pro-search": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_output_tokens": 32768, + "input_cost_per_token": 7.8e-06, + "output_cost_per_token": 3.9e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-pro-search" + }, + "empiriolabs/perplexity-sonar": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919327, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 2.4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-sonar" + }, + "empiriolabs/perplexity-sonar-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919328, + "max_output_tokens": 32768, + "input_cost_per_token": 7.2e-06, + "output_cost_per_token": 3.6e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-sonar-pro" + }, + "empiriolabs/perplexity-sonar-reasoning-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919330, + "max_output_tokens": 131072, + "input_cost_per_token": 4.8e-06, + "output_cost_per_token": 1.9e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-sonar-reasoning-pro" + }, + "empiriolabs/qwen3-5-122b-a10b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 1.15e-07, + "output_cost_per_token": 9.17e-07, + "input_cost_per_token_above_128k_tokens": 2.87e-07, + "output_cost_per_token_above_128k_tokens": 2.294e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" + }, + "empiriolabs/qwen3-5-27b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 8.6e-08, + "output_cost_per_token": 6.88e-07, + "input_cost_per_token_above_128k_tokens": 2.58e-07, + "output_cost_per_token_above_128k_tokens": 2.064e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-27b" + }, + "empiriolabs/qwen3-5-35b-a3b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 5.7e-08, + "output_cost_per_token": 4.59e-07, + "input_cost_per_token_above_128k_tokens": 2.29e-07, + "output_cost_per_token_above_128k_tokens": 1.835e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" + }, + "empiriolabs/qwen3-5-397b-a17b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 1.72e-07, + "output_cost_per_token": 1.032e-06, + "input_cost_per_token_above_128k_tokens": 4.3e-07, + "output_cost_per_token_above_128k_tokens": 2.58e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" + }, + "empiriolabs/qwen3-5-4b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 7e-08, + "cache_read_input_token_cost": 2e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-4b" + }, + "empiriolabs/qwen3-5-9b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 4.5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-9b" + }, + "empiriolabs/qwen3-5-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 995904, + "max_output_tokens": 32768, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3.68e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-flash" + }, + "empiriolabs/qwen3-5-flash:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 995904, + "max_output_tokens": 65536, + "input_cost_per_token": 2.9e-08, + "output_cost_per_token": 2.87e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-flash" + }, + "empiriolabs/qwen3-5-omni-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 196608, + "max_output_tokens": 32768, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" + }, + "empiriolabs/qwen3-5-omni-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 196608, + "max_output_tokens": 32768, + "input_cost_per_token": 2.8e-06, + "output_cost_per_token": 1.66e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" + }, + "empiriolabs/qwen3-5-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 3.6e-07, + "output_cost_per_token": 2.21e-06, + "input_cost_per_token_above_256k_tokens": 1.08e-06, + "output_cost_per_token_above_256k_tokens": 6.62e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-plus" + }, + "empiriolabs/qwen3-5-plus:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 1.15e-07, + "output_cost_per_token": 6.88e-07, + "input_cost_per_token_above_128k_tokens": 2.87e-07, + "output_cost_per_token_above_128k_tokens": 1.72e-06, + "input_cost_per_token_above_256k_tokens": 5.73e-07, + "output_cost_per_token_above_256k_tokens": 3.44e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-plus" + }, + "empiriolabs/qwen3-6-27b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229376, + "max_output_tokens": 64000, + "input_cost_per_token": 4.12564e-07, + "output_cost_per_token": 2.475384e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-27b" + }, + "empiriolabs/qwen3-6-35b-a3b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 3.5e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b" + }, + "empiriolabs/qwen3-6-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-06, + "input_cost_per_token_above_256k_tokens": 1e-06, + "output_cost_per_token_above_256k_tokens": 4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-flash" + }, + "empiriolabs/qwen3-6-flash:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 1.65e-07, + "output_cost_per_token": 9.9e-07, + "input_cost_per_token_above_256k_tokens": 6.6e-07, + "output_cost_per_token_above_256k_tokens": 3.961e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-flash" + }, + "empiriolabs/qwen3-6-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "input_cost_per_token_above_256k_tokens": 2e-06, + "output_cost_per_token_above_256k_tokens": 6e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.026, + "search_context_size_medium": 0.026, + "search_context_size_high": 0.026 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-plus" + }, + "empiriolabs/qwen3-6-plus:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.76e-07, + "output_cost_per_token": 1.651e-06, + "input_cost_per_token_above_256k_tokens": 1.101e-06, + "output_cost_per_token_above_256k_tokens": 6.602e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-plus" + }, + "empiriolabs/qwen3-7-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 1.3e-07, + "input_cost_per_token_above_32k_tokens": 1e-07, + "output_cost_per_token_above_32k_tokens": 4e-07, + "cache_read_input_token_cost_above_32k_tokens": 2e-08, + "input_cost_per_token_above_256k_tokens": 2e-07, + "output_cost_per_token_above_256k_tokens": 8e-07, + "cache_read_input_token_cost_above_256k_tokens": 4e-08, + "cache_read_input_token_cost": 6e-09, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.03, + "search_context_size_medium": 0.03, + "search_context_size_high": 0.03 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-flash" + }, + "empiriolabs/qwen3-7-flash:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.8e-08, + "output_cost_per_token": 1.1e-07, + "input_cost_per_token_above_32k_tokens": 8.3e-08, + "output_cost_per_token_above_32k_tokens": 3.3e-07, + "cache_read_input_token_cost_above_32k_tokens": 1.7e-08, + "input_cost_per_token_above_256k_tokens": 1.65e-07, + "output_cost_per_token_above_256k_tokens": 6.6e-07, + "cache_read_input_token_cost_above_256k_tokens": 3.3e-08, + "cache_read_input_token_cost": 6e-09, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-flash" + }, + "empiriolabs/qwen3-7-max": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 7.5e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-max" + }, + "empiriolabs/qwen3-7-max:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 4.951e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-max" + }, + "empiriolabs/qwen3-7-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "input_cost_per_token_above_256k_tokens": 1.2e-06, + "output_cost_per_token_above_256k_tokens": 4.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.03, + "search_context_size_medium": 0.03, + "search_context_size_high": 0.03 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-plus" + }, + "empiriolabs/qwen3-7-plus:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.76e-07, + "output_cost_per_token": 1.101e-06, + "input_cost_per_token_above_256k_tokens": 8.26e-07, + "output_cost_per_token_above_256k_tokens": 3.301e-06, + "cache_read_input_token_cost_above_256k_tokens": 1.66e-07, + "cache_read_input_token_cost": 5.6e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-plus" + }, + "empiriolabs/qwen3-8-27b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 1.7e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 8e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-27b" + }, + "empiriolabs/qwen3-8-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 1.6e-07, + "output_cost_per_token": 4.7e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-flash" + }, + "empiriolabs/qwen3-8-max": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-max" + }, + "empiriolabs/qwen3-8-max-0902": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-max-0902" + }, + "empiriolabs/qwen3-8-max:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 4.951e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-max" + }, + "empiriolabs/qwen3-8-omni-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1006592, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 9.4e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash" + }, + "empiriolabs/seed-2-0-code": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229096, + "max_output_tokens": 128000, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "input_cost_per_token_above_128k_tokens": 8e-07, + "output_cost_per_token_above_128k_tokens": 4.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "source": "https://empiriolabs.ai/models/seed-2-0-code" + }, + "empiriolabs/seed-2-0-lite": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229096, + "max_output_tokens": 128000, + "input_cost_per_token": 3.1e-07, + "output_cost_per_token": 2.5e-06, + "input_cost_per_token_above_128k_tokens": 6.2e-07, + "output_cost_per_token_above_128k_tokens": 5e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-0-lite" + }, + "empiriolabs/seed-2-0-mini": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229156, + "max_output_tokens": 128000, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 5e-07, + "input_cost_per_token_above_128k_tokens": 2.4e-07, + "output_cost_per_token_above_128k_tokens": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-0-mini" + }, + "empiriolabs/seed-2-0-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229096, + "max_output_tokens": 128000, + "input_cost_per_token": 6.3e-07, + "output_cost_per_token": 3.79e-06, + "input_cost_per_token_above_128k_tokens": 1.26e-06, + "output_cost_per_token_above_128k_tokens": 7.58e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-0-pro" + }, + "empiriolabs/seed-2-1-turbo": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229093, + "max_output_tokens": 65536, + "input_cost_per_token": 6.3e-07, + "output_cost_per_token": 3.13e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-1-turbo" + }, + "empiriolabs/skylark-embedding-vision": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 8000, + "input_cost_per_token": 2.5e-07, + "supports_embedding_image_input": true, + "source": "https://empiriolabs.ai/models/skylark-embedding-vision" + }, + "empiriolabs/step-3-5-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 2e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/step-3-5-flash" + }, + "empiriolabs/step-3-5-flash-2603": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 2e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/step-3-5-flash-2603" + }, + "empiriolabs/step-3-7-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.15e-06, + "cache_read_input_token_cost": 4e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/step-3-7-flash" + }, + "empiriolabs/step-5-preview": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1024000, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/step-5-preview" + }, + "empiriolabs/stepaudio-2-5-chat": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 3e-07, + "supports_system_messages": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/stepaudio-2-5-chat" + }, + "empiriolabs/stepaudio-3-chat": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 3e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/stepaudio-3-chat" + }, + "empiriolabs/text-embedding-v4": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 8192, + "input_cost_per_token": 7e-08, + "source": "https://empiriolabs.ai/models/text-embedding-v4" + }, + "empiriolabs/tongyi-embedding-vision-flash": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 1024, + "input_cost_per_token": 9e-08, + "supports_embedding_image_input": true, + "source": "https://empiriolabs.ai/models/tongyi-embedding-vision-flash" + }, + "empiriolabs/tongyi-embedding-vision-plus": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 1024, + "input_cost_per_token": 9e-08, + "supports_embedding_image_input": true, + "source": "https://empiriolabs.ai/models/tongyi-embedding-vision-plus" + }, "wally/glm-5.3-flash": { "cache_read_input_token_cost": 7.6e-08, "input_cost_per_token": 2.72e-07, @@ -82920,5 +84758,1354 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "bedrock/us-gov-east-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_assistant_prefill": true, + "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 + }, + "bedrock/us-gov-west-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_assistant_prefill": true, + "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 + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gmi/MiniMaxAI/MiniMax-M2": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.5": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.7": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3": { + "cache_read_input_token_cost": 4.8e-08, + "input_cost_per_token": 2.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9.6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3.1-Flash": { + "cache_read_input_token_cost": 1.125e-08, + "input_cost_per_token": 1.125e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-FP8": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.09e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Instruct-2507-FP8": { + "cache_read_input_token_cost": 1.75e-08, + "input_cost_per_token": 8.75e-08, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Thinking-2507-FP8": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262141, + "mode": "chat", + "output_cost_per_token": 4.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-30B-A3B": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 40946, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-32B-FP8": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 9e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262128, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-122B-A10B": { + "input_cost_per_token": 4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-27B": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-35B-A3B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-35B-A3B": { + "input_cost_per_token": 2.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.485e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus-2026-04-02": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.7-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-27B": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Flash": { + "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 1.6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max-0902": { + "cache_read_input_token_cost": 2.06e-07, + "input_cost_per_token": 1.65e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.951e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5": { + "cache_read_input_token_cost": 2.38e-09, + "input_cost_per_token": 1.19e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.38e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5-Pro": { + "cache_read_input_token_cost": 2.52e-09, + "input_cost_per_token": 3.045e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.09e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-0528": { + "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3.9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1-Terminus": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Exp": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4.1e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Speciale": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 1.82e-08, + "input_cost_per_token": 9.1e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.82e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash-0731": { + "cache_read_input_token_cost": 9.1e-09, + "input_cost_per_token": 2.86e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.58e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Pro-0813": { + "cache_read_input_token_cost": 3.52e-08, + "input_cost_per_token": 1.056e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.168e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-pro-preview": { + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.7-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.8-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-4-argon": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1.3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-31b-it": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-3.3-70B-Instruct": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2-Instruct-0905": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.6": { + "cache_read_input_token_cost": 1.44e-07, + "input_cost_per_token": 8.55e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code-highspeed": { + "cache_read_input_token_cost": 3.8e-07, + "input_cost_per_token": 1.9e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3.1-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/nemotron-3-ultra-550b-a55b": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-pro": { + "input_cost_per_token": 3e-05, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 0.00018, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-luna": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-sol": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-terra": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-astra": { + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-sol": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6.1-sol": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-oss-120b": { + "input_cost_per_token": 5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-3.7-Flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.15e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-5-Preview": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.7e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/HunyuanOCR": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/Hy3": { + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 5.8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy3-preview": { + "cache_read_input_token_cost": 6e-08, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 8.34e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/x-ai/grok-4.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/xai/grok4.6": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-Air-FP8": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131076, + "mode": "chat", + "output_cost_per_token": 1.1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-FP8": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.6": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 204800, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.7-Flash": { + "input_cost_per_token": 7e-08, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5-FP8": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 1.92e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.1-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.2-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3-Flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 7.5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "fireworks_ai/accounts/fireworks/models/kimi-k3p1-pro": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 1.5e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "baseten/moonshotai/Kimi-K3.1-Pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1000000, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "gpt-image-2.6-sunburst": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "gpt-image-2.6-flare": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64.0, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock/us-gov-west-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64.0, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0": { + "deprecation_date": "2026-09-10", + "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-3-haiku.md", + "input_cost_per_token": 3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 195904, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "successor_model": "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/xai.grok-4.7": { + "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_batches": 1.32e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 7.92e-06, + "output_cost_per_token_batches": 3.96e-06, + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "use_openai_responses_path": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0": { + "deprecation_date": "2026-09-10", + "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-3-haiku.md", + "input_cost_per_token": 3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 195904, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "successor_model": "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-west-1/xai.grok-4.7": { + "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_batches": 1.32e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 7.92e-06, + "output_cost_per_token_batches": 3.96e-06, + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "use_openai_responses_path": true + }, + "bedrock_mantle/xai.grok-4.7": { + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_flex": 2.75e-07, + "cache_read_input_token_cost_priority": 9.625e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_batches": 1.1e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 3.85e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "output_cost_per_token_batches": 3.3e-06, + "output_cost_per_token_flex": 3.3e-06, + "output_cost_per_token_priority": 1.155e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "use_openai_responses_path": true } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 672a2739964..00fee0a0d5a 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -5277,13 +5277,13 @@ } }, "empiriolabs": { - "display_name": "EmpirioLabs (`empiriolabs`)", + "display_name": "EmpirioLabs AI (`empiriolabs`)", "url": "https://docs.litellm.ai/docs/providers/empiriolabs", "endpoints": { - "chat_completions": false, - "responses": false, - "text_completion": false, - "embeddings": false, + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": true, "image_generations": false, "image_edits": false, "audio_transcriptions": false, @@ -5295,11 +5295,11 @@ "realtime": false, "video_generations": false, "batches": false, - "messages": false, - "interactions": false, + "messages": true, + "interactions": true, "systemone": false, - "a2a": false, - "generateContent": false + "a2a": true, + "generateContent": true } }, "openinfer": { diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b290adc47ce..ab85b6f8842 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -20233,7 +20233,7 @@ "type": "null" } ], - "description": "Regex pattern applied to the tool's function name", + "description": "Regex pattern applied to the tool's name, or to its type when the tool has no name", "title": "Tool Name" }, "tool_type": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index de1096701af..1864c3ae82e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -38,7 +38,6 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_no_callback_env_reference, ) from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled -from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.types.agents import ( AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export ) @@ -2481,20 +2480,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ - accounting_capture: UsageCaptureSettings | None = Field( - None, - description="Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes.", - ) - accounting_sync_ack: bool = Field( + use_spend_logs_ledger: bool = Field( False, - description="Wait for usage commit on foreground non-stream chat, text, embedding and Responses requests. " - "Other requests remain supported with asynchronous usage capture.", - ) - accounting_sync_ack_timeout_seconds: float = Field( - 20.0, gt=0, description="Maximum response wait for usage commit, including the batch flush interval" - ) - accounting_sync_ack_failure_policy: Literal["fail_closed", "log"] = Field( - "fail_closed", description="Return 503 or log the persistence failure after the response wait expires" + description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.", ) request_tag_mappings: Sequence[RequestTagMapping] | None = Field( @@ -3817,6 +3805,9 @@ class SpendLogsMetadata(TypedDict): class SpendLogsPayload(TypedDict): + accounting_payload: NotRequired[Mapping[str, object]] # writable-ok: capture before buffering + accounting_log_disabled: NotRequired[bool] # writable-ok: capture logging mode before buffering + accounting_delivery_id: NotRequired[str] # writable-ok: assigned once before buffering request_id: str call_type: str api_key: str diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 3dfaa7f213e..f0a4d47f1e4 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header +from litellm.llms.azure.common_utils import caused_by_entra_rejection from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy @@ -1202,6 +1203,8 @@ async def invoke_agent_a2a( pass if isinstance(e, litellm.BadRequestError): return _jsonrpc_error(body.get("id"), -32602, e.message, 400) + if isinstance(e, litellm.AuthenticationError) and caused_by_entra_rejection(e): + return _jsonrpc_error(body.get("id"), -32000, e.message, 401) return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) diff --git a/litellm/proxy/anthropic_endpoints/answered_status.py b/litellm/proxy/anthropic_endpoints/answered_status.py index 8a3469f9119..26eeb08c8c9 100644 --- a/litellm/proxy/anthropic_endpoints/answered_status.py +++ b/litellm/proxy/anthropic_endpoints/answered_status.py @@ -1,10 +1,29 @@ from typing import Final, TypeVar -from litellm.exceptions import InternalServerError, MidStreamFallbackError +from litellm.exceptions import APIError, InternalServerError, MidStreamFallbackError + +_PROVIDER_OVERLOAD_STATUS_CODE: Final = 529 +_OPENAI_OVERLOAD_ANSWERED_STATUS_CODE: Final = 503 _Raised = TypeVar("_Raised", bound=Exception) +def is_provider_overload(exc: Exception) -> bool: + return _provider_status_code(exc) == _PROVIDER_OVERLOAD_STATUS_CODE + + +def with_openai_overload_status(exc: _Raised) -> _Raised: + return ( + _copy_carrying( + _copy_carrying(exc, "provider_status_code", _PROVIDER_OVERLOAD_STATUS_CODE), + "status_code", + _OPENAI_OVERLOAD_ANSWERED_STATUS_CODE, + ) + if is_provider_overload(exc) + else exc + ) + + def _provider_status_code(exc: Exception) -> int | None: unwrapped = exc # rebind-ok: the recursion gate requires iterative wrapper traversal while isinstance(unwrapped, MidStreamFallbackError) and unwrapped.original_exception is not None: @@ -12,6 +31,10 @@ def _provider_status_code(exc: Exception) -> int | None: match unwrapped: case InternalServerError(provider_status_code=int() as provider_status): return provider_status + case APIError(provider_status_code=int() as provider_status): + return provider_status + case APIError(status_code=int() as provider_status): + return provider_status case _: return None diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 5df73230b6c..b8c90fe7e93 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -26,7 +26,6 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, get_args -from urllib.parse import urlsplit from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from fastapi.responses import JSONResponse, Response @@ -34,6 +33,7 @@ from pydantic import Json, TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.url_utils import is_https_url_with_host from litellm.models.plugin_upload import PluginUpload from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( @@ -377,14 +377,6 @@ _VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a- _VALID_SHA256_RE: Final = re.compile(r"^[0-9a-fA-F]{64}$") -def _is_https_url_with_host(url: str) -> bool: - try: - parts: Final = urlsplit(url) - except ValueError: - return False - return parts.scheme == "https" and bool(parts.hostname) - - def _validate_plugin_source(source: Mapping[str, str]) -> None: """Validate plugin source format, raising HTTPException on invalid input.""" source_type: Final = source.get("source") @@ -421,7 +413,7 @@ def _validate_plugin_source(source: Mapping[str, str]) -> None: }, ) elif source_type == "archive": - if not _is_https_url_with_host(source.get("url", "")): + if not is_https_url_with_host(source.get("url", "")): raise HTTPException( status_code=400, detail={ diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5ae6c628324..e812b03385d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -95,6 +95,7 @@ from litellm.litellm_core_utils.streaming_handler import ( ) from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.answered_status import is_provider_overload, with_openai_overload_status from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, is_model_cost_zero, @@ -115,6 +116,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from litellm.proxy.common_utils.openai_error_payload import ( LITELLM_CALL_ID_HEADER, + OverloadedProxyException, attribute_of, error_status_code, openai_error_param, @@ -140,7 +142,6 @@ from litellm.proxy.route_llm_request import ( route_request, ) from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost -from litellm.proxy.spend_tracking.usage_capture import capture_proxy_response from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ProxyLogging, @@ -165,6 +166,7 @@ _ChunkT: Final = TypeVar("_ChunkT") KNOWN_PROXY_ROUTES: Final = frozenset( route for member in LiteLLMRoutes for route in member.value if route.startswith("/") ) +_OPENAI_OVERLOAD_ROUTE_TYPES: Final = frozenset({"acompletion", "atext_completion", "aresponses"}) ProxyRouteType: TypeAlias = Literal[ "acompletion", @@ -3397,9 +3399,6 @@ class ProxyBaseLLMRequestProcessing: ) ), ) - await capture_proxy_response( - logging_obj, response, background=_REQUEST_FIELDS.validate_python(self.data).get("background") is True - ) except (Exception, asyncio.CancelledError) as exc: _exception_raised = True if logging_obj.prepared_cost is not None: @@ -4402,28 +4401,28 @@ class ProxyBaseLLMRequestProcessing: version: str | None = None, ) -> Never: """Raises ProxyException (OpenAI API compatible) if an exception is raised""" - log_llm_api_exception(e, self.litellm_call_id) + answered: Final = with_openai_overload_status(e) if self.route_type in _OPENAI_OVERLOAD_ROUTE_TYPES else e + log_llm_api_exception(answered, self.litellm_call_id) billing_logging_obj: Final = _REQUEST_FIELDS.validate_python(self.data).get("litellm_logging_obj") failed_in_agentic_follow_up: Final = is_agentic_followup_failure(e) # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, - original_exception=e, + original_exception=answered, request_data=self.data, dispatched_call_type=self.route_type, ) # Use transformed exception if callback returned one, otherwise use original - if transformed_exception is not None: - e = transformed_exception - litellm_debug_info: Final = getattr(e, "litellm_debug_info", "") + failure: Final = transformed_exception if transformed_exception is not None else answered + litellm_debug_info: Final = getattr(failure, "litellm_debug_info", "") verbose_proxy_logger.debug( "An error occurred: %s %s", - e, + failure, litellm_debug_info, ) timeout: Final = getattr( - e, "timeout", None + failure, "timeout", None ) # returns the timeout set by the wrapper. Used for testing if model-specific timeout are set correctly _litellm_logging_obj: Final = ( billing_logging_obj if isinstance(billing_logging_obj, LiteLLMLoggingObj) else None @@ -4455,11 +4454,11 @@ class ProxyBaseLLMRequestProcessing: # a failed request reports no timing, matching /v1/chat/completions read_timing_from_logging_obj=False, ) - headers = getattr(e, "headers", None) or {} + headers = getattr(failure, "headers", None) or {} if not headers: - _response: Final = attribute_of(e, "response") + _response: Final = attribute_of(failure, "response") _response_headers: Final = getattr(_response, "headers", None) if _response is not None else None - _provider_headers: Final = _response_headers or getattr(e, "litellm_response_headers", None) + _provider_headers: Final = _response_headers or getattr(failure, "litellm_response_headers", None) if _provider_headers: headers = get_response_headers(dict(_provider_headers)) headers.update(custom_headers) @@ -4479,21 +4478,21 @@ class ProxyBaseLLMRequestProcessing: safe_headers: Final = {k: v for k, v in headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} - self._apply_router_cooldown_retry_after(safe_headers, e) + self._apply_router_cooldown_retry_after(safe_headers, failure) - if isinstance(e, ProxyException): - e.headers = { - **{k: v for k, v in e.headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS}, + if isinstance(failure, ProxyException): + failure.headers = { + **{k: v for k, v in failure.headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS}, **{k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items()}, } - raise e + raise failure - if isinstance(e, HTTPException): - raise proxy_exception_from_http_exception(e, safe_headers) - elif isinstance(e, httpx.HTTPStatusError): + if isinstance(failure, HTTPException): + raise proxy_exception_from_http_exception(failure, safe_headers) + elif isinstance(failure, httpx.HTTPStatusError): # Handle httpx.HTTPStatusError - extract actual error from response # This matches the original behavior before the refactor in commit 511d435f6f - http_status_error: Final[httpx.HTTPStatusError] = e + http_status_error: Final[httpx.HTTPStatusError] = failure error_body: Final = await http_status_error.response.aread() error_text: Final = error_body.decode("utf-8") @@ -4503,12 +4502,12 @@ class ProxyBaseLLMRequestProcessing: detail={"error": error_text}, headers=error_headers, ) - error_msg: Final = f"{e}" + error_msg: Final = f"{failure}" # Check for AttributeError in the exception chain. # The AttributeError may be wrapped in multiple layers # (e.g. AttributeError -> OpenAIException -> APIConnectionError), # so walk __cause__, __context__, and original_exception recursively. - has_attribute_error: Final = _has_attribute_error_in_chain(e) + has_attribute_error: Final = _has_attribute_error_in_chain(failure) if has_attribute_error: raise ProxyException( @@ -4523,12 +4522,12 @@ class ProxyBaseLLMRequestProcessing: # VertexAIError, etc.) all have a status_code attribute reflecting # the upstream API response. Use it to return the correct HTTP code # instead of defaulting to 500. - _exc_status_code: Final = getattr(e, "status_code", None) + _exc_status_code: Final = getattr(failure, "status_code", None) if _exc_status_code is not None and isinstance(_exc_status_code, int) and 400 <= _exc_status_code <= 599: _code = _exc_status_code else: _code = status.HTTP_500_INTERNAL_SERVER_ERROR - if should_report_bug(e): + if should_report_bug(failure): proxy_server_request: Final = self.data.get("proxy_server_request") request_url: Final = ( proxy_server_request.get("url") if isinstance(proxy_server_request, Mapping) else None @@ -4537,24 +4536,33 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.error( bug_report_notice( build_proxy_bug_report( - e, + failure, call_type=allowlisted(request_path, KNOWN_PROXY_ROUTES), custom_llm_provider=self.data.get("custom_llm_provider"), stream=self.data.get("stream"), ) ) ) - client_message: Final = getattr(e, "message", error_msg) + client_message: Final = getattr(failure, "message", error_msg) + redacted_client_message: Final = redact_internal_details_from_client_message( + strip_bug_report_notice(client_message) if isinstance(client_message, str) else error_msg + ) + provider_specific_fields: Final = getattr( + failure, "provider_specific_fields", None + ) or relayed_provider_specific_fields(failure) + if self.route_type in _OPENAI_OVERLOAD_ROUTE_TYPES and is_provider_overload(failure): + raise OverloadedProxyException( + message=redacted_client_message, + headers=cast(dict[str, str], safe_headers), # cast-ok: ProxyException coerces header values to strings + provider_specific_fields=provider_specific_fields, + ) raise ProxyException( - message=redact_internal_details_from_client_message( - strip_bug_report_notice(client_message) if isinstance(client_message, str) else error_msg - ), - type=openai_error_type(e, _code), - param=openai_error_param(e), - openai_code=getattr(e, "code", None), + message=redacted_client_message, + type=openai_error_type(failure, _code), + param=openai_error_param(failure), + openai_code=getattr(failure, "code", None), code=_code, - provider_specific_fields=getattr(e, "provider_specific_fields", None) - or relayed_provider_specific_fields(e), + provider_specific_fields=provider_specific_fields, headers=safe_headers, ) diff --git a/litellm/proxy/common_utils/model_deprecation.py b/litellm/proxy/common_utils/model_deprecation.py index b86ebd830f9..832ebe175b5 100644 --- a/litellm/proxy/common_utils/model_deprecation.py +++ b/litellm/proxy/common_utils/model_deprecation.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timezone @@ -166,14 +167,20 @@ def _classify(days_until: int, warn_within_days: int) -> DeprecationStatus: return "upcoming" -def _checked_successor(model_name: str, model_info: Mapping[str, object], llm_router: Router) -> str | None: +def _checked_successor( + model_name: str, + model_info: Mapping[str, object], + llm_router: Router, + successor_log_level: int, +) -> str | None: """The deployment's successor_model as set, with a warning when no model in its team scope matches it""" successor: Final = _successor_model(model_info.get("successor_model")) if successor is None: return None team_id: Final = model_info.get("team_id") if not llm_router.registry().deployments(successor, team_id=team_id if isinstance(team_id, str) else None): - verbose_logger.warning( + verbose_logger.log( + successor_log_level, "model_deprecation: %s names successor %s, which matches no model group, alias, or wildcard route here", model_name, successor, @@ -187,6 +194,7 @@ def _build_info( today: date, warn_within_days: int, catalog: Mapping[str, object], + successor_log_level: int, ) -> _DeploymentInfo | None: model_name: Final = deployment.get("model_name") if not isinstance(model_name, str) or not model_name: @@ -197,10 +205,18 @@ def _build_info( return None days_until: Final = (resolved.deprecation_date - today).days - configured: Final = _checked_successor(model_name, _mapping_field(deployment, "model_info"), llm_router) + model_info: Final = _mapping_field(deployment, "model_info") + configured: Final = _checked_successor(model_name, model_info, llm_router, successor_log_level) + team_public_model_name_value: Final = model_info.get("team_public_model_name") + team_public_model_name: Final = ( + team_public_model_name_value + if isinstance(team_public_model_name_value, str) and team_public_model_name_value.strip() + else None + ) return _DeploymentInfo( info=ModelDeprecationInfo( model_name=model_name, + team_public_model_name=team_public_model_name, litellm_model=resolved.litellm_model, deprecation_date=resolved.deprecation_date, days_until_deprecation=days_until, @@ -218,7 +234,7 @@ class _DeploymentInfo: configured_successor: bool -def _merge_group(group: Sequence[_DeploymentInfo]) -> ModelDeprecationInfo: +def _merge_group(group: Sequence[_DeploymentInfo], successor_log_level: int) -> ModelDeprecationInfo: """Keep the group's first deployment in config order, naming the first successor any of its deployments names, a successor configured in model_info ahead of one the cost map supplies""" first: Final = group[0].info @@ -234,7 +250,8 @@ def _merge_group(group: Sequence[_DeploymentInfo]) -> ModelDeprecationInfo: if not successors: return first if len(frozenset(successors)) > 1: - verbose_logger.warning( + verbose_logger.log( + successor_log_level, "model_deprecation: %s names several successors %s, reporting %s", first.model_name, sorted(frozenset(successors)), @@ -245,11 +262,12 @@ def _merge_group(group: Sequence[_DeploymentInfo]) -> ModelDeprecationInfo: def _dedupe( models: Sequence[_DeploymentInfo], + successor_log_level: int, ) -> tuple[ModelDeprecationInfo, ...]: """Report a model group carrying the same date on several deployments once""" ordered: Final = sorted(models, key=lambda model: (model.info.model_name, model.info.deprecation_date)) return tuple( - _merge_group(tuple(group)) + _merge_group(tuple(group), successor_log_level) for _, group in groupby(ordered, key=lambda model: (model.info.model_name, model.info.deprecation_date)) ) @@ -264,11 +282,18 @@ def _bucket(models: Sequence[ModelDeprecationInfo], status: DeprecationStatus) - def _collect_from_router( - llm_router: Router, today: date, warn_within_days: int, catalog: Mapping[str, object] + llm_router: Router, + today: date, + warn_within_days: int, + catalog: Mapping[str, object], + successor_log_level: int, ) -> tuple[ModelDeprecationInfo, ...]: deployments: Final = llm_router.registry().deployments() or () - infos: Final = (_build_info(deployment, llm_router, today, warn_within_days, catalog) for deployment in deployments) - deduped: Final = _dedupe(tuple(info for info in infos if info is not None)) + infos: Final = ( + _build_info(deployment, llm_router, today, warn_within_days, catalog, successor_log_level) + for deployment in deployments + ) + deduped: Final = _dedupe(tuple(info for info in infos if info is not None), successor_log_level) verbose_logger.debug( "model_deprecation: %d/%d deployments carry a deprecation date", @@ -283,12 +308,13 @@ def collect_model_deprecations( warn_within_days: int = DEFAULT_DEPRECATION_WARN_DAYS, today: date | None = None, loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, + successor_log_level: int = logging.WARNING, ) -> ModelDeprecationResponse: """Bucket every deployment carrying a deprecation date by how urgent it is""" snapshot_time: Final = datetime.now(timezone.utc) effective_today: Final = today or snapshot_time.date() deduped: Final = ( - _collect_from_router(llm_router, effective_today, warn_within_days, loaded_catalog()) + _collect_from_router(llm_router, effective_today, warn_within_days, loaded_catalog(), successor_log_level) if llm_router is not None else () ) diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index a948c390375..26369bf3f87 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -17,6 +17,7 @@ from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id" +_SERVER_IS_OVERLOADED_CODE: Final = "server_is_overloaded" _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { @@ -44,6 +45,31 @@ _FAILOVER_REFUSALS: Final = ( ) +class OverloadedProxyException(ProxyException): + def __init__( + self, + message: str, + headers: dict[str, str] | None, # mutable-ok: matches ProxyException's constructor + provider_specific_fields: dict[str, object] | None, # mutable-ok: matches ProxyException's constructor + ) -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # ProxyException omits generic arguments + message=message, + type="service_unavailable_error", + param=None, + code=status.HTTP_503_SERVICE_UNAVAILABLE, + headers=None if headers is None else dict(headers), + provider_specific_fields=provider_specific_fields, + ) + + def to_dict(self) -> dict[str, object]: # mutable-ok: matches ProxyException.to_dict + return { + **_ERROR_BODY.validate_python( + super().to_dict() # pyright: ignore[reportUnknownMemberType] # ProxyException omits generic arguments + ), + "code": _SERVER_IS_OVERLOADED_CODE, + } + + def attribute_of(value: object, name: str, default: object = None) -> object: return getattr(value, name, default) diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py deleted file mode 100644 index 0a9c5db1619..00000000000 --- a/litellm/proxy/db/accounting_ledger.py +++ /dev/null @@ -1,352 +0,0 @@ -from __future__ import annotations - -import hashlib -import json -from collections.abc import AsyncGenerator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager -from dataclasses import dataclass -from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, Literal, TypeAlias - -from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter -from typing_extensions import assert_never - -from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent -from litellm.proxy.db.db_span import db_span -from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper - -if TYPE_CHECKING: - from prisma import Prisma - - from litellm.proxy.utils import PrismaClient - -AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent -_JSON: Final = TypeAdapter(dict[str, JsonValue]) -_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime) -_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...]) -_TIMEOUT: Final = timedelta(seconds=10) - - -class AccountingConflictError(ValueError): - pass - - -@dataclass(frozen=True, slots=True) -class UsageWrite: - event: InboundUsageEvent | OutboundUsageEvent - envelope: str - targets: tuple[DeliveryTarget, ...] = () - - -_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] - - -@dataclass(frozen=True, slots=True) -class _AcceptedEvent: - event_id: str - kind: EventKind - source_event_id: str | None - occurred_at: datetime - payload: str - payload_hash: str - targets: tuple[str, ...] - schema_version: int - version: int - - -def _json(value: object) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - - -def _prepare( - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...], - envelope: str, - schema_version: int, - version: int, -) -> _AcceptedEvent: - if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1: - raise ValueError("Accounting versions must be positive integers") - validated: Final = type(event).model_validate_json(event.model_dump_json()) - payload: Final = _JSON.validate_json(envelope) - if payload.get("event") != validated.model_dump(mode="json"): - raise ValueError("Accounting envelope must contain the matching public event") - destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets))) - serialized: Final = _json(payload) - identity: Final = _json((schema_version, version, destinations, payload)) - kind: Final = ( - EventKind.COST - if isinstance(validated, CostEvent) - else EventKind.INBOUND_USAGE - if isinstance(validated, InboundUsageEvent) - else EventKind.OUTBOUND_USAGE - ) - return _AcceptedEvent( - event_id=validated.event_id, - kind=kind, - source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None, - occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at, - payload=serialized, - payload_hash=hashlib.sha256(identity.encode()).hexdigest(), - targets=destinations, - schema_version=schema_version, - version=version, - ) - - -_INSERT_EVENT: Final = """ -INSERT INTO "LiteLLM_AccountingEvent" - (event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload) -VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb) -ON CONFLICT (event_id) DO NOTHING RETURNING event_id -""" -_INSERT_DELIVERIES: Final = """ -INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) -SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb) -""" -_LOCK_SOURCE: Final = """ -SELECT event_id FROM "LiteLLM_AccountingEvent" -WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE -""" -_EXISTING_COST: Final = """ -SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1 -""" -_LOCK_EVENT: Final = """ -SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE -""" -_INSERT_USAGE_BATCH: Final = """ -INSERT INTO "LiteLLM_AccountingEvent" - (event_id, kind, schema_version, payload_hash, occurred_at, payload) -SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload -FROM jsonb_to_recordset($1::jsonb) AS records( - event_id text, kind text, payload_hash text, occurred_at text, payload jsonb) -ORDER BY event_id -ON CONFLICT (event_id) DO NOTHING RETURNING event_id -""" -_LOCK_USAGE_BATCH: Final = """ -SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent" -WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) -ORDER BY event_id FOR UPDATE -""" -_LOAD_USAGE_BATCH: Final = """ -SELECT event.payload, ARRAY( - SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery - WHERE delivery.event_id = event.event_id ORDER BY target::text -) AS targets -FROM "LiteLLM_AccountingEvent" event -WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) - AND event.kind = 'outbound_usage' AND event.schema_version = 1 -""" -_INSERT_USAGE_DELIVERIES: Final = """ -INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) -SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1 -FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb) -CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination -""" -_COMPLETED_FAMILY: Final = """ -root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') -AND root.recorded_at < $1::timestamp -AND (root.kind = 'inbound_usage' OR EXISTS ( - SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost' -)) -AND NOT EXISTS ( - SELECT 1 FROM "LiteLLM_AccountingEvent" AS member - WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id)) - AND (member.recorded_at >= $1::timestamp - OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery - WHERE delivery.event_id = member.event_id AND delivery.state != 'applied')) -) -""" -_RETENTION_CANDIDATES: Final = ( - 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE ' - + _COMPLETED_FAMILY - + " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED" -) -_ELIGIBLE_GROUP: Final = ( - 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY -) -_DELETE_DELIVERIES: Final = """ -DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN ( - SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) -) -""" -_DELETE_EVENTS: Final = """ -DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) -""" - - -_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"] - - -async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict: - if record.source_event_id is not None: - source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id)) - if not source: - return "missing_source" - costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id))) - if any(cost["event_id"] != record.event_id for cost in costs): - return "competing_cost" - inserted: Final = _STRINGS.validate_python( - await db.query_raw( - _INSERT_EVENT, - record.event_id, - record.kind.value, - record.source_event_id, - record.schema_version, - record.payload_hash, - record.occurred_at.astimezone(timezone.utc).isoformat(), - record.payload, - ) - ) - if not inserted: - existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id))) - if not existing or existing[0]["payload_hash"] != record.payload_hash: - return "identity_conflict" - return "duplicate" - await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version) - return "inserted" - - -async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int: - if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)): - return 0 - await db.execute_raw(_DELETE_DELIVERIES, event_id) - return await db.execute_raw(_DELETE_EVENTS, event_id) - - -class AccountingLedger: - def __init__(self, transaction: Callable[[], _Transaction]) -> None: - self._transaction: Final = transaction - - @classmethod - def for_client(cls, client: PrismaClient) -> AccountingLedger: - from prisma.client import TransactionManager - - managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True)) - - @asynccontextmanager - async def transaction() -> AsyncGenerator[Prisma]: - primary: Final = writer_wrapper(client.db) - manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT)) - async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db: - await db.execute_raw("SET LOCAL statement_timeout = 5000") - await db.execute_raw("SET LOCAL lock_timeout = 1000") - yield db - - return cls(transaction) - - async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None: - records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes) - if not records: - return - unique: Final = {record.event_id: record for record in records} - if any(record.payload_hash != unique[record.event_id].payload_hash for record in records): - raise AccountingConflictError("Batch contains conflicting copies of a usage identity") - batch: Final = _json( - tuple( - { - "event_id": record.event_id, - "kind": record.kind.value, - "payload_hash": record.payload_hash, - "occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(), - "payload": _JSON.validate_json(record.payload), - } - for record in unique.values() - ) - ) - async with self._transaction() as db: - inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch)) - accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique)))) - if len(accepted) != len(unique) or any( - row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted - ): - raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set") - deliveries: Final = _json( - tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted) - ) - await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries) - - async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]: - if not event_ids: - return () - async with self._transaction() as db: - rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( - await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids)) - ) - return tuple( - UsageWrite( - OutboundUsageEvent.model_validate_json(_json(_JSON.validate_python(row["payload"])["event"])), - _json(row["payload"]), - TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]), - ) - for row in rows - ) - - async def insert_event( - self, - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...], - *, - envelope: str, - schema_version: int = 1, - version: int = 1, - ) -> Literal["inserted", "duplicate"]: - """Accept T5/T6's serialized envelope, containing its public event under `event`. - - The immutable acceptance hash includes the target set and delivery version. - Changing either for an accepted event conflicts. T7 must provide a separate - add-target or new-version workflow that preserves receipts and prevents - duplicate contributions; redelivery cannot silently create or reset work. - - The Prisma adapter opens one interactive transaction per event, including - two SET LOCAL statements and a source-row lock for costs. I/O is async, - but awaiting this method inline adds commit latency per event, potentially - several times per request. T5 owns bounded asynchronous capture and an - explicit batching strategy; looping over this method does not batch commits. - Success acknowledges commit, independently of HTTP response acknowledgement. - """ - record: Final = _prepare(event, targets, envelope, schema_version, version) - async with self._transaction() as db: - result: Final = await _accept(db, record) - match result: - case "inserted" | "duplicate": - return result - case "missing_source": - raise AccountingConflictError("Cost source must be a durably accepted outbound usage event") - case "competing_cost": - raise AccountingConflictError("Usage event already has an authoritative cost") - case "identity_conflict": - raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set") - case _: - assert_never(result) - - async def purge_completed( - self, - *, - redelivery_before: datetime | None = None, - recovery_before: datetime | None = None, - limit: int = 100, - ) -> int: - """Delete at most `limit` completed families once both retention owners permit it. - - T8 must supply a coherent recovery watermark including outstanding holds. - Omitting either watermark disables deletion. Each usage row and its cost - remain together so a retained charge never loses its source or vice versa. - Applied receipts are deleted with the family, ending its retained dedup - protection. The caller must guarantee events recorded before - redelivery_before cannot be redelivered through supported paths. - T8's recovery_before must also protect checkpoints, - outstanding holds and unfinished settlements; age alone is insufficient. - """ - if type(limit) is not int or not 1 <= limit <= 1000: - raise ValueError("Accounting retention batch size must be between 1 and 1000") - if redelivery_before is None or recovery_before is None: - return 0 - cutoff: Final = min( - _INSTANT.validate_python(redelivery_before).astimezone(timezone.utc), - _INSTANT.validate_python(recovery_before).astimezone(timezone.utc), - ).isoformat() - async with self._transaction() as db: - roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit))) - deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots]) - return sum(deleted) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 99b49331932..d2b69f24397 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -14,11 +14,13 @@ interval of increments queued elsewhere: the same lag the SpendLogs aggregate it replaces (and every other spend column) already has. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Final, Protocol +from pydantic import TypeAdapter from typing_extensions import LiteralString, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -74,30 +76,44 @@ _OPEN_WINDOW_SPEND_SQL: Final = ( "WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')" ) +_HAS_DELIVERY: Final = 'EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id)' + +_SETTLED_SPEND_LOG_FILTER: Final = ( + ' AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery ' + 'WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id ' + "AND delivery.state IN ('pending', 'parked'))" +) + _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, " - "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + "COALESCE(SUM(spend) FILTER (WHERE " + + _HAS_DELIVERY + + " OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) " + "AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, " - "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + "COALESCE(SUM(spend) FILTER (WHERE " + + _HAS_DELIVERY + + " OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) " + "AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) @@ -114,6 +130,10 @@ class _ExistingKeyRow(TypedDict): window_duration: ReadOnly[str] +_SEED_TOTALS_ROWS: Final = TypeAdapter(tuple[_SeedTotalsRow, ...]) +_EXISTING_KEY_ROWS: Final = TypeAdapter(tuple[_ExistingKeyRow, ...]) + + class _SeedTotalsQuery(Protocol): async def __call__(self, query: LiteralString, *args: object) -> Sequence[_SeedTotalsRow]: ... @@ -135,6 +155,8 @@ class _UpsertBatchManager(Protocol): class _UpsertTransaction(Protocol): def batch_(self) -> _UpsertBatchManager: ... + async def query_raw(self, query: LiteralString, *args: object) -> Sequence[Mapping[str, object]]: ... + class _UpsertTransactionManager(Protocol): async def __aenter__(self) -> _UpsertTransaction: ... @@ -178,6 +200,7 @@ class WindowSpendLogsAggregate(Protocol): entity_id: str, window_start: datetime, batch_started_at: datetime | None, + query: _SeedTotalsQuery | None = None, ) -> WindowSeedTotals | None: ... @@ -187,6 +210,7 @@ async def spend_logs_seed_totals( entity_id: str, window_start: datetime, batch_started_at: datetime | None, + query: _SeedTotalsQuery | None = None, ) -> WindowSeedTotals | None: """LiteLLM_SpendLogs spend for one entity since window_start, both in full and up to the start of the batch being flushed, in one scan. @@ -210,11 +234,12 @@ async def spend_logs_seed_totals( bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL else: return None + read: Final = query or _seed_totals_query(prisma_client) async with db_span("seed_window_spend", "LiteLLM_SpendLogs"): rows: Final = ( - await _seed_totals_query(prisma_client)(unbounded_sql, entity_id, window_start) + await read(unbounded_sql, entity_id, window_start) if batch_started_at is None - else await _seed_totals_query(prisma_client)( + else await read( bounded_sql, entity_id, window_start, @@ -246,9 +271,10 @@ def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]: async def _existing_primary_keys( prisma_client: "PrismaClient", transactions: tuple[WindowSpendTransaction, ...], + query: _ExistingKeysQuery | None = None, ) -> frozenset[tuple[str, str, str]]: async with db_span("select_window_spend_rows", "LiteLLM_BudgetWindowSpend"): - rows: Final = await _existing_keys_query(prisma_client)( + rows: Final = await (query or _existing_keys_query(prisma_client))( _SELECT_EXISTING_ROWS_SQL, tuple(transaction["entity_type"] for transaction in transactions), tuple(transaction["entity_id"] for transaction in transactions), @@ -262,6 +288,7 @@ async def _seed_base_for_missing_row( transaction: WindowSpendTransaction, existing_primary_keys: frozenset[tuple[str, str, str]], spend_logs_aggregate: WindowSpendLogsAggregate, + query: _SeedTotalsQuery | None = None, ) -> float: """Spend already recorded for a window that has no row yet. @@ -272,12 +299,24 @@ async def _seed_base_for_missing_row( """ if _primary_key(transaction) in existing_primary_keys: return 0.0 - totals: Final = await spend_logs_aggregate( - prisma_client=prisma_client, - entity_type=transaction["entity_type"], - entity_id=transaction["entity_id"], - window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), - batch_started_at=_transaction_started_at(transaction), + window_start: Final = datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc) + totals: Final = ( + await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=window_start, + batch_started_at=_transaction_started_at(transaction), + ) + if query is None + else await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=window_start, + batch_started_at=_transaction_started_at(transaction), + query=query, + ) ) if totals is None: return 0.0 @@ -324,10 +363,30 @@ def _upsert_params( ) +def _existing_keys_in(transaction: _UpsertTransaction) -> _ExistingKeysQuery: + async def query_raw(query: LiteralString, *args: object) -> Sequence[_ExistingKeyRow]: + return _EXISTING_KEY_ROWS.validate_python(await transaction.query_raw(query, *args)) + + return query_raw + + +def _seed_totals_in(transaction: _UpsertTransaction) -> _SeedTotalsQuery: + async def query_raw(query: LiteralString, *args: object) -> Sequence[_SeedTotalsRow]: + return _SEED_TOTALS_ROWS.validate_python(await transaction.query_raw(query, *args)) + + return query_raw + + +def _transaction_reads(transaction: _UpsertTransaction) -> tuple[_ExistingKeysQuery, _SeedTotalsQuery]: + """Reads made inside an enclosing transaction share its connection, so they cannot wait on it.""" + return _existing_keys_in(transaction), _seed_totals_in(transaction) + + async def commit_window_spend_updates( prisma_client: "PrismaClient", transactions: Sequence[WindowSpendTransaction], spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_seed_totals, + outer_transaction: _UpsertTransaction | None = None, ) -> None: """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. @@ -343,9 +402,11 @@ async def commit_window_spend_updates( return ordered: Final = tuple(sorted(transactions, key=window_spend_group_key)) + reads: Final = _transaction_reads(outer_transaction) if outer_transaction is not None else (None, None) existing_primary_keys: Final = await _existing_primary_keys( prisma_client=prisma_client, transactions=ordered, + query=reads[0], ) seed_bases: Final = tuple( [ @@ -354,6 +415,7 @@ async def commit_window_spend_updates( transaction=transaction, existing_primary_keys=existing_primary_keys, spend_logs_aggregate=spend_logs_aggregate, + query=reads[1], ) for transaction in ordered ] @@ -367,7 +429,9 @@ async def commit_window_spend_updates( ) async with ( db_span("commit_window_spend_updates", "LiteLLM_BudgetWindowSpend"), - _upsert_transaction(prisma_client) as db_transaction, + ( + nullcontext(outer_transaction) if outer_transaction is not None else _upsert_transaction(prisma_client) + ) as db_transaction, db_transaction.batch_() as batcher, ): for transaction, seed_base in zip(ordered, seed_bases): diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index eb130a5196f..2d1770f4c1d 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -55,7 +55,7 @@ _ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType( "tag": "request_tags", } ) -_TAGS: Final = TypeAdapter(tuple[str, ...]) +_TAGS: Final = TypeAdapter(tuple[object, ...]) def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]: @@ -67,7 +67,7 @@ def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEnti if value is None: return () tags: Final = _TAGS.validate_json(value) if isinstance(value, str) else _TAGS.validate_python(value) - return tuple(dict.fromkeys(tags)) + return tuple(dict.fromkeys(tag for tag in tags if isinstance(tag, str))) if value is None: return (None,) if entity == "user" else () if not isinstance(value, str) or (entity == "end_user" and not value): diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d9fb236632b..80f09df6c86 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,6 +12,7 @@ import json import os import random import time +import uuid from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from contextvars import ContextVar @@ -34,6 +35,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.batch_utils import BATCH_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils.core_helpers import finite_cost, single_log_line from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -92,8 +94,10 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import CallTypes +from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES if TYPE_CHECKING: + from litellm.accounting.projection import Contribution from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -158,23 +162,73 @@ def _org_member_transaction_key(org_id: str, user_id: str) -> str: return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}" -def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: - return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" +ACCOUNTING_LEDGER_KWARG: Final = "accounting_ledger" -_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) +def _persistable_api_key(token: str | None, logged_api_key: str) -> str | None: + from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash + from litellm.proxy.spend_tracking.spend_tracking_utils import redact_logged_api_key + + if token is None or token == logged_api_key or is_valid_sha256_hash(token): + return token + return redact_logged_api_key(token) -def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: - """Reduce a batch's cost row to what tells the retrieves apart when logging is off. +async def _track_guardrail_usage_of_cost_row( + payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool +) -> None: + """A job's cost row skips the spend log queue, so its guardrail runs are counted here instead of at the flush.""" + if disable_spend_logs: + return + from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage - A proxy run with spend logs disabled still needs one row per batch to charge it once, + try: + await process_spend_logs_guardrail_usage(prisma_client=prisma_client, logs_to_process=(payload,)) + except Exception as e: # noqa: BLE001 # guardrail usage is best effort, as at the spend log flush + verbose_proxy_logger.warning( + "Guardrail usage tracking failed for spend row %s (non-fatal): %s", payload["request_id"], e + ) + + +def _is_settle_once_cost_row(payload: SpendLogsPayload) -> bool: + """A poll of a job whose cost is charged once, by the first poll that sees it, however often it is polled. + + A batch's cost is known only when a retrieve sees it finished (LIT-7048). A video's cost is + charged at create by the built-in providers, but a custom provider may report it on the + status poll that sees the video completed. Only a row carrying the job's cost suffix settles + the job, so a poll that pays only its guardrails, under an id that client call ids cannot + turn into the job's, is charged like any other request. + """ + if payload.get("status") != "success": + return False + call_type: Final = payload.get("call_type") + request_id: Final = payload["request_id"] + if call_type == CallTypes.aretrieve_batch.value: + return request_id.endswith(BATCH_COST_REQUEST_ID_SUFFIX) + return call_type in VIDEO_STATUS_CALL_TYPES and request_id.endswith(VIDEO_COST_REQUEST_ID_SUFFIX) + + +_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) + + +def _cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: + """Reduce a job's cost row to what tells its polls apart when logging is off. + + A legacy proxy run with spend logs disabled still needs one row per job to charge it once, so the row is written either way, but it carries no request of its own: no metadata, - no requester IP, no key, model, or token counts (LIT-7048). + no requester IP, no key, model, or token counts (LIT-7048). A ledger batch is deduplicated + by its delivery row instead, so it writes no local row when logging is off. """ if disable_spend_logs is False: return payload - return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) + if "accounting_payload" in payload: + return { + "request_id": payload["request_id"], + "call_type": payload["call_type"], + "accounting_payload": payload["accounting_payload"], + "accounting_log_disabled": True, + } + return MappingProxyType({field: value for field, value in payload.items() if field in _COST_CLAIM_FIELDS}) class _SpendIncrement(TypedDict): @@ -187,8 +241,17 @@ class _MemberSpendRow(TypedDict): cost: ReadOnly[float] +class SpendUpsertTable(Protocol): + def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... + + class _SpendBatch(Protocol): litellm_usertable: BatchTable + litellm_endusertable: SpendUpsertTable + litellm_dailymodelusage: SpendUpsertTable + + def execute_raw(self, query: LiteralString, *args: object) -> None: ... + litellm_verificationtoken: BatchTable litellm_teamtable: BatchTable litellm_teammembership: BatchTable @@ -219,14 +282,14 @@ def _entity_spend_table(batcher: _SpendBatch, table_accessor: _EntitySpendTable) return _ENTITY_SPEND_TABLES[table_accessor](batcher) -class _SpendBatchManager(Protocol): +class SpendBatchManager(Protocol): async def __aenter__(self) -> _SpendBatch: ... async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... class _SpendTransaction(Protocol): - def batch_(self) -> _SpendBatchManager: ... + def batch_(self) -> SpendBatchManager: ... async def execute_raw(self, query: LiteralString, *args: object) -> int: ... @@ -381,8 +444,14 @@ _ENTITY_SPEND_MODELS: Final[Mapping[_EntitySpendTable, str]] = MappingProxyType( @asynccontextmanager async def _spend_update_tx( - prisma_client: PrismaClient, table: str, call_type: str = "commit_spend_updates" + prisma_client: PrismaClient, + table: str, + call_type: str = "commit_spend_updates", + outer_transaction: _SpendTransaction | None = None, ) -> AsyncGenerator[_SpendTransaction]: + if outer_transaction is not None: + yield outer_transaction + return tx: Final[_SpendTransactionManager] = prisma_client.db.tx(timeout=timedelta(seconds=60)) async with db_span(call_type, table), tx as transaction: await apply_rollup_lock_timeout(transaction) @@ -569,6 +638,21 @@ def debitable_model_access_groups( return tuple(group for group in ordered if group in served) +def _billable_cost(response_cost: float | None, payload: SpendLogsPayload) -> float | None: + billable_cost: Final = finite_cost(response_cost) + if response_cost is not None and billable_cost is None: + verbose_proxy_logger.warning( + "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", + response_cost, + single_log_line(payload["request_id"]), + ) + return billable_cost + + +def _iso_timestamp(value: datetime | str) -> str: + return value.isoformat() if isinstance(value, datetime) else value + + class DBSpendUpdateWriter: """ Module responsible for @@ -613,10 +697,11 @@ class DBSpendUpdateWriter: ) -> bool: """Record the request's spend, answering whether its cost still needs charging. - False only for a batch retrieve whose cost row another retrieve already wrote, + False only for a batch or video poll whose cost row another poll already wrote, so the caller leaves the key, team, and user counters alone (LIT-7048). """ from litellm.proxy.proxy_server import ( + accounting_window_keys, disable_spend_logs, general_settings, # pyright: ignore[reportUnknownVariableType] # legacy proxy configuration is narrowed below litellm_proxy_budget_name, @@ -624,6 +709,13 @@ class DBSpendUpdateWriter: ) from litellm.proxy.utils import ProxyUpdateSpend, hash_token + proxy_settings: Final = cast( # cast-ok: proxy settings use string keys + Mapping[str, object], general_settings + ) + use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True + request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {}) + if kwargs is not None: + kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker try: verbose_proxy_logger.debug( "Enters prisma db call, response_cost: %s, token: %s; user_id: %s; team_id: %s", @@ -634,10 +726,9 @@ class DBSpendUpdateWriter: ) if ProxyUpdateSpend.disable_spend_updates() is True: return True - if token is not None and isinstance(token, str) and token.startswith("sk-"): - hashed_token = hash_token(token=token) - else: - hashed_token = token + hashed_token: Final = ( + hash_token(token=token) if isinstance(token, str) and token.startswith("sk-") else token + ) ## CREATE SPEND LOG PAYLOAD ## from litellm.proxy.spend_tracking.spend_tracking_utils import ( @@ -653,32 +744,42 @@ class DBSpendUpdateWriter: end_time=end_time, llm_router=llm_router, ) - billable_cost: Final = finite_cost(response_cost) - if response_cost is not None and billable_cost is None: - verbose_proxy_logger.warning( - "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", - response_cost, - single_log_line(payload["request_id"]), - ) + billable_cost: Final = _billable_cost(response_cost, payload) payload["spend"] = billable_cost or 0.0 - if isinstance(payload["startTime"], datetime): - payload["startTime"] = payload["startTime"].isoformat() - if isinstance(payload["endTime"], datetime): - payload["endTime"] = payload["endTime"].isoformat() - - if org_id is not None and org_id != "": + payload["startTime"] = _iso_timestamp(payload["startTime"]) + payload["endTime"] = _iso_timestamp(payload["endTime"]) + if org_id: payload["organization_id"] = org_id - - if team_id is not None and team_id != "": + if team_id: payload["team_id"] = team_id - error_log_settings: Final = cast( # cast-ok: proxy settings use string keys - Mapping[str, object], general_settings - ) + if use_ledger: + from litellm.accounting.snapshot import build_accounting_payload + + if not _is_settle_once_cost_row(payload): + payload["accounting_delivery_id"] = uuid.uuid4().hex + payload["accounting_payload"] = build_accounting_payload( + payload, + entities={ + "api_key": _persistable_api_key(hashed_token, payload.get("api_key") or ""), + "user": user_id, + "end_user": end_user_id, + "team_id": team_id, + "organization_id": org_id, + "project_id": project_id, + "agent_id": payload["agent_id"], + "proxy_budget_name": litellm_proxy_budget_name if litellm.max_budget > 0 else None, + }, + model_access_groups=debitable_model_access_groups( + get_request_model_access_groups(request_fields), payload["model_id"], llm_router + ), + windows=await accounting_window_keys(hashed_token, team_id) if billable_cost else (), + router=get_llm_router, + ) keep_spend_logs: Final = disable_spend_logs is False and not ( "status" in payload and payload["status"] == "failure" - and error_log_settings.get("disable_error_logs") is True + and proxy_settings.get("disable_error_logs") is True ) if not await self._record_spend_log( payload=payload, @@ -707,24 +808,26 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) - spawn_logging_task( - asyncio.get_running_loop(), - self._batch_database_updates( - response_cost=billable_cost, - user_id=user_id, - hashed_token=hashed_token, - team_id=team_id, - org_id=org_id, - project_id=project_id, - end_user_id=end_user_id, - prisma_client=prisma_client, - litellm_proxy_budget_name=litellm_proxy_budget_name, - payload=payload, - request_model_access_groups=get_request_model_access_groups(kwargs), - ), - task_name="spend rollups", - ) - + if not use_ledger: + spawn_logging_task( + asyncio.get_running_loop(), + self._batch_database_updates( + response_cost=billable_cost, + user_id=user_id, + hashed_token=hashed_token, + team_id=team_id, + org_id=org_id, + project_id=project_id, + end_user_id=end_user_id, + prisma_client=prisma_client, + litellm_proxy_budget_name=litellm_proxy_budget_name, + payload=payload, + request_model_access_groups=get_request_model_access_groups(request_fields), + ), + task_name="spend rollups", + ) + else: + self._record_request_error_safely(payload=payload, prisma_client=prisma_client) self._enqueue_tool_registry_upsert( kwargs=kwargs, completion_response=completion_response, @@ -755,24 +858,32 @@ class DBSpendUpdateWriter: disable_spend_logs: bool, llm_router: "Router | None" = None, ) -> bool: - if prisma_client is not None and _is_batch_cost_row(payload): - return await self._claim_batch_cost_spend_log( + if prisma_client is not None and _is_settle_once_cost_row(payload): + return await self._claim_settle_once_cost_row( payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs ) - if disable_spend_logs is False: + if "accounting_payload" in payload: + await self._insert_spend_log_to_db( + payload={**payload, "accounting_log_disabled": disable_spend_logs}, + prisma_client=prisma_client, + llm_router=llm_router, + ) + elif disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router) return True - async def _claim_batch_cost_spend_log( + async def _claim_settle_once_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool ) -> bool: - """Write the batch's cost row now, or learn that another retrieve already did. + """Write the job's cost row now, or learn that another poll already did. - Every retrieve of one batch shares this row, so the insert that lands first owns + Every poll of one batch or video shares this row, so the insert that lands first owns the charge and every later one finds the row and charges nothing (LIT-7048). Only - a row that recorded a charge counts: a failed retrieve, a request whose client - picked the batch id as its call id, and the $0 row an older proxy left behind - while the batch was still running all leave the charge to be made. + a row that recorded a charge counts: a failed poll, a request whose client + picked the job id as its call id, and the $0 row a poll left behind + while the job was still running all leave the charge to be made. A ledger job + is owned by its delivery row, keyed by its cost row id, in every logging mode, and a + $0 poll records no delivery so it cannot block the poll that carries the cost. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -783,43 +894,59 @@ class DBSpendUpdateWriter: ) row: Final = spend_log_row_with_retained_metadata( - _batch_cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields() + _cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields() ) - spend_logs: Final = SpendLogsRepository(prisma_client).table + repository: Final = SpendLogsRepository(prisma_client) try: - claimed: Final = await spend_logs.create_many( + if "accounting_payload" in payload: + if payload["spend"] <= 0: + return True + accepted: Final = await repository.insert_with_deliveries([prisma_client.jsonify_object(row)]) == 1 + if accepted: + await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs) + return accepted + if await repository.has_delivery(request_id): + return False + claimed: Final = await repository.table.create_many( data=[prisma_client.jsonify_object(row)], skip_duplicates=True, ) - if claimed == 1: - return True - existing: Final = await spend_logs.find_unique(where={"request_id": request_id}) + existing: Final = ( + None if claimed == 1 else await repository.table.find_unique(where={"request_id": request_id}) + ) except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log verbose_proxy_logger.warning( - "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e + "Could not claim spend row %s for a job's cost, queueing it: %s", request_id, e ) await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) return True - if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": + if claimed == 1: + await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs) + return True + if existing is None or existing.call_type != payload["call_type"] or existing.status != "success": verbose_proxy_logger.warning( - "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + "Spend row %s belongs to a %s request, so this job's cost is charged without a row of its own", request_id, getattr(existing, "call_type", None), ) return True if existing.spend > 0: - verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this job", request_id) return False - return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) + taken_over: Final = await self._take_over_uncharged_cost_row( + payload=payload, prisma_client=prisma_client, row=row + ) + if taken_over: + await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs) + return taken_over - async def _take_over_uncharged_batch_cost_row( + async def _take_over_uncharged_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: - """Take the batch's cost row over from the poll that left it charging nothing. + """Take the job's cost row over from the poll that left it charging nothing. - A pre-upgrade proxy wrote that row every time it polled the batch while it was still - running, so the charge is still to be made and the row still has to end up carrying - it. The row stops matching the moment it carries a charge, so it is one retrieve that + A poll that saw the job still running wrote that row at $0, so the charge is still to be + made and the row still has to end up carrying it. The row stops matching the moment it carries a charge, so it is one poll that takes it over and charges, and every later one reads the charge and charges nothing. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -827,7 +954,7 @@ class DBSpendUpdateWriter: request_id: Final = payload["request_id"] if payload["spend"] <= 0: verbose_proxy_logger.debug( - "Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id + "Cost tracking skipped: this job costs nothing and spend row %s says so", request_id ) return False try: @@ -837,18 +964,18 @@ class DBSpendUpdateWriter: ), where={ "request_id": request_id, - "call_type": CallTypes.aretrieve_batch.value, + "call_type": payload["call_type"], "status": "success", "spend": 0.0, }, ) except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over verbose_proxy_logger.warning( - "Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e + "Could not take over spend row %s, leaving this job's cost to the next poll: %s", request_id, e ) return False if taken_over == 0: - verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this job", request_id) return False return True @@ -1231,10 +1358,7 @@ class DBSpendUpdateWriter: self.add_spend_log_transaction_to_daily_team_transaction(payload=payload_copy, prisma_client=prisma_client), payload=payload_copy, ) - try: - self._record_request_error(payload=payload_copy, prisma_client=prisma_client) - except Exception as e: # noqa: BLE001 # the rollup must never skip the sibling spend writes - _report_spend_update_step_failure("_record_request_error", payload=payload_copy, exc=e) + self._record_request_error_safely(payload=payload_copy, prisma_client=prisma_client) await self._run_spend_update_step( "add_spend_log_transaction_to_daily_org_transaction", self.add_spend_log_transaction_to_daily_org_transaction( @@ -1564,7 +1688,7 @@ class DBSpendUpdateWriter: async def _insert_spend_log_to_db( self, - payload: dict | SpendLogsPayload, + payload: Mapping[str, object], prisma_client: PrismaClient | None = None, spend_logs_url: str | None = os.getenv("SPEND_LOGS_URL"), llm_router: "Router | None" = None, @@ -1608,6 +1732,13 @@ class DBSpendUpdateWriter: else: - Regular flow of this method """ + from litellm.accounting.worker import flush_accounting_deliveries + + try: + await flush_accounting_deliveries(self, prisma_client, proxy_logging_obj) + except Exception: # noqa: BLE001 # pending deliveries survive outages while legacy queues must still drain + verbose_proxy_logger.exception("Accounting delivery flush failed; pending work will retry") + if RedisUpdateBuffer.should_commit_spend_updates_to_redis(): await self._commit_spend_updates_to_db_with_redis( prisma_client=prisma_client, @@ -2134,6 +2265,39 @@ class DBSpendUpdateWriter: _handle_spend_update_failure = handle_spend_update_failure + @staticmethod + async def invalidate_accounting_entities( + contributions: tuple["Contribution", ...], proxy_logging_obj: ProxyLogging + ) -> None: + await DBSpendUpdateWriter._invalidate_project_caches( + tuple(dict(item.key)["project_id"] for item in contributions if item.table == "LiteLLM_ProjectTable"), + proxy_logging_obj, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + details: Final = _SPEND_METADATA_ADAPTER.validate_python( + proxy_logging_obj.call_details # pyright: ignore[reportUnknownMemberType] # legacy logging metadata is validated at this boundary + ) + user_api_key_cache: Final = details.get("user_api_key_cache") + if isinstance(user_api_key_cache, UserApiKeyCache): + with service_target(AUTH_OBJECTS_TARGET): + for item in contributions: + if item.table == "LiteLLM_TeamMembership": + await user_api_key_cache.async_delete_cache( + key=f"team_membership:{dict(item.key)['user_id']}:{dict(item.key)['team_id']}" + ) + + async def apply_accounting_entities( + self, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + transaction: _SpendTransaction, + updates: DBSpendUpdateTransactions, + ) -> None: + await self._commit_spend_updates_to_db( + prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction + ) + async def _commit_spend_updates_to_db( self, prisma_client: PrismaClient, @@ -2141,6 +2305,7 @@ class DBSpendUpdateWriter: proxy_logging_obj: ProxyLogging, db_spend_update_transactions: DBSpendUpdateTransactions, on_table_committed: Callable[[_SpendTableName], None] | None = None, + outer_transaction: _SpendTransaction | None = None, ): """ Commits all the spend `UPDATE` transactions to the Database @@ -2155,7 +2320,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_UserTable") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_UserTable", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by ID for consistent lock ordering across pods to prevent deadlocks. # batch_() issues statements sequentially within the tx, so iteration @@ -2167,6 +2334,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2186,6 +2355,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, end_user_list_transactions=end_user_list_transactions, + outer_transaction=outer_transaction, ) if on_table_committed is not None: on_table_committed("end_user_list_transactions") @@ -2196,7 +2366,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_VerificationToken") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by token for consistent lock ordering across pods to prevent deadlocks. for token, response_cost in sorted(key_list_transactions.items()): @@ -2211,6 +2383,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2228,7 +2402,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_TeamTable") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. for team_id, response_cost in sorted(team_list_transactions.items()): @@ -2241,6 +2417,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2266,11 +2444,15 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_TeamMembership") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamMembership", outer_transaction=outer_transaction + ) as transaction: await _write_team_member_spend(transaction, team_member_list_transactions) # Transaction succeeded, break out of retry loop break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2283,7 +2465,7 @@ class DBSpendUpdateWriter: # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database - if team_memberships_to_invalidate and proxy_logging_obj is not None: + if outer_transaction is None and team_memberships_to_invalidate and proxy_logging_obj is not None: user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache") if user_api_key_cache is not None: with service_target(AUTH_OBJECTS_TARGET): @@ -2303,7 +2485,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_OrganizationTable") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. for org_id, response_cost in sorted(org_list_transactions.items()): @@ -2313,6 +2497,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2330,7 +2516,9 @@ class DBSpendUpdateWriter: start_time = time.time() try: async with ( - _spend_update_tx(prisma_client, "LiteLLM_OrganizationMembership") as transaction, + _spend_update_tx( + prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction + ) as transaction, transaction.batch_() as batcher, ): for key, response_cost in sorted(org_member_list_transactions.items()): @@ -2341,6 +2529,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2356,6 +2546,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Project", transactions=project_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_projecttable", where_field="project_id", n_retry_times=n_retry_times, @@ -2364,16 +2555,18 @@ class DBSpendUpdateWriter: ) if on_table_committed is not None: on_table_committed("project_list_transactions") - await DBSpendUpdateWriter._invalidate_project_caches( - project_ids=tuple(project_list_transactions or ()), - proxy_logging_obj=proxy_logging_obj, - ) + if outer_transaction is None: + await DBSpendUpdateWriter._invalidate_project_caches( + project_ids=tuple(project_list_transactions or ()), + proxy_logging_obj=proxy_logging_obj, + ) ### UPDATE TAG TABLE ### tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Tag", transactions=tag_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_tagtable", where_field="tag_name", n_retry_times=n_retry_times, @@ -2390,6 +2583,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Model access group", transactions=model_access_group_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_modelaccessgroupbudgettable", where_field="access_group_name", n_retry_times=n_retry_times, @@ -2404,6 +2598,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Agent", transactions=agent_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_agentstable", where_field="agent_id", n_retry_times=n_retry_times, @@ -2433,6 +2628,7 @@ class DBSpendUpdateWriter: n_retry_times: int, prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, + outer_transaction: _SpendTransaction | None = None, ): """ Helper function to update spend for any entity type (team, org, tag, etc). @@ -2445,13 +2641,16 @@ class DBSpendUpdateWriter: n_retry_times: Number of retries on failure prisma_client: Prisma client instance proxy_logging_obj: Proxy logging object + outer_transaction: Existing ledger transaction owning the complete delivery update """ verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions) if transactions is not None and len(transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, _ENTITY_SPEND_MODELS[table_accessor]) as transaction: + async with _spend_update_tx( + prisma_client, _ENTITY_SPEND_MODELS[table_accessor], outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by entity_id for consistent lock ordering across pods to prevent deadlocks. for entity_id, response_cost in sorted(transactions.items()): @@ -2468,6 +2667,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await DBSpendUpdateWriter.handle_spend_update_failure( e=e, attempt=i, @@ -2801,6 +3002,12 @@ class DBSpendUpdateWriter: entity_id_field="tag", ) + def _record_request_error_safely(self, *, payload: SpendLogsPayload, prisma_client: PrismaClient | None) -> None: + try: + self._record_request_error(payload=payload, prisma_client=prisma_client) + except Exception as error: # noqa: BLE001 # error rollup failure must not skip sibling writes + _report_spend_update_step_failure("_record_request_error", payload=payload, exc=error) + @staticmethod def _record_request_error(*, payload: SpendLogsPayload, prisma_client: PrismaClient | None) -> None: if prisma_client is None: diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index e9ec040023c..abbeab7f89d 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -55,3 +55,10 @@ class ToolDiscoveryQueue: """Return and clear all pending rows.""" items, self._pending = list(self._pending.values()), {} return items + + def requeue(self, items: Iterable[ToolDiscoveryQueueItem]) -> None: + """Put rows from a failed write back, ahead of and merged with anything queued since the flush.""" + since_flush: Final = self._pending + self._pending = {} + for item in (*items, *since_flush.values()): + self.add_update(item) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 8fbc8708f8c..c2d58512fea 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -31,10 +31,29 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( "too many connections for database", "remaining connection slots are reserved", ) +_NOT_ACCEPTING_CONNECTIONS_SQLSTATE: Final = "57P03" +_NOT_ACCEPTING_CONNECTIONS_PHRASES: Final = ( + "the database system is starting up", + "the database system is shutting down", + "the database system is not yet accepting connections", + "the database system is not accepting connections", + "the database system is in recovery mode", +) +_SERVER_TERMINATED_SQLSTATES: Final = frozenset({"57P01", "57P02", "57P05"}) +_SERVER_TERMINATED_PHRASES: Final = ( + "terminating connection due to administrator command", + "terminating connection due to immediate shutdown command", + "terminating connection because of unexpected sigquit signal", + "terminating connection because of crash of another server process", + "terminating connection due to idle-session timeout", +) +_UNREACHABLE_PRISMA_CODE: Final = "P1001" +_UNREACHABLE_PHRASES: Final = ("can't reach database server", "cannot reach database server") _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" _STATEMENT_TIMEOUT_SQLSTATE: Final = "57014" +_REJECTED_ROW_SQLSTATE_CLASSES: Final = frozenset({"22", "23", "54"}) _STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout" @@ -139,6 +158,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @staticmethod @@ -201,6 +222,20 @@ class PrismaDBExceptionHandler: return type(e) is prisma.errors.DataError + @staticmethod + def is_rejected_row_error(e: Exception) -> bool: + """True iff Postgres refused the rows themselves: a base ``DataError``, or a raw query failing with a data + exception (SQLSTATE class 22), integrity violation (class 23) or a value too large for an index (class 54). Systemic raw failures such as a missing + table (42P01) or a deadlock (40P01) stay False so the batch is retried instead of dropped.""" + import prisma + + if PrismaDBExceptionHandler.is_prisma_data_error(e): + return True + if not isinstance(e, _exception_types(prisma.errors.RawQueryError)): + return False + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is not None and sqlstate[:2] in _REJECTED_ROW_SQLSTATE_CLASSES + @staticmethod def is_database_transport_error(e: Exception) -> bool: """ @@ -289,6 +324,39 @@ class PrismaDBExceptionHandler: error_message: Final = str(e).lower() return any(phrase in error_message for phrase in _CONNECTION_CAPACITY_PHRASES) + @staticmethod + def is_database_not_accepting_connections_error(e: Exception) -> bool: + """True iff Postgres refused the connection because it is starting up, shutting + down or in crash recovery (SQLSTATE 57P03). The refusal comes before any statement + runs, so nothing the write carried was applied.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) == _NOT_ACCEPTING_CONNECTIONS_SQLSTATE: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _NOT_ACCEPTING_CONNECTIONS_PHRASES) + + @staticmethod + def is_database_lifecycle_error(e: Exception) -> bool: + """True iff Postgres failed the statement because of a server lifecycle event: + a refused connection during startup, shutdown or recovery (57P03), or a session + the server terminated (57P01, 57P02, 57P05). The server comes back on its own, + so the rows are not at fault. A terminated session may have committed before it + was cut, so only ``is_database_not_accepting_connections_error`` proves a + non-idempotent write never applied.""" + import prisma + + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) in _SERVER_TERMINATED_SQLSTATES: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _SERVER_TERMINATED_PHRASES) + @staticmethod def is_lock_timeout_error(e: Exception) -> bool: """True iff Postgres cancelled a statement for exceeding ``lock_timeout`` @@ -406,6 +474,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_lifecycle_error(e): + return True if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): return True if "cached plan must not change result type" in str(e).lower(): @@ -680,3 +750,18 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool: return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + +def db_write_never_sent(e: Exception) -> bool: + """Whether the database could not take the write at all, because it was unreachable (P1001), + refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again.""" + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if PrismaDBExceptionHandler.is_database_capacity_error(e): + return True + if _prisma_error_code(e) == _UNREACHABLE_PRISMA_CODE: + return True + if not PrismaDBExceptionHandler.is_prisma_error(e): + return False + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _UNREACHABLE_PHRASES) diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index 4566b10b260..f8cdabe9296 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -32,14 +32,15 @@ class _UpsertTable(Protocol): def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... -class _ModelUsageBatch(Protocol): - litellm_dailymodelusage: _UpsertTable +class ModelUsageBatch(Protocol): + @property + def litellm_dailymodelusage(self) -> _UpsertTable: ... def execute_raw(self, query: LiteralString, *args: object) -> None: ... -class _ModelUsageBatchManager(Protocol): - async def __aenter__(self) -> _ModelUsageBatch: ... +class ModelUsageBatchManager(Protocol): + async def __aenter__(self) -> ModelUsageBatch: ... async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... @@ -88,14 +89,14 @@ def _internal_call_origin(metadata: str) -> object: return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) -def _date_from_start_time(start_time: datetime | str) -> str | None: +def date_from_start_time(start_time: datetime | str) -> str | None: if isinstance(start_time, datetime): return start_time.date().isoformat() return start_time[:10] if len(start_time) >= 10 else None def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None: - date: Final = _date_from_start_time(payload["startTime"]) + date: Final = date_from_start_time(payload["startTime"]) origin: Final = _internal_call_origin(payload["metadata"]) if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): return None @@ -116,8 +117,8 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa ) -def _model_usage_batch(prisma_client: PrismaClient) -> _ModelUsageBatchManager: - batch: Final[_ModelUsageBatchManager] = prisma_client.db.batch_() +def _model_usage_batch(prisma_client: PrismaClient) -> ModelUsageBatchManager: + batch: Final[ModelUsageBatchManager] = prisma_client.db.batch_() return batch @@ -156,26 +157,16 @@ async def flush_model_usage_transactions( "custom_llm_provider": key.custom_llm_provider, "task_type": key.task_type, } - batcher.litellm_dailymodelusage.upsert( - where={"date_model_group_model_custom_llm_provider_task_type": key_fields}, - data={ - "create": { - **key_fields, - "spend": spend, - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "request_count": len(requests), - "successful_requests": successful, - "failed_requests": failed, - }, - "update": { - "spend": {"increment": spend}, - "prompt_tokens": {"increment": prompt_tokens}, - "completion_tokens": {"increment": completion_tokens}, - "request_count": {"increment": len(requests)}, - "successful_requests": {"increment": successful}, - "failed_requests": {"increment": failed}, - }, + queue_model_usage_upsert( + batcher, + key_fields, + { + "spend": spend, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "request_count": len(requests), + "successful_requests": successful, + "failed_requests": failed, }, ) return @@ -183,3 +174,15 @@ async def flush_model_usage_transactions( if attempt >= n_retry_times: raise await asyncio.sleep(2.0**attempt + random.uniform(0, 1)) + + +def queue_model_usage_upsert( + batcher: ModelUsageBatch, key_fields: Mapping[str, str], increments: Mapping[str, int | float] +) -> None: + batcher.litellm_dailymodelusage.upsert( + where={"date_model_group_model_custom_llm_provider_task_type": key_fields}, + data={ + "create": {**key_fields, **increments}, + "update": {name: {"increment": value} for name, value in increments.items()}, + }, + ) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index f7ec88945a1..9600fb40ad3 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -23,7 +23,11 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, aggregate_tool_discoveries, ) -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.exception_handler import ( + call_with_db_reconnect_retry, + db_write_failure_is_requeue_safe, + db_write_never_sent, +) from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions @@ -134,8 +138,19 @@ async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discove return try: await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail - verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + except Exception as e: # noqa: BLE001 # the rows go back to the queue or are dropped, the scheduler job must not fail + if db_write_never_sent(e): + tool_discovery_queue.requeue(items) + verbose_proxy_logger.warning( + "Tool discovery - %d tool rows went back to the queue for the next tick. Error: %s", len(items), str(e) + ) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) async def commit_tool_discoveries( diff --git a/litellm/proxy/enterprise_billing/billing_metrics.py b/litellm/proxy/enterprise_billing/billing_metrics.py index 166ece28b83..b6f69e8233c 100644 --- a/litellm/proxy/enterprise_billing/billing_metrics.py +++ b/litellm/proxy/enterprise_billing/billing_metrics.py @@ -2,29 +2,33 @@ Push-based OTLP metering for enterprise litellm deployments. Owns a dedicated OpenTelemetry meter provider and an OTLP/HTTP exporter -authenticated to our global collector with a TLS client certificate. The -collector front end terminates mutual TLS: the client certificate presented -here is validated against our CA at the edge, and the verified subject is -what identifies the deployment. It is intentionally isolated from the global -meter provider so the customer's own OTEL metrics are untouched and ours -never leak into their backend. +authenticated to our global collector with a Cognito bearer token. It is +intentionally isolated from the global meter provider so the customer's own +OTEL metrics are untouched and ours never leak into their backend. -The deployment's identity rides on the TLS client certificate, not on the -payload; the secret license key is never sent as an attribute or header. +The secret license key is never sent as an attribute or header. """ import os -import tempfile -from dataclasses import dataclass -from typing import TYPE_CHECKING, Final, Optional +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypedDict +import requests from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter from opentelemetry.metrics import Counter from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader from opentelemetry.sdk.resources import Resource +from pydantic import BaseModel, ConfigDict +from requests.auth import AuthBase +from typing_extensions import ReadOnly, Unpack from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.url_utils import is_https_url_with_host from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableCategory, ) @@ -33,23 +37,15 @@ if TYPE_CHECKING: from litellm.proxy._types import EnterpriseLicenseData ENDPOINT_ENV: Final = "LITELLM_BILLING_METRICS_ENDPOINT" -CLIENT_CERT_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_CERT" -CLIENT_KEY_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_KEY" -CA_CERT_ENV: Final = "LITELLM_BILLING_METRICS_CA_CERT" +CLIENT_ID_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_ID" +CLIENT_SECRET_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_SECRET" +TOKEN_URL_ENV: Final = "LITELLM_BILLING_METRICS_TOKEN_URL" EXPORT_INTERVAL_ENV: Final = "LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS" DEFAULT_EXPORT_INTERVAL_MS: Final = 60_000 SHUTDOWN_FLUSH_TIMEOUT_MS: Final = 5_000 +_TOKEN_REFRESH_SKEW_SECONDS: Final = 60 _METRICS_PATH: Final = "/v1/metrics" -# The cert env vars take a path or the PEM itself. Secret stores that inject -# values as env content cannot mount them as files, so inline PEM is written out. -_PEM_PREFIX: Final = "-----BEGIN" -_PEM_DIR_PREFIX: Final = "litellm-billing-mtls-" -_PEM_FILE_MODE: Final = 0o600 -_CLIENT_CERT_FILENAME: Final = "client.crt" -_CLIENT_KEY_FILENAME: Final = "client.key" -_CA_CERT_FILENAME: Final = "ca.crt" - METRIC_NAME: Final = "litellm.enterprise.billable_requests" METER_NAME: Final = "litellm.enterprise.billing" @@ -59,21 +55,28 @@ AttributeValue = str | int @dataclass(frozen=True, slots=True) class BillingMetricsConfig: endpoint: str - client_cert_path: str - client_key_path: str - ca_cert_path: str | None + token_url: str + client_id: str + client_secret: str = field(repr=False) export_interval_ms: int litellm_version: str license_id: str | None +class _BillingMetricsOAuthToken(BaseModel): + model_config = ConfigDict(frozen=True, strict=True) + + access_token: str + expires_in: int + + def _metrics_endpoint(endpoint: str) -> str: """The OTLP/HTTP metric exporter wants the full URL including the signal path.""" trimmed: Final = endpoint.rstrip("/") return trimmed if trimmed.endswith(_METRICS_PATH) else f"{trimmed}{_METRICS_PATH}" -def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]: +def resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]: base: Final[dict[str, AttributeValue]] = { "service.name": "litellm-proxy", "litellm.version": config.litellm_version, @@ -82,7 +85,7 @@ def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeVal return {**base, **license_attr} -def _billable_attributes( +def billable_attributes( category: BillableCategory, route: str, status_code: int, model_id: str | None ) -> dict[str, AttributeValue]: base: Final[dict[str, AttributeValue]] = { @@ -94,23 +97,95 @@ def _billable_attributes( return {**base, **model_attr} -def build_mtls_meter_provider(config: BillingMetricsConfig) -> MeterProvider: - """OTLP/HTTP exporter presenting a TLS client certificate. +class BillingMetricsBearerAuth(AuthBase): + def __init__( + self, + *, + client_id: str, + client_secret: str, + token_url: str, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._client_id: Final = client_id + self._client_secret: Final = client_secret + self._token_url: Final = token_url + self._clock: Final = clock + self._lock: Final = threading.Lock() + self._cached_token: str | None = None + self._refresh_at: float = 0.0 - The collector's load balancer terminates mutual TLS and validates the client - certificate against our CA. Server verification uses the system trust store - (the collector presents a public web-PKI certificate); ca_cert_path overrides - it only for private/test collectors. - """ - exporter: Final = OTLPMetricExporter( - endpoint=_metrics_endpoint(config.endpoint), - # None -> exporter falls back to the system trust store. - certificate_file=config.ca_cert_path, - client_certificate_file=config.client_cert_path, - client_key_file=config.client_key_path, + def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: + token: Final = self._get_token() + authenticated_request: Final = request.copy() + authenticated_request.headers["Authorization"] = f"Bearer {token}" + return authenticated_request + + def invalidate(self) -> None: + with self._lock: + self._cached_token = None + self._refresh_at = 0.0 + + def _get_token(self) -> str: + with self._lock: + now: Final = self._clock() + if self._cached_token is not None and now < self._refresh_at: + return self._cached_token + token, cache_seconds = self._fetch_token() + self._cached_token = token + self._refresh_at = now + cache_seconds + return token + + def _fetch_token(self) -> tuple[str, float]: + try: + # The OTel periodic exporter's background thread makes blocking token I/O safe here. + response: Final = requests.post( + self._token_url, + auth=(self._client_id, self._client_secret), + data={"grant_type": "client_credentials", "scope": "telemetry/write"}, + timeout=10, + ) + response.raise_for_status() + token_response: Final = _BillingMetricsOAuthToken.model_validate(response.json()) + return token_response.access_token, max(0.0, float(token_response.expires_in) - _TOKEN_REFRESH_SKEW_SECONDS) + except (requests.RequestException, ValueError): + raise requests.RequestException("Billing metrics token request failed") from None + + +class _SessionHookKwargs(TypedDict, total=False): + stream: ReadOnly[object] + timeout: ReadOnly[object] + verify: ReadOnly[object] + cert: ReadOnly[object] + proxies: ReadOnly[object] + + +class _BillingMetricsBearerSession(requests.Session): + def __init__(self, auth: BillingMetricsBearerAuth) -> None: + super().__init__() + self._bearer_auth: Final = auth + self.auth = auth + self.hooks = {"response": [self._invalidate_token_on_unauthorized]} + + def _invalidate_token_on_unauthorized( + self, + response: requests.Response, + **kwargs: Unpack[_SessionHookKwargs], # kwargs-ok: requests hooks forward send options + ) -> requests.Response: + if response.status_code == 401: + self._bearer_auth.invalidate() + return response + + +def build_bearer_meter_provider(config: BillingMetricsConfig) -> MeterProvider: + auth: Final = BillingMetricsBearerAuth( + client_id=config.client_id, + client_secret=config.client_secret, + token_url=config.token_url, ) + session: Final = _BillingMetricsBearerSession(auth) + exporter: Final = OTLPMetricExporter(endpoint=_metrics_endpoint(config.endpoint), session=session) reader: Final = PeriodicExportingMetricReader(exporter, export_interval_millis=config.export_interval_ms) - return MeterProvider(metric_readers=[reader], resource=Resource.create(_resource_attributes(config))) + return MeterProvider(metric_readers=[reader], resource=Resource.create(resource_attributes(config))) class BillingMetricsRecorder: @@ -125,7 +200,7 @@ class BillingMetricsRecorder: ) def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: str | None) -> None: - self._counter.add(1, _billable_attributes(category, route, status_code, model_id)) + self._counter.add(1, billable_attributes(category, route, status_code, model_id)) def shutdown(self) -> None: """Final flush + exporter-thread stop. Without this, up to one export @@ -146,115 +221,34 @@ def _export_interval_ms() -> int: return DEFAULT_EXPORT_INTERVAL_MS -@dataclass(frozen=True, slots=True) -class _CredentialPaths: - client_cert_path: str - client_key_path: str - ca_cert_path: str | None - - -def _is_pem_content(value: str) -> bool: - return value.lstrip().startswith(_PEM_PREFIX) - - -def _write_pem(directory: str, filename: str, pem: str) -> str: - path: Final = os.path.join(directory, filename) - with open(path, "w", encoding="utf-8") as handle: - handle.write(pem if pem.endswith("\n") else f"{pem}\n") - os.chmod(path, _PEM_FILE_MODE) - return path - - -def _resolve_credential_paths(*, client_cert: str, client_key: str, ca_cert: str | None) -> _CredentialPaths: - """ - Accept either a filesystem path or inline PEM content for each credential. - - Secret stores that inject values as environment content rather than mounted - files (ECS tasks reading AWS Secrets Manager, Cloud Run reading Secret - Manager) can only deliver the certificate as a string. The OTLP exporter - takes paths, so inline PEM is written to a private directory once, when the - recorder is built. Raises OSError if that write fails; the caller disables - metering rather than propagating. - """ - inline: Final = tuple(value for value in (client_cert, client_key, ca_cert) if value and _is_pem_content(value)) - if not inline: - return _CredentialPaths(client_cert, client_key, ca_cert) - - # mkdtemp is 0o700, so the 0o600 key file it holds is unreachable by other users. - directory: Final = tempfile.mkdtemp(prefix=_PEM_DIR_PREFIX) - return _CredentialPaths( - client_cert_path=( - _write_pem(directory, _CLIENT_CERT_FILENAME, client_cert) if _is_pem_content(client_cert) else client_cert - ), - client_key_path=( - _write_pem(directory, _CLIENT_KEY_FILENAME, client_key) if _is_pem_content(client_key) else client_key - ), - ca_cert_path=( - _write_pem(directory, _CA_CERT_FILENAME, ca_cert) if ca_cert and _is_pem_content(ca_cert) else ca_cert - ), - ) - - def load_billing_metrics_config( - *, license_data: Optional["EnterpriseLicenseData"], litellm_version: str + *, license_data: "EnterpriseLicenseData | None", litellm_version: str ) -> BillingMetricsConfig | None: - endpoint: Final = os.getenv(ENDPOINT_ENV) - client_cert: Final = os.getenv(CLIENT_CERT_ENV) - client_key: Final = os.getenv(CLIENT_KEY_ENV) - # Optional: only for private/test collectors whose server cert is not on the - # public web PKI. The production collector needs no CA override. - ca_cert: Final = os.getenv(CA_CERT_ENV) - - missing: Final = [ - name - for name, value in ( - (ENDPOINT_ENV, endpoint), - (CLIENT_CERT_ENV, client_cert), - (CLIENT_KEY_ENV, client_key), - ) - if not value - ] - if not endpoint or not client_cert or not client_key: + env: Final = MappingProxyType( + {name: os.getenv(name) or "" for name in (ENDPOINT_ENV, CLIENT_ID_ENV, CLIENT_SECRET_ENV, TOKEN_URL_ENV)} + ) + missing: Final = tuple(name for name, value in env.items() if not value) + if missing: verbose_proxy_logger.warning( "Enterprise billing metrics disabled: licensed deployment missing config (%s)", ", ".join(missing), ) return None - try: - paths: Final = _resolve_credential_paths(client_cert=client_cert, client_key=client_key, ca_cert=ca_cert) - except OSError as exc: - verbose_proxy_logger.warning( - "Enterprise billing metrics disabled: could not write inline certificate content to disk: %s", exc - ) - return None - - # Report the variable names, never their values. A value that is neither a - # readable path nor recognizable PEM is still secret material, and this - # warning would otherwise copy a client key straight into the proxy logs. - unreadable: Final = [ + insecure_urls: Final = tuple( env_name - for env_name, path in ( - (CLIENT_CERT_ENV, paths.client_cert_path), - (CLIENT_KEY_ENV, paths.client_key_path), - (CA_CERT_ENV, paths.ca_cert_path), - ) - if path and not os.path.isfile(path) - ] - if unreadable: - verbose_proxy_logger.warning( - "Enterprise billing metrics disabled: %s did not resolve to a readable certificate file. " - "Set each to a file path, or to inline PEM content beginning with '%s'.", - ", ".join(unreadable), - _PEM_PREFIX, - ) + for env_name, url in ((ENDPOINT_ENV, env[ENDPOINT_ENV]), (TOKEN_URL_ENV, env[TOKEN_URL_ENV])) + if not is_https_url_with_host(url) + ) + if insecure_urls: + verbose_proxy_logger.warning("Enterprise billing metrics disabled: %s must use HTTPS", insecure_urls[0]) return None return BillingMetricsConfig( - endpoint=endpoint, - client_cert_path=paths.client_cert_path, - client_key_path=paths.client_key_path, - ca_cert_path=paths.ca_cert_path, + endpoint=env[ENDPOINT_ENV], + client_id=env[CLIENT_ID_ENV], + client_secret=env[CLIENT_SECRET_ENV], + token_url=env[TOKEN_URL_ENV], export_interval_ms=_export_interval_ms(), litellm_version=litellm_version, license_id=(license_data or {}).get("user_id"), @@ -282,7 +276,7 @@ _ACTIVE_RECORDER: Final = _ActiveRecorderRegistry() def build_billing_metrics_recorder( - *, premium: bool, license_data: Optional["EnterpriseLicenseData"], litellm_version: str + *, premium: bool, license_data: "EnterpriseLicenseData | None", litellm_version: str ) -> BillingMetricsRecorder | None: """Build the recorder, or None when the deployment is not licensed or metering is unconfigured.""" if not premium: @@ -296,14 +290,12 @@ def build_billing_metrics_recorder( return None try: - recorder: Final = BillingMetricsRecorder(build_mtls_meter_provider(config)) + provider: Final = build_bearer_meter_provider(config) + recorder: Final = BillingMetricsRecorder(provider) except Exception as exc: # noqa: BLE001 -- metering must never break proxy startup verbose_proxy_logger.warning("Enterprise billing metrics disabled: failed to initialize exporter: %s", exc) return None _ACTIVE_RECORDER.set(recorder) - # The only positive signal that this component meters. Without it, a silent - # return above is indistinguishable from a working exporter in the logs, and - # a component that carries the cert but no license would look healthy. verbose_proxy_logger.info( "Enterprise billing metrics enabled: exporting to %s every %d ms", config.endpoint, diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 5f9dc4a7a31..5f99b1a4611 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1,13 +1,16 @@ import json import re -from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence -from typing import Any, Final, Literal, TypedDict +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from functools import partial, reduce +from itertools import chain, groupby +from types import MappingProxyType +from typing import Any, Final, Literal, TypeAlias, TypedDict, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter from typing_extensions import ReadOnly, Required -from litellm import ChatCompletionToolParam from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.completion_extras.litellm_responses_transformation.transformation import ( @@ -19,6 +22,9 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.tool_fields import declared_function_name, qualified_member_name +from litellm.llms.base_llm.guardrail_translation.utils import anthropic_tool_names from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -30,7 +36,10 @@ from litellm.proxy.guardrails.anthropic_sse import ( is_raw_sse_stream, is_sse_error_stream, ) +from litellm.responses.additional_tools import ADDITIONAL_TOOLS_INPUT_ITEM_TYPE from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.llms.anthropic import ANTHROPIC_HOSTED_TOOLS +from litellm.types.llms.vertex_ai import VertexToolName from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ToolPermissionRule, @@ -48,8 +57,48 @@ from litellm.types.utils import ( ) GUARDRAIL_NAME: Final = "tool_permission" +_MCP_TOOLSET_TYPE: Final = "mcp_toolset" +_HOSTED_SERVER_NAME_FIELDS: Final = MappingProxyType({"mcp": "server_label", _MCP_TOOLSET_TYPE: "mcp_server_name"}) +_TOOL_CARRYING_INPUT_TYPES: Final = frozenset({"tool_search_output", ADDITIONAL_TOOLS_INPUT_ITEM_TYPE}) +_DECLARATION_KEYS: Final = ("functionDeclarations", "function_declarations") +_NON_TOOL_PART_KEYS: Final = frozenset({*_DECLARATION_KEYS, "cachePoint", "type"}) +_TOOL_CONFIG_KEYS: Final = ("toolConfig", "tool_config") +_FUNCTION_CALLING_CONFIG_KEYS: Final = ("functionCallingConfig", "function_calling_config") +_ALLOWED_FUNCTION_NAMES_KEYS: Final = ("allowedFunctionNames", "allowed_function_names") +_FORCED_FUNCTION_CALLING_MODES: Final = frozenset({"ANY", "VALIDATED"}) +_NAME_IDENTIFIED_TYPES: Final = frozenset({"function", "custom", "namespace", *_HOSTED_SERVER_NAME_FIELDS}) +_TEXT_EDITOR_TOOL_NAMES: Final = frozenset({"str_replace_editor", "str_replace_based_edit_tool"}) +_TYPE_VERSION_SUFFIX: Final = re.compile(r"_\d{8}$") +_GEMINI_TYPED_HOSTED_TOOLS: Final = ("computer_use", "web_search", "web_search_preview") +_NAMED_CHOICE_TYPES: Final = frozenset({"tool", "function", "custom", "allowed_tools", "any", "auto", "none"}) _DECODED_ARGUMENTS: Final = TypeAdapter(object) _REQUEST_DATA: Final = TypeAdapter(object) +_ToolTarget: TypeAlias = tuple[str | None, str | None] +_JudgedCalls: TypeAlias = Mapping[str, tuple[tuple[str, str], ...]] +_GroupedT: Final = TypeVar("_GroupedT") + + +def _snake_case(name: str) -> str: + return re.sub(r"(? str: + head, *rest = name.split("_") + return head + "".join(part.title() for part in rest) + + +_GEMINI_HOSTED_TOOL_KEYS: Final = frozenset( + chain.from_iterable((_snake_case(tool.value), _camel_case(tool.value)) for tool in VertexToolName) +) + + +def _name_identifies_tool(name: str, tool_type: str) -> bool: + stem: Final = _TYPE_VERSION_SUFFIX.sub("", tool_type) + return ( + tool_type in _NAME_IDENTIFIED_TYPES + or name == stem + or (stem == ANTHROPIC_HOSTED_TOOLS.TEXT_EDITOR.value and name in _TEXT_EDITOR_TOOL_NAMES) + ) def _object_mapping(value: object) -> Mapping[str, object] | None: @@ -68,6 +117,356 @@ def _function_tool_name(tool: object) -> str | None: return name if isinstance(name, str) else None +def _non_empty_string(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _request_tool_names(tool: object) -> tuple[str, ...]: + fields: Final = _object_mapping(tool) + custom: Final = ( + _object_mapping(fields.get("custom")) if fields is not None and fields.get("type") == "custom" else None + ) + custom_name: Final = _non_empty_string(custom.get("name")) if custom is not None else None + server_field: Final = _HOSTED_SERVER_NAME_FIELDS.get(str(fields.get("type"))) if fields is not None else None + server_name: Final = ( + _non_empty_string(fields.get(server_field)) if fields is not None and server_field is not None else None + ) + extra_names: Final = tuple(name for name in (_function_name(tool), custom_name, server_name) if name is not None) + return tuple(dict.fromkeys(anthropic_tool_names(tool) + extra_names)) + + +def _function_name(tool: object) -> str | None: + fields: Final = _object_mapping(tool) + function: Final = _object_mapping(fields.get("function")) if fields is not None else None + return _non_empty_string(function.get("name")) if function is not None else None + + +def _namespace_function_targets(tool: object) -> tuple[_ToolTarget, ...]: + function_name: Final = _function_name(tool) + return ((function_name, "function"),) if function_name is not None else () + + +def _gemini_hosted_key_targets(tool: object) -> tuple[_ToolTarget, ...]: + fields: Final = _object_mapping(tool) + if fields is None or fields.get("type") in _GEMINI_TYPED_HOSTED_TOOLS: + return () + keys: Final = tuple(key for key in fields if key != "type") + return (_part_target(fields, keys[0]),) if len(keys) == 1 and keys[0] in _GEMINI_HOSTED_TOOL_KEYS else () + + +def _is_namespace(tool: object) -> bool: + fields: Final = _object_mapping(tool) + return fields is not None and fields.get("type") == "namespace" + + +def _namespace_members(tool: object) -> Sequence[object]: + fields: Final = _object_mapping(tool) + members: Final = _object_list(fields.get("tools")) if fields is not None and _is_namespace(tool) else None + return members or () + + +def _returned_namespace(tool_call: ChatCompletionMessageToolCall) -> str | None: + tags: Final = as_str_mapping(attribute_or_none(tool_call, "provider_specific_fields")) + return _non_empty_string(tags.get("namespace")) if tags is not None else None + + +def _judged_call_names( + tool_call: ChatCompletionMessageToolCall, returned_name: str, judged_names: _JudgedCalls +) -> tuple[tuple[str, str], ...]: + namespace: Final = _returned_namespace(tool_call) + if namespace is not None: + return ((qualified_member_name(namespace, returned_name), returned_name),) + return judged_names.get(returned_name, ()) + + +def _pair_namespace(qualified: str, bare: str) -> str | None: + suffix: Final = qualified_member_name("", bare) + if qualified == bare or not qualified.endswith(suffix): + return None + return _non_empty_string(qualified.removesuffix(suffix)) + + +def _qualified_target(namespace: str | None, target: _ToolTarget) -> _ToolTarget: + name, tool_type = target + return (qualified_member_name(namespace, name), tool_type) if namespace is not None and name is not None else target + + +def _member_name(member: object) -> str | None: + fields: Final = _object_mapping(member) + return _non_empty_string(fields.get("name")) if fields is not None else None + + +def _declaration_targets(declaration: object) -> tuple[_ToolTarget, ...]: + name: Final = _member_name(declaration) + return ((name, "function"),) if name is not None else () + + +@dataclass(frozen=True, slots=True) +class _MemberGroup: + key: str + members: Sequence[object] + targets_of: Callable[[object], tuple[_ToolTarget, ...]] + + +def _member_group( + tool: object, namespace_member_targets: Callable[[str | None, object], tuple[_ToolTarget, ...]] +) -> _MemberGroup | None: + fields: Final = _object_mapping(tool) + if fields is None: + return None + members: Final = _namespace_members(tool) + if members: + return _MemberGroup("tools", members, partial(namespace_member_targets, _member_name(tool))) + key: Final = next((key for key in _DECLARATION_KEYS if _object_list(fields.get(key))), None) + declarations: Final = _object_list(fields.get(key)) if key is not None else None + return _MemberGroup(key, declarations, _declaration_targets) if key is not None and declarations else None + + +def _keyed_tool_fields(tool: object) -> Mapping[str, object] | None: + fields: Final = _object_mapping(tool) + if fields is None or isinstance(fields.get("type"), str) or _request_tool_names(tool): + return None + return fields + + +def _part_keys(fields: Mapping[str, object]) -> tuple[str, ...]: + return tuple(key for key in fields if key not in _NON_TOOL_PART_KEYS) + + +def _part_target(fields: Mapping[str, object], key: str) -> _ToolTarget: + return (None if key in _GEMINI_HOSTED_TOOL_KEYS else _member_name(fields.get(key)), key) + + +def _input_item_tools(item: object) -> Sequence[object] | None: + fields: Final = _object_mapping(item) + carries_tools: Final = fields is not None and _non_empty_string(fields.get("type")) in _TOOL_CARRYING_INPUT_TYPES + return _object_list(fields.get("tools")) if fields is not None and carries_tools else None + + +def _request_list(data: object, key: str) -> Sequence[object] | None: + fields: Final = _object_mapping(data) + return _object_list(fields.get(key)) if fields is not None else None + + +def _input_items(data: object) -> Sequence[object]: + return _request_list(data, "input") or () + + +def _tools_in_items(items: Sequence[object]) -> tuple[object, ...]: + return tuple(chain.from_iterable(_input_item_tools(item) or () for item in items)) + + +def _carries_tools(items: Sequence[object]) -> bool: + return any(_input_item_tools(item) is not None for item in items) + + +def _pair_key(pair: tuple[str, object]) -> str: + return pair[0] + + +def _grouped_values(pairs: Iterable[tuple[str, _GroupedT]]) -> Mapping[str, tuple[_GroupedT, ...]]: + groups: Final = groupby(sorted(pairs, key=_pair_key), key=_pair_key) + return MappingProxyType({key: tuple(dict.fromkeys(value for _, value in group)) for key, group in groups}) + + +def _request_layers(data: object) -> tuple[Mapping[str, object], ...]: + fields: Final = _object_mapping(data) + extra_body: Final = _object_mapping(fields.get("extra_body")) if fields is not None else None + return tuple(layer for layer in (fields, extra_body) if layer is not None) + + +def _layer_sent_tools(layer: Mapping[str, object]) -> tuple[object, ...]: + return ( + *(_request_list(layer, "tools") or ()), + *(_request_list(layer, "functions") or ()), + *_tools_in_items(_input_items(layer)), + ) + + +def _sent_tools(data: object) -> tuple[object, ...]: + return tuple(chain.from_iterable(map(_layer_sent_tools, _request_layers(data)))) + + +def _sent_fields(data: object) -> Mapping[str, object]: + return dict(chain.from_iterable(layer.items() for layer in _request_layers(data))) + + +@dataclass(frozen=True, slots=True) +class _KeptLayer: + fields: Mapping[str, object] + tools: tuple[object, ...] | None + functions: tuple[object, ...] | None + mcp_servers: tuple[object, ...] | None + input: tuple[object, ...] | None + + def sent_tools(self) -> tuple[object, ...]: + kept_input: Final = _input_items(self.fields) if self.input is None else self.input + return (*(self.tools or ()), *(self.functions or ()), *_tools_in_items(kept_input)) + + def kept_lists(self) -> dict[str, object]: + lists: Final = ( + ("tools", self.tools), + ("functions", self.functions), + ("mcp_servers", self.mcp_servers), + ("input", self.input), + ) + return {key: list(kept) for key, kept in lists if kept is not None} + + def overlaid_by(self, over: "_KeptLayer") -> "_KeptLayer": + return _KeptLayer( + fields={**self.fields, **over.fields}, + tools=over.tools if "tools" in over.fields else self.tools, + functions=over.functions if "functions" in over.fields else self.functions, + mcp_servers=over.mcp_servers if "mcp_servers" in over.fields else self.mcp_servers, + input=over.input if "input" in over.fields else self.input, + ) + + +def _mcp_server_target(server: object) -> _ToolTarget: + fields: Final = _object_mapping(server) + return (_non_empty_string(fields.get("name")) if fields is not None else None, _MCP_TOOLSET_TYPE) + + +def _reference_name(part: object) -> str | None: + fields: Final = _object_mapping(part) + is_reference: Final = fields is not None and fields.get("type") == "tool_reference" + return _non_empty_string(fields.get("tool_name")) if fields is not None and is_reference else None + + +def _references_one_of(part: object, names: frozenset[str]) -> bool: + name: Final = _reference_name(part) + return name is not None and name in names + + +def _without_references_to(parts: Sequence[object], names: frozenset[str]) -> Sequence[object]: + removed: Final = [name for name in map(_reference_name, parts) if name is not None and name in names] + kept: Final = [part for part in parts if not _references_one_of(part, names)] + if kept or not removed: + return kept + return [{"type": "text", "text": f"Permission denied: tool {', '.join(removed)} is not available"}] + + +def _tool_result_without_references(fields: Mapping[str, object], names: frozenset[str]) -> object: + parts: Final = _object_list(fields.get("content")) + return {**fields, "content": _without_references_to(parts, names)} if parts else fields + + +def _search_result_without_references(fields: Mapping[str, object], names: frozenset[str]) -> object: + result: Final = _object_mapping(fields.get("content")) + references: Final = _object_list(result.get("tool_references")) if result is not None else None + if result is None or references is None: + return fields + kept: Final = [reference for reference in references if not _references_one_of(reference, names)] + return {**fields, "content": {**result, "tool_references": kept}} + + +def _without_references_in_block(block: object, names: frozenset[str]) -> object: + fields: Final = _object_mapping(block) + if fields is None: + return block + match fields.get("type"): + case "tool_result": + return _tool_result_without_references(fields, names) + case "tool_search_tool_result": + return _search_result_without_references(fields, names) + case _: + return block + + +def _without_references_in_message(message: object, names: frozenset[str]) -> object: + fields: Final = _object_mapping(message) + content: Final = _object_list(fields.get("content")) if fields is not None else None + if fields is None or content is None: + return message + if fields.get("role") == "tool": + return {**fields, "content": _without_references_to(content, names)} + return {**fields, "content": [_without_references_in_block(block, names) for block in content]} + + +def _allowed_tool_entries(tool_choice: Mapping[str, object]) -> Sequence[object] | None: + if tool_choice.get("type") != "allowed_tools": + return None + nested: Final = _object_mapping(tool_choice.get("allowed_tools")) + return _object_list((nested if nested is not None else tool_choice).get("tools")) + + +def _is_removed_entry(entry: object, removed_names: frozenset[str], removed_types: frozenset[str]) -> bool: + names: Final = _request_tool_names(entry) + if names: + return any(name in removed_names for name in names) + fields: Final = _object_mapping(entry) + entry_type: Final = _non_empty_string(fields.get("type")) if fields is not None else None + return entry_type is not None and entry_type in removed_types + + +def _with_allowed_tool_entries(tool_choice: Mapping[str, object], entries: list[object]) -> dict[str, object]: + nested: Final = _object_mapping(tool_choice.get("allowed_tools")) + if nested is None: + return {**tool_choice, "tools": entries} + return {**tool_choice, "allowed_tools": {**nested, "tools": entries}} + + +def _is_denied_target( + target: _ToolTarget, denied_targets: frozenset[_ToolTarget], denied_names: frozenset[str] +) -> bool: + return target in denied_targets or target[0] in denied_names + + +def _kept_members(group: _MemberGroup | None, is_denied: Callable[[_ToolTarget], bool]) -> tuple[object, ...]: + if group is None: + return () + return tuple(member for member in group.members if not any(map(is_denied, group.targets_of(member)))) + + +def _none_tool_choice(tool_choice: object) -> object: + choice: Final = _object_mapping(tool_choice) + return {"type": "none"} if choice is not None and choice.get("type") in ("tool", "any") else "none" + + +def _requires_a_tool(tool_choice: object) -> bool: + choice: Final = _object_mapping(tool_choice) + return tool_choice == "required" or (choice is not None and choice.get("type") == "any") + + +def _kept_allowed_function_names(config: Mapping[str, object], removed_names: frozenset[str]) -> Mapping[str, object]: + listed: Final = tuple((key, _object_list(config.get(key))) for key in _ALLOWED_FUNCTION_NAMES_KEYS) + return { + key: [name for name in names if _non_empty_string(name) not in removed_names] + for key, names in listed + if names is not None + } + + +def _rewritten_function_calling_config( + config: Mapping[str, object], removed_names: frozenset[str], functions_remain: bool +) -> Mapping[str, object]: + kept_names: Final = _kept_allowed_function_names(config, removed_names) + had_names: Final = any(_object_list(config.get(key)) for key in _ALLOWED_FUNCTION_NAMES_KEYS) + every_listed_name_removed: Final = had_names and not any(kept_names.values()) + forces_a_function: Final = str(config.get("mode", "")).upper() in _FORCED_FUNCTION_CALLING_MODES + if not every_listed_name_removed and (functions_remain or not forces_a_function): + return {**config, **kept_names} + return {**{key: value for key, value in config.items() if key not in _ALLOWED_FUNCTION_NAMES_KEYS}, "mode": "NONE"} + + +def _rewritten_tool_config_field( + key: str, value: object, removed_names: frozenset[str], functions_remain: bool +) -> object: + config: Final = _object_mapping(value) + if key not in _FUNCTION_CALLING_CONFIG_KEYS or config is None: + return value + return _rewritten_function_calling_config(config, removed_names, functions_remain) + + +def _rewritten_tool_config(tool_config: object, removed_names: frozenset[str], functions_remain: bool) -> object: + fields: Final = _object_mapping(tool_config) + if fields is None: + return tool_config + return { + key: _rewritten_tool_config_field(key, value, removed_names, functions_remain) for key, value in fields.items() + } + + class _ToolPermissionRuleFields(TypedDict, total=False): """The config-file shape a :class:`ToolPermissionRule` is built from.""" @@ -258,7 +657,7 @@ class ToolPermissionGuardrail(CustomGuardrail): name_required: Final = rule.tool_name is not None type_required: Final = rule.tool_type is not None - name_matched: Final = self._matches_regex(name_pattern, tool_name) if name_required else True + name_matched: Final = self._matches_regex(name_pattern, tool_name or tool_type) if name_required else True type_matched: Final = self._matches_regex(type_pattern, tool_type) if type_required else True overall_match: Final = name_matched and type_matched @@ -320,6 +719,72 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug(message) return is_allowed, None, message + def _is_name_rule_for(self, rule: ToolPermissionRule, tool_name: str, tool_type: str | None) -> bool: + return rule.tool_name is not None and self._rule_matches_tool(rule, tool_name=tool_name, tool_type=tool_type)[0] + + def _decides_call( + self, rule: ToolPermissionRule, tool_name: str, tool_call: ChatCompletionMessageToolCall | None + ) -> bool: + if tool_call is None or not rule.allowed_param_patterns: + return True + arguments, parse_error = self._parse_tool_call_arguments(tool_call) + if parse_error or not arguments: + return True + return self._patterns_match_for_rule(arguments=arguments, rule=rule, tool_name=tool_name)[0] + + def _first_name_rule( + self, tool_name: str, tool_type: str | None, tool_call: ChatCompletionMessageToolCall | None + ) -> ToolPermissionRule | None: + return next( + ( + rule + for rule in self.rules + if self._is_name_rule_for(rule, tool_name, tool_type) and self._decides_call(rule, tool_name, tool_call) + ), + None, + ) + + def _narrower_name( + self, + name: str, + other: str, + tool_type: str | None, + tool_call: ChatCompletionMessageToolCall | None, + ) -> str: + rule: Final = self._first_name_rule(name, tool_type, tool_call) + other_rule: Final = self._first_name_rule(other, tool_type, tool_call) + if other_rule is None or other_rule is rule: + return name + if rule is None or self._is_name_rule_for(rule, other, tool_type): + return other + if self._is_name_rule_for(other_rule, name, tool_type): + return name + return other if other_rule.decision == "deny" else name + + def _deciding_member_name( + self, + qualified: str, + bare: str, + tool_type: str | None, + tool_call: ChatCompletionMessageToolCall | None = None, + ) -> str: + member: Final = self._narrower_name(qualified, bare, tool_type, tool_call) + namespace: Final = _pair_namespace(qualified, bare) + return member if namespace is None else self._narrower_name(member, namespace, tool_type, tool_call) + + def _check_target_permission( + self, target: _ToolTarget, bare_names: Mapping[str, tuple[str, ...]] + ) -> tuple[bool, str | None, str | None]: + name, tool_type = target + candidates: Final = tuple(sorted(bare_names.get(name, ()))) if name is not None else () + if name is None or not candidates: + return self._check_tool_permission(name, tool_type) + results: Final = tuple( + self._check_tool_permission(self._deciding_member_name(name, bare_name, tool_type), tool_type) + for bare_name in candidates + ) + return next((result for result in results if not result[0]), results[0]) + def _parse_tool_call_arguments( self, tool_call: ChatCompletionMessageToolCall ) -> tuple[Mapping[str, object] | None, str | None]: @@ -411,9 +876,9 @@ class ToolPermissionGuardrail(CustomGuardrail): return True, None def _get_permission_for_tool_call( - self, tool_call: ChatCompletionMessageToolCall + self, tool_call: ChatCompletionMessageToolCall, judged_name: str | None = None ) -> tuple[bool, str | None, str | None]: - tool_name: Final = tool_call.function.name if tool_call.function else None + tool_name: Final = judged_name or (tool_call.function.name if tool_call.function else None) tool_type: Final = getattr(tool_call, "type", None) if not tool_name and not tool_type: return self.default_action == "allow", None, None @@ -556,10 +1021,46 @@ class ToolPermissionGuardrail(CustomGuardrail): tool_call for block in content if (tool_call := self._anthropic_tool_use_to_tool_call(block)) is not None ) + def _returned_call_names(self, data: Mapping[str, object]) -> _JudgedCalls: + sent_tools: Final = _sent_tools(_REQUEST_DATA.validate_python(data)) + member_names: Final = tuple(chain.from_iterable(map(self._bare_member_names, sent_tools))) + if not member_names: + return MappingProxyType({}) + top_level_tools: Final = (tool for tool in sent_tools if not _is_namespace(tool)) + top_level_names: Final = frozenset(name for name in map(declared_function_name, top_level_tools) if name) + listed: Final = frozenset(name for name, _ in self._collect_request_tools(data) if name is not None) + own_names: Final = frozenset(bare for _, bare in member_names if bare in listed) | top_level_names + returned_members: Final = _grouped_values( + chain( + ((bare, (name, bare)) for name, bare in member_names if bare not in top_level_names), + ((name, (name, bare)) for name, bare in member_names if name not in top_level_names), + ) + ) + return MappingProxyType( + { + name: (*members, *(((name, name),) if name in own_names else ())) + for name, members in returned_members.items() + } + ) + + def _judged_name(self, tool_call: ChatCompletionMessageToolCall, qualified: str, bare: str) -> str: + return qualified if qualified == bare else self._deciding_member_name(qualified, bare, "function", tool_call) + + def _judge_tool_call( + self, tool_call: ChatCompletionMessageToolCall, judged_names: _JudgedCalls + ) -> tuple[bool, str | None, str | None]: + returned_name: Final = tool_call.function.name if tool_call.function else None + names: Final = _judged_call_names(tool_call, returned_name, judged_names) if returned_name is not None else () + if returned_name is None or not names: + return self._get_permission_for_tool_call(tool_call) + judged: Final = dict.fromkeys(self._judged_name(tool_call, qualified, bare) for qualified, bare in names) + results: Final = tuple(self._get_permission_for_tool_call(tool_call, name) for name in judged) + return next((result for result in results if not result[0]), results[0]) + def _evaluate_tool_calls( - self, tool_calls: Sequence[ChatCompletionMessageToolCall] + self, tool_calls: Sequence[ChatCompletionMessageToolCall], judged_names: _JudgedCalls ) -> tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...]: - checked: Final = tuple((tool_call, *self._get_permission_for_tool_call(tool_call)) for tool_call in tool_calls) + checked: Final = tuple((tool_call, *self._judge_tool_call(tool_call, judged_names)) for tool_call in tool_calls) for _tool_call, is_allowed, _rule_id, message in checked: if not is_allowed and message is not None: @@ -619,112 +1120,311 @@ class ToolPermissionGuardrail(CustomGuardrail): if not any(_is_tool_use_block(block) for block in kept_blocks): response["stop_reason"] = "end_turn" # rebind-ok: dropping every tool_use ends the turn - def _get_request_tool_name(self, tool: object) -> tuple[str | None, str | None]: - tool_type: Final = self._get_mapping_value(tool, "type") - if tool_type != "function": - return None, tool_type + def _unqualified_member_targets(self, member: object) -> tuple[_ToolTarget, ...]: + fields: Final = _object_mapping(member) + if fields is None: + return () + targets: Final = self._get_request_tool_targets(member) + own_name: Final = _member_name(member) + judged_names: Final = frozenset(name for name, _ in targets) + own_target: Final = ( + ((own_name, _non_empty_string(fields.get("type"))),) + if own_name is not None and own_name not in judged_names + else () + ) + return targets + own_target - function: Final = self._get_mapping_value(tool, "function") - tool_name: Final = self._get_mapping_value(function, "name") - return tool_name, tool_type + def _namespace_member_targets(self, namespace: str | None, member: object) -> tuple[_ToolTarget, ...]: + return tuple(_qualified_target(namespace, target) for target in self._unqualified_member_targets(member)) - def _get_legacy_function_name(self, function: object) -> str | None: - return self._get_mapping_value(function, "name") + def _bare_member_names(self, tool: object) -> tuple[tuple[str, str], ...]: + namespace: Final = _member_name(tool) + members: Final = _namespace_members(tool) + if namespace is None or not members: + return () + targets: Final = chain.from_iterable(map(self._unqualified_member_targets, members)) + return tuple((qualified_member_name(namespace, name), name) for name, _ in targets if name is not None) - def _get_named_tool_choice(self, data: dict) -> str | None: - tool_choice: Final = data.get("tool_choice") + def _get_request_tool_targets(self, tool: object) -> tuple[_ToolTarget, ...]: + group: Final = _member_group(tool, self._namespace_member_targets) + member_targets: Final = ( + tuple(chain.from_iterable(map(group.targets_of, group.members))) if group is not None else () + ) + keyed: Final = _keyed_tool_fields(tool) + part_targets: Final = tuple(_part_target(keyed, key) for key in _part_keys(keyed)) if keyed is not None else () + return member_targets + part_targets + self._own_tool_targets(tool) + + def _own_tool_targets(self, tool: object) -> tuple[_ToolTarget, ...]: + if _keyed_tool_fields(tool) is not None: + return () + named: Final = _namespace_function_targets(tool) if _namespace_members(tool) else self._named_tool_targets(tool) + return named + _gemini_hosted_key_targets(tool) + + def _named_tool_targets(self, tool: object) -> tuple[_ToolTarget, ...]: + raw_type: Final = self._get_mapping_value(tool, "type") + tool_type: Final = raw_type if isinstance(raw_type, str) else None + anthropic_custom: Final = ( + tool_type == "custom" + and self._get_mapping_value(tool, "input_schema") is not None + and _object_mapping(self._get_mapping_value(tool, "custom")) is None + ) + normalized_type: Final = "function" if tool_type is None or anthropic_custom else tool_type + names: Final = _request_tool_names(tool) + named_targets: Final = tuple((name, normalized_type) for name in names) + if names and all(_name_identifies_tool(name, normalized_type) for name in names): + return named_targets + if names: + return (*named_targets, (None, normalized_type)) + return ((None, tool_type),) if tool_type is not None else () + + def _get_named_tool_choice(self, tool_choice: object) -> str | None: if not tool_choice or tool_choice in ("auto", "none", "required"): return None if isinstance(tool_choice, str): return tool_choice - if self._get_mapping_value(tool_choice, "type") != "function": + choice: Final = _object_mapping(tool_choice) + if choice is None: return None - return self._get_mapping_value(self._get_mapping_value(tool_choice, "function"), "name") + choice_type: Final = choice.get("type") + server_field: Final = _HOSTED_SERVER_NAME_FIELDS.get(str(choice_type)) + if server_field is not None: + return _non_empty_string(choice.get(server_field)) + if choice_type not in ("tool", "function", "custom"): + return None + nested: Final = _object_mapping(choice.get(str(choice_type))) + nested_name: Final = _non_empty_string(nested.get("name")) if nested is not None else None + return nested_name or _non_empty_string(choice.get("name")) - def _get_named_function_call(self, data: dict) -> str | None: + def _get_named_function_call(self, data: Mapping[str, object]) -> str | None: function_call: Final = data.get("function_call") if not function_call or function_call in ("auto", "none"): return None if isinstance(function_call, str): return function_call - return self._get_mapping_value(function_call, "name") + return _non_empty_string(self._get_mapping_value(function_call, "name")) - def _collect_request_tools(self, data: dict) -> list[tuple[str, str | None]]: - request_tools: Final[list[tuple[str, str | None]]] = [] + def _listed_layer_targets(self, layer: Mapping[str, object]) -> tuple[_ToolTarget, ...]: + return ( + *chain.from_iterable(map(self._get_request_tool_targets, _request_list(layer, "tools") or ())), + *map(_mcp_server_target, _request_list(layer, "mcp_servers") or ()), + *chain.from_iterable(map(self._legacy_function_targets, _request_list(layer, "functions") or ())), + ) - for tool in data.get("tools") or []: - tool_name, tool_type = self._get_request_tool_name(tool) - if tool_name is not None: - request_tools.append((tool_name, tool_type)) + def _forced_layer_names(self, layer: Mapping[str, object]) -> tuple[str | None, ...]: + return (self._get_named_tool_choice(layer.get("tool_choice")), self._get_named_function_call(layer)) - for function in data.get("functions") or []: - function_name = self._get_legacy_function_name(function) - if function_name is not None: - request_tools.append((function_name, "function")) + def _collect_request_tools(self, data: Mapping[str, object]) -> tuple[_ToolTarget, ...]: + request: Final = _REQUEST_DATA.validate_python(data) + layers: Final = _request_layers(request) + listed: Final = tuple(chain.from_iterable(map(self._listed_layer_targets, layers))) + listed_names: Final = frozenset(name for name, _ in listed) + forced: Final = chain.from_iterable(map(self._forced_layer_names, layers)) + replayed: Final = map(_function_tool_name, replayed_tool_search_tools(request)) + input_tools: Final = _tools_in_items(tuple(chain.from_iterable(map(_input_items, layers)))) + return ( + *listed, + *((name, "function") for name in forced if name is not None and name not in listed_names), + *((name, "function") for name in replayed if name is not None), + *chain.from_iterable(map(self._get_request_tool_targets, input_tools)), + ) - for forced_tool_name in ( - self._get_named_tool_choice(data), - self._get_named_function_call(data), - ): - if forced_tool_name is not None: - request_tools.append((forced_tool_name, "function")) + def _legacy_function_targets(self, function: object) -> tuple[_ToolTarget, ...]: + targets: Final = self._get_request_tool_targets(function) + as_functions: Final = tuple((name, "function") for name, _ in targets if name is not None) + return tuple(dict.fromkeys(targets + as_functions)) - replayed: Final = map(_function_tool_name, replayed_tool_search_tools(_REQUEST_DATA.validate_python(data))) - request_tools.extend((name, "function") for name in replayed if name is not None) + def _kept_layer( + self, + layer: Mapping[str, object], + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> _KeptLayer: + is_denied: Final = partial(_is_denied_target, denied_targets=denied_targets, denied_names=denied_names) + tools: Final = _request_list(layer, "tools") + functions: Final = _request_list(layer, "functions") + mcp_servers: Final = _request_list(layer, "mcp_servers") + return _KeptLayer( + fields=layer, + tools=None if tools is None else self._kept_tool_list(tools, denied_targets, denied_names), + functions=None if functions is None else self._kept_tool_list(functions, denied_targets, denied_names), + mcp_servers=( + None + if mcp_servers is None + else tuple(server for server in mcp_servers if not is_denied(_mcp_server_target(server))) + ), + input=self._kept_input_list(_input_items(layer), denied_targets, denied_names), + ) - return request_tools + def _rewritten_layer_fields( + self, + kept: _KeptLayer, + rewrite_choice: Callable[[object], object], + removed_names: frozenset[str], + functions_remain: bool, + ) -> dict[str, object]: + layer: Final = kept.fields + choice: Final = layer.get("tool_choice") + relaxes_function_call: Final = self._get_named_function_call(layer) in removed_names + return { + **kept.kept_lists(), + **({"tool_choice": rewrite_choice(choice)} if choice is not None else {}), + **({"function_call": "none"} if relaxes_function_call else {}), + **{ + key: _rewritten_tool_config(layer[key], removed_names, functions_remain) + for key in _TOOL_CONFIG_KEYS + if key in layer + }, + } + + def _tool_names_in(self, tools: Sequence[object]) -> frozenset[str]: + targets: Final = chain.from_iterable(self._get_request_tool_targets(tool) for tool in tools) + return frozenset(name for name, _ in targets if name is not None) + + def _referable_names_in(self, tools: Sequence[object]) -> frozenset[str]: + member_names: Final = chain.from_iterable(map(self._bare_member_names, tools)) + namespace_names: Final = (_member_name(tool) for tool in tools if _namespace_members(tool)) + return ( + self._tool_names_in(tools) + | frozenset(bare for _, bare in member_names) + | frozenset(name for name in namespace_names if name is not None) + ) + + def _kept_tools( + self, + tool: object, + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...]: + is_denied: Final = partial(_is_denied_target, denied_targets=denied_targets, denied_names=denied_names) + group: Final = _member_group(tool, self._namespace_member_targets) + keyed: Final = _keyed_tool_fields(tool) + fields: Final = _object_mapping(tool) + if fields is None or (group is None and keyed is None): + return () if any(map(is_denied, self._get_request_tool_targets(tool))) else (tool,) + own_targets: Final = self._own_tool_targets(tool) + if any(map(is_denied, own_targets)): + return () + kept_members: Final = _kept_members(group, is_denied) + denied_keys: Final = ( + frozenset(key for key in _part_keys(keyed) if is_denied(_part_target(keyed, key))) + if keyed is not None + else frozenset[str]() + ) + if not denied_keys and len(kept_members) == len(group.members if group is not None else ()): + return (tool,) + emptied_group: Final = frozenset({group.key}) if group is not None and not kept_members else frozenset[str]() + dropped_keys: Final = denied_keys | emptied_group + kept_fields: Final = { + key: list(kept_members) if group is not None and key == group.key else value + for key, value in fields.items() + if key not in dropped_keys + } + own_name_left: Final = any(name is not None for name, _ in own_targets) + parts_left: Final = bool(kept_members) or own_name_left or (keyed is not None and bool(_part_keys(kept_fields))) + return (kept_fields,) if parts_left else () + + def _kept_tool_list( + self, + tools: Sequence[object], + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...]: + return tuple(chain.from_iterable(self._kept_tools(tool, denied_targets, denied_names) for tool in tools)) + + def _kept_input_list( + self, + items: Sequence[object], + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...] | None: + if not _carries_tools(items): + return None + return tuple(chain.from_iterable(self._kept_input_items(item, denied_targets, denied_names) for item in items)) + + def _kept_input_items( + self, + item: object, + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...]: + fields: Final = _object_mapping(item) + tools: Final = _input_item_tools(item) + if fields is None or tools is None: + return (item,) + kept: Final = self._kept_tool_list(tools, denied_targets, denied_names) + emptied: Final = bool(tools) and not kept and fields.get("type") == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE + return () if emptied else ({**fields, "tools": list(kept)},) + + def _tool_types_in(self, tools: Sequence[object]) -> frozenset[str]: + targets: Final = chain.from_iterable(self._get_request_tool_targets(tool) for tool in tools) + return frozenset(tool_type for _, tool_type in targets if tool_type is not None) + + def _rewritten_tool_choice( + self, tool_choice: object, removed_names: frozenset[str], removed_types: frozenset[str], tools_remain: bool + ) -> object: + named: Final = self._get_named_tool_choice(tool_choice) + if named is not None and named in removed_names: + return _none_tool_choice(tool_choice) + choice: Final = _object_mapping(tool_choice) + hosted_type: Final = _non_empty_string(choice.get("type")) if choice is not None else None + if hosted_type not in _NAMED_CHOICE_TYPES and hosted_type in removed_types: + return _none_tool_choice(tool_choice) + allowed: Final = _allowed_tool_entries(choice) if choice is not None else None + if choice is not None and allowed is not None: + kept: Final = [entry for entry in allowed if not _is_removed_entry(entry, removed_names, removed_types)] + return _with_allowed_tool_entries(choice, kept) if kept and tools_remain else _none_tool_choice(choice) + if not tools_remain and (named is not None or _requires_a_tool(tool_choice)): + return _none_tool_choice(tool_choice) + return tool_choice def _modify_request_with_permission_errors( self, data: dict, - denied_tool_names: list[str], + denied_targets: frozenset[_ToolTarget], ): """ - Modify the request to replace denied tool_calls blocks with error results - - Args: - data: The model request to modify - denied_tools: List of (tool_use, error) tuples for denied tools + Remove denied tools from the request, and relax any tool_choice that would still require one """ - if not denied_tool_names: + if not denied_targets: return data - verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_tool_names)) + verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_targets)) - # Create a mapping of tool_use_id to error result - error_tool_names: Final = set() - for tool_use in denied_tool_names: - error_tool_names.add(tool_use) - - tools: Final[list[ChatCompletionToolParam] | None] = data.get("tools") - if tools is not None: - new_tools: Final = [] - for tool in tools: - tool_name, tool_type = self._get_request_tool_name(tool) - if tool_type == "function" and tool_name in error_tool_names: - continue - new_tools.append(tool) - data["tools"] = new_tools - - functions: Final = data.get("functions") - if functions is not None: - data["functions"] = [ - function for function in functions if self._get_legacy_function_name(function) not in error_tool_names - ] - - named_tool_choice: Final = self._get_named_tool_choice(data) - if named_tool_choice in error_tool_names: - data["tool_choice"] = "none" - - named_function_call: Final = self._get_named_function_call(data) - if named_function_call in error_tool_names: - data["function_call"] = "none" - - messages: Final = without_tool_search_items_loading( - _REQUEST_DATA.validate_python(data), frozenset(denied_tool_names) + denied_names: Final = frozenset(name for name, _ in denied_targets if name is not None) + request: Final = _REQUEST_DATA.validate_python(data) + kept_layers: Final = tuple( + self._kept_layer(layer, denied_targets, denied_names) for layer in _request_layers(request) ) + + sent_tools: Final = _layer_sent_tools(_sent_fields(request)) + kept_sent: Final = reduce(_KeptLayer.overlaid_by, kept_layers) + kept_sent_tools: Final = kept_sent.sent_tools() + removed_names: Final = denied_names | ( + self._referable_names_in(sent_tools) - self._referable_names_in(kept_sent_tools) + ) + kept_types: Final = self._tool_types_in(kept_sent_tools) + removed_types: Final = self._tool_types_in(sent_tools) - kept_types + + tools_remain: Final = bool(kept_sent_tools) or bool(kept_sent.mcp_servers) + rewrite_choice: Final = partial( + self._rewritten_tool_choice, + removed_names=removed_names, + removed_types=removed_types, + tools_remain=tools_remain, + ) + rewritten: Final = tuple( + self._rewritten_layer_fields(kept, rewrite_choice, removed_names, "function" in kept_types) + for kept in kept_layers + ) + data.update(rewritten[0]) + if len(rewritten) > 1 and rewritten[1]: + data["extra_body"] = {**kept_layers[1].fields, **rewritten[1]} # rebind-ok: rewrite mode edits in place + + messages: Final = without_tool_search_items_loading(_REQUEST_DATA.validate_python(data), denied_names) if messages is not None: - data["messages"] = messages # rebind-ok: rewrite mode edits the request in place + data["messages"] = [ # rebind-ok: rewrite mode edits the request in place + _without_references_in_message(message, removed_names) for message in messages + ] return data @@ -843,25 +1543,29 @@ class ToolPermissionGuardrail(CustomGuardrail): ) return data - # Check permissions for each tool - denied_tool_names: Final = [] - for tool_name, tool_type in new_tools: - is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) + member_names: Final = tuple( + chain.from_iterable(map(self._bare_member_names, _sent_tools(_REQUEST_DATA.validate_python(data)))) + ) + bare_names: Final = _grouped_values(member_names) + decisions: Final = tuple((target, self._check_target_permission(target, bare_names)) for target in new_tools) + denied: Final = tuple( + (target, message) + for target, (is_allowed, _, message) in decisions + if not is_allowed and message is not None + ) + for _, message in denied: + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) + if denied and self.on_disallowed_action == "block": + raise HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "detection_message": denied[0][1], + }, + ) - if not is_allowed and message is not None: - verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) - if self.on_disallowed_action == "block": - raise HTTPException( - status_code=400, - detail={ - "error": "Violated guardrail policy", - "detection_message": message, - }, - ) - denied_tool_names.append(tool_name) - - if denied_tool_names: - data = self._modify_request_with_permission_errors(data, denied_tool_names) + if denied: + self._modify_request_with_permission_errors(data, frozenset(target for target, _ in denied)) verbose_proxy_logger.debug("Tool Permission Guardrail Pre-Call Hook: All tools allowed") @@ -908,7 +1612,7 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: Found %s tool calls", len(tool_calls)) - denied_tools: Final = self._evaluate_tool_calls(tool_calls) + denied_tools: Final = self._evaluate_tool_calls(tool_calls, self._returned_call_names(data)) if not denied_tools: verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") @@ -932,7 +1636,7 @@ class ToolPermissionGuardrail(CustomGuardrail): Args: user_api_key_dict: User API key information (unused but required by interface) response: The model response to check - request_data: The model request (unused but required by interface) + request_data: The model request """ # Import here to avoid circular imports @@ -949,7 +1653,9 @@ class ToolPermissionGuardrail(CustomGuardrail): stream_chunk_builder(chunks=all_chunks) if not is_raw_sse_stream(all_chunks) else None ) if isinstance(assembled_model_response, ModelResponse): - denied_tools = self._check_assembled_stream(assembled_model_response) + denied_tools = self._check_assembled_stream( + assembled_model_response, self._returned_call_names(request_data) + ) if denied_tools: self._modify_response_with_permission_errors(assembled_model_response, denied_tools) @@ -975,18 +1681,23 @@ class ToolPermissionGuardrail(CustomGuardrail): yield chunk return - for sse_chunk in self._anthropic_stream_to_deliver(anthropic_response, all_chunks): + for sse_chunk in self._anthropic_stream_to_deliver( + anthropic_response, all_chunks, self._returned_call_names(request_data) + ): yield sse_chunk def _anthropic_stream_to_deliver( - self, anthropic_response: ModelResponse, source_chunks: Sequence[object] + self, + anthropic_response: ModelResponse, + source_chunks: Sequence[object], + judged_names: _JudgedCalls, ) -> Sequence[object]: """The frames the client gets: on a failed upstream only what the check read, never the upstream's own. Assembly reads past an upstream error, and a client can read a frame of a broken stream that assembly does not, so forwarding those frames would pass a tool call nobody checked. """ - denials: Final = self._check_assembled_stream(anthropic_response) + denials: Final = self._check_assembled_stream(anthropic_response, judged_names) if denials: self._modify_response_with_permission_errors(anthropic_response, denials) elif not anthropic_sse_error_events(source_chunks): @@ -994,7 +1705,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return anthropic_sse_chunks_from_response(anthropic_response, source_chunks, keep_explanations=True) def _check_assembled_stream( - self, assembled: ModelResponse + self, assembled: ModelResponse, judged_names: _JudgedCalls ) -> tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...]: verbose_proxy_logger.debug("Tool Permission Guardrail: Checking response") tool_calls: Final = self._extract_tool_calls_from_response(assembled) @@ -1002,7 +1713,7 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: No tool uses found") return () verbose_proxy_logger.debug("Tool Permission Guardrail: Found %s tool calls", len(tool_calls)) - denied_tools: Final = self._evaluate_tool_calls(tool_calls) + denied_tools: Final = self._evaluate_tool_calls(tool_calls, judged_names) if not denied_tools: verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") return denied_tools diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 1ddbc139564..55e2e129e36 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -11,7 +11,7 @@ from datetime import datetime, timezone from functools import partial from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guard from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES from litellm.proxy.db.spend_log_batching import spend_log_write_batches from litellm.proxy.utils import PrismaClient +from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, DailyGuardrailUsageUnitsRepository, @@ -40,7 +41,7 @@ _RowKey = TypeVar("_RowKey") _RowValue = TypeVar("_RowValue") -class _UsageUnitKey(NamedTuple): +class UsageUnitKey(NamedTuple): guardrail_id: str date: str team_id: str @@ -48,7 +49,7 @@ class _UsageUnitKey(NamedTuple): usage_unit: str -class _UsageUnitIncrement(NamedTuple): +class UsageUnitIncrement(NamedTuple): units: int cost: float """USD for the priced share of units.""" @@ -56,14 +57,14 @@ class _UsageUnitIncrement(NamedTuple): """Units recorded with no known price, the share cost leaves out.""" -def _usage_unit_key(entry: tuple[_UsageUnitKey, _UsageUnitIncrement]) -> _UsageUnitKey: +def _usage_unit_key(entry: tuple[UsageUnitKey, UsageUnitIncrement]) -> UsageUnitKey: return entry[0] -def _usage_unit_increment(units: int, cost: float | None) -> _UsageUnitIncrement: +def _usage_unit_increment(units: int, cost: float | None) -> UsageUnitIncrement: if cost is None: - return _UsageUnitIncrement(units=units, cost=0.0, untracked_units=units) - return _UsageUnitIncrement(units=units, cost=cost, untracked_units=0) + return UsageUnitIncrement(units=units, cost=0.0, untracked_units=units) + return UsageUnitIncrement(units=units, cost=cost, untracked_units=0) class _MetricsKey(NamedTuple): @@ -89,22 +90,22 @@ class PendingRollups: def __init__(self) -> None: self.lock: Final = asyncio.Lock() self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({}) - self.units: Mapping[_UsageUnitKey, _UsageUnitIncrement] = MappingProxyType({}) + self.units: Mapping[UsageUnitKey, UsageUnitIncrement] = MappingProxyType({}) _PENDING_ROLLUPS: Final = PendingRollups() _NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({}) -_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0, untracked_units=0) +_NO_INCREMENT: Final = UsageUnitIncrement(units=0, cost=0.0, untracked_units=0) def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]: return (*base, *(key for key in extra if key not in base)) -def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitIncrement: +def _summed_increments(increments: Iterable[UsageUnitIncrement]) -> UsageUnitIncrement: materialized: Final = tuple(increments) - return _UsageUnitIncrement( + return UsageUnitIncrement( units=sum(i.units for i in materialized), cost=sum(i.cost for i in materialized), untracked_units=sum(i.untracked_units for i in materialized), @@ -112,8 +113,8 @@ def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitI def _merged_unit_rows( - base: Mapping[_UsageUnitKey, _UsageUnitIncrement], extra: Mapping[_UsageUnitKey, _UsageUnitIncrement] -) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]: + base: Mapping[UsageUnitKey, UsageUnitIncrement], extra: Mapping[UsageUnitKey, UsageUnitIncrement] +) -> Mapping[UsageUnitKey, UsageUnitIncrement]: return MappingProxyType( { key: _summed_increments((base.get(key, _NO_INCREMENT), extra.get(key, _NO_INCREMENT))) @@ -250,7 +251,7 @@ def _parse_payload_start_time(payload: Mapping[str, object]) -> datetime | None: def _iter_usage_unit_increments( logs_to_process: Sequence[Mapping[str, object]], -) -> Iterator[tuple[_UsageUnitKey, _UsageUnitIncrement]]: +) -> Iterator[tuple[UsageUnitKey, UsageUnitIncrement]]: for payload in logs_to_process: start_time = _parse_payload_start_time(payload) if not payload.get("request_id") or start_time is None: @@ -266,14 +267,14 @@ def _iter_usage_unit_increments( cost_by_unit = billed_guardrail_cost_by_unit(entry) for unit_name, units in usage.items(): if isinstance(units, int) and not isinstance(units, bool) and units > 0: - key = _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)) + key = UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)) cost = cost_by_unit.get(str(unit_name)) if cost_by_unit is not None else None yield key, _usage_unit_increment(units=units, cost=cost) def _sum_usage_unit_increments( logs_to_process: Sequence[Mapping[str, object]], -) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]: +) -> Mapping[UsageUnitKey, UsageUnitIncrement]: ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=_usage_unit_key) return MappingProxyType( { @@ -283,8 +284,16 @@ def _sum_usage_unit_increments( ) -async def _upsert_usage_unit_row( - prisma_client: PrismaClient, key: _UsageUnitKey, increment: _UsageUnitIncrement +class GuardrailUsageDatabase(Protocol): + @property + def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ... + + +async def upsert_usage_unit_row( + prisma_client: PrismaClient, + key: UsageUnitKey, + increment: UsageUnitIncrement, + outer_transaction: GuardrailUsageDatabase | None = None, ) -> None: row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = { "guardrail_id": key.guardrail_id, @@ -314,7 +323,12 @@ async def _upsert_usage_unit_row( "untracked_units": {"increment": increment.untracked_units}, }, } - await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) + table: Final = ( + outer_transaction.litellm_dailyguardrailusageunits + if outer_transaction is not None + else DailyGuardrailUsageUnitsRepository(prisma_client).table + ) + await table.upsert(where=where, data=data) async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None: @@ -363,7 +377,7 @@ async def process_spend_logs_guardrail_usage( ) index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {} - for payload in logs_to_process: + for payload in (row for row in logs_to_process if row.get("accounting_log_disabled") is not True): request_id = payload.get("request_id") start_time = _parse_payload_start_time(payload) if not isinstance(request_id, str) or not request_id or start_time is None: @@ -416,7 +430,10 @@ async def process_spend_logs_guardrail_usage( {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0} ) metrics_rows: Final = _merged_metric_rows(pending_metrics, evaluated_metrics) - unit_rows: Final = _merged_unit_rows(pending_units, _sum_usage_unit_increments(logs_to_process)) + unit_rows: Final = _merged_unit_rows( + pending_units, + _sum_usage_unit_increments(tuple(row for row in logs_to_process if "accounting_payload" not in row)), + ) if not metrics_rows and not index_rows and not unit_rows: return @@ -435,7 +452,7 @@ async def process_spend_logs_guardrail_usage( metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep ) failed_units: Final = await _upsert_rows_with_retry( - unit_rows, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep + unit_rows, partial(upsert_usage_unit_row, prisma_client), "usage unit", sleep ) if failed_metrics or failed_units: async with pending.lock: diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 002aeb145a5..636345a7ef7 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -26,6 +26,14 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import log_redis_failure from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import ( + ChargeOnceJob, + charged_poll_cost, + defers_batch_cost, + job_charged_once_marker_key, + job_to_charge_once, +) +from litellm.litellm_core_utils.core_helpers import finite_cost from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -44,6 +52,7 @@ _CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float]] = TypeAdapt _OPTIONAL_CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float | None]] = TypeAdapter( str | bytes | int | float | None, config=ConfigDict(strict=True, hide_input_in_errors=True) ) +_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) # Redis Lua script for atomic float increment with TTL. @@ -63,6 +72,30 @@ end return new_val """ +MAX_BUDGET_SESSION_CHARGE_JOB_ONCE_SCRIPT: Final = """ +local spend_key = KEYS[1] +local marker_key = KEYS[2] +local ttl = tonumber(ARGV[3]) + +local polls = redis.call('INCR', marker_key) +redis.call('EXPIRE', marker_key, ttl) +local amount = ARGV[2] +if polls == 1 then + amount = ARGV[1] +end +if tonumber(amount) <= 0 then + return '0' +end + +local existed = redis.call('EXISTS', spend_key) +redis.call('INCRBYFLOAT', spend_key, amount) +if existed == 0 then + redis.call('EXPIRE', spend_key, ttl) +end + +return amount +""" + # Default TTL for session budget counters (1 hour) DEFAULT_MAX_BUDGET_PER_SESSION_TTL: Final = 3600 @@ -87,12 +120,15 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) ) - if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - MAX_BUDGET_SESSION_INCREMENT_SCRIPT - ) - else: - self.increment_script = None + redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache + self.increment_script = ( + redis_cache.async_register_script(MAX_BUDGET_SESSION_INCREMENT_SCRIPT) if redis_cache is not None else None + ) + self.charge_job_once_script = ( + redis_cache.async_register_script(MAX_BUDGET_SESSION_CHARGE_JOB_ONCE_SCRIPT) + if redis_cache is not None + else None + ) @with_service_target("session_budgets") async def async_pre_call_hook( @@ -140,7 +176,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return None @with_service_target("session_budgets") - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: """ After a successful LLM call, increment the session spend by the response cost. """ @@ -168,15 +204,23 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): if max_budget is None: return - response_cost: Final = kwargs.get("response_cost") or 0.0 + callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs) + response_cost: Final = finite_cost(charged_poll_cost(callback_fields)) or 0.0 if response_cost <= 0: return + if defers_batch_cost(callback_fields, response_obj): + return + cache_key: Final = self._make_cache_key(str(session_id)) - await self._increment_spend(cache_key, float(response_cost)) + job: Final = job_to_charge_once(callback_fields, response_obj, response_cost) + if job is None: + await self._increment_spend(cache_key, response_cost) + else: + await self._charge_job_once(cache_key, job, response_cost) verbose_proxy_logger.debug( - "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", + "MaxBudgetPerSessionHandler: charged session %s for a call that cost %.6f", session_id, response_cost, ) @@ -221,6 +265,31 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): def _make_cache_key(self, session_id: str) -> str: return f"{{session_budget:{session_id}}}:spend" + async def _charge_job_once(self, cache_key: str, job: ChargeOnceJob, response_cost: float) -> None: + marker_key: Final = job_charged_once_marker_key(spend_key=cache_key, job=job) + if self.charge_job_once_script is not None: + try: + await self.charge_job_once_script( + keys=[cache_key, marker_key], + args=[str(response_cost), str(job.repeat_poll_cost), self.ttl], + ) + except Exception as e: + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "MaxBudgetPerSessionHandler: Redis job charge failed, falling back to in-memory", + e, + ) + else: + return + + polls: Final = await self.internal_usage_cache.dual_cache.in_memory_cache.async_increment( + key=marker_key, value=1, ttl=self.ttl + ) + charge: Final = response_cost if polls == 1 else job.repeat_poll_cost + if charge > 0: + await self._in_memory_increment_spend(cache_key, charge) + @with_service_target("session_budgets") async def _get_current_spend(self, cache_key: str) -> float: """Read current accumulated spend for a session.""" diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 1c10826124e..889fb9fb373 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,7 +19,12 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once +from litellm.litellm_core_utils.batch_utils import ( + ChargeOnceJob, + defers_batch_cost, + job_charged_once_marker_key, + job_to_charge_once, +) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions @@ -1208,18 +1213,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) return - batch_id: Final = batch_id_to_charge_once( - call_type=kwargs.get("call_type"), - response_obj=response_obj, - response_cost=response_cost, - ) + if defers_batch_cost(kwargs, response_obj): + return + + job: Final = job_to_charge_once(kwargs=kwargs, response_obj=response_obj, response_cost=response_cost) for entity_type, entity_id, resolved in resolved_budgets: await self._charge_entity( entity_type=entity_type, entity_id=entity_id, resolved=resolved, response_cost=response_cost, - batch_id=batch_id, + job=job, ) if self.dual_cache.redis_cache is not None: @@ -1236,7 +1240,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): entity_id: str | None, resolved: ResolvedModelBudget, response_cost: float, - batch_id: str | None, + job: ChargeOnceJob | None, ) -> None: budget_duration: Final = resolved.budget_config.budget_duration if budget_duration is None: @@ -1247,11 +1251,19 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): budget_model=resolved.budget_model, budget_duration=budget_duration, ) - if batch_id is not None and not await self._claim_batch_charge( + if job is not None and not await self._claim_job_charge( spend_key=spend_key, - batch_id=batch_id, + job=job, ttl_seconds=duration_in_seconds(budget_duration), ): + if job.repeat_poll_cost > 0: + await self._charge_entity( + entity_type=entity_type, + entity_id=entity_id, + resolved=resolved, + response_cost=job.repeat_poll_cost, + job=None, + ) return start_time_key: Final = model_budget_start_time_cache_key( entity_type=entity_type, @@ -1362,8 +1374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): spend=response_cost, ) - async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: - marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) + async def _claim_job_charge(self, spend_key: str, job: ChargeOnceJob, ttl_seconds: int) -> bool: + marker_key: Final = job_charged_once_marker_key(spend_key=spend_key, job=job) polls: Final = await self.dual_cache.async_increment_cache( key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 031bdfa19cd..ff67f2cb7c6 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,17 +4,19 @@ from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime, timezone +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast from pydantic import TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger -from litellm.batches.batch_utils import batch_cost_is_final from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import defers_batch_cost, repeat_poll_cost from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, finite_cost, @@ -24,7 +26,6 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, - UsageCaptureObservation, UsageObservation, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost @@ -38,6 +39,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import ( + ACCOUNTING_LEDGER_KWARG, DBSpendUpdateWriter, debitable_model_access_groups, get_llm_router, @@ -67,7 +69,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export should_store_prompts_and_responses_in_spend_logs, ) -from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, capture_callback, captured_callback_kwargs from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, @@ -77,7 +78,6 @@ from litellm.proxy.utils import ( from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, - LiteLLMBatch, StandardLoggingPayload, ) from litellm.utils import get_end_user_id_for_cost_tracking @@ -375,26 +375,13 @@ def _attempt_spend_row_kwargs( "combined_usage_object": attempt.usage, "response_cost": attempt.response_cost, "_accounting_usage_observed": UsageObservation.ESTIMATED, - "_accounting_usage_snapshot": UsageCaptureObservation.from_usage( - attempt.usage, UsageObservation.ESTIMATED, attempt.response_cost - ), - "_accounting_prior_captures": (), } row.pop("billed_failed_attempts", None) - row.pop("_accounting_capture", None) call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id") if isinstance(call_id, str): row["litellm_call_id"] = f"{call_id}{suffix}" row["parent_request_id"] = call_id metadata["parent_request_id"] = call_id - captured_attempts: Final = request_data.get("_accounting_billed_attempts") - if isinstance(captured_attempts, tuple): - for captured in TypeAdapter(tuple[object, ...]).validate_python(captured_attempts): - if ( - isinstance(captured, CapturedUsage) - and captured.payload["litellm_call_id"] == row["litellm_call_id"] - ): - row["_accounting_capture"] = captured standard_logging_object: Final = request_data.get("standard_logging_object") if isinstance(standard_logging_object, dict): payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict @@ -430,7 +417,7 @@ class _ProxyDBLogger(CustomLogger): async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: - if self.spend_event_producer is None or not is_offloadable_success(response_obj): + if self.spend_event_producer is None or not is_offloadable_success(response_obj, kwargs.get("call_type")): await self._PROXY_track_cost_callback(dict(kwargs), response_obj, start_time, end_time) return event: Final = build_spend_event( @@ -480,9 +467,6 @@ class _ProxyDBLogger(CustomLogger): "Failed to invalidate budget reservation counters after failure release failed" ) - capture_metadata: Final = await self._try_capture_failure_usage( - request_data, original_exception, user_api_key_dict - ) if completed_cost and tracking_disabled: return @@ -519,10 +503,8 @@ class _ProxyDBLogger(CustomLogger): _error_information.pop("traceback", None) _metadata["error_information"] = _error_information - _metadata = ( # rebind-ok: enriched identity is shared with the earlier immutable capture - {**_metadata, **capture_metadata} - if capture_metadata is not None - else await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( + _metadata = ( # rebind-ok: enrich identity before building failure spend rows + await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( metadata=_metadata, original_exception=original_exception ) ) @@ -629,106 +611,6 @@ class _ProxyDBLogger(CustomLogger): ), ) - async def _failure_capture_metadata( - self, request_data: Mapping[str, object], exception: Exception, auth: UserAPIKeyAuth - ) -> Mapping[str, object] | None: - from litellm.proxy.proxy_server import usage_capture - - if usage_capture.meter is None: - return None - captured_identity: Final = ( - _batch_accounting_for(request_data.get("litellm_call_id")) is not None - and request_data.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is not REJECTED_BEFORE_REQUEST_SETUP - ) - current_identity: Final = { - **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth), - "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth), - } - metadata: Final = ( - { - **dict.fromkeys(current_identity), - **{ - key: value - for key, value in _REQUEST_DATA.validate_python( - get_litellm_metadata_from_kwargs(dict(request_data)) - ).items() - if key in current_identity or key == "tags" - }, - } - if captured_identity - else current_identity - ) - return await self._enrich_failure_metadata_unless_db_stalled( - metadata=metadata, original_exception=exception, resolve_missing_key_identity=not captured_identity - ) - - async def _try_capture_failure_usage( - self, - request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer - original_exception: Exception, - auth: UserAPIKeyAuth, - ) -> Mapping[str, object] | None: - try: - metadata: Final = await self._failure_capture_metadata(request_data, original_exception, auth) - self._capture_failure_usage(request_data, original_exception, auth, metadata) - except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing or reservation release - verbose_proxy_logger.exception("Usage capture failed on the request failure path") - return None - return metadata - - def _capture_failure_usage( - self, - request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer - original_exception: Exception, - auth: UserAPIKeyAuth, - metadata: Mapping[str, object] | None = None, - ) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.db.exception_handler import exception_chain - from litellm.proxy.proxy_server import usage_capture - from litellm.proxy.spend_tracking.usage_capture import UsageAcknowledgementError - - if any(isinstance(error, UsageAcknowledgementError) for error in exception_chain(original_exception)): - return - usage_meter: Final = usage_capture.meter - - if usage_meter is None: - return - logging_obj: Final = request_data.get("litellm_logging_obj") - details: Final[Mapping[str, object]] = ( - _str_keyed_fields(logging_obj.model_call_details) if isinstance(logging_obj, Logging) else {} - ) - params: Final = _str_keyed_fields(request_data.get("litellm_params") or details.get("litellm_params")) - failure_data: Final = { - **details, - **request_data, - "litellm_params": { - **params, - "metadata": { - **_str_keyed_fields(params.get("metadata")), - **( - metadata - if metadata is not None - else { - **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth), - "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth), - } - ), - "status": "failure", - }, - }, - } - start: Final = logging_obj.start_time if isinstance(logging_obj, Logging) else request_data.get("start_time") - if not isinstance(start, datetime): - return - end: Final = datetime.now(timezone.utc) - captured: Final = tuple( - capture_callback(usage_meter, row, original_exception, start, end) - for _, row in _failure_spend_rows(failure_data, 0.0, None, self._deployment_lookup) - ) - request_data["_accounting_capture"] = captured[0] # rebind-ok: share one receipt with the legacy writer - request_data["_accounting_billed_attempts"] = captured[1:] # rebind-ok: share attempt receipts - async def _write_failure_spend_rows( self, user_api_key_dict: UserAPIKeyAuth, @@ -777,6 +659,7 @@ class _ProxyDBLogger(CustomLogger): token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id + accounting_rows: Final = tuple((cost, dict(row)) for cost, row in spend_rows) charged: Final = [ await self._spend_writer().update_database( token=token, @@ -784,17 +667,17 @@ class _ProxyDBLogger(CustomLogger): user_id=user_api_key_dict.user_id, end_user_id=end_user_id, team_id=user_api_key_dict.team_id, - kwargs=dict(row_kwargs), + kwargs=row_kwargs, completion_response=original_exception, start_time=start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, ) - for row_cost, row_kwargs in spend_rows + for row_cost, row_kwargs in accounting_rows ] counted_rows: Final = tuple( - row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0 + row for row, row_charged in zip(accounting_rows, charged, strict=True) if row_charged and row[0] > 0 ) if counted_rows and not ProxyUpdateSpend.disable_spend_updates(): await self._increment_failure_spend_counters( @@ -873,6 +756,7 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), project_id=user_api_key_dict.project_id, + **_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") @@ -914,20 +798,11 @@ class _ProxyDBLogger(CustomLogger): resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES, ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) - kwargs.update(await captured_callback_kwargs(kwargs, completion_response, start_time, end_time)) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) if accounting is not None and (budget_reservation is not None or litellm.apply_search_free_tier): accounting.write_started = True - if ( - isinstance(completion_response, LiteLLMBatch) - and kwargs.get("call_type") == CallTypes.aretrieve_batch.value - and not batch_cost_is_final(completion_response) - ): - verbose_proxy_logger.debug( - "Cost tracking deferred for batch %s still in status %s", - completion_response.id, - completion_response.status, - ) + if defers_batch_cost(kwargs, completion_response): + verbose_proxy_logger.debug("Cost tracking deferred for a batch whose cost is not final yet") await _release_budget_reservation(budget_reservation=budget_reservation) return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) @@ -989,25 +864,28 @@ class _ProxyDBLogger(CustomLogger): agent_id=metadata.get("billing_agent_id") or metadata.get("agent_id"), ): ## UPDATE DATABASE - charged: Final = await _update_database_and_spend_counters( - proxy_logging_obj=proxy_logging_obj, - increment_spend_counters=increment_spend_counters, - user_api_key=user_api_key, - user_id=user_id, - end_user_id=end_user_id, - team_id=team_id, - org_id=org_id, + charged_cost: Final = await _charge_request_cost( + partial( + _update_database_and_spend_counters, + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + team_id=team_id, + org_id=org_id, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + request_tags=tags, + model_access_groups=model_access_groups, + project_id=project_id, + ), kwargs=kwargs, - completion_response=completion_response, - start_time=start_time, - end_time=end_time, response_cost=response_cost, budget_reservation=budget_reservation, - request_tags=tags, - model_access_groups=model_access_groups, - project_id=project_id, ) - if not charged: + if charged_cost is None: return # update cache (fire-and-forget for backward compat: @@ -1017,7 +895,7 @@ class _ProxyDBLogger(CustomLogger): token=user_api_key, user_id=user_id, end_user_id=end_user_id, - response_cost=response_cost, + response_cost=charged_cost, team_id=team_id, parent_otel_span=parent_otel_span, tags=tags, @@ -1028,7 +906,7 @@ class _ProxyDBLogger(CustomLogger): token=user_api_key, key_alias=key_alias, end_user_id=end_user_id, - response_cost=response_cost, + response_cost=charged_cost, max_budget=end_user_max_budget, ) elif budget_reservation is not None: @@ -1322,6 +1200,58 @@ def _get_request_tags_for_cost_tracking( return None +class _LedgerCounterOptions(TypedDict): + windows_from_ledger: ReadOnly[NotRequired[bool]] + + +def _ledger_counter_options(ledger: object) -> _LedgerCounterOptions: + return {"windows_from_ledger": True} if ledger is True else {} + + +class _ChargeSpend(Protocol): + async def __call__( + self, + *, + kwargs: dict[str, object], # mutable-ok: update_database takes the logging kwargs dict itself + response_cost: float, + budget_reservation: dict[str, object] | None, # mutable-ok: reconcile stamps the shared reservation dict + ) -> bool: ... + + +async def _charge_request_cost( + charge_spend: _ChargeSpend, + kwargs: dict[str, object], # mutable-ok: the first charge passes the caller's logging kwargs dict on unchanged + response_cost: float, + budget_reservation: dict[str, object] | None, # mutable-ok: reconcile stamps the shared reservation dict +) -> float | None: + """What the request charged, since a poll of a job an earlier poll charged still pays for its own guardrails.""" + if await charge_spend(kwargs=kwargs, response_cost=response_cost, budget_reservation=budget_reservation): + return response_cost + standard_logging_payload: Final = kwargs.get("standard_logging_object") + cost_after_job: Final = repeat_poll_cost(kwargs.get("call_type"), response_cost, standard_logging_payload) + if cost_after_job <= 0: + return None + charged: Final = await charge_spend( + kwargs={ + **kwargs, + "response_cost": cost_after_job, + **( + { + "standard_logging_object": { + **_OBJECT_MAPPING.validate_python(standard_logging_payload), + "response_cost": cost_after_job, + } + } + if isinstance(standard_logging_payload, Mapping) + else {} + ), + }, + response_cost=cost_after_job, + budget_reservation=None, + ) + return cost_after_job if charged else None + + class _IncrementSpendCounters(Protocol): """The ``increment_spend_counters`` coroutine :func:`_update_database_and_spend_counters` awaits.""" @@ -1338,6 +1268,7 @@ class _IncrementSpendCounters(Protocol): request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + windows_from_ledger: bool = False, ) -> None: ... @@ -1475,6 +1406,7 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, + **_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index c0d81abeba7..18ce18da936 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -21,8 +21,8 @@ from typing import ( from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import MAX_FILE_LIST_LIMIT +from litellm.litellm_core_utils.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 64cde072dc5..627164abb93 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -31,7 +31,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider -from litellm._logging import verbose_proxy_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, @@ -58,6 +58,7 @@ from litellm.llms.anthropic.common_utils import ( anthropic_pass_through_api_base, merge_anthropic_beta_headers, ) +from litellm.llms.azure.common_utils import entra_rejection_in from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, is_azure_body_model_inference_endpoint, @@ -2381,6 +2382,16 @@ async def _relay_router_model( ) except httpx.HTTPStatusError as upstream_error: return await _relay_upstream_response(upstream_error.response) + except Exception as error: + rejection: Final = entra_rejection_in(error) + if rejection is None: + raise + raise ProxyException( + message=redact_internal_details_from_client_message(rejection.message), + type=ProxyErrorTypes.auth_error, + param=None, + code=rejection.status_code, + ) from error if not is_streaming_request: return await _relay_upstream_response(cast(httpx.Response, result)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0c487de2b42..f67d89a6bf6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6,6 +6,7 @@ import importlib import inspect import io import itertools +import logging import os import random import re @@ -426,6 +427,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolv from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) +from litellm.proxy.anthropic_endpoints.answered_status import is_provider_overload from litellm.proxy.auth.auth_checks import ( ROLE_BASED_PERMISSIONS_ADAPTER, ExperimentalUIJWTToken, @@ -877,8 +879,6 @@ from litellm.proxy.spend_tracking.spend_event_producer import ( SpendEventProducer, build_spend_event_producer, ) -from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, UsageCaptureRuntime -from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.proxy.telemetry.endpoints import router as telemetry_router from litellm.proxy.telemetry.middleware import TelemetryMiddleware from litellm.proxy.telemetry.runtime import TelemetryRuntime @@ -1392,7 +1392,6 @@ async def _flush_spend_around_logging_drain_on_shutdown( finally: await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock)) await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock)) - await usage_capture.stop(GracefulShutdownManager.remaining(clock=clock)) await flush_spend() @@ -3309,39 +3308,9 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): spend_event_producer: SpendEventProducer | None = None -usage_capture: Final = UsageCaptureRuntime() - - -def _configure_usage_capture() -> None: - settings: Final = _GENERAL_SETTINGS_VIEW.validate_python(general_settings) - capture_settings: Final = settings.get("accounting_capture") - if settings.get("accounting_sync_ack") and capture_settings is None: - raise ValueError("accounting_sync_ack requires accounting_capture") - if capture_settings is None: - if usage_capture.meter is not None: - raise ValueError("Disabling accounting_capture requires a worker restart") - return - if ( - prisma_client is None - or CollectorSettings().enabled - or os.getenv("SPEND_LOGS_URL") - or litellm.apply_search_free_tier - ): - raise ValueError( - "Usage capture requires local Prisma writes and disabled collector, remote writer and search free tier" - ) - from litellm.proxy.db.accounting_ledger import AccountingLedger - - ledger: Final = AccountingLedger.for_client(prisma_client) - writer: Final = UsageBatchWriter(ledger.insert_usage_batch, ledger.load_usage_batch) - usage_capture.start( - writer.write, UsageCaptureSettings.model_validate(capture_settings), load=ledger.load_usage_batch - ) - def cost_tracking(): global prisma_client, spend_event_producer - _configure_usage_capture() if prisma_client is not None: from litellm.integrations.shadow_eval_logger import ShadowEvalLogger @@ -3649,6 +3618,7 @@ async def increment_spend_counters( request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + windows_from_ledger: bool = False, ): """ Atomically increment spend counters for budget enforcement. @@ -3685,6 +3655,7 @@ async def increment_spend_counters( request_started_at=request_started_at, model_access_groups=model_access_groups, project_id=project_id, + windows_from_ledger=windows_from_ledger, ) @@ -3714,6 +3685,50 @@ def _finish_pending_read(absent_key: str, pending_read: str, remember_absent: bo user_api_key_cache.budget_windows_absent_cache.delete_cache(absent_key) +async def accounting_window_keys(token: str | None, team_id: str | None) -> tuple[Mapping[str, object], ...]: + scopes: Final = tuple( + (entity_type, entity_id, cache_key) + for entity_type, entity_id, cache_key in ( + (Litellm_EntityType.KEY, token, token), + (Litellm_EntityType.TEAM, team_id, f"team_id:{team_id}"), + ) + if entity_id is not None and cache_key is not None + ) + snapshots: Final = tuple( + [ + await _accounting_scope_windows(entity_type, entity_id, cache_key) + for entity_type, entity_id, cache_key in scopes + ] + ) + from itertools import chain + + return tuple(chain.from_iterable(snapshots)) + + +async def _accounting_scope_windows( + entity_type: Litellm_EntityType, entity_id: str, cache_key: str +) -> tuple[Mapping[str, object], ...]: + from litellm.proxy.spend_tracking.budget_reservation import fixed_budget_window + + cached: Final[object] = _UNTYPED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) + limits: Final = ( + await _stored_budget_limits(entity_type, entity_id) if cached is None else _cached_budget_limits(cached) + ) + if not isinstance(limits, list): + return () + windows: Final = TypeAdapter(tuple[object, ...]).validate_python(limits) + return tuple( + { + "entity_type": entity_type.value, + "entity_id": entity_id, + "window_duration": fixed[0], + "window_start": fixed[1].isoformat(), + } + for window in windows + if (fixed := fixed_budget_window(window)) is not None + ) + + async def _stored_budget_limits(entity_type: Litellm_EntityType, entity_id: str) -> object: """Read only: caching the row here could race a delete and put a just-deleted key back into the auth cache.""" if prisma_client is None: @@ -3755,6 +3770,7 @@ async def _increment_spend_counters_batched( request_started_at: datetime | None, model_access_groups: Sequence[str] | None, project_id: str | None = None, + windows_from_ledger: bool = False, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments.""" @@ -3796,15 +3812,16 @@ async def _increment_spend_counters_batched( if window_counter not in reserved_counter_keys else None ) - await _enqueue_window_spend_row_update( - entity_type=entity_type, - entity_id=entity_id, - reset_at=window_reset_at, - window_duration=duration, - window_start=window_start, - increment=cost, - request_started_at=request_started_at, - ) + if not windows_from_ledger: + await _enqueue_window_spend_row_update( + entity_type=entity_type, + entity_id=entity_id, + reset_at=window_reset_at, + window_duration=duration, + window_start=window_start, + increment=cost, + request_started_at=request_started_at, + ) return pending_window async def _window_increments( @@ -13408,6 +13425,13 @@ async def completion( _response.choices[0].text = e.message return _response except Exception as e: + if is_provider_overload(e): + raise await base_llm_response_processor.handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -18639,7 +18663,11 @@ async def model_deprecations( -H "Authorization: Bearer $LITELLM_MASTER_KEY" ``` """ - return collect_model_deprecations(llm_router=llm_router, warn_within_days=warn_within_days) + return collect_model_deprecations( + llm_router=llm_router, + warn_within_days=warn_within_days, + successor_log_level=logging.DEBUG, + ) def get_model_group_info( diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 9d4ede7b90d..5ded2fa5f5e 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1553,6 +1553,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "EMPIRIOLABS", + "provider_display_name": "EmpirioLabs AI", + "litellm_provider": "empiriolabs", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.empiriolabs.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "empiriolabs/aplomb-1" + }, { "provider": "EMPOWER", "provider_display_name": "Empower", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3f76bcf69ad..c4cd7c2d411 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2162,7 +2162,7 @@ _WS_FRAME: Final = TypeAdapter(dict[str, object]) _GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]]) _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" -_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") +_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools", "tool_choice", "extra_body") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") _METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c007b047087..3aaa74d4fc2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery { @@index([target, state, next_attempt_at]) } +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) +} + model LiteLLM_SubscriptionAccountTable { subscription_account_id String @id @default(uuid()) custom_llm_provider String diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index dcfd6148f54..af7eb742aa7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1264,6 +1264,15 @@ def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float return default_reserved_cost +def fixed_budget_window(window: object) -> tuple[str, datetime] | None: + fields: Final = _coerce_window(window) + start: Final = get_budget_window_start(window) + duration: Final = fields.get("budget_duration") + if not fields.get("reset_at") or not isinstance(duration, str) or start is None: + return None + return duration, start + + def get_budget_window_start(window: object) -> datetime | None: window_dict: Final = _coerce_window(window) budget_duration: Final = window_dict.get("budget_duration") diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index 031e84896b2..d054519cb9c 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -28,6 +28,7 @@ from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import LiteLLMBatch, Usage +from litellm.types.videos.utils import is_video_status_call_type SPEND_EVENT_VERSION: Final = 1 CACHE_OFF_KEY: Final = "Cache OFF" @@ -205,9 +206,10 @@ class SpendEventDecodeError: reason: str -def is_offloadable_success(response_obj: object) -> bool: - """Batch retrieves and interaction polls branch on the concrete response class, so they stay in-process.""" - return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) +def is_offloadable_success(response_obj: object, call_type: object) -> bool: + """Batch retrieves and interaction polls branch on the concrete response class, and a video status poll on + the status and video id the event leaves out, so they stay in-process.""" + return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) and not is_video_status_call_type(call_type) def _json_fallback(value: object) -> str: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 696f4e1ec48..fc54d1da0b2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -3,7 +3,7 @@ import json import os import re import secrets -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt from functools import reduce @@ -38,6 +38,12 @@ from litellm.constants import ( ) from litellm.integrations.otel.model.metadata import caller_session_id from litellm.integrations.otel.model.trace_controls import TraceControls +from litellm.litellm_core_utils.batch_utils import ( + BATCH_COST_REQUEST_ID_SUFFIX, + charged_poll_cost, + charged_video_id, + pays_only_its_guardrail, +) from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, @@ -58,7 +64,6 @@ from litellm.litellm_core_utils.internal_call_metadata import ( ) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, - UsageObservation, coerce_model_access_groups, is_valid_sha256_hash, request_model_access_groups_from_litellm_params, @@ -99,6 +104,7 @@ from litellm.types.utils import ( VectorStoreSearchResponse, batch_chat_usage, ) +from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: @@ -155,9 +161,7 @@ def _is_non_secret_key_value(value: str, *, key_alias: object = None) -> bool: ) -def _redact_logged_api_key( - value: str | None, *, already_redacted: bool = False, key_alias: object = None -) -> str | None: +def redact_logged_api_key(value: str | None, *, already_redacted: bool = False, key_alias: object = None) -> str | None: if not isinstance(value, str) or not value: return None stripped: Final = re.sub(r"(?i)^bearer ", "", value) @@ -321,7 +325,7 @@ def _get_spend_logs_metadata( and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias) and _trusted_hash == _raw_key ) - clean_metadata["user_api_key"] = _redact_logged_api_key( + clean_metadata["user_api_key"] = redact_logged_api_key( _raw_key, already_redacted=_already_redacted, key_alias=_key_alias ) clean_metadata["applied_guardrails"] = applied_guardrails @@ -343,7 +347,7 @@ def _get_spend_logs_metadata( return clean_metadata -BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" +JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX: Final = "_guardrail_cost" def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: @@ -371,6 +375,13 @@ def get_spend_logs_id( ) if resolved_id is None: return None + response_cost: Final = charged_poll_cost(kwargs) + charged_video: Final = charged_video_id(call_type, kwargs, resolved_id, response_cost) + if charged_video is not None: + return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" + call_id: Final = kwargs.get("litellm_call_id") + if pays_only_its_guardrail(call_type, kwargs, response_cost) and isinstance(call_id, str) and call_id: + return f"{call_id}{JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX}" return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" @@ -592,7 +603,6 @@ def get_logging_payload( start_time: datetime, end_time: datetime, llm_router: "Router | None" = None, - capture_usage: Callable[[SpendLogsPayload, str | None], None] | None = None, ) -> SpendLogsPayload: if kwargs is None: kwargs = {} @@ -680,14 +690,14 @@ def get_logging_payload( and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias) and _trusted_hash == api_key ) - api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or "" + api_key = redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or "" if ( standard_logging_payload is not None ): # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data api_key = ( api_key - or _redact_logged_api_key( + or redact_logged_api_key( standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True, key_alias=standard_logging_payload["metadata"].get("user_api_key_alias"), @@ -982,14 +992,6 @@ def get_logging_payload( payload.get("spend"), ) - if capture_usage is not None: - observed_usage: Final = ( - StandardLoggingPayloadSetup.get_usage_as_dict({"usage": usage}) - if usage and kwargs.get("_accounting_usage_observed") is not UsageObservation.UNKNOWN - else None - ) - capture_usage(payload, safe_dumps(observed_usage) if observed_usage is not None else None) - # Explicitly clear large intermediate objects to reduce memory pressure del response_obj_dict, usage, clean_metadata, additional_usage_values diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py deleted file mode 100644 index 19480a66678..00000000000 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ /dev/null @@ -1,1084 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import re -from collections.abc import Awaitable, Callable, Iterator, Mapping -from dataclasses import dataclass, replace -from datetime import datetime, timezone -from decimal import Decimal -from hashlib import sha256 -from itertools import chain, groupby, islice, repeat -from typing import TYPE_CHECKING, Final - -from pydantic import BaseModel, JsonValue, TypeAdapter - -from litellm._logging import verbose_proxy_logger -from litellm.accounting import ( - Dimension, - DimensionName, - InboundUsageEvent, - Metric, - MetricName, - OutboundUsageEvent, - Outcome, -) -from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY -from litellm.proxy.spend_tracking.usage_meter import ( - DurableUsageMeter, - UsageCaptureError, - UsageCaptureSettings, - usage_delivery_targets, -) -from litellm.types.utils import CallTypes, CustomPricingLiteLLMParams - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy._types import SpendLogsPayload - - -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite - - -class UsageAcknowledgementError(UsageCaptureError): - def __init__(self, cause: Exception) -> None: - super().__init__("Usage persistence did not acknowledge commit") - self.__cause__ = cause - - -@dataclass(frozen=True, slots=True) -class CapturedUsage: - payload: SpendLogsPayload - writes: tuple[UsageWrite, ...] - receipt: asyncio.Future[None] - - -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) - - -_FIELDS: Final = TypeAdapter(dict[str, object]) - - -_INSTANT: Final = TypeAdapter(datetime) - - -_DIMENSIONS: Final = ( - ("api_key", DimensionName.KEY), - ("user", DimensionName.USER), - ("team_id", DimensionName.TEAM), - ("organization_id", DimensionName.ORGANIZATION), - ("end_user", DimensionName.END_USER), - ("model", DimensionName.MODEL), - ("model_group", DimensionName.MODEL_GROUP), - ("model_id", DimensionName.DEPLOYMENT), - ("custom_llm_provider", DimensionName.PROVIDER), - ("mcp_namespaced_tool_name", DimensionName.TOOL), - ("agent_id", DimensionName.AGENT), -) - - -_USAGE_METRICS: Final = tuple( - metric - for metric in MetricName - if metric not in (MetricName.SPEND, MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS) -) - - -def _fields(value: object) -> Mapping[str, object]: - return _FIELDS.validate_python(value) if isinstance(value, Mapping) else {} - - -def _text(value: object) -> str | None: - return value if isinstance(value, str) and value else None - - -def _utc_time(value: datetime | str) -> datetime: - parsed: Final = _INSTANT.validate_python(value) - return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) - - -def _json(value: object) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - - -def _usage_metrics(usage: Mapping[str, JsonValue] | None) -> tuple[Metric, ...] | None: - if usage is None: - return None - return tuple( - Metric(name=metric, value=Decimal(str(value))) - for metric in _USAGE_METRICS - if isinstance(value := usage.get(metric.value), (int, float)) and not isinstance(value, bool) - ) - - -@dataclass(frozen=True, slots=True) -class UsageCaptureContext: - caller_id: str | None - attempt_id: str - parent_request_id: str | None - retry_of_attempt_id: str | None - job_id: str | None - charge_eligible: bool - request_count_eligible: bool - ownership: str - pricing: str - finalized_cost: str | None = None - allocation: str | None = None - outcome: Outcome | None = None - started_at: datetime | None = None - ended_at: datetime | None = None - phase: str = "execution" - extra_dimensions: tuple[Dimension, ...] = () - provider_timestamps: str = "{}" - usage_source: str = "measured" - included_attempt_ids: tuple[str, ...] = () - related_attempt_ids: tuple[str, ...] = () - - -def _database_usage_write(event: InboundUsageEvent | OutboundUsageEvent, envelope: Mapping[str, object]) -> UsageWrite: - from sys import maxsize - - from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure - - normalized: Final = _FIELDS.validate_python(safe_json_structure(envelope, max_depth=maxsize)) - if normalized.get("event") != event.model_dump(mode="json"): - raise UsageCaptureError("Usage event text cannot be normalized without changing identity or attribution") - serialized: Final = _json(normalized) - _JSON_OBJECT.validate_json(serialized) - return UsageWrite(event, serialized, usage_delivery_targets(event.dimensions)) - - -def usage_writes( - payload: SpendLogsPayload, usage_json: str | None, context: UsageCaptureContext -) -> tuple[UsageWrite, ...]: - fields: Final = _FIELDS.validate_python(payload) - tags: Final = TypeAdapter(tuple[str, ...]).validate_json(payload["request_tags"]) - dimensions: Final = ( - *(Dimension(name=name, value=value) for field, name in _DIMENSIONS if (value := _text(fields.get(field)))), - *(Dimension(name=DimensionName.TAG, value=tag) for tag in dict.fromkeys(tags)), - *context.extra_dimensions, - ) - outcome: Final = context.outcome or (Outcome.FAILED if payload["status"] == "failure" else Outcome.SUCCEEDED) - started: Final = context.started_at or _utc_time(payload["startTime"]) - ended: Final = context.ended_at or _utc_time(payload["endTime"]) - usage: Final = None if usage_json is None else _JSON_OBJECT.validate_json(usage_json) - outbound: Final = OutboundUsageEvent( - event_id=context.attempt_id, - attempt_id=context.attempt_id, - request_id=context.caller_id, - retry_of_attempt_id=context.retry_of_attempt_id, - call_type=CallTypes(payload["call_type"]), - started_at=started, - ended_at=ended, - outcome=outcome, - dimensions=dimensions, - metrics=_usage_metrics(None if context.usage_source == "unknown" else usage), - ) - envelope: Final = { - "event": outbound.model_dump(mode="json"), - "usage": usage, - "usage_source": "unknown" if usage is None else context.usage_source, - "included_attempt_ids": context.included_attempt_ids, - "related_attempt_ids": context.related_attempt_ids, - "pricing": _JSON_OBJECT.validate_json(context.pricing), - "charge_eligible": context.charge_eligible, - "request_count_eligible": context.request_count_eligible, - "ownership": context.ownership, - "parent_request_id": context.parent_request_id, - "job_id": context.job_id, - "phase": context.phase, - "provider_timestamps": _JSON_OBJECT.validate_json(context.provider_timestamps), - "finalized_cost": None - if context.finalized_cost is None - else _JSON_OBJECT.validate_json(context.finalized_cost), - "allocation": None if context.allocation is None else _JSON_OBJECT.validate_json(context.allocation), - } - outbound_write: Final = _database_usage_write(outbound, envelope) - if not context.request_count_eligible or context.caller_id is None: - return (outbound_write,) - inbound: Final = InboundUsageEvent( - event_id=f"{context.caller_id}:inbound", - request_id=context.caller_id, - call_type=outbound.call_type, - started_at=started, - ended_at=ended, - outcome=outcome, - dimensions=dimensions, - metrics=( - Metric(name=MetricName.REQUESTS, value=Decimal(1)), - Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(int(outcome is Outcome.SUCCEEDED))), - Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(int(outcome is not Outcome.SUCCEEDED))), - ), - ) - return ( - _database_usage_write(inbound, {"event": inbound.model_dump(mode="json")}), - outbound_write, - ) - - -_INCLUDED_IN_PARENT: Final = object() - - -_PRICING_FIELDS: Final = frozenset( - { - "service_tier", - "data_residency", - "vertex_location", - "region_name", - "base_model", - "custom_pricing", - "input_cost_per_token", - "output_cost_per_token", - "input_cost_per_second", - "output_cost_per_second", - "input_cost_per_character", - "output_cost_per_character", - "size", - "quality", - "n", - "duration", - "seconds", - "search_context_size", - "additional_response_cost", - } -) - - -def _selected_pricing(fields: Mapping[str, object]) -> Mapping[str, JsonValue]: - return _JSON_OBJECT.validate_python( - { - key: value - for key, value in fields.items() - if key in _PRICING_FIELDS or key in CustomPricingLiteLLMParams.model_fields - } - ) - - -def _response_fields(response: object) -> Mapping[str, object]: - if isinstance(response, BaseModel): - return _FIELDS.validate_python(response.model_dump()) - return _fields(response) - - -def _job_time(response: Mapping[str, object], key: str) -> datetime | None: - stamp: Final = response.get(key) - if isinstance(stamp, (int, float)) and not isinstance(stamp, bool): - return datetime.fromtimestamp(stamp, tz=timezone.utc) - return None - - -def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], response: object) -> UsageCaptureContext: - from dataclasses import asdict - - from litellm.litellm_core_utils.hidden_params import get_hidden_params - from litellm.litellm_core_utils.internal_call_metadata import is_background_response, is_unbilled_non_inference_call - from litellm.litellm_core_utils.litellm_logging import ( - Logging, - UsageCaptureObservation, - UsageObservation, - mask_api_base_credentials, - ) - from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost - from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation - from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups - - metadata: Final = _JSON_OBJECT.validate_json(payload["metadata"]) - params: Final = _fields(kwargs.get("litellm_params")) - request_metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) - standard: Final = _fields(kwargs.get("standard_logging_object")) - logging_obj: Final = kwargs.get("litellm_logging_obj") - raw_response: Final = _response_fields(response) - hidden: Final = _fields(get_hidden_params(response)) - tool: Final = _fields(kwargs.get("mcp_tool_call_metadata")) - tool_pricing: Final = _fields(tool.get("mcp_server_cost_info")) - caller: Final = payload["litellm_call_id"] - cached: Final = payload["cache_hit"].lower() == "true" - included_in_parent: Final = kwargs.get("_accounting_internal_aggregate") is _INCLUDED_IN_PARENT - billed_id: Final = payload["request_id"] - terminal: Final = raw_response.get("status") in { - "completed", - "complete", - "failed", - "cancelled", - "expired", - "incomplete", - } - deferred: Final = billed_id.endswith(("_batch_cost", "_background_cost")) and terminal - management: Final = ( - (billed_id.endswith(("_batch_cost", "_background_cost")) and not terminal) - or (is_background_response(response) and not terminal) - or is_unbilled_non_inference_call(payload["call_type"], request_metadata, response) - ) - attempt: Final = ( - f"{caller}_cache_hit" - if cached and caller - else f"{caller}:status" - if management - else f"{billed_id}:included" - if included_in_parent - else billed_id - ) - if not attempt or attempt == "None" or (cached and not caller): - raise UsageCaptureError("Usage capture requires a stable execution identity") - retry: Final = metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is True - origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) - parent: Final = _text(kwargs.get("parent_request_id")) or _text(request_metadata.get("parent_request_id")) - finalized_value: Final = ( - logging_obj.finalized_cost if isinstance(logging_obj, Logging) else kwargs.get("_accounting_finalized_cost") - ) - finalized: Final = finalized_value if isinstance(finalized_value, FinalizedCost) else None - allocation_value: Final = ( - logging_obj.cost_allowance_allocation - if isinstance(logging_obj, Logging) - else kwargs.get("_accounting_allocation") - ) - allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None - pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot") - observation: Final = kwargs.get("_accounting_usage_snapshot") - captured_pricing: Final = { - "model_map_information": ( - _JSON_OBJECT.validate_json(pricing_snapshot.model_map_information) - if isinstance(pricing_snapshot, PricingCapture) - else metadata.get("model_map_information") - ), - "cost_breakdown": metadata.get("cost_breakdown"), - **( - _JSON_OBJECT.validate_json(pricing_snapshot.inputs) - if isinstance(pricing_snapshot, PricingCapture) - else _pricing_inputs(kwargs, logging_obj, response) - ), - "tool": { - "name": tool.get("name"), - "cost_per_query": _fields(tool_pricing.get("tool_name_to_cost_per_query")).get( - _text(tool.get("name")) or "", tool_pricing.get("default_cost_per_query") - ), - } - if tool - else None, - "response_inputs": { - "audio_transcription_duration": hidden.get("audio_transcription_duration", raw_response.get("duration")), - "billed_units": _fields(raw_response.get("meta")).get("billed_units"), - "batch_models": hidden.get("batch_models"), - "batch_successful_requests": hidden.get("batch_successful_requests"), - "batch_failed_requests": hidden.get("batch_failed_requests"), - "service_tier": raw_response.get("service_tier"), - "region_name": hidden.get("region_name"), - }, - "reported_cost": observation.response_cost - if isinstance(observation, UsageCaptureObservation) - else standard.get( - "response_cost", kwargs.get("response_cost", _fields(get_hidden_params(response)).get("response_cost")) - ), - "adjustments": ( - _JSON_OBJECT.validate_json(pricing_snapshot.adjustments) - if isinstance(pricing_snapshot, PricingCapture) - else None - ), - } - return UsageCaptureContext( - caller_id=None if deferred else parent if retry or included_in_parent else caller, - attempt_id=attempt, - parent_request_id=parent, - retry_of_attempt_id=_text(request_metadata.get("retry_of_attempt_id")), - job_id=_text(raw_response.get("id")) if deferred or management else None, - charge_eligible=not cached and not included_in_parent and not management, - request_count_eligible=bool(caller) and not origin and not retry and not deferred and not included_in_parent, - ownership="included_in_parent" - if included_in_parent - else "aggregate" - if kwargs.get("additional_response_cost") is not None - else "attempt", - pricing=_json(captured_pricing), - finalized_cost=_json( - { - "response_cost": finalized.amount, - "cost_breakdown": finalized.logging_breakdown(), - "allowance_discount": finalized.allowance_discount, - } - ) - if finalized is not None - else ( - _json({"response_cost": standard.get("response_cost"), "cost_breakdown": standard.get("cost_breakdown")}) - if standard.get("cost_finalized") is True - else None - ), - allocation=_json(asdict(allocation)) if allocation is not None else None, - started_at=_job_time(raw_response, "created_at") if deferred else None, - ended_at=( - _job_time(raw_response, "completed_at") - or _job_time(raw_response, "failed_at") - or _job_time(raw_response, "cancelled_at") - or _job_time(raw_response, "expired_at") - ) - if deferred - else None, - outcome=( - Outcome.CANCELLED - if raw_response.get("status") == "cancelled" - or isinstance(response, asyncio.CancelledError) - or getattr(response, "status_code", None) == 499 - else Outcome.FAILED - if raw_response.get("status") in {"failed", "expired", "incomplete"} or isinstance(response, BaseException) - else None - ), - phase="terminal" if deferred else "management" if management else "execution", - usage_source="estimated" - if kwargs.get("_accounting_usage_observed") is UsageObservation.ESTIMATED - else "measured", - included_attempt_ids=TypeAdapter(tuple[str, ...]).validate_python( - kwargs.get("_accounting_included_attempt_ids", ()) - ), - related_attempt_ids=tuple(write.event.event_id for write in _prior_writes(kwargs)), - provider_timestamps=_json( - { - key: raw_response.get(key) - for key in ("created_at", "completed_at", "failed_at", "cancelled_at", "expired_at") - if raw_response.get(key) is not None - } - ) - if deferred - else "{}", - extra_dimensions=( - *( - (Dimension(name=DimensionName.TOOL, value=tool_name),) - if not payload.get("mcp_namespaced_tool_name") - and (tool_name := _text(tool.get("namespaced_tool_name"))) - else () - ), - *( - (Dimension(name=DimensionName.API_BASE, value=mask_api_base_credentials(api_base)),) - if (api_base := _text(payload["api_base"])) - else () - ), - *( - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value=group) - for group in get_request_model_access_groups(kwargs) - ), - *( - (Dimension(name=DimensionName.PROJECT, value=project),) - if (project := _text(request_metadata.get("user_api_key_project_id"))) - else () - ), - *( - (Dimension(name=DimensionName.SERVICE_TIER, value=tier),) - if (tier := _text(raw_response.get("service_tier") or kwargs.get("service_tier"))) - else () - ), - ), - ) - - -class _CaptureBuilder: - def __init__(self, kwargs: Mapping[str, object], response: object) -> None: - self.kwargs: Final = kwargs - self.response: Final = response - self.writes: tuple[UsageWrite, ...] = () - - def observe(self, payload: SpendLogsPayload, usage: str | None) -> None: - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation, UsageObservation - - snapshot: Final = self.kwargs.get("_accounting_usage_snapshot") - if isinstance(snapshot, UsageCaptureError): - raise snapshot - effective: Final = snapshot.usage_json if isinstance(snapshot, UsageCaptureObservation) else usage - parsed: Final = _JSON_OBJECT.validate_json(effective) if effective is not None else {} - batch: Final = _fields(parsed.get("batch_usage_by_model")) - unknown: Final = (bool(batch) or bool(parsed.get("batch_unknown_usage_requests"))) and all( - value is None for value in batch.values() - ) - observed: Final = ( - _json( - { - "batch_usage_by_model": batch, - "batch_unknown_usage_requests": parsed.get("batch_unknown_usage_requests"), - **( - {"batch_usage_by_request": parsed["batch_usage_by_request"]} - if "batch_usage_by_request" in parsed - else {} - ), - } - ) - if unknown - else effective - ) - context: Final = capture_context(payload, self.kwargs, self.response) - provenance: Final = ( - "unknown" - if unknown - else "estimated" - if isinstance(snapshot, UsageCaptureObservation) and snapshot.provenance is UsageObservation.ESTIMATED - else context.usage_source - ) - self.writes = usage_writes(payload, observed, replace(context, usage_source=provenance)) - - -def _prior_captures(kwargs: Mapping[str, object]) -> tuple[CapturedUsage, ...]: - prior: Final = kwargs.get("_accounting_prior_captures", ()) - return ( - tuple( - item for item in TypeAdapter(tuple[object, ...]).validate_python(prior) if isinstance(item, CapturedUsage) - ) - if isinstance(prior, tuple) - else () - ) - - -def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]: - return tuple(chain.from_iterable(item.writes for item in _prior_captures(kwargs))) - - -def capture_callback( - meter: DurableUsageMeter, - kwargs: Mapping[str, object], - response: object, - start_time: datetime, - end_time: datetime, -) -> CapturedUsage: - from litellm.proxy.db.db_spend_update_writer import get_llm_router - from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - - existing: Final = kwargs.get("_accounting_capture") - if isinstance(existing, CapturedUsage): - return existing - builder: Final = _CaptureBuilder(kwargs, response) - payload: Final = get_logging_payload( - dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe - ) - prerequisites: Final = tuple(item.receipt for item in _prior_captures(kwargs)) - return CapturedUsage(payload, builder.writes, meter.submit(builder.writes, prerequisites=prerequisites)) - - -def _batch_rows(response: object) -> tuple[Mapping[str, JsonValue], ...]: - usage: Final = _response_fields(_response_fields(response).get("usage")) - rows: Final = usage.get("batch_usage_by_request") - return TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(rows) if rows else () - - -def _batch_root(write: UsageWrite, rows: tuple[Mapping[str, JsonValue], ...]) -> UsageWrite: - event: Final = write.event.model_copy(update={"metrics": None}) - envelope: Final = _JSON_OBJECT.validate_json(write.envelope) - return _database_usage_write( - event, - { - **envelope, - "event": event.model_dump(mode="json"), - "usage": { - "batch_row_count": len(rows), - "batch_rows_sha256": sha256(_json(rows).encode()).hexdigest(), - "batch_charge_eligible": envelope["charge_eligible"], - }, - "usage_source": "unknown", - "charge_eligible": False, - "ownership": "aggregate", - }, - ) - - -def _batch_row_write(root: UsageWrite, row: Mapping[str, JsonValue], index: int) -> UsageWrite: - envelope: Final = _JSON_OBJECT.validate_json(root.envelope) - pricing: Final = _JSON_OBJECT.validate_python(envelope["pricing"]) - model: Final = _text(row["model"]) - usage: Final = _JSON_OBJECT.validate_python(row["usage"]) if row["usage"] is not None else None - identity: Final = f"{root.event.event_id}:row:{index}" - event: Final = OutboundUsageEvent.model_validate( - { - **root.event.model_dump(), - "event_id": identity, - "attempt_id": identity, - "metrics": _usage_metrics(usage), - "dimensions": ( - *(dimension for dimension in root.event.dimensions if dimension.name is not DimensionName.MODEL), - *((Dimension(name=DimensionName.MODEL, value=model),) if model else ()), - ), - } - ) - return _database_usage_write( - event, - { - **envelope, - "event": event.model_dump(mode="json"), - "usage": usage, - "usage_source": "unknown" if usage is None else "measured", - "charge_eligible": _fields(envelope["usage"])["batch_charge_eligible"], - "ownership": "attempt", - "parent_request_id": root.event.event_id, - "finalized_cost": None, - "allocation": None, - "pricing": { - **pricing, - "model_map_information": _fields(pricing.get("batch_models")).get(model or ""), - "batch_models": {}, - "reported_cost": None, - "cost_breakdown": None, - "response_inputs": { - key: value - for key, value in _fields(pricing.get("response_inputs")).items() - if key not in {"batch_models", "batch_successful_requests"} - }, - }, - }, - ) - - -@dataclass(frozen=True, slots=True) -class _BatchCapture: - payload: SpendLogsPayload - root: UsageWrite - rows: bytes - prerequisites: tuple[asyncio.Future[None], ...] - - -def _prepare_batch_capture( - kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime -) -> _BatchCapture: - from litellm.proxy.db.db_spend_update_writer import get_llm_router - from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - - rows: Final = tuple(sorted(_batch_rows(response), key=_json)) - builder: Final = _CaptureBuilder(kwargs, response) - payload: Final = get_logging_payload( - dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe - ) - if len(builder.writes) != 1 or _terminal_facts(builder.writes[0]) is None: - raise UsageCaptureError("Batch rows require one terminal job observation") - return _BatchCapture( - payload, - _batch_root(builder.writes[0], rows), - b"\n".join(_json((row, sum(1 for _ in group))).encode("utf-8") for row, group in groupby(map(_json, rows))), - tuple(prior.receipt for prior in _prior_captures(kwargs)), - ) - - -def _batch_groups(rows: bytes) -> Iterator[tuple[str, int]]: - return (TypeAdapter(tuple[str, int]).validate_json(match.group()) for match in re.finditer(rb"[^\n]+", rows)) - - -def _batch_writes(root: UsageWrite, rows: bytes) -> Iterator[UsageWrite]: - expanded: Final = chain.from_iterable(repeat(row, count) for row, count in _batch_groups(rows)) - return (_batch_row_write(root, _JSON_OBJECT.validate_json(row), index) for index, row in enumerate(expanded)) - - -def _batch_write_bound(root: UsageWrite, rows: bytes) -> int: - last_index: Final = sum(count for _, count in _batch_groups(rows)) - 1 - return max( - len(root.envelope.encode("utf-8")), - max( - len(_batch_row_write(root, _JSON_OBJECT.validate_json(row), last_index).envelope.encode("utf-8")) - for row, _ in _batch_groups(rows) - ), - ) - - -def _batch_capacity( - meter: DurableUsageMeter, root: UsageWrite, rows: bytes, *, extra_bytes: int = 0 -) -> tuple[int, int, int]: - retained: Final = len(rows) + len(root.envelope.encode("utf-8")) + extra_bytes - bound: Final = _batch_write_bound(root, rows) - chunk_bytes: Final = max( - bound, - min( - bound * min(meter.settings.batch_size, sum(count for _, count in _batch_groups(rows))), - meter.settings.max_pending_bytes // 4, - (meter.settings.max_pending_bytes - retained) // 2, - ), - ) - return retained + chunk_bytes, chunk_bytes, bound - - -async def _persist_batch_capture( - meter: DurableUsageMeter, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], - root: UsageWrite, - rows: bytes, - prerequisites: tuple[asyncio.Future[None], ...], - chunk_bytes: int, - reservation: object | None = None, -) -> None: - for receipt in prerequisites: - await asyncio.shield(receipt) - await meter.record_writes((root,)) - accepted: Final = await load((root.event.event_id,)) - if len(accepted) != 1: - raise UsageCaptureError("Committed batch snapshot could not be loaded") - frozen: Final = _reuse_terminal_snapshot(root, accepted[0]) - retained, capacity, bound = _batch_capacity(meter, frozen, rows, extra_bytes=len(root.envelope.encode("utf-8"))) - if reservation is not None: - meter.resize_producer(reservation, retained, capacity) - write_capacity: Final = capacity if reservation is not None else chunk_bytes - if bound > write_capacity: - raise UsageCaptureError("Batch observation exceeds reserved producer write capacity") - group_size: Final = min(meter.settings.batch_size, max(1, write_capacity // bound)) - remaining: Final = _batch_writes(frozen, rows) - while group := tuple(islice(remaining, group_size)): - await meter.record_writes(group) - - -async def capture_batch_callback( - meter: DurableUsageMeter, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None, - kwargs: Mapping[str, object], - response: object, - start_time: datetime, - end_time: datetime, -) -> CapturedUsage: - if load is None: - raise UsageCaptureError("Batch capture requires the existing ledger snapshot reader") - capture: Final = _prepare_batch_capture(kwargs, response, start_time, end_time) - await _persist_batch_capture( - meter, load, capture.root, capture.rows, capture.prerequisites, meter.settings.max_pending_bytes - ) - receipt: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - receipt.set_result(None) - return CapturedUsage(capture.payload, (capture.root,), receipt) - - -async def captured_callback_kwargs( - kwargs: Mapping[str, object], response: object, start_time: datetime | None, end_time: datetime | None -) -> Mapping[str, object]: - from litellm.proxy.proxy_server import usage_capture - - usage_meter: Final = usage_capture.meter - - if usage_meter is None: - return kwargs - if start_time is None or end_time is None: - verbose_proxy_logger.error("Usage capture requires callback start and end timestamps") - return kwargs - try: - captured: Final = ( - usage_capture.submit_batch(kwargs, response, start_time, end_time) - if _batch_rows(response) - else capture_callback(usage_meter, kwargs, response, start_time, end_time) - ) - except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing - verbose_proxy_logger.exception("Usage capture failed before legacy cost tracking") - return kwargs - return {**kwargs, "_accounting_capture": captured} - - -async def capture_proxy_response(logging_obj: object, response: object, *, background: bool = False) -> None: - from fastapi import HTTPException - - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.proxy_server import general_settings, usage_capture - - settings: Final = _FIELDS.validate_python(general_settings) - usage_meter: Final = usage_capture.meter - - if usage_meter is None or not isinstance(logging_obj, Logging): - return - sync_ack: Final = ( - settings.get("accounting_sync_ack") is True - and logging_obj.stream is not True - and not background - and logging_obj.call_type in {"acompletion", "atext_completion", "aembedding", "aresponses"} - ) - try: - capture_logging_usage(usage_meter, logging_obj, response, logging_obj.start_time, datetime.now(timezone.utc)) - captured: Final = logging_obj.model_call_details.get("_accounting_capture") - if sync_ack: - if not isinstance(captured, CapturedUsage): - raise UsageCaptureError("Usage capture was not admitted") - timeout: Final = TypeAdapter(float).validate_python( - settings.get("accounting_sync_ack_timeout_seconds", 20.0) - ) - await asyncio.wait_for(asyncio.shield(captured.receipt), timeout=timeout) - except Exception as error: # noqa: BLE001 # the configured response policy owns every capture failure - verbose_proxy_logger.exception("Usage capture failed before the response boundary") - if sync_ack and settings.get("accounting_sync_ack_failure_policy", "fail_closed") == "fail_closed": - acknowledgement: Final = UsageAcknowledgementError(error) - raise HTTPException(503, str(acknowledgement)) from acknowledgement - - -def capture_logging_usage( - meter: DurableUsageMeter, - logging_obj: object, - response: object, - start_time: datetime, - end_time: datetime, - freeze_only: bool = False, -) -> CapturedUsage | UsageCaptureError | None: - from litellm._internal_context import internal_parent_request_id, is_internal_call - from litellm.batches.batch_utils import batch_cost_is_final - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import LiteLLMBatch - - if not isinstance(logging_obj, Logging): - return - internal: Final = is_internal_call.get() - if isinstance(response, LiteLLMBatch) and batch_cost_is_final(response) and response.usage is None: - return - try: - if not internal and isinstance(logging_obj.model_call_details.get("_accounting_capture"), CapturedUsage): - return capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time) - normalized_result: Final = logging_obj.usage_capture_response(response) - prepared: Final = logging_obj.prepared_cost - normalized: Final = ( - {**_response_fields(normalized_result), "usage": prepared.usage.model_dump()} - if prepared is not None - else normalized_result - ) - if normalized is None: - return - details: Final = _FIELDS.validate_python(logging_obj.model_call_details) - params: Final = _fields(details.get("litellm_params")) - metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) - if metadata.get("batch_ignore_default_logging") is True: - return - existing_pricing: Final = details.get("_accounting_pricing_snapshot") - pricing: Final = ( - existing_pricing - if isinstance(existing_pricing, PricingCapture) and not internal - else _capture_pricing(logging_obj, normalized, details) - ) - if freeze_only or _batch_rows(normalized): - if not internal: - logging_obj.record_usage_pricing(pricing) - return - captured: Final = capture_callback( - meter, - { - **details, - "litellm_logging_obj": logging_obj, - "_accounting_capture": None if internal else details.get("_accounting_capture"), - "_accounting_pricing_snapshot": pricing, - "_accounting_internal_aggregate": _INCLUDED_IN_PARENT if internal else None, - "parent_request_id": internal_parent_request_id.get() if internal else details.get("parent_request_id"), - }, - normalized, - start_time, - end_time, - ) - if not internal: - logging_obj.record_usage_capture(captured) - except Exception as error: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging - verbose_proxy_logger.exception("Usage capture failed at the execution boundary") - return UsageCaptureError(str(error)) - else: - return captured - - -@dataclass(frozen=True, slots=True) -class PricingCapture: - model_map_information: str - adjustments: str - inputs: str - - -def _capture_pricing(logging_obj: Logging, normalized: object, details: Mapping[str, object]) -> PricingCapture: - import litellm - from litellm.litellm_core_utils.cost_discount import resolve_cost_discount - from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup - - params: Final = _fields(details.get("litellm_params")) - prepared: Final = logging_obj.prepared_cost - provider: Final = _text(details.get("custom_llm_provider")) - builtins: Final = logging_obj.standard_built_in_tools_params - pricing: Final = ( - prepared.payload["model_map_information"] - if prepared is not None - else StandardLoggingPayloadSetup.get_model_cost_information( - base_model=logging_obj.get_deployment_model_for_cost(), - custom_pricing=bool(params.get("custom_pricing")), - custom_llm_provider=_text(details.get("custom_llm_provider")), - init_response_obj=normalized, - api_base=_text(params.get("api_base")), - router_model_id=logging_obj.get_router_model_id(), - ) - ) - adjustments: Final = _json( - { - "discount": resolve_cost_discount( - litellm.cost_discount_config, - provider, - logging_obj.get_deployment_model_for_cost(), - _text(params.get("region_name") or params.get("vertex_location")), - ), - "margin": litellm.cost_margin_config.get(provider or "", litellm.cost_margin_config.get("global")), - "web_search_options": ( - _selected_pricing(_fields(builtins.get("web_search_options"))) - if builtins.get("web_search_options") is not None - else None - ), - "file_search": builtins.get("file_search") is not None, - } - ) - batch_usage: Final = _fields(_fields(_response_fields(normalized).get("usage")).get("batch_usage_by_model")) - return PricingCapture( - _json(pricing), - adjustments, - _json( - { - **_pricing_inputs(details, logging_obj, normalized), - "batch_models": { - model: StandardLoggingPayloadSetup.get_model_cost_information( - base_model=model, - custom_pricing=bool(params.get("custom_pricing")), - custom_llm_provider=provider, - init_response_obj={"model": model}, - api_base=_text(params.get("api_base")), - router_model_id=logging_obj.get_router_model_id(), - ) - for model in batch_usage - }, - } - ), - ) - - -def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: object) -> Mapping[str, object]: - from litellm.litellm_core_utils.litellm_logging import Logging - - return { - "request": _selected_pricing(kwargs), - "deployment": _selected_pricing(_fields(kwargs.get("litellm_params"))), - "optional_params": _selected_pricing(_fields(kwargs.get("optional_params"))), - "logging": logging_obj.usage_capture_pricing(response) if isinstance(logging_obj, Logging) else None, - "model_info": _selected_pricing( - _fields( - _fields( - _fields(kwargs.get("litellm_params")).get("litellm_metadata") - or _fields(kwargs.get("litellm_params")).get("metadata") - ).get("model_info") - ) - ), - } - - -def _terminal_facts(write: UsageWrite) -> Mapping[str, JsonValue] | None: - payload: Final = _JSON_OBJECT.validate_json(write.envelope) - if payload.get("phase") != "terminal": - return None - event: Final = _JSON_OBJECT.validate_python(payload["event"]) - return { - **{ - key: value - for key, value in payload.items() - if key not in {"event", "pricing", "finalized_cost", "allocation"} - }, - "event": {key: value for key, value in event.items() if key not in {"started_at", "ended_at"}}, - } - - -def _reuse_terminal_snapshot(candidate: UsageWrite, accepted: UsageWrite) -> UsageWrite: - facts: Final = _terminal_facts(candidate) - if facts is None or facts != _terminal_facts(accepted) or set(candidate.targets) != set(accepted.targets): - raise AccountingConflictError("Repeated terminal job disagrees with its accepted usage or attribution") - return accepted - - -class UsageBatchWriter: - def __init__( - self, - insert: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], - ) -> None: - self._insert: Final = insert - self._load: Final = load - - async def _restore(self, writes: tuple[UsageWrite, ...]) -> tuple[UsageWrite, ...]: - terminal: Final = tuple(write for write in writes if _terminal_facts(write) is not None) - if not terminal: - return writes - identities: Final = tuple(dict.fromkeys(write.event.event_id for write in terminal)) - stored: Final = await self._load(identities) - # Re-read the first accepted observation and pricing before the ledger's unchanged id/hash check. - snapshots: Final = {write.event.event_id: write for write in (*reversed(terminal), *stored)} - return tuple( - _reuse_terminal_snapshot(write, snapshots[write.event.event_id]) - if write.event.event_id in snapshots - else write - for write in writes - ) - - async def write(self, writes: tuple[UsageWrite, ...]) -> None: - restored: Final = await self._restore(writes) - try: - await self._insert(restored) - except AccountingConflictError: - if not any(_terminal_facts(write) is not None for write in writes): - raise - await self._insert(await self._restore(writes)) - - -class UsageCaptureRuntime: - def __init__(self) -> None: - self._batches: Final[set[asyncio.Task[None]]] = set() # mutable-ok: owned producer lifecycle - self._stopping = False - self.meter: DurableUsageMeter | None = None - self.load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None - - def start( - self, - write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - settings: UsageCaptureSettings, - *, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None, - ) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - - if self.meter is not None: - if self.meter.settings != settings: - raise ValueError("Changing accounting_capture requires a worker restart") - return - meter: Final = DurableUsageMeter(write_batch, settings) - self.meter = meter - self.load = load - self._stopping = False - Logging.usage_capture_callback = lambda logger, response, start, end, freeze: capture_logging_usage( - meter, logger, response, start, end, freeze - ) - - def submit_batch( - self, kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime - ) -> CapturedUsage: - meter: Final = self.meter - load: Final = self.load - if meter is None or load is None or self._stopping: - raise UsageCaptureError("Batch capture requires an active ledger writer and snapshot reader") - if len(self._batches) >= meter.settings.max_pending_records: - raise UsageCaptureError("Batch capture producer capacity exhausted") - capture: Final = _prepare_batch_capture(kwargs, response, start_time, end_time) - retained, chunk_bytes, _ = _batch_capacity( - meter, capture.root, capture.rows, extra_bytes=len(capture.root.envelope.encode("utf-8")) - ) - reservation: Final = meter.reserve_producer(retained, chunk_bytes) - task: Final = asyncio.create_task( - _persist_batch_capture( - meter, load, capture.root, capture.rows, capture.prerequisites, chunk_bytes, reservation - ), - name="accounting batch capture", - ) - task.add_done_callback(lambda _: meter.release_producer(reservation)) - self._batches.add(task) - task.add_done_callback(self._batch_finished) - return CapturedUsage(capture.payload, (capture.root,), task) - - def _batch_finished(self, task: asyncio.Task[None]) -> None: - self._batches.discard(task) - if task.cancelled(): - verbose_proxy_logger.error("Batch usage capture stopped before complete persistence") - elif (error := task.exception()) is not None: - verbose_proxy_logger.error("Batch usage capture failed", exc_info=error) - - async def stop(self, timeout: float) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - - if self.meter is None: - return - self._stopping = True - deadline: Final = asyncio.get_running_loop().time() + timeout - try: - if self._batches: - await asyncio.wait(tuple(self._batches), timeout=max(0, timeout)) - finally: - remaining: Final = tuple(self._batches) - for task in remaining: - task.cancel() - await asyncio.gather(*remaining, return_exceptions=True) - try: - await self.meter.stop(max(0, deadline - asyncio.get_running_loop().time())) - finally: - Logging.usage_capture_callback = None - self.meter = None - self.load = None diff --git a/litellm/proxy/spend_tracking/usage_meter.py b/litellm/proxy/spend_tracking/usage_meter.py deleted file mode 100644 index 431e4c2c3c3..00000000000 --- a/litellm/proxy/spend_tracking/usage_meter.py +++ /dev/null @@ -1,330 +0,0 @@ -from __future__ import annotations - -import asyncio -from collections import deque -from collections.abc import Awaitable, Callable, Mapping -from concurrent.futures import Future -from dataclasses import dataclass -from itertools import accumulate, chain, count, takewhile -from threading import Lock -from typing import Final - -from pydantic import BaseModel, ConfigDict, Field - -from litellm._logging import verbose_proxy_logger -from litellm.accounting import ( - DeliveryTarget, - Dimension, - DimensionName, - InboundUsageEvent, - OutboundUsageEvent, -) -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite - -_DAILY_TARGETS: Final = ( - (DimensionName.KEY, DeliveryTarget.DAILY_USER_SPEND), - (DimensionName.USER, DeliveryTarget.DAILY_USER_SPEND), - (DimensionName.TEAM, DeliveryTarget.DAILY_TEAM_SPEND), - (DimensionName.ORGANIZATION, DeliveryTarget.DAILY_ORGANIZATION_SPEND), - (DimensionName.END_USER, DeliveryTarget.DAILY_END_USER_SPEND), - (DimensionName.AGENT, DeliveryTarget.DAILY_AGENT_SPEND), - (DimensionName.TAG, DeliveryTarget.DAILY_TAG_SPEND), - (DimensionName.MODEL, DeliveryTarget.DAILY_MODEL_USAGE), - (DimensionName.TOOL, DeliveryTarget.DAILY_TOOL_SPEND), -) - - -def usage_delivery_targets(dimensions: tuple[Dimension, ...]) -> tuple[DeliveryTarget, ...]: - names: Final = frozenset(dimension.name for dimension in dimensions) - return tuple(dict.fromkeys(target for dimension, target in _DAILY_TARGETS if dimension in names)) - - -class UsageCaptureSettings(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - batch_size: int = Field(default=1000, ge=2, le=1000) - flush_interval_seconds: float = Field(default=5.0, gt=0, le=10) - max_pending_bytes: int = Field(default=16 * 1024 * 1024, ge=1024) - max_pending_records: int = Field(default=10000, ge=2) - - -class UsageCaptureError(RuntimeError): - pass - - -@dataclass(frozen=True, slots=True) -class _Pending: - writes: tuple[UsageWrite, ...] - size: int - committed: Future[None] - finished: asyncio.Event - prerequisites: tuple[asyncio.Future[None], ...] - - -async def _wait_for_batch(wakeup: asyncio.Event, timeout: float) -> None: - try: - await asyncio.wait_for(wakeup.wait(), timeout) - except asyncio.TimeoutError: - return - - -def _report_failure(receipt: Future[None]) -> None: - if not receipt.cancelled() and receipt.exception() is not None: - verbose_proxy_logger.error("Usage capture did not acknowledge commit", exc_info=receipt.exception()) - - -def _retryable_persistence_failure(error: Exception) -> bool: - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain - - outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(error) - return (outage is not None and not PrismaDBExceptionHandler.is_permanent_database_fault(outage)) or any( - isinstance(cause, Exception) - and ( - PrismaDBExceptionHandler.is_deadlock_error(cause) - or PrismaDBExceptionHandler.is_lock_timeout_error(cause) - or PrismaDBExceptionHandler.postgres_sqlstate(cause) == "40001" - ) - for cause in exception_chain(error) - ) - - -def _isolatable_persistence_failure(error: Exception) -> bool: - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain - - return isinstance(error, AccountingConflictError) or any( - isinstance(cause, Exception) and (PrismaDBExceptionHandler.postgres_sqlstate(cause) or "").startswith("22") - for cause in exception_chain(error) - ) - - -class DurableUsageMeter: - def __init__( - self, - write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - settings: UsageCaptureSettings, - *, - wait_for_batch: Callable[[asyncio.Event, float], Awaitable[None]] = _wait_for_batch, - retry_delay: Callable[[float], Awaitable[None]] = asyncio.sleep, - ) -> None: - self._write_batch: Final = write_batch - self.settings: Final = settings - self._wait_for_batch: Final = wait_for_batch - self._retry_delay: Final = retry_delay - self._loop: Final = asyncio.get_running_loop() - self._admission: Final = Lock() - self._queue: Final[deque[_Pending]] = deque() # mutable-ok: one bounded producer/consumer queue - self._wakeup: Final = asyncio.Event() - self._capacity_changed: Final = asyncio.Event() - self._reservations: Mapping[object, tuple[int, int]] = {} - self._reserved_bytes = 0 - self._task: asyncio.Task[None] | None = None - self._closed = False - self.pending_bytes = 0 - self.pending_records = 0 - - def submit( - self, writes: tuple[UsageWrite, ...], *, prerequisites: tuple[asyncio.Future[None], ...] = () - ) -> asyncio.Future[None]: - size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes) - with self._admission: - if ( - self._closed - or not writes - or len(writes) > self.settings.batch_size - or self.pending_bytes + size > self.settings.max_pending_bytes - or self.pending_records + len(writes) > self.settings.max_pending_records - ): - raise UsageCaptureError("Usage capture admission rejected: closed, empty or capacity exhausted") - receipt: Final = Future[None]() - receipt.set_running_or_notify_cancel() - receipt.add_done_callback(_report_failure) - self._queue.append(_Pending(writes, size, receipt, asyncio.Event(), prerequisites)) - self.pending_bytes += size - self.pending_records += len(writes) - self._loop.call_soon_threadsafe(self._ensure_worker) - wrapped: Final = asyncio.wrap_future(receipt, loop=self._loop) - wrapped.add_done_callback(lambda done: None if done.cancelled() else done.exception()) - return wrapped - - def reserve_producer(self, size: int, headroom: int) -> object: - token: Final = object() - self.resize_producer(token, size, headroom) - return token - - def resize_producer(self, token: object, size: int, headroom: int) -> None: - with self._admission: - previous: Final = self._reservations.get(token, (0, 0))[0] - delta: Final = size - previous - others: Final = (value[1] for key, value in self._reservations.items() if key is not token) - required: Final = max((headroom, *others)) - if ( - self._closed - or self.pending_bytes + delta > self.settings.max_pending_bytes - or self._reserved_bytes + delta + required > self.settings.max_pending_bytes - ): - raise UsageCaptureError("Batch capture producer byte capacity exhausted") - self._reservations = {**self._reservations, token: (size, headroom)} - self._reserved_bytes += delta - self.pending_bytes += delta - self._capacity_changed.set() - - def release_producer(self, token: object) -> None: - with self._admission: - size, _ = self._reservations[token] - self._reservations = {key: value for key, value in self._reservations.items() if key is not token} - self._reserved_bytes -= size - self.pending_bytes -= size - self._capacity_changed.set() - - def _ensure_worker(self) -> None: - if self.pending_records >= self.settings.batch_size or self._queue and self._queue[0].prerequisites: - self._wakeup.set() - if self._task is None or self._task.done(): - self._task = asyncio.create_task(self._run(), name="accounting usage capture") - - async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: - envelope: Final = '{"event":' + event.model_dump_json() + "}" - await asyncio.shield(self.submit((UsageWrite(event, envelope, usage_delivery_targets(event.dimensions)),))) - - async def record_writes(self, writes: tuple[UsageWrite, ...]) -> None: - size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes) - capacity: Final = self.settings.max_pending_bytes - self._reserved_bytes - if len(writes) > min(self.settings.batch_size, self.settings.max_pending_records) or ( - size > capacity and len(writes) > 1 - ): - middle: Final = len(writes) // 2 - await self.record_writes(writes[:middle]) - await self.record_writes(writes[middle:]) - return - if size > capacity: - raise UsageCaptureError("One usage observation exceeds capture capacity") - while (receipt := self._submit_available(writes)) is None: - self._wakeup.set() - await self._capacity_changed.wait() - self._wakeup.set() - await asyncio.shield(receipt) - - def _submit_available(self, writes: tuple[UsageWrite, ...]) -> asyncio.Future[None] | None: - self._capacity_changed.clear() - try: - receipt: Final = self.submit(writes) - except UsageCaptureError: - if self._closed or not self.pending_records or not writes: - raise - return None - return receipt - - async def flush(self) -> None: - self._ensure_worker() - self._wakeup.set() - if self._task is not None: - await asyncio.shield(self._task) - - async def close(self) -> None: - with self._admission: - self._closed = True - await self.flush() - - async def stop(self, timeout: float) -> None: - with self._admission: - self._closed = True - try: - await asyncio.wait_for(self.close(), timeout) - except asyncio.TimeoutError: - if self._task is not None: - self._task.cancel() - await asyncio.gather(self._task, return_exceptions=True) - while self._queue: - self._finish( - self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") - ) - verbose_proxy_logger.error("Usage capture shutdown deadline exhausted; commit acknowledgement incomplete") - - def _take_batch(self) -> tuple[_Pending, ...]: - with self._admission: - if self._queue[0].prerequisites: - return (self._queue.popleft(),) - independent: Final = takewhile(lambda item: not item.prerequisites, self._queue) - sizes: Final = accumulate(len(item.writes) for item in independent) - count: Final = sum(1 for _ in takewhile(lambda size: size <= self.settings.batch_size, sizes)) - return tuple(self._queue.popleft() for _ in range(count)) - - async def _persist(self, writes: tuple[UsageWrite, ...]) -> None: - for attempt in count(): - try: - await self._write_batch(writes) - except Exception as error: # noqa: BLE001 # retry only classified database failures, preserving frozen writes - if not _retryable_persistence_failure(error): - raise - verbose_proxy_logger.warning("Usage persistence unavailable; retaining admitted batch for retry") - await self._retry_delay(min(10.0, 0.2 * 2 ** min(attempt, 6))) - else: - return - - async def _commit(self, batch: tuple[_Pending, ...]) -> None: - writes: Final = tuple(chain.from_iterable(item.writes for item in batch)) - try: - if batch[0].prerequisites: - outcomes: Final = await asyncio.gather( - *(asyncio.shield(receipt) for receipt in batch[0].prerequisites), return_exceptions=True - ) - for outcome in outcomes: - if isinstance(outcome, BaseException): - raise UsageCaptureError("Usage dependency did not acknowledge commit") from outcome - await self._persist(writes) - except asyncio.CancelledError: - self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) - raise - except Exception as error: # noqa: BLE001 # isolate data faults, surface systemic failures to all producers - if len(batch) > 1 and _isolatable_persistence_failure(error): - try: - for item in batch: - await self._commit((item,)) - except asyncio.CancelledError: - self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) - raise - return - self._finish( - batch, - UsageCaptureError("Usage identity conflicts with accepted content") - if isinstance(error, AccountingConflictError) - else error, - ) - else: - self._finish(batch, None) - - def _finish(self, batch: tuple[_Pending, ...], error: Exception | None) -> None: - for item in batch: - if item.finished.is_set(): - continue - item.finished.set() - with self._admission: - self.pending_bytes -= item.size - self.pending_records -= len(item.writes) - self._capacity_changed.set() - if item.committed.done(): - continue - if error is None: - item.committed.set_result(None) - else: - item.committed.set_exception(error) - - async def _run(self) -> None: - try: - while self._queue: - await self._wait_for_batch(self._wakeup, self.settings.flush_interval_seconds) - self._wakeup.clear() - await self._commit(self._take_batch()) - if ( - self._closed - or self.pending_records >= self.settings.batch_size - or self._queue - and self._queue[0].prerequisites - ): - self._wakeup.set() - except (Exception, asyncio.CancelledError): # noqa: BLE001 # reject queued receipts if the writer exits abnormally - while self._queue: - self._finish( - self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") - ) - raise diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e68d348c16c..71ecae74058 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -168,7 +168,6 @@ from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int from litellm.litellm_core_utils.litellm_logging import ( BilledRoundsDispatch, Logging, - UsageObservation, deployment_pricing_model_info, ) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -214,6 +213,7 @@ from litellm.proxy.db.db_url_settings import ( from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, + db_write_never_sent, ) from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, @@ -278,6 +278,7 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.spend_log_insert import RejectedSpendLogRow from litellm.repositories.table_repositories import ( EndUserRepository, HealthCheckRepository, @@ -1805,39 +1806,6 @@ def _failure_fields_to_lift( if _prompt_estimate is not None else () ) - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation - - _recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed") - _current_usage: Final = ( - _current_failure_usage_to_lift( - _model_call_details, - request_data, - _model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, - _billed, - _deployment_id, - _base_model, - _estimate_call_type, - ) - if Logging.usage_capture_callback is not None - and (_billed or _model_call_details.get("combined_usage_object") is not None) - else None - ) - _current_observation: Final = ( - _recorded_observation - if isinstance(_recorded_observation, UsageObservation) - else UsageObservation.UNKNOWN - if _current_usage is None - else UsageObservation.KNOWN - if _model_call_details.get("combined_usage_object") is not None - else UsageObservation.ESTIMATED - ) - _capture_observation: Final = UsageCaptureObservation.from_usage( - _current_usage[0] if _current_usage is not None and isinstance(_current_usage[0], Usage) else None, - _current_observation, - float(_current_usage[1]) - if _current_usage is not None and isinstance(_current_usage[1], (int, float)) - else None, - ) _entries: Final = ( ("start_time", _model_call_details.get("start_time")), ("first_api_call_start_time", _first_handoff), @@ -1846,16 +1814,12 @@ def _failure_fields_to_lift( ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), ("billed_failed_attempts", _earlier_attempts or None), - ("_accounting_capture", _model_call_details.get("_accounting_capture")), ("_accounting_pricing_snapshot", _model_call_details.get("_accounting_pricing_snapshot")), ("_accounting_finalized_cost", _logging_obj.finalized_cost if isinstance(_logging_obj, Logging) else None), ( "_accounting_allocation", _logging_obj.cost_allowance_allocation if isinstance(_logging_obj, Logging) else None, ), - ("_accounting_usage_observed", _current_observation if Logging.usage_capture_callback is not None else None), - ("_accounting_usage_snapshot", _capture_observation if Logging.usage_capture_callback is not None else None), - ("_accounting_prior_captures", _model_call_details.get("_accounting_prior_captures")), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) @@ -8455,14 +8419,17 @@ async def enqueue_spend_logs( """ added: Final = sum(spend_log_row_bytes(row) for row in logs) async with prisma_client._spend_log_transactions_lock: - queued: Final = ( + queued: Final[tuple[Mapping[str, object], ...]] = ( tuple(logs) + tuple(prisma_client.spend_log_transactions) if at_head else tuple(prisma_client.spend_log_transactions) + tuple(logs) ) kept, kept_bytes = spend_log_queue_within_budget(queued, PrismaClient.spend_log_queue_bytes + added, max_bytes) - prisma_client.spend_log_transactions[:] = kept - PrismaClient.spend_log_queue_bytes = kept_bytes + rescued: Final = tuple( + accounting_only_spend_log(row) for row in queued[: len(queued) - len(kept)] if _carries_a_charge(row) + ) + prisma_client.spend_log_transactions[:] = (*rescued, *kept) + PrismaClient.spend_log_queue_bytes = kept_bytes + sum(spend_log_row_bytes(row) for row in rescued) if len(kept) < len(queued): verbose_proxy_logger.error( "Spend tracking - spend log queue is at its %d byte budget; dropped the %d oldest spend logs", @@ -8517,13 +8484,18 @@ class ProxyUpdateSpend: prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, end_user_list_transactions: dict[str, float], + outer_transaction: _EndUserSpendTransaction | None = None, ): for i in range(n_retry_times + 1): start_time = time.time() try: async with ( db_span("update_end_user_spend", "LiteLLM_EndUserTable"), - _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) as transaction, + ( + contextlib.nullcontext(outer_transaction) + if outer_transaction is not None + else _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) + ) as transaction, ): batcher: _EndUserSpendBatch async with transaction.batch_() as batcher: @@ -8545,6 +8517,8 @@ class ProxyUpdateSpend: break except Exception as e: + if outer_transaction is not None: + raise await DBSpendUpdateWriter.handle_spend_update_failure( e=e, attempt=i, @@ -8584,11 +8558,18 @@ class ProxyUpdateSpend: for i in range(n_retry_times + 1): try: base_url = os.getenv("SPEND_LOGS_URL", None) - if len(logs_to_process) > 0 and base_url is not None and db_writer_client is not None: + await _commit_export_accounting( + prisma_client, rows_to_write, base_url if db_writer_client is not None else None + ) + if ( + base_url is not None + and db_writer_client is not None + and any(row.get("accounting_log_disabled") is not True for row in rows_to_write) + ): if not base_url.endswith("/"): base_url += "/" verbose_proxy_logger.debug("base_url: %s", base_url) - json_data = json.dumps(rows_to_write) + json_data = json.dumps(_spend_log_export_rows(rows_to_write)) response = await db_writer_client.post( url=base_url + "spend/update", data=json_data, @@ -8666,6 +8647,28 @@ class ProxyUpdateSpend: return False +def _spend_log_export_rows(rows: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple( + { + key: value + for key, value in row.items() + if key not in {"accounting_payload", "accounting_log_disabled", "accounting_delivery_id"} + } + for row in rows + if row.get("accounting_log_disabled") is not True + ) + + +async def _commit_export_accounting( + prisma_client: PrismaClient, rows: Sequence[Mapping[str, object]], base_url: str | None +) -> None: + if base_url is None: + return + await SpendLogsRepository(prisma_client).insert_with_deliveries( + tuple(row for row in rows if "accounting_payload" in row), write_spend_logs=False + ) + + async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend.""" writer: Final = proxy_logging_obj.db_spend_update_writer @@ -8841,9 +8844,20 @@ async def update_daily_tag_spend( def _rollup_batch_never_applied(e: Exception) -> bool: """True when a drained rollup batch provably never reached its rows: Postgres had no - connection for it (53300) or cancelled it under ``lock_timeout`` before it took the row - lock (55P03). Either way re-sending the batch cannot double-count, so it is requeued.""" - return PrismaDBExceptionHandler.is_database_capacity_error(e) or PrismaDBExceptionHandler.is_lock_timeout_error(e) + connection for it (53300), was unreachable or refusing connections (P1001, 57P03), or + cancelled it under ``lock_timeout`` before it took the row lock (55P03). In each case + re-sending the batch cannot double-count, so it is requeued.""" + return ( + PrismaDBExceptionHandler.is_database_capacity_error(e) + or db_write_never_sent(e) + or PrismaDBExceptionHandler.is_lock_timeout_error(e) + ) + + +def _rollup_failure_stops_job(e: Exception) -> bool: + """True when the database cannot take any write, so the job stops instead of draining + the queues after this one into the same failure.""" + return PrismaDBExceptionHandler.is_database_capacity_error(e) or db_write_never_sent(e) async def update_spend_logs_job( @@ -8934,12 +8948,12 @@ async def _run_spend_logs_job( *prisma_client.tool_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during tool usage " + "Spend tracking - database unavailable or rollup row lock timed out during tool usage " "flush, requeued %s tool usage transactions for the next flush: %s", len(tool_usage_to_process), tool_tracking_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(tool_tracking_err): + if _rollup_failure_stops_job(tool_tracking_err): raise else: verbose_proxy_logger.error( @@ -8966,12 +8980,12 @@ async def _run_spend_logs_job( *prisma_client.model_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during model usage " + "Spend tracking - database unavailable or rollup row lock timed out during model usage " "flush, requeued %s model usage transactions for the next flush: %s", len(model_usage_to_process), model_usage_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(model_usage_err): + if _rollup_failure_stops_job(model_usage_err): raise else: verbose_proxy_logger.error( @@ -9148,10 +9162,45 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool: return ( PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) + or PrismaDBExceptionHandler.is_database_lifecycle_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) ) +_ACCOUNTING_ONLY_KEYS: Final = frozenset({"request_id", "call_type", "accounting_payload", "accounting_delivery_id"}) + + +def _carries_a_charge(row: Mapping[str, object]) -> bool: + return "accounting_payload" in row + + +def accounting_only_spend_log(row: Mapping[str, object]) -> Mapping[str, object]: + return { + **{key: value for key, value in row.items() if key in _ACCOUNTING_ONLY_KEYS}, + "accounting_log_disabled": True, + } + + +async def _keep_charges_of_dropped_logs(repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]]) -> None: + charged: Final = tuple(accounting_only_spend_log(row) for row in rows if _carries_a_charge(row)) + if not charged: + return + try: + await repo.insert_with_deliveries(charged) + except Exception as e: # noqa: BLE001 # the request log is already lost; keep the flush going for other rows + if not ( + PrismaDBExceptionHandler.is_rejected_row_error(e) + or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow)) + ): + raise + spend_log_error( + "Spend tracking - could not keep the charges of %d dropped spend log rows. error=%s", + len(charged), + str(e), + exc=e, + ) + + async def _create_spend_logs_with_poison_isolation( repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]], @@ -9180,10 +9229,13 @@ async def _create_spend_logs_with_poison_isolation( persist. Returns the budget left after this subtree. """ try: - await repo.table.create_many(data=rows, skip_duplicates=True) + await repo.insert_with_deliveries(rows) return failure_budget except Exception as e: - if not (PrismaDBExceptionHandler.is_prisma_data_error(e) or isinstance(e, UnicodeEncodeError)): + if not ( + PrismaDBExceptionHandler.is_rejected_row_error(e) + or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow)) + ): raise if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): raise @@ -9198,6 +9250,7 @@ async def _create_spend_logs_with_poison_isolation( str(e), exc=e, ) + await _keep_charges_of_dropped_logs(repo, rows) return budget_left if budget_left <= 0: spend_log_error( @@ -9205,6 +9258,7 @@ async def _create_spend_logs_with_poison_isolation( "isolation failure budget exhausted for this flush", len(rows), ) + await _keep_charges_of_dropped_logs(repo, rows) return 0 mid: Final = len(rows) // 2 remaining: Final = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], budget_left) @@ -9214,6 +9268,7 @@ async def _create_spend_logs_with_poison_isolation( "isolation failure budget exhausted for this flush", len(rows) - mid, ) + await _keep_charges_of_dropped_logs(repo, rows[mid:]) return 0 return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining) diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py new file mode 100644 index 00000000000..7dd92aecb2d --- /dev/null +++ b/litellm/repositories/accounting_delivery_repository.py @@ -0,0 +1,193 @@ +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import timedelta +from itertools import chain +from types import TracebackType +from typing import TYPE_CHECKING, Final, Literal, Protocol + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from litellm.proxy.db.db_spend_update_writer import SpendBatchManager + + +class LedgerTransaction(Protocol): + @property + def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ... + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + async def execute_raw(self, query: str, *args: object) -> int: ... + + def batch_(self) -> "SpendBatchManager": ... + + +class _TransactionManager(Protocol): + async def __aenter__(self) -> LedgerTransaction: ... + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None + ) -> None: ... + + +class LedgerDatabase(Protocol): + def tx(self, *, timeout: timedelta) -> _TransactionManager: ... + + +class _Delivery(BaseModel): + model_config = ConfigDict(frozen=True) + + delivery_id: str + attempts: int + accounting_payload: JsonValue + + +@dataclass(frozen=True, slots=True) +class DeliveryProjection: + row: Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class DeliveryBatchResult: + applied: int = 0 + failed: int = 0 + parked: tuple[str, ...] = () + error: str | None = None + + +class DeliveryMetrics(BaseModel): + model_config = ConfigDict(frozen=True) + + pending_count: int + parked_count: int + oldest_pending_seconds: float + batch_count: int + table_bytes: int + + +_DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...]) +_PAYLOAD: Final = TypeAdapter(dict[str, object]) +_TIMEOUT: Final = timedelta(seconds=60) +_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery" +WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP + AND (attempts > 0) = $2::boolean +ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED""" +_APPLIED: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL +WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """ +_FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET attempts = attempts + 1, last_error = $3, + next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer), + state = $4::"AccountingDeliveryState" +WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer +RETURNING delivery_id, attempts, accounting_payload""" + + +def _is_permanent_failure(error: Exception) -> bool: + """A delivery whose own data failed fails the same way on every retry, while any other error is the database's + and clears once it recovers.""" + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + + if PrismaDBExceptionHandler.is_database_service_unavailable_error(error): + return False + return isinstance(error, (ValueError, TypeError, KeyError)) or PrismaDBExceptionHandler.is_rejected_row_error(error) + + +class AccountingDeliveryRepository: + def __init__(self, database: LedgerDatabase) -> None: + self._database = database + + async def metrics(self) -> DeliveryMetrics: + async with self._database.tx(timeout=_TIMEOUT) as transaction: + rows: Final = await transaction.query_raw( + "SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, " + "COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, " + "COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) " + "AS oldest_pending_seconds, COUNT(*) FILTER (WHERE is_batch) AS batch_count, " + "pg_total_relation_size('\"LiteLLM_SpendLogDelivery\"') AS table_bytes " + 'FROM "LiteLLM_SpendLogDelivery"' + ) + return DeliveryMetrics.model_validate(rows[0]) + + async def cleanup_applied(self, *, retention_days: int = 1, batch_size: int = 1000) -> int: + if retention_days < 1 or batch_size < 1: + raise ValueError("Delivery retention and cleanup batch size must be positive") + async with self._database.tx(timeout=_TIMEOUT) as transaction: + return await transaction.execute_raw( + 'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id IN (' + 'SELECT delivery_id FROM "LiteLLM_SpendLogDelivery" ' + "WHERE state = 'applied' AND NOT is_batch AND applied_at < CURRENT_TIMESTAMP - make_interval(days => $1::integer) " + "ORDER BY applied_at, delivery_id LIMIT $2::integer FOR UPDATE SKIP LOCKED)", + retention_days, + batch_size, + ) + + async def apply_pending( + self, + apply: Callable[[LedgerTransaction, tuple[DeliveryProjection, ...]], Awaitable[None]], + *, + batch_size: int = 100, + max_attempts: int = 5, + retry_only: bool | None = None, + ) -> DeliveryBatchResult: + if batch_size < 1 or max_attempts < 1: + raise ValueError("Delivery batch size and maximum attempts must be positive") + claimed: tuple[_Delivery, ...] = () # rebind-ok: preserve claims for retry accounting after rollback + try: + async with self._database.tx(timeout=_TIMEOUT) as transaction: + retries: Final = ( + _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True)) + if retry_only is not False + else () + ) + claimed = retries or ( # rebind-ok: retain claimed identities after rollback + _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) + if retry_only is not True + else () + ) + if not claimed: + return DeliveryBatchResult() + await apply( + transaction, + tuple( + DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload)) + for delivery in claimed + ), + ) + await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed)) + return DeliveryBatchResult(applied=len(claimed)) + except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry + if not claimed: + raise + return await self._record_failure(claimed, error, max_attempts) + + async def _record_failure( + self, claimed: tuple[_Delivery, ...], error: Exception, max_attempts: int + ) -> DeliveryBatchResult: + parks: Final = len(claimed) == 1 and _is_permanent_failure(error) + async with self._database.tx(timeout=_TIMEOUT) as transaction: + updates: Final = tuple( + [ + await self._record_row_failure( + transaction, + delivery, + type(error).__name__, + "parked" if parks and delivery.attempts + 1 >= max_attempts else "pending", + ) + for delivery in claimed + ] + ) + return DeliveryBatchResult( + failed=len(claimed), parked=tuple(chain.from_iterable(updates)), error=type(error).__name__ + ) + + @staticmethod + async def _record_row_failure( + transaction: LedgerTransaction, delivery: _Delivery, error: str, state: Literal["pending", "parked"] + ) -> tuple[str, ...]: + updated: Final = _DELIVERIES.validate_python( + await transaction.query_raw(_FAILED, delivery.delivery_id, delivery.attempts, error, state) + ) + return tuple(row.delivery_id for row in updated) if state == "parked" else () diff --git a/litellm/repositories/spend_log_insert.py b/litellm/repositories/spend_log_insert.py new file mode 100644 index 00000000000..a958edb9f8a --- /dev/null +++ b/litellm/repositories/spend_log_insert.py @@ -0,0 +1,88 @@ +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import LiteralString + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + +_COLUMNS: Final[tuple[LiteralString, ...]] = ( + "request_id", + "call_type", + "api_key", + "spend", + "total_tokens", + "prompt_tokens", + "completion_tokens", + "startTime", + "endTime", + "request_duration_ms", + "completionStartTime", + "model", + "model_id", + "model_group", + "custom_llm_provider", + "api_base", + "user", + "metadata", + "cache_hit", + "cache_key", + "request_tags", + "team_id", + "organization_id", + "end_user", + "requester_ip_address", + "messages", + "response", + "session_id", + "status", + "mcp_namespaced_tool_name", + "billing_agent_id", + "agent_id", + "proxy_server_request", + "litellm_call_id", + "created_at", + "updated_at", +) +_JSON_COLUMNS: Final = frozenset({"metadata", "request_tags", "messages", "response", "proxy_server_request"}) + + +class RejectedSpendLogRow(ValueError): + """A row the raw insert cannot encode, so it is isolated like a row Postgres refuses.""" + + +def _json_value(column: str, value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if column not in _JSON_COLUMNS or not isinstance(value, str): + return value + try: + return _JSON.validate_json(value) + except ValidationError as error: + raise RejectedSpendLogRow(f"Spend log column {column} is not valid JSON") from error + + +def spend_log_insert_statement(rows: Sequence[Mapping[str, object]]) -> tuple[LiteralString, str]: + if not rows: + raise ValueError("Spend log insert requires at least one row") + fields: Final = frozenset(rows[0]) | {"created_at", "updated_at"} + if not fields.issubset(_COLUMNS) or any(frozenset(row) | {"created_at", "updated_at"} != fields for row in rows): + raise RejectedSpendLogRow("Spend log insert requires matching known columns") + columns: Final = ", ".join('"' + column + '"' for column in _COLUMNS if column in fields) + now: Final = datetime.now(timezone.utc).isoformat() + values: Final = tuple( + {"created_at": now, "updated_at": now, **{column: _json_value(column, value) for column, value in row.items()}} + for row in rows + ) + sql: Final = ( + 'INSERT INTO "LiteLLM_SpendLogs" (' + + columns + + ") SELECT " + + columns + + ' FROM jsonb_populate_recordset(NULL::"LiteLLM_SpendLogs", $1::jsonb) ' + "ON CONFLICT DO NOTHING RETURNING request_id" + ) + return sql, safe_dumps(values, non_finite_as_null=True) diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 2d3139f7a2c..10db81eca4b 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -7,8 +7,15 @@ These are thin wrappers for tables that do not (yet) need domain-specific query methods; richer repositories live in their own modules. """ +from collections.abc import Mapping, Sequence +from datetime import timedelta +from itertools import groupby from typing import TYPE_CHECKING, Any, Final, Generic +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions if TYPE_CHECKING: @@ -80,9 +87,117 @@ class OrganizationMembershipRepository(PrismaTableRepository["prisma_models.Lite table_name = "litellm_organizationmembership" +_SPEND_LOG_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def _accounting_payload(value: object) -> Mapping[str, object]: + return ( + _SPEND_LOG_FIELDS.validate_json(value) if isinstance(value, str) else _SPEND_LOG_FIELDS.validate_python(value) + ) + + +_ACCOUNTING_ONLY_FIELDS: Final = frozenset({"accounting_payload", "accounting_log_disabled", "accounting_delivery_id"}) + + +def _settles_once(row: Mapping[str, object]) -> bool: + """A batch or video cost row: update_database gives every other ledger row its own delivery id.""" + return "accounting_delivery_id" not in row + + +def _delivery_id(row: Mapping[str, object]) -> str: + return str(row["accounting_delivery_id"]) if not _settles_once(row) else str(row["request_id"]) + + class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]): table_name = "litellm_spendlogs" + async def insert_with_deliveries( + self, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool = True + ) -> int: + legacy: Final = tuple(row for row in rows if "accounting_payload" not in row) + legacy_count: Final = await self.table.create_many(data=legacy, skip_duplicates=True) if legacy else 0 + ledger: Final = tuple( + sorted( + {_delivery_id(row): row for row in reversed(rows) if "accounting_payload" in row}.values(), + key=_delivery_id, + ) + ) + if not ledger: + return legacy_count + database: Final[LedgerDatabase] = self.prisma_client.writer_db + async with database.tx(timeout=timedelta(seconds=60)) as transaction: + accepted: Final = await self._insert_ledger(transaction, ledger, write_spend_logs=write_spend_logs) + return legacy_count + accepted + + @staticmethod + async def _insert_ledger( + transaction: LedgerTransaction, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool + ) -> int: + from litellm.proxy.db.db_span import db_span + from litellm.repositories.spend_log_insert import spend_log_insert_statement + + payloads: Final = tuple( + { + "delivery_id": _delivery_id(row), + "request_id": row["request_id"], + "accounting_payload": _accounting_payload(row["accounting_payload"]), + "is_batch": _settles_once(row), + } + for row in rows + ) + inserted_deliveries: Final = await transaction.query_raw( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch) ' + "SELECT delivery_id, request_id, accounting_payload, is_batch FROM jsonb_to_recordset($1::jsonb) " + "AS inputs(delivery_id text, request_id text, accounting_payload jsonb, is_batch boolean) " + "ON CONFLICT (delivery_id) DO NOTHING RETURNING delivery_id", + safe_dumps(payloads, non_finite_as_null=True), + ) + accepted_ids: Final = frozenset(str(row["delivery_id"]) for row in inserted_deliveries) + charged_by_legacy: Final = await SpendLogsRepository._jobs_charged_by_legacy( + transaction, + tuple(_delivery_id(row) for row in rows if _settles_once(row) and _delivery_id(row) in accepted_ids), + ) + if charged_by_legacy: + await transaction.execute_raw( + 'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = ANY($1::text[])', list(charged_by_legacy) + ) + accepted: Final = accepted_ids - charged_by_legacy + local_rows: Final = tuple( + {key: value for key, value in row.items() if key not in _ACCOUNTING_ONLY_FIELDS} + for row in rows + if _delivery_id(row) in accepted and row.get("accounting_log_disabled") is not True + ) + if write_spend_logs and local_rows: + groups: Final = tuple( + tuple(group) + for _, group in groupby( + sorted(local_rows, key=lambda row: tuple(sorted(row))), key=lambda row: tuple(sorted(row)) + ) + ) + async with db_span("insert_spend_logs", "LiteLLM_SpendLogs"): + for group in groups: + await transaction.query_raw(*spend_log_insert_statement(group)) + return len(accepted) + + @staticmethod + async def _jobs_charged_by_legacy(transaction: LedgerTransaction, cost_row_ids: Sequence[str]) -> frozenset[str]: + """Batches and videos a pre-ledger poll already charged through their SpendLogs cost row.""" + if not cost_row_ids: + return frozenset() + rows: Final = await transaction.query_raw( + 'SELECT DISTINCT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1::text[]) ' + "AND status = 'success' AND spend > 0", + list(cost_row_ids), + ) + return frozenset(str(row["request_id"]) for row in rows) + + async def has_delivery(self, delivery_id: str) -> bool: + database: Final[LedgerTransaction] = self.prisma_client.writer_db + rows: Final = await database.query_raw( + 'SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = $1', delivery_id + ) + return bool(rows) + class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]): table_name = "litellm_budgetwindowspend" diff --git a/litellm/router_api.py b/litellm/router_api.py index 3007b5345c7..f22e9b1f351 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1402,6 +1402,9 @@ class RoutingContext: def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]: return self._router._params_sent_to_deployment(request, deployment) + def default_litellm_params(self) -> Mapping[str, object]: + return self._router._active_default_litellm_params() + def settings(self) -> RouterSettings: return self._router.settings() diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 83a1251d3c0..0a02408a717 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -27,11 +27,12 @@ import inspect import sys import warnings from collections import Counter -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload +from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet +from litellm.router_strategy.least_busy import picks_reserve if TYPE_CHECKING: from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged @@ -58,6 +59,7 @@ if TYPE_CHECKING: from litellm.types.utils import ModelInfo as ModelMapInfo _T: Final = TypeVar("_T") +_P: Final = ParamSpec("_P") DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: the live tally _note increments @@ -164,6 +166,18 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib return _DeprecatedAttribute(replacement, read) +def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coroutine[object, object, _T]]: + @functools.wraps(lookup) + async def looked_up( + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: forwards the wrapped lookup's ParamSpec unchanged + ) -> _T: + with picks_reserve(False): + return await lookup(*args, **kwargs) + + return looked_up + + # (id(router), member name) pairs whose subclass override is running in this context. _OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar( "router_overrides_running", default=frozenset() @@ -333,7 +347,7 @@ class DeprecatedRouterMembers: deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map) async_get_available_deployment = _attr( "select(model, request_kwargs)", - lambda r: r._core_async_get_available_deployment, + lambda r: _without_reserving(r._core_async_get_available_deployment), ) # -- observers ------------------------------------------------------------------------------ diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index e28ebf439b5..4b5ecc82581 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -282,7 +282,15 @@ from litellm.router_strategy.complexity_router.context_compaction import ( reject_recursive_compactor, surface_for_call, ) -from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import ( + CALL_IDENTITY, + SELECTION_RESERVATION_MAX_KEYS, + LeastBusyLoggingHandler, + call_id_of, + owning_picks, + pick_reserves, + picks_reserve, +) from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler @@ -306,6 +314,7 @@ from litellm.router_utils.cooldown_handlers import ( get_cooldown_deployment_info, get_cooldown_deployments, is_background_response_poll_not_found, + is_caller_entra_credential_failure, is_caller_scoped_auth_failure, is_caller_timeout_408, is_failure_outside_selected_deployment, @@ -1461,13 +1470,14 @@ class PythonCore(DeprecatedRouterMembers): (``pass_through``). """ if kind == "pick": - return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) + with picks_reserve(False): + return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts + model=model, + request_kwargs=request_kwargs if request_kwargs is not None else {}, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) if kind == "healthy": return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts model=model, @@ -1479,13 +1489,14 @@ class PythonCore(DeprecatedRouterMembers): health_check_probe=health_check_probe, ) if kind == "pass_through": - return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) + with picks_reserve(False): + return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts + model=model, + messages=messages, + input=input, + specific_deployment=specific_deployment, + request_kwargs=request_kwargs, + ) raise ValueError(f"unknown select kind {kind!r}") def apply(self, changes: ChangeSet) -> ApplyResult: @@ -1767,6 +1778,9 @@ class PythonCore(DeprecatedRouterMembers): self._core_cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._core_least_busy_reservations: InMemoryCache = InMemoryCache( + max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS + ) self._claude_code_session_router_cache: DualCache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache(), @@ -2270,7 +2284,9 @@ class PythonCore(DeprecatedRouterMembers): selector: RouterStrategySelector | None = None match self._normalize_strategy(strategy): case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler(router_cache=self._core_cache) + selector = LeastBusyLoggingHandler( + router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations + ) case RoutingStrategy.USAGE_BASED_ROUTING.value: selector = LowestTPMLoggingHandler( router_cache=self._core_cache, @@ -2735,6 +2751,7 @@ class PythonCore(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, + reserve: bool, ) -> Any | None: """ Asks the strategy selector for a deployment. Caller handles @@ -2748,6 +2765,8 @@ class PythonCore(DeprecatedRouterMembers): return await selector.async_get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, + reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), ) case "usage-based-routing": # `LowestTPMLoggingHandler` (v1) only exposes the sync @@ -2788,6 +2807,7 @@ class PythonCore(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, + reserve: bool, ) -> Any | None: """ Sync sibling of `_select_deployment_async`. Caller handles @@ -2804,6 +2824,8 @@ class PythonCore(DeprecatedRouterMembers): return selector.get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, + reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), ) case "usage-based-routing" | "usage-based-routing-v2": return selector.get_available_deployments( @@ -3523,12 +3545,13 @@ class PythonCore(DeprecatedRouterMembers): ) # pick the one that is available (lowest TPM/RPM) - deployment = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._drop_unsupported_classifier_reasoning_effort( deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment model=model, @@ -3842,11 +3865,6 @@ class PythonCore(DeprecatedRouterMembers): priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" } - if LiteLLMLogging.usage_capture_callback is not None: - from litellm._internal_context import internal_sub_call - - with internal_sub_call(cut_logging.litellm_call_id): - cut_logging.capture_usage_before_logging(priced_partial) return ContinuedPartial( usage=BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[partial_usage, inherited_usage] @@ -5216,8 +5234,19 @@ class PythonCore(DeprecatedRouterMembers): # _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping # the static shared deployment's id instead would let one tenant's bad credentials # cool down the deployment every other tenant sharing this config relies on. + if getattr(exception, "failed_deployment_id", None): + return effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({}) self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info})) + configured_deployment_id: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info).get( + "original_model_id" + ) + if isinstance(configured_deployment_id, str) and configured_deployment_id: + exception.failed_configured_deployment_id = configured_deployment_id # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: the exception carries the failure to _maybe_run_weighted_failover + + def _stamp_failed_attempt(self, exception: Exception, kwargs: Mapping[str, object], *, picked: bool) -> None: + if picked: + self._stamp_failed_deployment_id_with_effective_model_info(exception, MappingProxyType({}), kwargs) @staticmethod def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: @@ -6094,14 +6123,16 @@ class PythonCore(DeprecatedRouterMembers): def _image_generation(self, prompt: str, model: str, **kwargs): model_name: Final = "" + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - deployment: Final = self.get_available_deployment( - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init + model=model, + messages=[{"role": "user", "content": "prompt"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() @@ -6132,6 +6163,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aimage_generation(self, prompt: str, model: str, **kwargs): @@ -6157,11 +6191,12 @@ class PythonCore(DeprecatedRouterMembers): async def _aimage_generation(self, prompt: str, model: str, **kwargs): model_name = model start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) try: - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6222,6 +6257,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def atranscription(self, file: FileTypes, model: str, **kwargs): @@ -6268,10 +6306,11 @@ class PythonCore(DeprecatedRouterMembers): async def _atranscription(self, file: FileTypes, model: str, **kwargs): model_name: Final = model + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _atranscription()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6305,6 +6344,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): @@ -6361,10 +6403,11 @@ class PythonCore(DeprecatedRouterMembers): async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): model_name: Final = model + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _aspeech()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "prompt"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6398,6 +6441,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def arerank(self, model: str, **kwargs): @@ -6423,9 +6469,10 @@ class PythonCore(DeprecatedRouterMembers): async def _arerank(self, model: str, **kwargs): model_name = None + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _rerank()- model: %s", model) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, @@ -6457,6 +6504,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e def text_completion( @@ -6469,37 +6519,38 @@ class PythonCore(DeprecatedRouterMembers): **kwargs, ): messages: Final = [{"role": "user", "content": prompt}] - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - kwargs.setdefault("metadata", {}).update({"model_group": model}) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + kwargs["model"] = model + kwargs["prompt"] = prompt + kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) + kwargs.setdefault("metadata", {}).update({"model_group": model}) - # pick the one that is available (lowest TPM/RPM) - with endpoint_selection("text_completion", model=model): - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, + # pick the one that is available (lowest TPM/RPM) + with endpoint_selection("text_completion", model=model), picks_reserve(True): + deployment: Final = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + + data: Final = deployment["litellm_params"].copy() + for k, v in self._core_default_litellm_params.items(): + if k not in kwargs: # prioritize model-specific params > default router params + kwargs[k] = v + elif k == "metadata": + kwargs[k].update(v) + pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) + if pinned_model_id is not None: + kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults + + # call via litellm.completion() + return self._host.invoke("text_completion")( + **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} ) - - data: Final = deployment["litellm_params"].copy() - for k, v in self._core_default_litellm_params.items(): - if k not in kwargs: # prioritize model-specific params > default router params - kwargs[k] = v - elif k == "metadata": - kwargs[k].update(v) - pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults - - # call via litellm.completion() - return self._host.invoke("text_completion")( - **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} - ) - except Exception as e: - raise e + except Exception as e: + raise e async def atext_completion( self, @@ -6538,10 +6589,11 @@ class PythonCore(DeprecatedRouterMembers): raise e async def _atext_completion(self, model: str, prompt: str, **kwargs): + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": prompt}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6577,6 +6629,9 @@ class PythonCore(DeprecatedRouterMembers): if model is not None: self.fail_calls[model] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aadapter_completion( @@ -6608,10 +6663,11 @@ class PythonCore(DeprecatedRouterMembers): raise e async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs): + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _aadapter_completion()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, messages=[{"role": "user", "content": "default text"}], specific_deployment=kwargs.pop("specific_deployment", None), @@ -6647,6 +6703,9 @@ class PythonCore(DeprecatedRouterMembers): if model is not None: self.fail_calls[model] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs) -> object: @@ -7635,61 +7694,70 @@ class PythonCore(DeprecatedRouterMembers): """ handler_name: Final = original_function.__name__ metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") - try: - verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): - deployment: Final = self.get_available_deployment( + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) + self._update_kwargs_before_fallbacks( model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, + kwargs=kwargs, + metadata_variable_name=metadata_variable_name, + ) + with ( + self._video_call_pin(handler_name, kwargs), + self._route_capability_pin(handler_name), + picks_reserve(True), + ): + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init + model=model, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="generic_api_call", + splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, ) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="generic_api_call", - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] - self.total_calls[model_name] += 1 + self.total_calls[model_name] += 1 - # For passthrough routes, use the actual model from deployment - # and swap model name in endpoint if present - if "endpoint" in kwargs and kwargs["endpoint"]: - kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) - kwargs["model"] = model_name + # For passthrough routes, use the actual model from deployment + # and swap model name in endpoint if present + if "endpoint" in kwargs and kwargs["endpoint"]: + kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) + kwargs["model"] = model_name - # Perform pre-call checks for routing strategy - self.routing_strategy_pre_call_checks(deployment=deployment) + # Perform pre-call checks for routing strategy + self.routing_strategy_pre_call_checks(deployment=deployment) - custom_llm_provider: Final = provider_for_generic_call(data) + custom_llm_provider: Final = provider_for_generic_call(data) - response: Final = original_function( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - **kwargs, - } - ) + response: Final = original_function( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + **kwargs, + } + ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) - return response - except Exception as e: - verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) - if model is not None: - self.fail_calls[model] += 1 - raise e + self.success_calls[model_name] += 1 + verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) + return response + except Exception as e: + verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) + if model is not None: + self.fail_calls[model] += 1 + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) + raise e def embedding( self, @@ -7710,10 +7778,11 @@ class PythonCore(DeprecatedRouterMembers): def _embedding(self, input: str | list[object], model: str, **kwargs): model_name = None + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside embedding()- model: %s", model) - with embedding_deployment_selection(): - deployment: Final = self.get_available_deployment( + with embedding_deployment_selection(), picks_reserve(True): + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init model=model, input=input, specific_deployment=kwargs.pop("specific_deployment", None), @@ -7770,6 +7839,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e async def aembedding( @@ -7799,11 +7871,12 @@ class PythonCore(DeprecatedRouterMembers): async def _aembedding(self, input: str | list[object], model: str, **kwargs): model_name = None + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _aembedding()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) with embedding_deployment_selection(): - deployment: Final = await self._core_async_get_available_deployment( + deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init model=model, input=input, specific_deployment=kwargs.pop("specific_deployment", None), @@ -7850,6 +7923,9 @@ class PythonCore(DeprecatedRouterMembers): if model_name is not None: self.fail_calls[model_name] += 1 self._stamp_retry_skip_deployment_id(e, kwargs) + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) raise e #### FILES API #### @@ -8014,68 +8090,71 @@ class PythonCore(DeprecatedRouterMembers): Returns: VectorStoreCreateResponse """ - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + # If model is None, use the factory function approach (direct SDK call) + if model is None: + from litellm.vector_stores.main import acreate - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) + # Use the factory function to handle the call + factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") + return await factory_fn(**kwargs) - from litellm.vector_stores import acreate as avector_store_create_sdk + from litellm.vector_stores import acreate as avector_store_create_sdk - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } + parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "vector-store-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="avector_store_create", ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + model_client: Final = self._get_async_openai_model_client( + deployment=deployment, + kwargs=kwargs, + ) + self.total_calls[model_name] += 1 - return response - except Exception as e: - verbose_router_logger.exception( - "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e + # Get custom provider from deployment params + custom_llm_provider = data.get("custom_llm_provider") + _, inferred_custom_llm_provider, _, _ = get_llm_provider( + model=data["model"], + custom_llm_provider=custom_llm_provider, + ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider + + async with self._deployment_slot( + deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span + ): + response = await avector_store_create_sdk( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + "client": model_client, + **kwargs, + } + ) + + self.success_calls[model_name] += 1 + verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + + return response + except Exception as e: + verbose_router_logger.exception( + "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e + ) + if model is not None: + self.fail_calls[model] += 1 + raise e def _override_vector_store_methods_for_router(self): """ @@ -8323,12 +8402,13 @@ class PythonCore(DeprecatedRouterMembers): try: verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "batch-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(False): + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "batch-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id") data: Final = deployment["litellm_params"].copy() @@ -9088,8 +9168,9 @@ class PythonCore(DeprecatedRouterMembers): kwargs[metadata_variable_name] = meta if not isinstance(meta, dict): return None + configured_failed_id: Final[str | None] = getattr(exception, "failed_configured_deployment_id", None) prev_excluded: Final = set(meta.get("_failover_excluded_ids") or []) - excluded: Final = prev_excluded | {failed_id} + excluded: Final = prev_excluded | {configured_failed_id or failed_id} all_ids: Final = { (d.get("model_info") or {}).get("id") @@ -9472,38 +9553,39 @@ class PythonCore(DeprecatedRouterMembers): mock_timeout: Final = kwargs.pop("mock_timeout", None) - try: - self._handle_mock_testing_fallbacks( - kwargs=kwargs, - model_group=model_group, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + self._handle_mock_testing_fallbacks( + kwargs=kwargs, + model_group=model_group, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + ) - if mock_timeout is not None: - response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) - else: - response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Async Response: %s", response) - response = add_fallback_headers_to_response( - response=response, - attempted_fallbacks=0, - ) - return response - except Exception as e: - return await self.async_function_with_fallbacks_common_utils( - e, - disable_fallbacks, - fallbacks, - context_window_fallbacks, - content_policy_fallbacks, - model_group, - args, - kwargs, - include_fallback_errors=include_fallback_errors, - ) + if mock_timeout is not None: + response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) + else: + response = await self.async_function_with_retries(*args, **kwargs) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug("Async Response: %s", response) + response = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=0, + ) + return response + except Exception as e: + return await self.async_function_with_fallbacks_common_utils( + e, + disable_fallbacks, + fallbacks, + context_window_fallbacks, + content_policy_fallbacks, + model_group, + args, + kwargs, + include_fallback_errors=include_fallback_errors, + ) def _handle_mock_testing_fallbacks( self, @@ -10262,8 +10344,9 @@ class PythonCore(DeprecatedRouterMembers): return False exception_status: Final = getattr(exception, "status_code", "") + failure_call_details: Final = _CALL_KWARGS_ADAPTER.validate_python(kwargs) - if is_caller_timeout_408(kwargs, exception_status): + if is_caller_timeout_408(failure_call_details, exception_status): verbose_router_logger.debug( "Router: Exiting 'deployment_callback_on_failure' without cooldown. " "A timeout the caller set caused this 408, not the deployment's health." @@ -10287,6 +10370,18 @@ class PythonCore(DeprecatedRouterMembers): "Caller-scoped OAuth authentication failed, not the deployment." ) return False + if is_caller_entra_credential_failure( + caller_failure_deployment, + self._active_default_litellm_params(), + failure_call_details, + exception_status, + exception, + ): + verbose_router_logger.debug( + "Router: Exiting 'deployment_callback_on_failure' without cooldown. " + "Entra ID rejected a credential the caller sent, not the deployment's own." + ) + return False exception_headers: Final = exception_mapping_utils._get_response_headers(original_exception=exception) @@ -10795,6 +10890,7 @@ class PythonCore(DeprecatedRouterMembers): messages=None, input=None, request_kwargs=routing_kwargs, + reserve=False, ) if selected is None: no_deployment: Final = await async_raise_no_deployment_exception( @@ -16153,6 +16249,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, + reserve=pick_reserves(default=True), ) async def _async_get_available_deployment( @@ -16162,6 +16259,7 @@ class PythonCore(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, specific_deployment: bool | None, + reserve: bool, ): if ( self._core_routing_strategy != "usage-based-routing-v2" @@ -16258,6 +16356,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=reserve and request_kwargs not in (None, {}), ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( @@ -16427,6 +16526,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17304,6 +17404,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17484,6 +17585,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 1d502dd8d95..98dd39b40e3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,7 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once +from litellm.litellm_core_utils.batch_utils import ( + ChargeOnceJob, + defers_batch_cost, + job_charged_once_marker_key, + job_to_charge_once, +) from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) @@ -52,7 +57,7 @@ from litellm.types.utils import BudgetConfig as GenericBudgetInfo DEFAULT_REDIS_SYNC_INTERVAL: Final = 1 _CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) -_BATCH_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) +_JOB_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) class _LiteLLMParamsDictView: @@ -545,13 +550,16 @@ class RouterBudgetLimiting(CustomLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is required") - batch_response, response_cost = _BATCH_CHARGE_VALUES.validate_python( + polled_response, response_cost = _JOB_CHARGE_VALUES.validate_python( (response_obj, _CALLBACK_FIELDS.validate_python(standard_logging_payload).get("response_cost", 0)) ) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") - batch_id: Final = batch_id_to_charge_once(callback_fields.get("call_type"), batch_response, response_cost) + if defers_batch_cost(callback_fields, polled_response): + return + + job: Final = job_to_charge_once(callback_fields, polled_response, response_cost) budget_config: Final = ( self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None @@ -565,7 +573,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=spend_key, start_time_key=start_time_key, response_cost=response_cost, - batch_id=batch_id, + job=job, ) deployment_budget_config: Final = self._get_budget_config_for_deployment(model_id) @@ -578,7 +586,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=deployment_spend_key, start_time_key=deployment_start_time_key, response_cost=response_cost, - batch_id=batch_id, + job=job, ) request_tags: Final = get_tags_from_request_kwargs( @@ -596,7 +604,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=_tag_spend_key, start_time_key=_tag_start_time_key, response_cost=response_cost, - batch_id=batch_id, + job=job, ) async def _increment_spend_for_key( @@ -605,14 +613,21 @@ class RouterBudgetLimiting(CustomLogger): spend_key: str, start_time_key: str, response_cost: float, - batch_id: str | None = None, + job: ChargeOnceJob | None = None, ): if budget_config.budget_duration is None: return current_time: Final = self._clock() new_window_seconds: Final = budget_window_seconds(current_time, budget_config.budget_duration) - if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, new_window_seconds): + if job is not None and not await self._claim_job_charge(spend_key, job, new_window_seconds): + if job.repeat_poll_cost > 0: + await self._increment_spend_for_key( + budget_config=budget_config, + spend_key=spend_key, + start_time_key=start_time_key, + response_cost=job.repeat_poll_cost, + ) return budget_start: Final = await self._get_or_set_budget_start_time( @@ -639,8 +654,8 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost) - async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: - marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) + async def _claim_job_charge(self, spend_key: str, job: ChargeOnceJob, ttl_seconds: int) -> bool: + marker_key: Final = job_charged_once_marker_key(spend_key=spend_key, job=job) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is not None: polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 8e0c3d6a7e7..22159079f98 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,6 +1,11 @@ +import contextlib import logging -from collections.abc import Mapping, Sequence -from typing import Final +import threading +import time +import uuid +from collections.abc import Callable, Generator, Mapping, Sequence +from contextvars import ContextVar +from typing import Final, Protocol, TypeAlias from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -8,11 +13,31 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type +from litellm.types.utils import CallTypes IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 +SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 +SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 +_RESERVATIONS_LOCK: Final = threading.Lock() +_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves") +_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None) + + +@contextlib.contextmanager +def picks_reserve(reserve: bool) -> Generator[None, None, None]: + token: Final = _PICK_RESERVES.set(reserve) + try: + yield + finally: + _PICK_RESERVES.reset(token) + + +def pick_reserves(default: bool) -> bool: + return _PICK_RESERVES.get(default) class _ModelInfo(TypedDict, total=False): @@ -21,10 +46,12 @@ class _ModelInfo(TypedDict, total=False): class _Metadata(TypedDict, total=False): model_group: ReadOnly[str | None] + model_info: ReadOnly[_ModelInfo | None] class _LitellmParams(TypedDict, total=False): metadata: ReadOnly[_Metadata | None] + litellm_metadata: ReadOnly[_Metadata | None] model_info: ReadOnly[_ModelInfo | None] @@ -32,6 +59,10 @@ class _CallKwargs(TypedDict, total=False): litellm_params: ReadOnly[_LitellmParams | None] +class CallIdentity(TypedDict, total=False): + litellm_call_id: ReadOnly[object] + + class _DeploymentModelInfo(TypedDict): id: ReadOnly[str | int] @@ -41,38 +72,85 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +CALL_IDENTITY: Final = TypeAdapter(CallIdentity) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) -_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) +_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None) +_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None) +_Pick: TypeAlias = tuple[float, str | None] +_PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None) +_PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: return f"{model_group}_request_count:{deployment_id}" +def _model_group(metadata: _Metadata | None) -> str | None: + return metadata.get("model_group") if metadata else None + + +def _deployment_id(model_info: _ModelInfo | None) -> str | int | None: + return model_info.get("id") if model_info else None + + +def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None: + return _deployment_id(metadata.get("model_info")) if metadata else None + + def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: - if is_batch_retrieve_call_type(kwargs.get("call_type")): + call_type: Final = kwargs.get("call_type") + if is_batch_retrieve_call_type(call_type) or call_type == CallTypes.arealtime.value: return None try: call: Final = _CALL_KWARGS.validate_python(kwargs) except ValidationError: return None - litellm_params: Final = call.get("litellm_params") - metadata: Final = litellm_params.get("metadata") if litellm_params else None - model_info: Final = litellm_params.get("model_info") if litellm_params else None - model_group: Final = metadata.get("model_group") if metadata else None - deployment_id: Final = model_info.get("id") if model_info else None + litellm_params: Final = call.get("litellm_params") or _LitellmParams() + routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata")) + model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None) + deployment_id: Final = next( + ( + deployment_id + for deployment_id in ( + _deployment_id(litellm_params.get("model_info")), + *map(_routed_deployment_id, routing_metadata), + ) + if deployment_id is not None + ), + None, + ) if model_group is None or deployment_id is None: return None return model_group, str(deployment_id) -def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: +def _reservation_key(deployment_id: str) -> str: + return f"least_busy_reservations:{deployment_id}" + + +def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str: + return f"least_busy_started:{call_id}:{_request_count_key(*ref)}" + + +def _picked_deployments_key(call_id: str) -> str: + return f"least_busy_picked:{call_id}" + + +def call_id_of(kwargs: Mapping[str, object]) -> str | None: + call_id: Final = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: return tuple( - _request_count_key(model_group, str(deployment["model_info"]["id"])) - for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) ) +def _request_count_keys(model_group: str, deployment_ids: tuple[str, ...]) -> tuple[str, ...]: + return tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids) + + def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: return tuple(0 if value is None else int(value) for value in values) @@ -84,12 +162,21 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) -def _least_busy( - healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] -) -> Mapping[str, object] | None: - if not healthy_deployments: - return None - return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] +def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]: + return tuple(pick for pick in picks if now - pick[0] < SELECTION_RESERVATION_TTL_SECONDS) + + +def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]: + tags: Final = tuple(tag for _, tag in picks) + owner: Final = next((tag for tag in (call_id, _PICK_OWNER.get()) if tag is not None and tag in tags), None) + if owner not in tags: + return picks + index: Final = tags.index(owner) + return (*picks[:index], *picks[index + 1 :]) + + +def _least_busy_index(counts: tuple[int, ...]) -> int: + return min(range(len(counts)), key=lambda index: counts[index]) def _warn_unreadable(model_group: str, error: Exception) -> None: @@ -111,24 +198,156 @@ def _warn_unwritable(key: str, error: Exception) -> None: ) +class _ReservationStore(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + +def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]: + return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or () + + +def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]: + return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or () + + +def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]: + return _PICKED_DEPLOYMENTS.validate_python(store.get_cache(_picked_deployments_key(call_id))) or () + + +def _put(store: _ReservationStore, key: str, values: tuple[object, ...], ttl: float) -> None: + # InMemoryCache keeps a live key's first TTL, so drop the key to restart it + store.delete_cache(key) + if values: + store.set_cache(key, values, ttl=ttl) + + +def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str) -> None: + kept: Final = tuple(pick for pick in _picks(store, deployment_id) if pick[1] != call_id) + _put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS) + + +def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None: + if call_id is None: + return + if not _picked_deployments(store, call_id): + return + with _RESERVATIONS_LOCK: + for deployment_id in frozenset(_picked_deployments(store, call_id)): + _release_call_on(store, deployment_id, call_id) + store.delete_cache(_picked_deployments_key(call_id)) + + +@contextlib.contextmanager +def owning_picks(store: _ReservationStore, call_id: str | None) -> Generator[None, None, None]: + owner: Final = call_id or _PICK_OWNER.get() or uuid.uuid4().hex + token: Final = _PICK_OWNER.set(owner) + try: + yield + finally: + _PICK_OWNER.reset(token) + release_unstarted_picks(store, owner) + + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 logged_failure: int = 0 - def __init__(self, router_cache: DualCache): + def __init__( + self, + router_cache: DualCache, + clock: Callable[[], float] = time.monotonic, + reservation_store: _ReservationStore | None = None, + ): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) + self._clock = clock + self._reservations: Final[_ReservationStore] = ( + reservation_store + if reservation_store is not None + else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock) + ) + + def _live_reservations(self, deployment_id: str, now: float) -> tuple[_Pick, ...]: + return _unexpired(_picks(self._reservations, deployment_id), now) + + def _store_reservations(self, deployment_id: str, picks: tuple[_Pick, ...]) -> None: + _put(self._reservations, _reservation_key(deployment_id), picks, SELECTION_RESERVATION_TTL_SECONDS) + + def _take_reservation(self, deployment_id: str, call_id: str | None, now: float) -> None: + self._store_reservations(deployment_id, _without_one_pick(self._live_reservations(deployment_id, now), call_id)) + + def _record_pick(self, call_id: str, deployment_id: str) -> None: + picked: Final = (*_picked_deployments(self._reservations, call_id), deployment_id) + _put(self._reservations, _picked_deployments_key(call_id), picked, SELECTION_RESERVATION_TTL_SECONDS) + + def _select( + self, + healthy_deployments: Sequence[Mapping[str, object]], + deployment_ids: tuple[str, ...], + keys: tuple[str, ...], + shared: tuple[int, ...] | None, + reserve: bool, + call_id: str | None, + ) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + with _RESERVATIONS_LOCK: + now: Final = self._clock() + counts: Final = ( + self._this_workers_counts(keys) if shared is None else self._at_least_this_workers(keys, shared) + ) + live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids) + index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True))) + if reserve: + tag: Final = call_id if call_id is not None else _PICK_OWNER.get() + self._store_reservations(deployment_ids[index], (*live[index], (now, tag))) + if tag is not None: + self._record_pick(tag, deployment_ids[index]) + return healthy_deployments[index] + + def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None: + with _RESERVATIONS_LOCK: + now: Final = self._clock() + self._take_reservation(ref[1], call_id, now) + self._count_locally(_request_count_key(*ref), 1) + if call_id is not None: + key: Final = _started_attempts_key(call_id, ref) + started: Final = (*_attempt_times(self._reservations, key), now) + _put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS) + + def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool: + if call_id is None: + with _RESERVATIONS_LOCK: + self._count_locally(_request_count_key(*ref), -1) + return True + key: Final = _started_attempts_key(call_id, ref) + with _RESERVATIONS_LOCK: + started: Final = _attempt_times(self._reservations, key) + if started: + _put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) + self._count_locally(_request_count_key(*ref), -1) + else: + self._take_reservation(ref[1], call_id, self._clock()) + return bool(started) @with_service_target("router_usage") def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: - self._increment(kwargs, 1) + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + self._start_attempt(ref, call_id_of(kwargs)) + self._count_shared(_request_count_key(*ref), 1) @with_service_target("router_usage") def log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self._increment(kwargs, -1) + self._settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -136,7 +355,7 @@ class LeastBusyLoggingHandler(CustomLogger): def log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self._increment(kwargs, -1) + self._settle(kwargs) if self.test_flag: self.logged_failure += 1 @@ -144,7 +363,7 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -152,15 +371,20 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_failure += 1 @with_service_target("router_usage") def get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + self, + model_group: str, + healthy_deployments: Sequence[Mapping[str, object]], + reserve: bool = False, + call_id: str | None = None, ) -> Mapping[str, object] | None: - keys: Final = _request_count_keys(model_group, healthy_deployments) + deployment_ids: Final = _deployment_ids(healthy_deployments) + keys: Final = _request_count_keys(model_group, deployment_ids) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: @@ -168,15 +392,19 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) - local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) - return _least_busy(healthy_deployments, local) + return self._select(healthy_deployments, deployment_ids, keys, shared, reserve, call_id) + return self._select(healthy_deployments, deployment_ids, keys, None, reserve, call_id) @with_service_target("router_usage") async def async_get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + self, + model_group: str, + healthy_deployments: Sequence[Mapping[str, object]], + reserve: bool = False, + call_id: str | None = None, ) -> Mapping[str, object] | None: - keys: Final = _request_count_keys(model_group, healthy_deployments) + deployment_ids: Final = _deployment_ids(healthy_deployments) + keys: Final = _request_count_keys(model_group, deployment_ids) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: @@ -184,42 +412,51 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) - local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) - return _least_busy(healthy_deployments, local) + return self._select(healthy_deployments, deployment_ids, keys, shared, reserve, call_id) + return self._select(healthy_deployments, deployment_ids, keys, None, reserve, call_id) - def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + def _settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None: + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): return - key: Final = _request_count_key(*ref) - redis_cache: Final = self.router_cache.redis_cache + self._count_shared(_request_count_key(*ref), -1) + + async def _async_settle(self, kwargs: Mapping[str, object]) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): + return + await self._async_count_shared(_request_count_key(*ref), -1) + + def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]: + return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + + def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]: + return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True)) + + def _count_locally(self, key: str, delta: int) -> None: try: local: Final = self.router_cache.increment_cache( key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS ) if local < 0: self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: - return + except TypeError as e: + _warn_unwritable(key, e) + + def _count_shared(self, key: str, delta: int) -> None: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return + try: redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) - async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: - ref: Final = _deployment_ref(kwargs) - if ref is None: - return - key: Final = _request_count_key(*ref) + async def _async_count_shared(self, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return try: - local: Final = await self.router_cache.async_increment_cache( - key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS - ) - if local is not None and local < 0: - await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: - return await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e50d5bb47a6..7a8a645ad13 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -847,6 +847,45 @@ def is_stream_stall(exception: BaseException) -> bool: return isinstance(exception, litellm.Timeout) and isinstance(exception.__cause__, MidStreamFallbackError) +def is_caller_entra_credential_failure( + deployment: Deployment | None, + router_defaults: Mapping[str, object], + model_call_details: Mapping[str, object], + exception_status: str | int, + exception: BaseException, +) -> bool: + if cast_exception_status_to_int(exception_status) != 401 or deployment is None: + return False + + from litellm.llms.azure.common_utils import caused_by_entra_rejection + + if not caused_by_entra_rejection(exception): + return False + from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, azure_entra_environment_defaults + + call_params: Final = _CREDENTIAL_VALUES_ADAPTER.validate_python(model_call_details.get("litellm_params") or {}) + + credential_name: Final = deployment.litellm_params.litellm_credential_name + configured_by_precedence: Final[tuple[Mapping[str, object], ...]] = ( + router_defaults, + _CREDENTIAL_VALUES_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)), + _CREDENTIAL_VALUES_ADAPTER.validate_python( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} + ), + ) + environment: Final = azure_entra_environment_defaults() + return any( + call_params.get(key) is not None + and (call_params.get(key) or environment.get(key)) + != (_first_configured(configured_by_precedence, key) or environment.get(key)) + for key in AZURE_ENTRA_LITELLM_PARAM_KEYS + ) + + +def _first_configured(sources: tuple[Mapping[str, object], ...], key: str) -> object: + return next((value for source in sources if (value := source.get(key)) is not None), None) + + def is_caller_scoped_auth_failure(deployment: Deployment | None, exception_status: str | int) -> bool: if cast_exception_status_to_int(exception_status) not in (401, 403) or deployment is None: return False diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index c929cca44be..e8f832b4b2d 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -28,6 +28,7 @@ from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, is_advisor_orchestration_failure, + is_caller_entra_credential_failure, is_caller_scoped_auth_failure, is_caller_timeout_408, is_stream_stall, @@ -132,6 +133,14 @@ def _trigger_cooldown_for_failed_deployment( deployment_id, ) return + if is_caller_entra_credential_failure( + deployment, ctx.default_litellm_params(), model_call_details, exception_status, exception + ): + verbose_router_logger.debug( + "Not triggering cooldown for fallback deployment %s: Entra ID rejected a credential the caller sent.", + deployment_id, + ) + return # Priority: deployment config > response header > router default, matching # Router.deployment_callback_on_failure's precedence for the primary path. diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py index 6f521d1fd1b..99ff7953ee2 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -16,7 +16,7 @@ class ToolPermissionRule(LiteLLMBaseModel): id: str = Field(description="Unique identifier for the rule") tool_name: str | None = Field( default=None, - description="Regex pattern applied to the tool's function name", + description="Regex pattern applied to the tool's name, or to its type when the tool has no name", ) tool_type: str | None = Field( default=None, diff --git a/litellm/types/proxy/model_deprecation.py b/litellm/types/proxy/model_deprecation.py index 3699f01a684..7e459904d59 100644 --- a/litellm/types/proxy/model_deprecation.py +++ b/litellm/types/proxy/model_deprecation.py @@ -18,6 +18,10 @@ DeprecationStatus = Literal["upcoming", "imminent", "deprecated"] class ModelDeprecationInfo(LiteLLMBaseModel): model_name: str = Field(description="The public name of the model on the proxy (model_group).") + team_public_model_name: str | None = Field( + default=None, + description="The model name visible to the team, when configured.", + ) litellm_model: str | None = Field( default=None, description="The underlying litellm model string the deprecation date is sourced from.", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dc71ba149da..4c47489e08f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4650,6 +4650,7 @@ class LlmProviders(str, Enum): STRANDS_DECIDER = "strands_decider" LAYA = "laya" BESPOKE = "bespoke" + EMPIRIOLABS = "empiriolabs" EU_JEV = "eu_jev" MISTRAL = "mistral" MILVUS = "milvus" diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 0fc2b8ede1b..240e36ecdf6 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -17,10 +17,11 @@ VIDEO_ID_PREFIX: Final = "video_" VIDEO_ID_SERVED_MODEL_PREFIX: Final = "served_model:" VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:" VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:" -VIDEO_ID_CALL_TYPES: Final = frozenset( +VIDEO_STATUS_CALL_TYPES: Final = frozenset({"avideo_status", "video_status"}) +VIDEO_COST_REQUEST_ID_SUFFIX: Final = "_video_cost" +POLLED_VIDEO_STATUS_KWARG: Final = "polled_video_status" +VIDEO_ID_CALL_TYPES: Final = VIDEO_STATUS_CALL_TYPES | frozenset( { - "avideo_status", - "video_status", "avideo_content", "video_content", "avideo_remix", @@ -31,6 +32,12 @@ VIDEO_ID_CALL_TYPES: Final = frozenset( "video_extension", } ) + + +def is_video_status_call_type(call_type: object) -> bool: + return isinstance(call_type, str) and call_type in VIDEO_STATUS_CALL_TYPES + + CHARACTER_ID_PREFIX: Final = "character_" CHARACTER_ID_TEMPLATE: Final = "litellm:custom_llm_provider:{};model_id:{};character_id:{}" @@ -94,6 +101,27 @@ def _trailing_part(parts: Sequence[str], prefix: str) -> str | None: return next((part.removeprefix(prefix) or None for part in parts[3:] if part.startswith(prefix)), None) +def _decoded_video_id_payload(encoded_video_id: str) -> str | None: + if not encoded_video_id.startswith(VIDEO_ID_PREFIX): + return None + try: + decoded_id: Final = base64.b64decode( + _add_base64_padding(encoded_video_id.removeprefix(VIDEO_ID_PREFIX)).encode("utf-8") + ).decode("utf-8") + except ValueError as e: + verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e) + return None + return decoded_id if ";" in decoded_id else None + + +def canonical_video_id(video_id: str) -> str: + """One spelling per video, since a client may poll a LiteLLM video id with or without its base64 padding.""" + payload: Final = _decoded_video_id_payload(video_id) + if payload is None: + return video_id + return f"{VIDEO_ID_PREFIX}{base64.b64encode(payload.encode('utf-8')).decode('utf-8')}" + + def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: """Decode provider and model_id from encoded video_id.""" if not encoded_video_id: @@ -103,26 +131,8 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: video_id=encoded_video_id, ) - if not encoded_video_id.startswith(VIDEO_ID_PREFIX): - return DecodedVideoId( - custom_llm_provider=None, - model_id=None, - video_id=encoded_video_id, - ) - - try: - cleaned_id = encoded_video_id.replace(VIDEO_ID_PREFIX, "") - cleaned_id = _add_base64_padding(cleaned_id) - decoded_id: Final = base64.b64decode(cleaned_id.encode("utf-8")).decode("utf-8") - except Exception as e: - verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e) - return DecodedVideoId( - custom_llm_provider=None, - model_id=None, - video_id=encoded_video_id, - ) - - if ";" not in decoded_id: + decoded_id: Final = _decoded_video_id_payload(encoded_video_id) + if decoded_id is None: return DecodedVideoId( custom_llm_provider=None, model_id=None, diff --git a/litellm/utils.py b/litellm/utils.py index b457da058ad..4add9d53cb3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1636,7 +1636,6 @@ def dispatch_success_logging( is_litellm_internal_call: bool, ) -> None: if is_litellm_internal_call: - logging_obj.capture_usage_before_logging(result, start_time, end_time) return _schedule_async_success_logging( @@ -1675,14 +1674,9 @@ def _schedule_async_success_logging( if not logging_obj.claim_async_success_log(): return - if logging_obj.defer_async_logging and not is_completion_with_fallbacks: - logging_obj.capture_usage_before_logging(result, start_time, end_time, freeze_only=True) - def _enqueue_async_logging() -> None: from litellm.responses.streaming_iterator import spawn_logging_task - if not is_completion_with_fallbacks: - logging_obj.capture_usage_before_logging(result, start_time, end_time) if logging_obj.defer_async_logging: logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, end_time=end_time @@ -2280,7 +2274,6 @@ def client(original_function): # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") - logging_obj.capture_usage_before_logging(TypeAdapter(object).validate_python(result), start_time, end_time) if not is_internal_call.get(): # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation @@ -10303,6 +10296,7 @@ class ProviderConfigManager: if provider == LlmProviders.BESPOKE: return litellm.BespokeDecisionsConfig() if provider in ( + LlmProviders.EMPIRIOLABS, LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI, @@ -10310,6 +10304,7 @@ class ProviderConfigManager: LlmProviders.PARALLEL_AI, ): configs: Final = { + LlmProviders.EMPIRIOLABS: litellm.EmpirioLabsDecisionsConfig, LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig, LlmProviders.AZURE: litellm.AzureDecisionsConfig, LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7419f48f3fb..4d576443b30 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9003,11 +9003,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9039,7 +9035,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, @@ -9055,11 +9055,7 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/chat/completions", @@ -9091,7 +9087,11 @@ "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, "input_cost_per_token_above_272k_tokens_flex": 1e-05, - "output_cost_per_token_above_272k_tokens_flex": 3.75e-05 + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "input_cost_per_token_flex": 5e-06, + "output_cost_per_token_flex": 2.5e-05, + "cache_read_input_token_cost_flex": 5e-07, + "cache_creation_input_token_cost_flex": 6.25e-06 }, "azure/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, @@ -9107,12 +9107,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9139,7 +9135,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-luna-2026-09-22": { "cache_creation_input_token_cost": 1.25e-07, @@ -9155,12 +9159,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, - "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -9187,7 +9187,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-08, + "cache_read_input_token_cost_flex": 5e-09, + "input_cost_per_token_above_272k_tokens_flex": 1e-07, + "input_cost_per_token_flex": 5e-08, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, + "output_cost_per_token_flex": 2.5e-07 }, "azure/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, @@ -9239,7 +9247,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6-sol-2026-09-22": { "cache_creation_input_token_cost": 2.5e-06, @@ -9291,7 +9307,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 1e-07, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.25e-06, @@ -9354,11 +9378,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9386,7 +9406,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/gpt-6.1-sol-2026-09-29": { "cache_creation_input_token_cost": 2.5e-06, @@ -9402,11 +9430,7 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.014, - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014 - }, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ "/v1/chat/completions", @@ -9434,7 +9458,15 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 5e-06, + "cache_read_input_token_cost_flex": 5e-08, + "cache_creation_input_token_cost_flex": 1.25e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "output_cost_per_token_above_272k_tokens_flex": 7.5e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, "azure/us/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, @@ -15566,7 +15598,7 @@ "output_vector_size": 1536 }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 3e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "mode": "embedding", @@ -15683,8 +15715,10 @@ "input_cost_per_image": 7.2e-05, "input_cost_per_video_per_second": 0.00084, "input_cost_per_audio_per_second": 0.000168, + "input_cost_per_token_batches": 8.1e-08, "output_cost_per_token": 0.0, "output_vector_size": 3072, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_embedding_image_input": true, "supports_image_input": true, "supports_video_input": true, @@ -15693,11 +15727,14 @@ "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.8e-08, "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -15708,11 +15745,14 @@ "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 1.05e-08, "input_cost_per_token": 4.2e-08, + "input_cost_per_token_batches": 2.1e-08, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.68e-07, + "output_cost_per_token_batches": 8.4e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -15720,12 +15760,21 @@ }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, "input_cost_per_token": 9.6e-07, + "input_cost_per_token_batches": 4.8e-07, + "input_cost_per_token_flex": 4.8e-07, + "input_cost_per_token_priority": 1.68e-06, "litellm_provider": "bedrock", "max_input_tokens": 300000, "max_output_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.84e-06, + "output_cost_per_token_batches": 1.92e-06, + "output_cost_per_token_flex": 1.92e-06, + "output_cost_per_token_priority": 6.72e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -22047,6 +22096,7 @@ "output_cost_per_token": 2.8e-07, "output_dbu_cost_per_token": 4e-06, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -22060,8 +22110,12 @@ "supports_vision": false }, "databricks/system.ai.deepseek-v4-flash-0731": { + "deprecation_date": "2026-11-05", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "litellm_provider": "databricks", "mode": "chat", + "source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { @@ -22081,6 +22135,7 @@ "output_cost_per_token": 3.95997e-06, "output_dbu_cost_per_token": 5.6571e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017", "supported_modalities": [ "text" ], @@ -23158,6 +23213,7 @@ "output_cost_per_token": 4.04999e-06, "output_dbu_cost_per_token": 5.7857e-05, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "successor_model": "databricks/databricks-glm-5-3", "supported_modalities": [ "text", "image" @@ -36252,6 +36308,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_image_token_batches": 1.6e-05, "output_cost_per_token_batches": 5e-06, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ @@ -36275,6 +36332,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3.2e-05, + "output_cost_per_image_token_batches": 1.6e-05, "output_cost_per_token_batches": 5e-06, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ @@ -36294,6 +36352,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", @@ -36329,6 +36388,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -36364,6 +36424,7 @@ "input_cost_per_image_token_batches": 4e-06, "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" @@ -38126,7 +38187,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38165,7 +38226,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-05, @@ -38204,7 +38265,7 @@ "input_cost_per_token": 1.25e-05, "input_cost_per_token_above_272k_tokens": 2.5e-05, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "mode": "responses", "output_cost_per_token": 7.5e-05, @@ -39718,6 +39779,7 @@ "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 8e-06, + "output_cost_per_image_token_batches": 4e-06, "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/images/generations", @@ -55961,6 +56023,7 @@ "openai/sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://platform.openai.com/docs/api-reference/videos", @@ -59823,6 +59886,7 @@ "sora-2-pro-high-res": { "litellm_provider": "openai", "deprecation_date": "2026-09-24", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, "source": "https://developers.openai.com/api/docs/pricing", @@ -59848,6 +59912,7 @@ "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image_token": 3.2e-05, + "output_cost_per_image_token_batches": 1.6e-05, "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "source": "https://developers.openai.com/api/docs/pricing", @@ -64661,7 +64726,8 @@ { "name": "claude-family-baseline", "pattern": "claude-[a-z]+-\\d+(?:[-.]\\d+)?(?:-\\d{8})?", - "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number. On mapped Claude entries from Anthropic, Azure Foundry, Bedrock and Vertex it only fills the keys an entry leaves out, so the entry's own values always win.", "model_info": { "mode": "chat", "max_input_tokens": 200000, @@ -69076,12 +69142,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69094,12 +69166,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69110,12 +69188,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69124,12 +69208,18 @@ }, "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69141,12 +69231,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69160,12 +69256,18 @@ "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69371,12 +69473,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 262144, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.88e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_native_structured_output": true, "supports_system_messages": true, @@ -69389,12 +69497,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_vision": true, "supports_audio_input": false, @@ -69405,12 +69519,18 @@ "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.76e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_system_messages": true, "supports_audio_input": false, "supports_function_calling": true, @@ -69419,12 +69539,18 @@ }, "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 256000, "max_output_tokens": 32768, "mode": "chat", "output_cost_per_token": 7.8e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, @@ -69436,12 +69562,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69455,12 +69587,18 @@ "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -69664,16 +69802,32 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-06, - "input_cost_per_token_above_272k_tokens": 5.4e-06, - "cache_creation_input_token_cost": 3.375e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-06, - "cache_read_input_token_cost": 2.7e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, - "output_cost_per_token": 1.62e-05, - "output_cost_per_token_above_272k_tokens": 2.43e-05, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-05, + "cache_creation_input_token_cost_flex": 1.65e-06, + "cache_creation_input_token_cost_priority": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-06, + "cache_read_input_token_cost_flex": 1.32e-07, + "cache_read_input_token_cost_priority": 5.28e-07, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-05, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 5.28e-06, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-05, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-05, + "output_cost_per_token_flex": 7.92e-06, + "output_cost_per_token_priority": 3.168e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { @@ -69701,16 +69855,32 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 2.7e-07, - "input_cost_per_token_above_272k_tokens": 5.4e-07, - "cache_creation_input_token_cost": 3.375e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.75e-07, - "cache_read_input_token_cost": 2.7e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, - "output_cost_per_token": 1.62e-06, - "output_cost_per_token_above_272k_tokens": 2.43e-06, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.32e-06, + "cache_creation_input_token_cost_flex": 1.65e-07, + "cache_creation_input_token_cost_priority": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.056e-07, + "cache_read_input_token_cost_flex": 1.32e-08, + "cache_read_input_token_cost_priority": 5.28e-08, + "input_cost_per_token_above_272k_tokens_flex": 2.64e-07, + "input_cost_per_token_above_272k_tokens_priority": 1.056e-06, + "input_cost_per_token_flex": 1.32e-07, + "input_cost_per_token_priority": 5.28e-07, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.188e-06, + "output_cost_per_token_above_272k_tokens_priority": 4.752e-06, + "output_cost_per_token_flex": 7.92e-07, + "output_cost_per_token_priority": 3.168e-06, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { @@ -69737,9 +69907,9 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "input_cost_per_token": 3.375e-06, - "cache_read_input_token_cost": 3.375e-07, - "output_cost_per_token": 2.025e-05, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05, "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", "supports_web_search": true @@ -69759,20 +69929,36 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "input_cost_per_token": 1.5e-06, "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 2.4e-07 + "cache_read_input_token_cost": 2.4e-07, + "cache_read_input_token_cost_batches": 1.2e-07, + "cache_read_input_token_cost_flex": 1.2e-07, + "cache_read_input_token_cost_priority": 4.2e-07, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.625e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.25e-06 }, "bedrock_mantle/us-gov-west-1/xai.grok-4.6": { "use_openai_responses_path": true, "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69786,11 +69972,17 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": { "input_cost_per_token": 4.8e-08, + "input_cost_per_token_flex": 2.4e-08, + "input_cost_per_token_priority": 8.4e-08, "output_cost_per_token": 9.6e-08, "litellm_provider": "bedrock_mantle", "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 4.8e-08, + "output_cost_per_token_flex": 4.8e-08, + "output_cost_per_token_priority": 1.68e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69804,11 +69996,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": { "input_cost_per_token": 1.56e-07, + "input_cost_per_token_batches": 7.8e-08, + "input_cost_per_token_flex": 7.8e-08, + "input_cost_per_token_priority": 2.73e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69822,11 +70021,18 @@ }, "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": { "input_cost_per_token": 1.68e-07, + "input_cost_per_token_batches": 8.4e-08, + "input_cost_per_token_flex": 8.4e-08, + "input_cost_per_token_priority": 2.94e-07, "output_cost_per_token": 4.8e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 256000, "max_output_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2.4e-07, + "output_cost_per_token_flex": 2.4e-07, + "output_cost_per_token_priority": 8.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -69841,12 +70047,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69860,12 +70072,18 @@ "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69981,10 +70199,17 @@ "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -69999,12 +70224,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": { "supports_vision": false, "input_cost_per_token": 8.4e-08, + "input_cost_per_token_batches": 4.2e-08, + "input_cost_per_token_flex": 4.2e-08, + "input_cost_per_token_priority": 1.47e-07, "output_cost_per_token": 3.6e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.8e-07, + "output_cost_per_token_flex": 1.8e-07, + "output_cost_per_token_priority": 6.3e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -70018,12 +70249,18 @@ "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": { "supports_vision": false, "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, "output_cost_per_token": 7.2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -81729,7 +81966,8 @@ "input_cost_per_token_priority": 5.25e-07, "output_cost_per_token_priority": 2.1e-06, "cache_read_input_token_cost_priority": 5.25e-08, - "cache_creation_input_token_cost_priority": 5.25e-07 + "cache_creation_input_token_cost_priority": 5.25e-07, + "successor_model": "databricks/databricks-deepseek-v4-1-flash-1017" }, "databricks/databricks-deepseek-v4-1-flash-1017": { "cache_creation_input_token_cost": 3.0002e-07, @@ -82181,6 +82419,8 @@ "databricks/databricks-kimi-k2-7": { "cache_creation_input_token_cost": 9.4997e-07, "cache_read_input_token_cost": 1.8998e-07, + "deprecation_date": "2026-10-30", + "deprecation_source": "https://docs.databricks.com/aws/en/machine-learning/retired-models-policy", "input_cost_per_token": 9.4997e-07, "input_dbu_cost_per_token": 1.3571e-05, "litellm_provider": "databricks", @@ -82193,6 +82433,7 @@ "output_cost_per_token": 4.00001e-06, "output_dbu_cost_per_token": 5.7143e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "successor_model": "databricks/databricks-kimi-k3", "supported_modalities": [ "text", "image" @@ -82594,6 +82835,1603 @@ "supports_function_calling": true, "source": "https://tokenfactory.nebius.com/api/public/models_info" }, + "empiriolabs/aplomb-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 0.0, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1/decisions" + ], + "supports_vision": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/aplomb-1" + }, + "empiriolabs/deepseek-v4-1-flash-1017": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-1-flash-1017" + }, + "empiriolabs/deepseek-v4-flash-0731": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 4.24e-07, + "output_cost_per_token": 1.272e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + }, + "empiriolabs/deepseek-v4-flash-0731:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 4.4e-07, + "output_cost_per_token": 1.32e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-flash-0731" + }, + "empiriolabs/deepseek-v4-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 3.3e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/deepseek-v4-pro-0813": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 1.32e-06, + "output_cost_per_token": 3.96e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro-0813" + }, + "empiriolabs/deepseek-v4-pro:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/deepseek-v4-pro:variant2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 3.301e-06, + "cache_read_input_token_cost": 1.38e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/deepseek-v4-pro:variant3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/deepseek-v4-pro" + }, + "empiriolabs/fugu-max": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 2.5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.007, + "search_context_size_medium": 0.007, + "search_context_size_high": 0.007 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-max" + }, + "empiriolabs/fugu-ultra-v1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-06, + "output_cost_per_token": 4.5e-05, + "input_cost_per_token_above_272k_tokens": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 6.75e-05, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 1.5e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-ultra-v1-0" + }, + "empiriolabs/fugu-ultra-v1-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 1e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost": 5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-ultra-v1-1" + }, + "empiriolabs/fugu-ultra-v2-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 1e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost": 5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/fugu-ultra-v2-0" + }, + "empiriolabs/gemma-4-26b-a4b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 5e-08, + "output_cost_per_token": 2.9e-07, + "cache_read_input_token_cost": 2.5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/gemma-4-26b-a4b" + }, + "empiriolabs/glm-4-6v-flashx": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 3.6e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-4-6v-flashx" + }, + "empiriolabs/glm-4-7-flashx": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 202752, + "max_output_tokens": 131072, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 6.3e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-4-7-flashx" + }, + "empiriolabs/glm-5-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 202745, + "input_cost_per_token": 8.25e-07, + "output_cost_per_token": 3.301e-06, + "input_cost_per_token_above_32k_tokens": 1.1e-06, + "output_cost_per_token_above_32k_tokens": 3.851e-06, + "cache_read_input_token_cost_above_32k_tokens": 2.2e-07, + "cache_read_input_token_cost": 1.65e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-1" + }, + "empiriolabs/glm-5-2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-2" + }, + "empiriolabs/glm-5-2:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 3.851e-06, + "cache_read_input_token_cost": 2.75e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-2" + }, + "empiriolabs/glm-5-2:variant2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-2" + }, + "empiriolabs/glm-5-3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-3" + }, + "empiriolabs/glm-5-3-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 2.5e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/glm-5-3-flash" + }, + "empiriolabs/kimi-k2-6": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 16000, + "input_cost_per_token": 8.939e-07, + "output_cost_per_token": 3.7131e-06, + "cache_read_input_token_cost": 1.788e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-6" + }, + "empiriolabs/kimi-k2-7-code": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 9.5e-07, + "output_cost_per_token": 4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + }, + "empiriolabs/kimi-k2-7-code-highspeed": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1.9e-06, + "output_cost_per_token": 8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-7-code-highspeed" + }, + "empiriolabs/kimi-k2-7-code:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "input_cost_per_token": 8.939e-07, + "output_cost_per_token": 3.7131e-06, + "cache_read_input_token_cost": 1.788e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k2-7-code" + }, + "empiriolabs/kimi-k3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/kimi-k3" + }, + "empiriolabs/mimo-v2-6-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/mimo-v2-6-flash" + }, + "empiriolabs/mimo-v2-6-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 4.35e-07, + "output_cost_per_token": 8.7e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro" + }, + "empiriolabs/mimo-v2-6-pro-ultraspeed": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 4.35e-06, + "output_cost_per_token": 8.7e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/mimo-v2-6-pro-ultraspeed" + }, + "empiriolabs/minimax-m2-7": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 3e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m2-7" + }, + "empiriolabs/minimax-m2-7-highspeed": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 3e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m2-7-highspeed" + }, + "empiriolabs/minimax-m3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 524288, + "input_cost_per_token": 2.25e-07, + "output_cost_per_token": 9e-07, + "input_cost_per_token_above_512k_tokens": 4.5e-07, + "output_cost_per_token_above_512k_tokens": 1.8e-06, + "cache_read_input_token_cost_above_512k_tokens": 9e-08, + "cache_read_input_token_cost": 4.5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m3" + }, + "empiriolabs/minimax-m3-1-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 524288, + "input_cost_per_token": 1.125e-07, + "output_cost_per_token": 4.5e-07, + "cache_read_input_token_cost": 1.125e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m3-1-flash" + }, + "empiriolabs/minimax-m3:priority": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048458, + "max_output_tokens": 524288, + "input_cost_per_token": 3.375e-07, + "output_cost_per_token": 1.35e-06, + "input_cost_per_token_above_512k_tokens": 6.75e-07, + "output_cost_per_token_above_512k_tokens": 2.7e-06, + "cache_read_input_token_cost_above_512k_tokens": 1.35e-07, + "cache_read_input_token_cost": 6.75e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/minimax-m3" + }, + "empiriolabs/mistral-small-4": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.084, + "search_context_size_medium": 0.084, + "search_context_size_high": 0.14 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/mistral-small-4" + }, + "empiriolabs/muse-glimmer-30b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 8e-07, + "cache_read_input_token_cost": 5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/muse-glimmer-30b" + }, + "empiriolabs/muse-spark-1-1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 4.25e-06, + "cache_read_input_token_cost": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/muse-spark-1-1" + }, + "empiriolabs/muse-spark-1-2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 4.25e-06, + "cache_read_input_token_cost": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/muse-spark-1-2" + }, + "empiriolabs/muse-spark-1-3": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 4.25e-06, + "cache_read_input_token_cost": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/muse-spark-1-3" + }, + "empiriolabs/nova-lite-1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 317379, + "max_output_tokens": 10000, + "input_cost_per_token": 6.9e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 3.86e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_prompt_caching": true, + "source": "https://empiriolabs.ai/models/nova-lite-1-0" + }, + "empiriolabs/nova-lite-2": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 3.16e-06, + "cache_read_input_token_cost": 2.128e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "source": "https://empiriolabs.ai/models/nova-lite-2" + }, + "empiriolabs/nova-micro-1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 133003, + "max_output_tokens": 10000, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, + "cache_read_input_token_cost": 2.24e-08, + "supports_system_messages": true, + "supports_prompt_caching": true, + "source": "https://empiriolabs.ai/models/nova-micro-1-0" + }, + "empiriolabs/nova-pro-1-0": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 317319, + "max_output_tokens": 10000, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 9.6e-06, + "supports_system_messages": true, + "supports_vision": true, + "source": "https://empiriolabs.ai/models/nova-pro-1-0" + }, + "empiriolabs/perplexity-advanced-deep-research": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_output_tokens": 16000, + "input_cost_per_token": 1.2e-05, + "output_cost_per_token": 6e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012, + "search_context_size_high": 0.012 + }, + "source": "https://empiriolabs.ai/models/perplexity-advanced-deep-research" + }, + "empiriolabs/perplexity-deep-research": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919327, + "input_cost_per_token": 4.8e-06, + "output_cost_per_token": 1.9e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-deep-research" + }, + "empiriolabs/perplexity-pro-search": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_output_tokens": 32768, + "input_cost_per_token": 7.8e-06, + "output_cost_per_token": 3.9e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-pro-search" + }, + "empiriolabs/perplexity-sonar": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919327, + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 2.4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-sonar" + }, + "empiriolabs/perplexity-sonar-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919328, + "max_output_tokens": 32768, + "input_cost_per_token": 7.2e-06, + "output_cost_per_token": 3.6e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-sonar-pro" + }, + "empiriolabs/perplexity-sonar-reasoning-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 919330, + "max_output_tokens": 131072, + "input_cost_per_token": 4.8e-06, + "output_cost_per_token": 1.9e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "source": "https://empiriolabs.ai/models/perplexity-sonar-reasoning-pro" + }, + "empiriolabs/qwen3-5-122b-a10b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 1.15e-07, + "output_cost_per_token": 9.17e-07, + "input_cost_per_token_above_128k_tokens": 2.87e-07, + "output_cost_per_token_above_128k_tokens": 2.294e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" + }, + "empiriolabs/qwen3-5-27b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 8.6e-08, + "output_cost_per_token": 6.88e-07, + "input_cost_per_token_above_128k_tokens": 2.58e-07, + "output_cost_per_token_above_128k_tokens": 2.064e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-27b" + }, + "empiriolabs/qwen3-5-35b-a3b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 5.7e-08, + "output_cost_per_token": 4.59e-07, + "input_cost_per_token_above_128k_tokens": 2.29e-07, + "output_cost_per_token_above_128k_tokens": 1.835e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" + }, + "empiriolabs/qwen3-5-397b-a17b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 258048, + "max_output_tokens": 64000, + "input_cost_per_token": 1.72e-07, + "output_cost_per_token": 1.032e-06, + "input_cost_per_token_above_128k_tokens": 4.3e-07, + "output_cost_per_token_above_128k_tokens": 2.58e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" + }, + "empiriolabs/qwen3-5-4b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 7e-08, + "cache_read_input_token_cost": 2e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-4b" + }, + "empiriolabs/qwen3-5-9b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 4.5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-9b" + }, + "empiriolabs/qwen3-5-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 995904, + "max_output_tokens": 32768, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3.68e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-flash" + }, + "empiriolabs/qwen3-5-flash:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 995904, + "max_output_tokens": 65536, + "input_cost_per_token": 2.9e-08, + "output_cost_per_token": 2.87e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-flash" + }, + "empiriolabs/qwen3-5-omni-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 196608, + "max_output_tokens": 32768, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4.4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" + }, + "empiriolabs/qwen3-5-omni-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 196608, + "max_output_tokens": 32768, + "input_cost_per_token": 2.8e-06, + "output_cost_per_token": 1.66e-05, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" + }, + "empiriolabs/qwen3-5-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 3.6e-07, + "output_cost_per_token": 2.21e-06, + "input_cost_per_token_above_256k_tokens": 1.08e-06, + "output_cost_per_token_above_256k_tokens": 6.62e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-plus" + }, + "empiriolabs/qwen3-5-plus:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 1.15e-07, + "output_cost_per_token": 6.88e-07, + "input_cost_per_token_above_128k_tokens": 2.87e-07, + "output_cost_per_token_above_128k_tokens": 1.72e-06, + "input_cost_per_token_above_256k_tokens": 5.73e-07, + "output_cost_per_token_above_256k_tokens": 3.44e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-5-plus" + }, + "empiriolabs/qwen3-6-27b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229376, + "max_output_tokens": 64000, + "input_cost_per_token": 4.12564e-07, + "output_cost_per_token": 2.475384e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-27b" + }, + "empiriolabs/qwen3-6-35b-a3b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 3.5e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-35b-a3b" + }, + "empiriolabs/qwen3-6-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-06, + "input_cost_per_token_above_256k_tokens": 1e-06, + "output_cost_per_token_above_256k_tokens": 4e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-flash" + }, + "empiriolabs/qwen3-6-flash:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 1.65e-07, + "output_cost_per_token": 9.9e-07, + "input_cost_per_token_above_256k_tokens": 6.6e-07, + "output_cost_per_token_above_256k_tokens": 3.961e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-flash" + }, + "empiriolabs/qwen3-6-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "input_cost_per_token_above_256k_tokens": 2e-06, + "output_cost_per_token_above_256k_tokens": 6e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.026, + "search_context_size_medium": 0.026, + "search_context_size_high": 0.026 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-plus" + }, + "empiriolabs/qwen3-6-plus:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.76e-07, + "output_cost_per_token": 1.651e-06, + "input_cost_per_token_above_256k_tokens": 1.101e-06, + "output_cost_per_token_above_256k_tokens": 6.602e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-6-plus" + }, + "empiriolabs/qwen3-7-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 3e-08, + "output_cost_per_token": 1.3e-07, + "input_cost_per_token_above_32k_tokens": 1e-07, + "output_cost_per_token_above_32k_tokens": 4e-07, + "cache_read_input_token_cost_above_32k_tokens": 2e-08, + "input_cost_per_token_above_256k_tokens": 2e-07, + "output_cost_per_token_above_256k_tokens": 8e-07, + "cache_read_input_token_cost_above_256k_tokens": 4e-08, + "cache_read_input_token_cost": 6e-09, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.03, + "search_context_size_medium": 0.03, + "search_context_size_high": 0.03 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-flash" + }, + "empiriolabs/qwen3-7-flash:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.8e-08, + "output_cost_per_token": 1.1e-07, + "input_cost_per_token_above_32k_tokens": 8.3e-08, + "output_cost_per_token_above_32k_tokens": 3.3e-07, + "cache_read_input_token_cost_above_32k_tokens": 1.7e-08, + "input_cost_per_token_above_256k_tokens": 1.65e-07, + "output_cost_per_token_above_256k_tokens": 6.6e-07, + "cache_read_input_token_cost_above_256k_tokens": 3.3e-08, + "cache_read_input_token_cost": 6e-09, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-flash" + }, + "empiriolabs/qwen3-7-max": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 7.5e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-max" + }, + "empiriolabs/qwen3-7-max:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 4.951e-06, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-max" + }, + "empiriolabs/qwen3-7-plus": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "input_cost_per_token_above_256k_tokens": 1.2e-06, + "output_cost_per_token_above_256k_tokens": 4.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.03, + "search_context_size_medium": 0.03, + "search_context_size_high": 0.03 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-plus" + }, + "empiriolabs/qwen3-7-plus:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 65536, + "input_cost_per_token": 2.76e-07, + "output_cost_per_token": 1.101e-06, + "input_cost_per_token_above_256k_tokens": 8.26e-07, + "output_cost_per_token_above_256k_tokens": 3.301e-06, + "cache_read_input_token_cost_above_256k_tokens": 1.66e-07, + "cache_read_input_token_cost": 5.6e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-7-plus" + }, + "empiriolabs/qwen3-8-27b": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "input_cost_per_token": 1.7e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 8e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-27b" + }, + "empiriolabs/qwen3-8-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 1.6e-07, + "output_cost_per_token": 4.7e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-flash" + }, + "empiriolabs/qwen3-8-max": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-max" + }, + "empiriolabs/qwen3-8-max-0902": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-max-0902" + }, + "empiriolabs/qwen3-8-max:variant1": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 983616, + "max_output_tokens": 131072, + "input_cost_per_token": 1.65e-06, + "output_cost_per_token": 4.951e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/qwen3-8-max" + }, + "empiriolabs/qwen3-8-omni-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1006592, + "max_output_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 9.4e-07, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.02 + }, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/qwen3-8-omni-flash" + }, + "empiriolabs/seed-2-0-code": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229096, + "max_output_tokens": 128000, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "input_cost_per_token_above_128k_tokens": 8e-07, + "output_cost_per_token_above_128k_tokens": 4.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "source": "https://empiriolabs.ai/models/seed-2-0-code" + }, + "empiriolabs/seed-2-0-lite": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229096, + "max_output_tokens": 128000, + "input_cost_per_token": 3.1e-07, + "output_cost_per_token": 2.5e-06, + "input_cost_per_token_above_128k_tokens": 6.2e-07, + "output_cost_per_token_above_128k_tokens": 5e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-0-lite" + }, + "empiriolabs/seed-2-0-mini": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229156, + "max_output_tokens": 128000, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 5e-07, + "input_cost_per_token_above_128k_tokens": 2.4e-07, + "output_cost_per_token_above_128k_tokens": 1e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-0-mini" + }, + "empiriolabs/seed-2-0-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229096, + "max_output_tokens": 128000, + "input_cost_per_token": 6.3e-07, + "output_cost_per_token": 3.79e-06, + "input_cost_per_token_above_128k_tokens": 1.26e-06, + "output_cost_per_token_above_128k_tokens": 7.58e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-0-pro" + }, + "empiriolabs/seed-2-1-turbo": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 229093, + "max_output_tokens": 65536, + "input_cost_per_token": 6.3e-07, + "output_cost_per_token": 3.13e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/seed-2-1-turbo" + }, + "empiriolabs/skylark-embedding-vision": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 8000, + "input_cost_per_token": 2.5e-07, + "supports_embedding_image_input": true, + "source": "https://empiriolabs.ai/models/skylark-embedding-vision" + }, + "empiriolabs/step-3-5-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 2e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/step-3-5-flash" + }, + "empiriolabs/step-3-5-flash-2603": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 2e-08, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/step-3-5-flash-2603" + }, + "empiriolabs/step-3-7-flash": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.15e-06, + "cache_read_input_token_cost": 4e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/step-3-7-flash" + }, + "empiriolabs/step-5-preview": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1024000, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 5e-08, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "source": "https://empiriolabs.ai/models/step-5-preview" + }, + "empiriolabs/stepaudio-2-5-chat": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 3e-07, + "supports_system_messages": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.013, + "search_context_size_medium": 0.013, + "search_context_size_high": 0.013 + }, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/stepaudio-2-5-chat" + }, + "empiriolabs/stepaudio-3-chat": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 3e-07, + "supports_system_messages": true, + "supports_reasoning": true, + "supports_prompt_caching": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_audio_input": true, + "source": "https://empiriolabs.ai/models/stepaudio-3-chat" + }, + "empiriolabs/text-embedding-v4": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 8192, + "input_cost_per_token": 7e-08, + "source": "https://empiriolabs.ai/models/text-embedding-v4" + }, + "empiriolabs/tongyi-embedding-vision-flash": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 1024, + "input_cost_per_token": 9e-08, + "supports_embedding_image_input": true, + "source": "https://empiriolabs.ai/models/tongyi-embedding-vision-flash" + }, + "empiriolabs/tongyi-embedding-vision-plus": { + "litellm_provider": "empiriolabs", + "mode": "embedding", + "max_input_tokens": 1024, + "input_cost_per_token": 9e-08, + "supports_embedding_image_input": true, + "source": "https://empiriolabs.ai/models/tongyi-embedding-vision-plus" + }, "wally/glm-5.3-flash": { "cache_read_input_token_cost": 7.6e-08, "input_cost_per_token": 2.72e-07, @@ -82920,5 +84758,1354 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "bedrock/us-gov-east-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "supports_assistant_prefill": true, + "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 + }, + "bedrock/us-gov-west-1/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "input_cost_per_token": 3.6e-06, + "input_cost_per_token_batches": 1.8e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "output_cost_per_token_batches": 9e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0": { + "cache_creation_input_token_cost": 4.5e-06, + "cache_read_input_token_cost": 3.6e-07, + "input_cost_per_token": 3.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.8e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "supports_assistant_prefill": true, + "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 + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "input_cost_per_token_batches": 1.2e-07, + "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4.2e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "output_cost_per_token_batches": 3.6e-07, + "output_cost_per_token_flex": 3.6e-07, + "output_cost_per_token_priority": 1.26e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 5.04e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-nano-9b-v2": { + "input_cost_per_token": 7.2e-08, + "input_cost_per_token_batches": 3.6e-08, + "input_cost_per_token_flex": 3.6e-08, + "input_cost_per_token_priority": 1.32e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.76e-07, + "output_cost_per_token_batches": 1.44e-07, + "output_cost_per_token_flex": 1.44e-07, + "output_cost_per_token_priority": 4.8e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": false + }, + "bedrock_mantle/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "input_cost_per_token_batches": 9e-08, + "input_cost_per_token_flex": 9e-08, + "input_cost_per_token_priority": 3.15e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "output_cost_per_token_batches": 3.9e-07, + "output_cost_per_token_flex": 3.9e-07, + "output_cost_per_token_priority": 1.365e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": false, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gmi/MiniMaxAI/MiniMax-M2": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.5": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M2.7": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 196608, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3": { + "cache_read_input_token_cost": 4.8e-08, + "input_cost_per_token": 2.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9.6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/MiniMaxAI/MiniMax-M3.1-Flash": { + "cache_read_input_token_cost": 1.125e-08, + "input_cost_per_token": 1.125e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-FP8": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.09e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Instruct-2507-FP8": { + "cache_read_input_token_cost": 1.75e-08, + "input_cost_per_token": 8.75e-08, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-235B-A22B-Thinking-2507-FP8": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262141, + "mode": "chat", + "output_cost_per_token": 4.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-30B-A3B": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 40946, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-32B-FP8": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 9e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262128, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-122B-A10B": { + "input_cost_per_token": 4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-27B": { + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-35B-A3B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-35B-A3B": { + "input_cost_per_token": 2.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.485e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.6-Plus-2026-04-02": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.7-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-27B": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Flash": { + "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 1.6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/Qwen/Qwen3.8-Max-0902": { + "cache_read_input_token_cost": 2.06e-07, + "input_cost_per_token": 1.65e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.951e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5": { + "cache_read_input_token_cost": 2.38e-09, + "input_cost_per_token": 1.19e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.38e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.5-Pro": { + "cache_read_input_token_cost": 2.52e-09, + "input_cost_per_token": 3.045e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.09e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.18e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-0528": { + "cache_read_input_token_cost": 2.3e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3.9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.1-Terminus": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Exp": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4.1e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V3.2-Speciale": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 1.82e-08, + "input_cost_per_token": 9.1e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.82e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Flash-0731": { + "cache_read_input_token_cost": 9.1e-09, + "input_cost_per_token": 2.86e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8.58e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/DeepSeek-V4-Pro-0813": { + "cache_read_input_token_cost": 3.52e-08, + "input_cost_per_token": 1.056e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 3.168e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { + "cache_read_input_token_cost": 1.4e-08, + "input_cost_per_token": 4.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.1-pro-preview": { + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 9e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.7-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-3.8-flash": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemini-4-argon": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1.3e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/google/gemma-4-31b-it": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-3.3-70B-Instruct": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "input_cost_per_token": 8e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2-Instruct-0905": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 5.7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.29e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/Kimi-K2.6": { + "cache_read_input_token_cost": 1.44e-07, + "input_cost_per_token": 8.55e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k2.7-code-highspeed": { + "cache_read_input_token_cost": 3.8e-07, + "input_cost_per_token": 1.9e-06, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/moonshotai/kimi-k3.1-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/nvidia/nemotron-3-ultra-550b-a55b": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-mini": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 4.5e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.4-pro": { + "input_cost_per_token": 3e-05, + "litellm_provider": "gmi", + "max_input_tokens": 409600, + "mode": "chat", + "output_cost_per_token": 0.00018, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-luna": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-sol": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-5.6-terra": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-astra": { + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6-sol": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-6.1-sol": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/openai/gpt-oss-120b": { + "input_cost_per_token": 5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-3.7-Flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.15e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/stepfun-ai/Step-5-Preview": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.7e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/HunyuanOCR": { + "input_cost_per_token": 5e-07, + "litellm_provider": "gmi", + "max_input_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/Hy3": { + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 5.8e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy3-preview": { + "cache_read_input_token_cost": 6e-08, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/tencent/hy4-preview": { + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 8.34e-07, + "litellm_provider": "gmi", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.501e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/x-ai/grok-4.5": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/xai/grok4.6": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "gmi", + "max_input_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-Air-FP8": { + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131076, + "mode": "chat", + "output_cost_per_token": 1.1e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.5-FP8": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.6": { + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 204800, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-4.7-Flash": { + "input_cost_per_token": 7e-08, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5-FP8": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 1.92e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.1-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 202752, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.2-FP8": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 7e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.2e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/zai-org/GLM-5.3-Flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 7.5e-08, + "litellm_provider": "gmi", + "max_input_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "fireworks_ai/accounts/fireworks/models/kimi-k3p1-pro": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 1.5e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "baseten/moonshotai/Kimi-K3.1-Pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1000000, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "gpt-image-2.6-sunburst": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "gpt-image-2.6-flare": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_pdf_input": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64.0, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock/us-gov-west-1/moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 4.95e-06, + "cache_creation_input_token_cost_flex": 2.475e-06, + "cache_creation_input_token_cost_priority": 8.6625e-06, + "cache_read_input_token_cost": 3.96e-07, + "cache_read_input_token_cost_flex": 1.98e-07, + "cache_read_input_token_cost_priority": 6.93e-07, + "input_cost_per_token": 3.96e-06, + "input_cost_per_token_batches": 1.98e-06, + "input_cost_per_token_flex": 1.98e-06, + "input_cost_per_token_priority": 6.93e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 868928, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.98e-05, + "output_cost_per_token_batches": 9.9e-06, + "output_cost_per_token_flex": 9.9e-06, + "output_cost_per_token_priority": 3.465e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_cache_breakpoint": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_regex_lookaround": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 64.0, + "token_count_tokenizer": "kimi_k3" + }, + "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0": { + "deprecation_date": "2026-09-10", + "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-3-haiku.md", + "input_cost_per_token": 3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 195904, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json", + "successor_model": "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-east-1/xai.grok-4.7": { + "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_batches": 1.32e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 7.92e-06, + "output_cost_per_token_batches": 3.96e-06, + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "use_openai_responses_path": true + }, + "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0": { + "deprecation_date": "2026-09-10", + "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-3-haiku.md", + "input_cost_per_token": 3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 195904, + "max_output_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", + "successor_model": "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0", + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/us-gov-west-1/xai.grok-4.7": { + "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_batches": 1.32e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 7.92e-06, + "output_cost_per_token_batches": 3.96e-06, + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "use_openai_responses_path": true + }, + "bedrock_mantle/xai.grok-4.7": { + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_flex": 2.75e-07, + "cache_read_input_token_cost_priority": 9.625e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_batches": 1.1e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 3.85e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "output_cost_per_token_batches": 3.3e-06, + "output_cost_per_token_flex": 3.3e-06, + "output_cost_per_token_priority": 1.155e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "use_openai_responses_path": true } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 672a2739964..00fee0a0d5a 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -5277,13 +5277,13 @@ } }, "empiriolabs": { - "display_name": "EmpirioLabs (`empiriolabs`)", + "display_name": "EmpirioLabs AI (`empiriolabs`)", "url": "https://docs.litellm.ai/docs/providers/empiriolabs", "endpoints": { - "chat_completions": false, - "responses": false, - "text_completion": false, - "embeddings": false, + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": true, "image_generations": false, "image_edits": false, "audio_transcriptions": false, @@ -5295,11 +5295,11 @@ "realtime": false, "video_generations": false, "batches": false, - "messages": false, - "interactions": false, + "messages": true, + "interactions": true, "systemone": false, - "a2a": false, - "generateContent": false + "a2a": true, + "generateContent": true } }, "openinfer": { diff --git a/pyproject.toml b/pyproject.toml index f2774e87b7b..2cc2d0f70a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,7 +155,8 @@ stt-vertex-chirp = [ stt-nvidia-riva = [ # NVIDIA Riva STT provider (gRPC). These are imported lazily inside the # provider handler so litellm core remains usable without them. - "nvidia-riva-client>=2.15.0", + "grpcio>=1.67.1,<2.0", + "protobuf>=5.26.1,<7.0", "soundfile>=0.12.1", "audioread>=3.0.1", "numpy>=1.26.0", @@ -241,6 +242,7 @@ dev = [ "pytest-timeout==2.4.0", "vcrpy==8.2.1", "pytest-recording==0.13.4", + "nvidia-riva-client==2.16.0", ] e2e-dev = [ "playwright==1.61.0", @@ -275,6 +277,9 @@ ci = [ # protobuf, Pillow is a compiled C extension). "tenacity==8.5.0", "google-generativeai==0.8.6", + # The Riva tests serve NVIDIA's own message definitions, which checks + # LiteLLM's hand-written copy of them on every run. + "nvidia-riva-client==2.16.0", "Pillow==12.3.0; python_version == '3.10'", "Pillow==13.0.0; python_version >= '3.11'", # Azure batch E2E tests still import psycopg2 directly. diff --git a/schema.prisma b/schema.prisma index c007b047087..3aaa74d4fc2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery { @@index([target, state, next_attempt_at]) } +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) +} + model LiteLLM_SubscriptionAccountTable { subscription_account_id String @id @default(uuid()) custom_llm_provider String diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index 898c65dff2c..b8e9ac3c07a 100644 --- a/scripts/check_schema_compatibility.py +++ b/scripts/check_schema_compatibility.py @@ -52,14 +52,20 @@ def check_breaking_changes(old_schema: Schema, new_schema: Schema) -> tuple[str, return tuple(changes()) -def main() -> int: - parser: Final = argparse.ArgumentParser() - parser.add_argument("--base-ref", required=True) - options: Final = parser.parse_args() - base_ref: Final = TypeAdapter(str).validate_python(options.base_ref) - base: Final = subprocess.run( +def _merge_base(base_ref: str) -> str: + return subprocess.run( ("git", "merge-base", "HEAD", base_ref), capture_output=True, text=True, check=True ).stdout.strip() + + +def main() -> int: + parser: Final = argparse.ArgumentParser() + source: Final = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--base-ref", help="compare against the merge base of HEAD and this ref (needs history)") + source.add_argument("--base-commit", help="compare against this commit as is (works in a shallow clone)") + options: Final = TypeAdapter(dict[str, str | None]).validate_python(vars(parser.parse_args())) + base_commit: Final = options["base_commit"] + base: Final = base_commit if base_commit else _merge_base(options["base_ref"] or "") previous: Final = subprocess.run( ("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True ).stdout diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index bedb6f4e98b..6f0450c0091 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -212,33 +212,23 @@ orthogonal to the OTLP variables above; their credentials still go in ### Enterprise billing metrics -License-gated request metering is opt-in and gated entirely on -`billing_metrics_endpoint`. Empty (default) and no billing env is added to -the container, so existing deployments are unchanged. Set it and both -gateway and backend export billable-request counts over OTLP/HTTP, -authenticating to the collector with the mTLS client certificate issued for -your deployment. - -The proxy accepts the certificate, key, and CA bundle as either a file path -or literal PEM content. This stack takes the PEM, writes each one to its own -Secrets Manager entry, grants the task-execution role -`secretsmanager:GetSecretValue` on them, and injects them as -`LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and `_CA_CERT` when -set), so no volume mount is needed on Fargate. +License-gated request metering uses Cognito bearer authentication. Set the +endpoint, client ID, and token URL as plain values, and provide the client +secret through the sensitive Terraform variable. The client secret is stored +in Secrets Manager and injected into both gateway and backend tasks ```hcl billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics" +billing_metrics_client_id = "cognito-client-id" +billing_metrics_token_url = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" ``` ```bash -export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)" -export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)" +export TF_VAR_billing_metrics_client_secret="cognito-client-secret" ``` -`billing_metrics_ca_cert_pem` is only for private or test collectors whose -CA is not in the system trust store; leave it empty against -`telemetry.litellm.ai`. Metering requires an enterprise license, so pair -this with `litellm_license`. To tune the export cadence, set +The client ID and client secret are required when the endpoint is set +Metering also requires an enterprise license. To tune the export cadence, set `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / `backend_extra_env` diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 412e6c423f5..dec7f599ced 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -79,32 +79,19 @@ locals { { name = "OTEL_HEADERS", valueFrom = var.otel_headers_secret_arn }, ] : [] - # Enterprise request metering, gated on billing_metrics_endpoint. The - # endpoint rides in as a plain env var; the mTLS material is stored in - # Secrets Manager (secrets.tf) and injected as PEM-valued env vars, which - # the proxy accepts in place of file paths. Each PEM is wired only when the - # operator supplied it, so an empty ca_cert_pem falls back to the system - # trust store. - billing_metrics_enabled = var.billing_metrics_endpoint != "" - billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != "" - billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != "" - billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != "" + # Enterprise request metering, gated on billing_metrics_endpoint. + billing_metrics_enabled = var.billing_metrics_endpoint != "" + billing_metrics_client_secret_enabled = local.billing_metrics_enabled && var.billing_metrics_client_secret != "" billing_metrics_env = local.billing_metrics_enabled ? [ { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint }, + { name = "LITELLM_BILLING_METRICS_CLIENT_ID", value = var.billing_metrics_client_id }, + { name = "LITELLM_BILLING_METRICS_TOKEN_URL", value = var.billing_metrics_token_url }, ] : [] - billing_metrics_secrets = concat( - local.billing_metrics_client_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_client_cert[0].arn }, - ] : [], - local.billing_metrics_client_key_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", valueFrom = aws_secretsmanager_secret.billing_metrics_client_key[0].arn }, - ] : [], - local.billing_metrics_ca_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CA_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_ca_cert[0].arn }, - ] : [], - ) + billing_metrics_secrets = local.billing_metrics_client_secret_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_SECRET", valueFrom = aws_secretsmanager_secret.billing_metrics_client_secret[0].arn }, + ] : [] managed_db_env = var.create_database ? [ { name = "IAM_TOKEN_DB_AUTH", value = "true" }, @@ -338,27 +325,12 @@ locals { # ---------- Gateway ---------- resource "aws_ecs_task_definition" "gateway" { - # Metering needs a client certificate AND its key. Each secret is created only - # when its own PEM is supplied, so an endpoint set with a missing key would - # otherwise apply cleanly and leave the proxy logging "missing config" and - # never exporting. ca_cert_pem stays optional: empty means fall back to the - # system trust store. - # - # The guard lives here, on an unconditional resource, rather than on the cert - # secret: that secret is count-gated on the cert itself, so it has zero - # instances in exactly the case this must catch. Adding count or for_each to - # this resource would silently stop the guard from evaluating. - # - # endpoint cert key -> result - # "" any any -> metering off, no secrets created - # set set set -> metering on - # set any-missing -> plan fails here lifecycle { precondition { condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != "" ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set." } precondition { @@ -465,9 +437,7 @@ resource "aws_ecs_service" "gateway" { aws_secretsmanager_secret_version.license, aws_secretsmanager_secret_version.database_url, aws_secretsmanager_secret_version.redis_url, - aws_secretsmanager_secret_version.billing_metrics_client_cert, - aws_secretsmanager_secret_version.billing_metrics_client_key, - aws_secretsmanager_secret_version.billing_metrics_ca_cert, + aws_secretsmanager_secret_version.billing_metrics_client_secret, ] tags = local.tags @@ -481,9 +451,9 @@ resource "aws_ecs_task_definition" "backend" { lifecycle { precondition { condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != "" ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set." } } @@ -566,9 +536,7 @@ resource "aws_ecs_service" "backend" { aws_secretsmanager_secret_version.ui_password, aws_secretsmanager_secret_version.database_url, aws_secretsmanager_secret_version.redis_url, - aws_secretsmanager_secret_version.billing_metrics_client_cert, - aws_secretsmanager_secret_version.billing_metrics_client_key, - aws_secretsmanager_secret_version.billing_metrics_ca_cert, + aws_secretsmanager_secret_version.billing_metrics_client_secret, ] tags = local.tags diff --git a/terraform/litellm/aws/iam.tf b/terraform/litellm/aws/iam.tf index 3c55f07b02a..a653bd82d31 100644 --- a/terraform/litellm/aws/iam.tf +++ b/terraform/litellm/aws/iam.tf @@ -53,9 +53,7 @@ data "aws_iam_policy_document" "secrets_access" { [aws_secretsmanager_secret.master_key.arn], aws_secretsmanager_secret.license[*].arn, aws_secretsmanager_secret.ui_password[*].arn, - aws_secretsmanager_secret.billing_metrics_client_cert[*].arn, - aws_secretsmanager_secret.billing_metrics_client_key[*].arn, - aws_secretsmanager_secret.billing_metrics_ca_cert[*].arn, + aws_secretsmanager_secret.billing_metrics_client_secret[*].arn, aws_secretsmanager_secret.database_url[*].arn, aws_secretsmanager_secret.redis_url[*].arn, local.extra_secret_arns, diff --git a/terraform/litellm/aws/secrets.tf b/terraform/litellm/aws/secrets.tf index 921bae4d827..6821b8b43dc 100644 --- a/terraform/litellm/aws/secrets.tf +++ b/terraform/litellm/aws/secrets.tf @@ -75,59 +75,21 @@ resource "aws_secretsmanager_secret_version" "ui_password" { secret_string = var.ui_password } -# Billing-metrics mTLS material — only created when metering is enabled -# (billing_metrics_endpoint non-empty) and the operator supplied the PEM. -# The task-execution role gets GetSecretValue via iam.tf, and gateway + -# backend pick the env vars up through shared_secrets in ecs.tf. -resource "aws_secretsmanager_secret" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "aws_secretsmanager_secret" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - name = "${local.name}-billing-metrics-client-cert" - description = "LITELLM_BILLING_METRICS_CLIENT_CERT for gateway + backend." + name = "${local.name}-billing-metrics-client-secret" + description = "LITELLM_BILLING_METRICS_CLIENT_SECRET for gateway + backend." recovery_window_in_days = 0 tags = local.tags } -resource "aws_secretsmanager_secret_version" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "aws_secretsmanager_secret_version" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret_id = aws_secretsmanager_secret.billing_metrics_client_cert[0].id - secret_string = var.billing_metrics_client_cert_pem -} - -resource "aws_secretsmanager_secret" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - name = "${local.name}-billing-metrics-client-key" - description = "LITELLM_BILLING_METRICS_CLIENT_KEY for gateway + backend." - recovery_window_in_days = 0 - - tags = local.tags -} - -resource "aws_secretsmanager_secret_version" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret_id = aws_secretsmanager_secret.billing_metrics_client_key[0].id - secret_string = var.billing_metrics_client_key_pem -} - -resource "aws_secretsmanager_secret" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - name = "${local.name}-billing-metrics-ca-cert" - description = "LITELLM_BILLING_METRICS_CA_CERT for gateway + backend." - recovery_window_in_days = 0 - - tags = local.tags -} - -resource "aws_secretsmanager_secret_version" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret_id = aws_secretsmanager_secret.billing_metrics_ca_cert[0].id - secret_string = var.billing_metrics_ca_cert_pem + secret_id = aws_secretsmanager_secret.billing_metrics_client_secret[0].id + secret_string = var.billing_metrics_client_secret } resource "aws_secretsmanager_secret" "db_master_password" { diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 9b7a279a907..be7656c2d48 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -764,13 +764,7 @@ variable "otel_headers_secret_arn" { # ---------- Enterprise billing metrics ---------- # -# License-gated request metering. Opt-in and gated entirely on -# billing_metrics_endpoint: leave it empty (the default) and nothing -# metering-related lands in the container env. Set it and gateway + backend -# export billable-request counts over OTLP/HTTP, authenticating to the -# collector with an mTLS client cert. The proxy accepts the cert, key, and CA -# as either a file path or literal PEM content, so on Fargate they are -# injected straight from Secrets Manager as env vars and no volume is needed. +# License-gated request metering. The client secret is stored in Secrets Manager. variable "billing_metrics_endpoint" { description = <<-EOT @@ -784,44 +778,34 @@ variable "billing_metrics_endpoint" { default = "" } -variable "billing_metrics_client_cert_pem" { +variable "billing_metrics_client_id" { description = <<-EOT - PEM content of the mTLS client certificate issued for this deployment. - When billing_metrics_endpoint is set, the stack stores this in a - `-litellm--billing-metrics-client-cert` Secrets Manager - entry, grants the task-execution role GetSecretValue on it, and exposes - it to gateway + backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required - whenever metering is enabled. + Cognito OAuth2 client ID exposed to the gateway and backend as + LITELLM_BILLING_METRICS_CLIENT_ID. Required when billing_metrics_endpoint + is set. + EOT + type = string + default = "" +} + +variable "billing_metrics_client_secret" { + description = <<-EOT + Cognito OAuth2 client secret stored in Secrets Manager and exposed to + the gateway and backend as LITELLM_BILLING_METRICS_CLIENT_SECRET. + Required when billing_metrics_endpoint is set. EOT type = string default = "" sensitive = true } -variable "billing_metrics_client_key_pem" { +variable "billing_metrics_token_url" { description = <<-EOT - PEM content of the private key matching - billing_metrics_client_cert_pem. Stored in a - `-litellm--billing-metrics-client-key` Secrets Manager - entry and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required - whenever metering is enabled. + Cognito OAuth2 client-credentials token endpoint exposed to gateway and + backend as LITELLM_BILLING_METRICS_TOKEN_URL. EOT type = string - default = "" - sensitive = true -} - -variable "billing_metrics_ca_cert_pem" { - description = <<-EOT - PEM content of the CA bundle used to verify the metering collector. - Only needed for private or test collectors whose CA is not in the - system trust store; telemetry.litellm.ai is publicly trusted, so leave - this empty for production. When set, it is exposed as - LITELLM_BILLING_METRICS_CA_CERT. - EOT - type = string - default = "" - sensitive = true + default = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" } # ---------- Collector sidecar ---------- diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 3227dcac179..9e21ba14543 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -210,33 +210,23 @@ Behavior matches the AWS stack 1:1; the only naming differences are ### Enterprise billing metrics -License-gated request metering is opt-in and gated entirely on -`billing_metrics_endpoint`. Empty (default) and no billing env is added to -the container, so existing deployments are unchanged. Set it and both -gateway and backend export billable-request counts over OTLP/HTTP, -authenticating to the collector with the mTLS client certificate issued for -your deployment. - -The proxy accepts the certificate, key, and CA bundle as either a file path -or literal PEM content. This stack takes the PEM, writes each one to its own -Secret Manager entry, grants the runtime service account -`roles/secretmanager.secretAccessor` on them, and injects them as Cloud Run -secret env vars `LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and -`_CA_CERT` when set), so no volume mount is needed. +License-gated request metering uses Cognito bearer authentication. Set the +endpoint, client ID, and token URL as plain values, and provide the client +secret through the sensitive Terraform variable. The client secret is stored +in Secret Manager and injected into both gateway and backend services ```hcl billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics" +billing_metrics_client_id = "cognito-client-id" +billing_metrics_token_url = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" ``` ```bash -export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)" -export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)" +export TF_VAR_billing_metrics_client_secret="cognito-client-secret" ``` -`billing_metrics_ca_cert_pem` is only for private or test collectors whose -CA is not in the system trust store; leave it empty against -`telemetry.litellm.ai`. Metering requires an enterprise license, so pair -this with `litellm_license`. To tune the export cadence, set +The client ID and client secret are required when the endpoint is set +Metering also requires an enterprise license. To tune the export cadence, set `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / `backend_extra_env` diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index d0b32a367d6..07bf68b3a62 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -62,32 +62,19 @@ locals { { name = "OTEL_HEADERS", secret = var.otel_headers_secret, version = "latest" }, ] : [] - # Enterprise request metering, gated on billing_metrics_endpoint. The - # endpoint rides in as a plain env var; the mTLS material lives in Secret - # Manager (secrets.tf) and is injected as PEM-valued env vars, which the - # proxy accepts in place of file paths. Each PEM is wired only when the - # operator supplied it, so an empty ca_cert_pem falls back to the system - # trust store. - billing_metrics_enabled = var.billing_metrics_endpoint != "" - billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != "" - billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != "" - billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != "" + # Enterprise request metering, gated on billing_metrics_endpoint. + billing_metrics_enabled = var.billing_metrics_endpoint != "" + billing_metrics_client_secret_enabled = local.billing_metrics_enabled && var.billing_metrics_client_secret != "" billing_metrics_env_kv = local.billing_metrics_enabled ? [ { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint }, + { name = "LITELLM_BILLING_METRICS_CLIENT_ID", value = var.billing_metrics_client_id }, + { name = "LITELLM_BILLING_METRICS_TOKEN_URL", value = var.billing_metrics_token_url }, ] : [] - billing_metrics_env_secrets = concat( - local.billing_metrics_client_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", secret = google_secret_manager_secret.billing_metrics_client_cert[0].id, version = "latest" }, - ] : [], - local.billing_metrics_client_key_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", secret = google_secret_manager_secret.billing_metrics_client_key[0].id, version = "latest" }, - ] : [], - local.billing_metrics_ca_cert_enabled ? [ - { name = "LITELLM_BILLING_METRICS_CA_CERT", secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id, version = "latest" }, - ] : [], - ) + billing_metrics_env_secrets = local.billing_metrics_client_secret_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_SECRET", secret = google_secret_manager_secret.billing_metrics_client_secret[0].id, version = "latest" }, + ] : [] # Cloud Run v2 secret env vars use value_source.secret_key_ref pointing at a # secret resource ID. Shared between gateway and backend (the migrations @@ -481,9 +468,7 @@ resource "google_cloud_run_v2_service" "gateway" { google_secret_manager_secret_iam_member.license, google_secret_manager_secret_iam_member.extras, google_secret_manager_secret_iam_member.otel_headers, - google_secret_manager_secret_iam_member.billing_metrics_client_cert, - google_secret_manager_secret_iam_member.billing_metrics_client_key, - google_secret_manager_secret_iam_member.billing_metrics_ca_cert, + google_secret_manager_secret_iam_member.billing_metrics_client_secret, google_secret_manager_secret_iam_member.metrics_run_monitoring, google_project_iam_member.runtime_metric_writer, google_storage_bucket_iam_member.proxy_config_runtime, @@ -603,9 +588,7 @@ resource "google_cloud_run_v2_service" "backend" { google_secret_manager_secret_iam_member.ui_password, google_secret_manager_secret_iam_member.extras, google_secret_manager_secret_iam_member.otel_headers, - google_secret_manager_secret_iam_member.billing_metrics_client_cert, - google_secret_manager_secret_iam_member.billing_metrics_client_key, - google_secret_manager_secret_iam_member.billing_metrics_ca_cert, + google_secret_manager_secret_iam_member.billing_metrics_client_secret, google_storage_bucket_iam_member.proxy_config_runtime, google_sql_user.app, terraform_data.migration, diff --git a/terraform/litellm/gcp/iam.tf b/terraform/litellm/gcp/iam.tf index 5c2bc38bc12..1e0a5fe9626 100644 --- a/terraform/litellm/gcp/iam.tf +++ b/terraform/litellm/gcp/iam.tf @@ -10,9 +10,9 @@ resource "google_service_account" "runtime" { lifecycle { precondition { condition = !var.create_runtime || var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != "" ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set and create_runtime is true." + error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set and create_runtime is true." } } } @@ -91,28 +91,10 @@ resource "google_secret_manager_secret_iam_member" "otel_headers" { member = "serviceAccount:${google_service_account.runtime.email}" } -# Billing-metrics mTLS accessors — only created when request metering is -# enabled and the matching PEM was supplied. -resource "google_secret_manager_secret_iam_member" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "google_secret_manager_secret_iam_member" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret_id = google_secret_manager_secret.billing_metrics_client_cert[0].id - role = "roles/secretmanager.secretAccessor" - member = "serviceAccount:${google_service_account.runtime.email}" -} - -resource "google_secret_manager_secret_iam_member" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret_id = google_secret_manager_secret.billing_metrics_client_key[0].id - role = "roles/secretmanager.secretAccessor" - member = "serviceAccount:${google_service_account.runtime.email}" -} - -resource "google_secret_manager_secret_iam_member" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret_id = google_secret_manager_secret.billing_metrics_ca_cert[0].id + secret_id = google_secret_manager_secret.billing_metrics_client_secret[0].id role = "roles/secretmanager.secretAccessor" member = "serviceAccount:${google_service_account.runtime.email}" } diff --git a/terraform/litellm/gcp/secrets.tf b/terraform/litellm/gcp/secrets.tf index b3b656721ab..f3636b6e8e8 100644 --- a/terraform/litellm/gcp/secrets.tf +++ b/terraform/litellm/gcp/secrets.tf @@ -64,59 +64,21 @@ resource "google_secret_manager_secret_version" "ui_password" { secret_data = var.ui_password } -# Billing-metrics mTLS material — only created when metering is enabled -# (billing_metrics_endpoint non-empty) and the operator supplied the PEM. -# The runtime SA gets accessor permission via iam.tf, and gateway + backend -# pick the env vars up through billing_metrics_env_secrets in cloudrun.tf. -resource "google_secret_manager_secret" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "google_secret_manager_secret" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret_id = "${local.name}-billing-metrics-client-cert" + secret_id = "${local.name}-billing-metrics-client-secret" labels = local.labels replication { auto {} } } -resource "google_secret_manager_secret_version" "billing_metrics_client_cert" { - count = local.billing_metrics_client_cert_enabled ? 1 : 0 +resource "google_secret_manager_secret_version" "billing_metrics_client_secret" { + count = local.billing_metrics_client_secret_enabled ? 1 : 0 - secret = google_secret_manager_secret.billing_metrics_client_cert[0].id - secret_data = var.billing_metrics_client_cert_pem -} - -resource "google_secret_manager_secret" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret_id = "${local.name}-billing-metrics-client-key" - labels = local.labels - replication { - auto {} - } -} - -resource "google_secret_manager_secret_version" "billing_metrics_client_key" { - count = local.billing_metrics_client_key_enabled ? 1 : 0 - - secret = google_secret_manager_secret.billing_metrics_client_key[0].id - secret_data = var.billing_metrics_client_key_pem -} - -resource "google_secret_manager_secret" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret_id = "${local.name}-billing-metrics-ca-cert" - labels = local.labels - replication { - auto {} - } -} - -resource "google_secret_manager_secret_version" "billing_metrics_ca_cert" { - count = local.billing_metrics_ca_cert_enabled ? 1 : 0 - - secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id - secret_data = var.billing_metrics_ca_cert_pem + secret = google_secret_manager_secret.billing_metrics_client_secret[0].id + secret_data = var.billing_metrics_client_secret } resource "google_secret_manager_secret" "metrics_run_monitoring" { diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 6161e893e02..e40ee36fa18 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -602,14 +602,7 @@ variable "gateway_metrics_collector_image" { # ---------- Enterprise billing metrics ---------- # -# License-gated request metering. Opt-in and gated entirely on -# billing_metrics_endpoint: leave it empty (the default) and nothing -# metering-related is added to the container env. Set it and gateway + -# backend export billable-request counts over OTLP/HTTP, authenticating to -# the collector with an mTLS client cert. The proxy accepts the cert, key, -# and CA as either a file path or literal PEM content, so on Cloud Run they -# are injected straight from Secret Manager as env vars and no volume is -# needed. +# License-gated request metering. The client secret is stored in Secret Manager. variable "billing_metrics_endpoint" { description = <<-EOT @@ -623,44 +616,34 @@ variable "billing_metrics_endpoint" { default = "" } -variable "billing_metrics_client_cert_pem" { +variable "billing_metrics_client_id" { description = <<-EOT - PEM content of the mTLS client certificate issued for this deployment. - When billing_metrics_endpoint is set, the stack stores this in a - `-litellm--billing-metrics-client-cert` Secret Manager - entry, grants the runtime SA accessor on it, and exposes it to gateway + - backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required whenever - metering is enabled. + Cognito OAuth2 client ID exposed to the gateway and backend as + LITELLM_BILLING_METRICS_CLIENT_ID. Required when billing_metrics_endpoint + is set. + EOT + type = string + default = "" +} + +variable "billing_metrics_client_secret" { + description = <<-EOT + Cognito OAuth2 client secret stored in Secret Manager and exposed to + the gateway and backend as LITELLM_BILLING_METRICS_CLIENT_SECRET. + Required when billing_metrics_endpoint is set. EOT type = string default = "" sensitive = true } -variable "billing_metrics_client_key_pem" { +variable "billing_metrics_token_url" { description = <<-EOT - PEM content of the private key matching - billing_metrics_client_cert_pem. Stored in a - `-litellm--billing-metrics-client-key` Secret Manager entry - and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required whenever - metering is enabled. + Cognito OAuth2 client-credentials token endpoint exposed to gateway and + backend as LITELLM_BILLING_METRICS_TOKEN_URL. EOT type = string - default = "" - sensitive = true -} - -variable "billing_metrics_ca_cert_pem" { - description = <<-EOT - PEM content of the CA bundle used to verify the metering collector. - Only needed for private or test collectors whose CA is not in the - system trust store; telemetry.litellm.ai is publicly trusted, so leave - this empty for production. When set, it is exposed as - LITELLM_BILLING_METRICS_CA_CERT. - EOT - type = string - default = "" - sensitive = true + default = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token" } # ---------- Collector sidecar ---------- diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 57ae9f1905f..42ee6178bf0 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -176,6 +176,15 @@ ignored_function_names = [ "_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py "_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py + "job_cost_of_poll", # Read by job_to_charge_once and charged_video_id, tested through the guardrail-priced poll tests in test_batch_utils.py + "charged_video_id", # Read by job_to_charge_once and get_spend_logs_id, tested through both in test_batch_utils.py and test_spend_tracking_utils.py + "_requested_video_id", # Read by charged_video_id, tested through the polled-id video test in test_batch_utils.py + "_guardrail_cost_of_poll", # Read by job_cost_of_poll and repeat_poll_cost, tested through the guardrail-priced poll tests + "repeat_poll_cost", # Read by job_to_charge_once and _charge_request_cost, tested through every limiter's guardrail poll test + "pays_only_its_guardrail", # Read by get_spend_logs_id, tested through its guardrail-only poll test in test_spend_tracking_utils.py + "charged_poll_cost", # Read by get_spend_logs_id, tested through the finished job guardrail poll tests in test_db_spend_update_writer.py + "_polled_batch", # Read by defers_batch_cost and _charged_batch, tested through the redacted batch retrieve tests + "_charged_batch", # Read by job_to_charge_once, tested through the batch charge identity tests in test_batch_utils.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 781d8fd051b..b4ca8bd3d50 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -40,11 +40,11 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below) +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below). The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). The memory tests (`test_reliability_memory_e2e.py`, run by the perf lane on its own fresh stack: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches, the check for the v1.100.0 retry-breadcrumb leak). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests @@ -100,7 +100,7 @@ Each suite provides its own `client` fixture (see `llm_translation/passthrough_c Request and response bodies are typed pydantic models in `models.py`. Only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass -Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. Add `@pytest.mark.quiet_stack` to a test that measures the proxy itself (RSS, latency): the shared stack lock in `stack_lock.py` then runs it while no other test on the host is hitting the stack, marked or not, so the reading depends only on the test's own traffic. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache +Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. A test that measures the proxy itself (RSS, latency) cannot share the stack with the rest of the suite, so it goes in `perf/`, marked `perf`. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache ## Record and replay fixtures diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index b2442d492cf..14201ef6846 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig): ) -def _register_deployment(proxy, deployment: CompatDeployment) -> str: - """Register one deployment and return its proxy-assigned model_id - once it is servable on the data plane. The aliases are shared by every - cell and, under xdist, by every worker, so no call to them belongs to - one test and none is cached: the matrix exists to reach real providers.""" - return proxy.create_model( - deployment.model_name, - deployment.litellm_params, +def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]): + """The aliases are shared by every cell and, under xdist, by every worker, + so no call to them belongs to one test and none is cached: the matrix + exists to reach real providers.""" + from models import ModelInfoBody, ModelNewBody + + return proxy.register_models( + tuple( + ModelNewBody( + model_name=deployment.model_name, + litellm_params=deployment.litellm_params, + model_info=ModelInfoBody(), + ) + for deployment in deployments + ), provider_live=True, ) @@ -632,18 +639,14 @@ def _compat_models_registered() -> Any: yield return + from proxy_client import ModelNotRegistered, ModelRegistered from requests import RequestException proxy = _build_control_plane_client(proxy_config) - registered_ids: list[str] = [] - failures: list[tuple[str, str]] = [] + registrations = _register_deployments(proxy, load_all_deployments()) + registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) + failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered)) try: - for deployment in load_all_deployments(): - try: - model_id = _register_deployment(proxy, deployment) - registered_ids.append(model_id) - except (AssertionError, RequestException) as exc: - failures.append((deployment.model_name, str(exc))) if failures: summary = "\n".join( f" - {name}: {reason}" for name, reason in failures diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4c7c57ca18f..8fb96bef2f9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -35,6 +35,7 @@ from e2e_config import ( PROMPT_CACHING_OPT_IN_ENV, PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, + PROXY_WIDE_CONFIG_OPT_IN_ENV, SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, SAGEMAKER_NOVA_OPT_IN_ENV, SECRET_MANAGER_OPT_IN_ENV, @@ -54,7 +55,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse from provider_cache_routing import LIVE_PROVIDER_REQUIRED from provider_edge import replay_leftover_error from proxy_client import ProxyClient, build_proxy_client -from stack_lock import stack_lock _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() @@ -77,6 +77,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "secret_manager": SECRET_MANAGER_OPT_IN_ENV, "sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV, "sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, + "proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV, } ) @@ -162,11 +163,6 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set", ) - config.addinivalue_line( - "markers", - "quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; " - "every other test waits for it to finish", - ) config.addinivalue_line( "markers", "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " @@ -200,6 +196,11 @@ def pytest_configure(config: pytest.Config) -> None: "sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; " "deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it", ) + config.addinivalue_line( + "markers", + "proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite " + "is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: @@ -269,7 +270,7 @@ def pytest_collection_finish(session: pytest.Session) -> None: ahead of xdist's own collection-finish report, and the controller schedules no test until every worker has reported, so this is the idle footprint of a stack that just passed its readiness gate. The fixture hands the capture to the - idle-budget test in router/test_reliability_memory_e2e.py.""" + idle-budget test in perf/test_reliability_memory_e2e.py.""" if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items): return session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS) @@ -299,12 +300,6 @@ def _proxy_fail_reason() -> str | None: return None -@pytest.hookimpl(wrapper=True) -def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]: - with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None): - return (yield) - - @pytest.hookimpl(tryfirst=True) def pytest_runtest_setup(item: pytest.Item) -> None: """Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe. diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 854958fc39e..457f03406f1 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -20,8 +20,8 @@ - {id: guardrail.ibm_guardrails.post_call.blocks, module: guardrail, tier: P1, hook_point: post_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/ibm_guardrails/ibm_detector.py", rationale: "Output policy validation"} - {id: guardrail.semantic_guard.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/semantic_guard", rationale: "Semantic policy compliance"} - {id: guardrail.block_code_execution.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/block_code_execution", rationale: "Code-injection prevention"} -- {id: guardrail.tool_permission.pre_call.allows, module: guardrail, tier: P1, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_hooks/tool_permission.py", rationale: "Grant allowed tools"} -- {id: guardrail.tool_permission.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/tool_permission.py", rationale: "Block unauthorized tools"} +- {id: guardrail.tool_permission.pre_call.allows, module: guardrail, tier: P1, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/tool_permission.py", rationale: "Grant allowed tools"} +- {id: guardrail.tool_permission.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/tool_permission.py", rationale: "Block unauthorized tools"} - {id: guardrail.microsoft_purview.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/microsoft_purview/purview_dlp.py", rationale: "DLP sensitive-data disclosure"} - {id: guardrail.headroom.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/headroom/headroom.py", rationale: "Anomaly detection threshold"} - {id: guardrail.generic_guardrail_api.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py", rationale: "Vendor-agnostic custom API"} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index cbaf46130d6..49997aef162 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -44,8 +44,6 @@ - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"} - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"} -- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"} -- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"} - {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"} - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"} - {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 724aa2362e7..6316d22d3fa 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -210,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT" SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER" SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT" SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT" +PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) @@ -234,7 +235,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES", MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1")) MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200")) MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768")) -MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64")) def ws_base_url() -> str: diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py index d5a6a9d9866..87545299169 100644 --- a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -19,10 +19,9 @@ verdict comes from the proxy itself. from __future__ import annotations -from typing import Final +from typing import Final, Literal import pytest - from e2e_config import unique_marker from e2e_http import StreamingResponse, UnknownApiError from e2e_metadata import Capability, Domain, Mode, Provider, Subject, meta @@ -31,14 +30,29 @@ from guardrails_client import ( ToolPermissionParamsBody, ToolPermissionRuleBody, poll_until_blocked, + poll_until_blocked_stream, poll_until_guardrail_applied, ) from lifecycle import ResourceManager -from models import ChatResponse, ChatTool, ChatToolFunction +from models import ( + AnthropicCustomTool, + AnthropicMessagesBody, + AnthropicMessagesResponse, + AnthropicToolChoice, + ChatMessage, + ChatResponse, + ChatTool, + ChatToolFunction, + JsonSchemaProperty, + ResponsesFunctionTool, + ResponsesToolBody, + ResponsesToolOutput, + ToolInputSchema, +) -pytestmark = pytest.mark.e2e +pytestmark: Final = pytest.mark.e2e -MODEL = "gemini-2.5-flash" +MODEL: Final = "gemini-2.5-flash" #: The one tool the guardrail permits, and one it does not. Both are declared by #: the caller in the request body; the guardrail reads them there. @@ -63,11 +77,15 @@ DENIED_TOOL: Final = ChatTool( TOOL_PROMPT: Final = "What is the weather in Paris right now?" +type FlatToolEndpoint = Literal["messages", "responses"] +ALLOWED_TOOL_SCHEMA: Final = ToolInputSchema(properties={"city": JsonSchemaProperty(type="string")}, required=["city"]) +DENIED_TOOL_SCHEMA: Final = ToolInputSchema() + def _register_tool_permission(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None: """Allow-list exactly one tool: everything else falls to `default_action=deny` and, with `on_disallowed_action=block`, is rejected outright.""" - guardrail_id = client.register( + guardrail_id: Final = client.register( name, ToolPermissionParamsBody( mode="pre_call", @@ -100,6 +118,57 @@ def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: ) +def _send_flat_tool( + client: GuardrailsClient, + key: str, + endpoint: FlatToolEndpoint, + *, + guardrail: str, + tool: ChatTool, + schema: ToolInputSchema, + require_tool: bool = False, +) -> StreamingResponse: + name: Final = tool.function.name + match endpoint: + case "messages": + return client.proxy.transport.send( + "/v1/messages", + headers=client.proxy.transport.bearer(key), + json=AnthropicMessagesBody( + model=MODEL, + messages=[ChatMessage(role="user", content=TOOL_PROMPT)], + max_tokens=128, + tools=[ + AnthropicCustomTool(name=name, description=tool.function.description or "", input_schema=schema) + ], + tool_choice=AnthropicToolChoice(type="any") if require_tool else None, + guardrails=[guardrail], + ), + ) + case "responses": + return client.proxy.transport.send( + "/v1/responses", + headers=client.proxy.transport.bearer(key), + json=ResponsesToolBody( + model=MODEL, + input=TOOL_PROMPT, + tools=[ResponsesFunctionTool(name=name, description=tool.function.description, parameters=schema)], + tool_choice="required" if require_tool else None, + guardrails=[guardrail], + ), + ) + + +def _flat_tool_call_names(endpoint: FlatToolEndpoint, body: str) -> tuple[str, ...]: + match endpoint: + case "messages": + content: Final = AnthropicMessagesResponse.model_validate_json(body).content or [] + return tuple(block.name for block in content if block.type == "tool_use" and block.name) + case "responses": + output: Final = ResponsesToolOutput.model_validate_json(body).output + return tuple(item.name for item in output if item.type == "function_call" and item.name) + + class TestToolPermissionPreCall: @pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["chat_completions"]) @meta( @@ -118,10 +187,10 @@ class TestToolPermissionPreCall: rejected with a 400 that names the denied tool. An unauthorized tool that merely reaches the model is the whole failure mode this guardrail exists to prevent, so a 200 here is a hard failure.""" - name = f"e2e-toolperm-block-{unique_marker()}" + name: Final = f"e2e-toolperm-block-{unique_marker()}" _register_tool_permission(client, resources, name=name) - result = poll_until_blocked( + result: Final = poll_until_blocked( lambda: client.chat( scoped_key, MODEL, @@ -161,10 +230,10 @@ class TestToolPermissionPreCall: and the model calls it. Without the header check a guardrail that never attached would pass this test for the wrong reason, so the 200 alone is not the contract.""" - name = f"e2e-toolperm-allow-{unique_marker()}" + name: Final = f"e2e-toolperm-allow-{unique_marker()}" _register_tool_permission(client, resources, name=name) - outcome = poll_until_guardrail_applied( + outcome: Final = poll_until_guardrail_applied( lambda: client.chat_raw( scoped_key, MODEL, @@ -178,13 +247,117 @@ class TestToolPermissionPreCall: ) assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" - applied = _applied_guardrails(outcome) + applied: Final = _applied_guardrails(outcome) assert name in applied, ( "the allowed call must carry x-litellm-applied-guardrails naming the guardrail; " f"without it the 200 only proves the guardrail never ran. Got {applied!r}" ) - called = _tool_call_names(ChatResponse.model_validate_json(outcome.body)) + called: Final = _tool_call_names(ChatResponse.model_validate_json(outcome.body)) + assert called == (ALLOWED_TOOL.function.name,), ( + f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}" + ) + + @pytest.mark.parametrize( + "endpoint", + [ + pytest.param( + "messages", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["messages"]), + ), + pytest.param( + "responses", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["responses"]), + ), + ], + ) + @meta( + Subject( + domain=Domain.GUARDRAILS, + providers=(Provider.GEMINI,), + models=(MODEL,), + capabilities=(Capability.FUNCTION_CALLING,), + mode=Mode.NONSTREAM, + ) + ) + def test_pre_call_blocks_flat_tool_outside_the_allow_list( + self, + client: GuardrailsClient, + resources: ResourceManager, + scoped_key: str, + endpoint: FlatToolEndpoint, + ) -> None: + name: Final = f"e2e-toolperm-block-{endpoint}-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + outcome: Final = poll_until_blocked_stream( + lambda: _send_flat_tool( + client, scoped_key, endpoint, guardrail=name, tool=DENIED_TOOL, schema=DENIED_TOOL_SCHEMA + ) + ) + + assert outcome.status_code == 400, ( + f"tool_permission let a tool outside the allow-list through /{endpoint}; " + f"got {outcome.status_code}: {outcome.body[:400]}" + ) + assert DENIED_TOOL.function.name in outcome.body, ( + f"the block must name the denied tool so the caller can fix the request; got: {outcome.body[:400]}" + ) + assert "Violated guardrail policy" in outcome.body, ( + f"the block must be a guardrail verdict, got: {outcome.body[:400]}" + ) + + @pytest.mark.parametrize( + "endpoint", + [ + pytest.param( + "messages", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["messages"]), + ), + pytest.param( + "responses", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["responses"]), + ), + ], + ) + @meta( + Subject( + domain=Domain.GUARDRAILS, + providers=(Provider.GEMINI,), + models=(MODEL,), + capabilities=(Capability.FUNCTION_CALLING,), + mode=Mode.NONSTREAM, + ) + ) + def test_pre_call_allows_permitted_flat_tool( + self, + client: GuardrailsClient, + resources: ResourceManager, + scoped_key: str, + endpoint: FlatToolEndpoint, + ) -> None: + name: Final = f"e2e-toolperm-allow-{endpoint}-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + outcome: Final = poll_until_guardrail_applied( + lambda: _send_flat_tool( + client, + scoped_key, + endpoint, + guardrail=name, + tool=ALLOWED_TOOL, + schema=ALLOWED_TOOL_SCHEMA, + require_tool=True, + ), + name, + ) + + assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" + applied: Final = _applied_guardrails(outcome) + assert name in applied, ( + f"the allowed call must carry x-litellm-applied-guardrails naming the guardrail; got {applied!r}" + ) + called: Final = _flat_tool_call_names(endpoint, outcome.body) assert called == (ALLOWED_TOOL.function.name,), ( f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}" ) diff --git a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py index 6e564f4625a..1ec2a6822e3 100644 --- a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py +++ b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py @@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t /config/update. A chat completion that attaches the server and asks to search logs must report, in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's -OPENAI_API_KEY for the embedding deployment. +OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is +marked `proxy_wide_config` and runs on a stack the main suite is not driving. """ from __future__ import annotations @@ -23,7 +24,7 @@ from mcp_client import McpClient from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool from pydantic import BaseModel -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config] TOP_K: Final = 3 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9c12e2146c9..86af30f413a 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -653,6 +653,35 @@ class ResponsesStreamBody(BaseModel): cache: dict[str, bool] | None = {"no-cache": True} +class ResponsesFunctionTool(BaseModel): + model_config = ConfigDict(frozen=True) + type: Literal["function"] = "function" + name: str + description: str | None = None + parameters: ToolInputSchema + + +class ResponsesToolBody(BaseModel): + model_config = ConfigDict(frozen=True) + model: str + input: str + tools: list[ResponsesFunctionTool] + tool_choice: Literal["required"] | None = None + guardrails: list[str] + max_output_tokens: int = 128 + + +class ResponsesOutputItem(BaseModel): + model_config = ConfigDict(frozen=True) + type: str + name: str | None = None + + +class ResponsesToolOutput(BaseModel): + model_config = ConfigDict(frozen=True) + output: list[ResponsesOutputItem] + + class CountTokensBody(BaseModel): """POST /v1/messages/count_tokens body: the /v1/messages shape minus max_tokens (the endpoint only counts the prompt).""" @@ -1262,27 +1291,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]): pass -class ConfigPatchResponse(BaseModel): - status: str - values: dict[str, float | dict[str, float]] - - -class CostDiscountConfig(RootModel[dict[str, float]]): - pass - - -class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]): - pass - - -class CostDiscountConfigResponse(BaseModel): - values: dict[str, float] - - -class CostMarginConfigResponse(BaseModel): - values: dict[str, float | dict[str, float]] - - class BedrockGuardrailParams(BaseModel): guardrail: Literal["bedrock"] = "bedrock" mode: Literal["pre_call"] = "pre_call" diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py deleted file mode 100644 index c9114ccbff2..00000000000 --- a/tests/e2e/perf/setup/upgrade.py +++ /dev/null @@ -1,223 +0,0 @@ -from __future__ import annotations - -import math -import time -from collections.abc import Callable -from dataclasses import replace -from typing import Final - -import psycopg -from cases.small_bodies_5rps import CASE -from e2e_config import MASTER_KEY -from e2e_http import send_text_external -from prometheus_client.parser import text_string_to_metric_families -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter - -from setup.performance_test_case import PerformanceTestCase -from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes -from setup.wire import WireRecord - -MAX_TRAFFIC_SECONDS: Final = 6 * 3600 -_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...]) -_ACTIVITY: Final = TypeAdapter(tuple[list[int], int]) -_COUNT: Final = TypeAdapter(tuple[int]) -_SERVICES: Final = ("postgres", "batch_write_to_db") - - -def upgrade_case() -> PerformanceTestCase: - return replace( - CASE, - name="database_upgrade_5rps", - traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)), - limits=tuple( - limit - for limit in CASE.limits - if limit.metric - in { - "unexpected_status_count", - "gateway_added_seconds", - "time_to_first_token_seconds", - "memory_growth_megabytes", - } - ), - ) - - -class MetricsHeaders(BaseModel): - model_config = ConfigDict(frozen=True) - authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization") - - -def postgres_failures(text: str) -> float: - families: Final = tuple(text_string_to_metric_families(text)) - names: Final = frozenset(family.name for family in families) - required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES) - samples: Final = tuple(sample for family in families for sample in family.samples) - has_database_activity: Final = any( - sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0 - for sample in samples - ) - assert required <= names or has_database_activity, f"Missing database service counters: {required - names}" - failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required}) - assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter" - return sum(failures) - - -def assert_database_metrics(urls: tuple[str, ...]) -> None: - assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint" - for url in urls: - assert_database_metrics_url(url) - - -def assert_database_metrics_url(url: str) -> None: - result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10) - assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}" - assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}" - - -def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]: - ids: Final = tuple(record.call_id for record in records if record.call_id) - if not ids: - return frozenset() - rows: Final = _ROWS.validate_python( - connection.execute( - 'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),) - ).fetchall() - ) - return matched_response_ids(records, rows) - - -def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]: - call_ids: Final = frozenset(call_id for _, call_id in rows if call_id) - return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids) - - -def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]: - return _ACTIVITY.validate_python( - connection.execute("""SELECT - ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()), - (SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock') - """).fetchone() - ) - - -def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None: - assert state.phase == "complete", f"Upgrade did not finish: {state.phase}" - assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid" - (missing,) = _COUNT.validate_python( - connection.execute( - """ - SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS ( - SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name - AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""", - (list(state.required_migrations),), - ).fetchone() - ) - assert missing == 0, f"{missing} target migrations were not applied" - (unfinished,) = _COUNT.validate_python( - connection.execute(""" - SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL - """).fetchone() - ) - assert unfinished == 0, f"{unfinished} unfinished migrations" - assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, ( - "Baseline rows changed or disappeared" - ) - - -class UpgradeMonitor: - def __init__( - self, - *, - state: Callable[[], UpgradeState], - check_metrics: Callable[[], None], - read_activity: Callable[[], tuple[list[int], int]], - read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]], - start_indexes: Callable[[], None], - clock: Callable[[], float] = time.monotonic, - wall_clock: Callable[[], float] = time.time, - ) -> None: - self.state: Final = state - self.check_metrics: Final = check_metrics - self.read_activity: Final = read_activity - self.read_persisted: Final = read_persisted - self.start_indexes: Final = start_indexes - self.clock: Final = clock - self.wall_clock: Final = wall_clock - self.started: Final = clock() - self.status_limit: Final = next( - limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count" and limit.max is not None - ) - self.pending: tuple[WireRecord, ...] = () - self.processed = 0 - self.successes = 0 - self.persisted_count = 0 - self.overlap_writes = 0 - self.max_lock_waiters = 0 - self.previous_builds: frozenset[int] = frozenset() - self.build_seen_at: float | None = None - self.released = False - self.completed_at: float | None = None - - def poll(self, records: tuple[WireRecord, ...]) -> bool: - self.check_metrics() - current: Final = self.state() - assert current.phase != "failed", "Target migration/index worker failed" - assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout" - successful: Final = tuple(record for record in records if record.status == 200) - self.processed += len(records) - self.successes += len(successful) - assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses" - assert all(record.ok and record.text for record in successful), "Failed or empty inference response" - assert all(record.response_id and record.call_id for record in successful), ( - "Response missing spend-log identifiers" - ) - builds, lock_waiters = self.read_activity() - if not self.previous_builds.intersection(builds): - self.build_seen_at = self.wall_clock() if builds else None - self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters) - self.pending += successful - found: Final = self.read_persisted(self.pending) - self.persisted_count += len(found) - still_building, _ = self.read_activity() - since: Final = self.build_seen_at - if since is not None and frozenset(builds).intersection(still_building): - self.overlap_writes += sum( - 1 for record in self.pending if record.response_id in found and record.started_at >= since - ) - self.previous_builds = frozenset(builds) - self.pending = tuple(record for record in self.pending if record.response_id not in found) - if self.successes and not self.released: - self.start_indexes() - self.released = True - if current.phase == "complete" and self.completed_at is None: - self.completed_at = self.clock() - return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300 - - def assert_complete(self, state: UpgradeState) -> None: - assert self.successes >= 300, f"Only {self.successes} requests completed" - assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log" - assert self.persisted_count == self.successes, "Spend-log count does not match successful requests" - assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work" - - -def upgrade_properties(state: UpgradeState, monitor: UpgradeMonitor) -> tuple[tuple[str, float], ...]: - """The upgrade's own numbers under the `perf_` prefix the dashboard ingests, next to the perf rows.""" - return ( - ("perf_upgrade_schema_seconds", state.schema_seconds), - ("perf_upgrade_index_seconds", state.index_seconds), - ("perf_upgrade_overlap_writes", float(monitor.overlap_writes)), - ("perf_upgrade_max_lock_waiters", float(monitor.max_lock_waiters)), - ("perf_upgrade_pending_indexes", float(len(state.pending_indexes))), - ("perf_upgrade_unmatched_requests", float(len(monitor.pending))), - ) - - -def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor: - return UpgradeMonitor( - state=lambda: read_state(connection), - check_metrics=lambda: assert_database_metrics(urls), - read_activity=lambda: activity(connection), - read_persisted=lambda records: persisted(connection, records), - start_indexes=lambda: release_indexes(connection), - ) diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py deleted file mode 100644 index fba6af30014..00000000000 --- a/tests/e2e/perf/test_database_upgrade.py +++ /dev/null @@ -1,102 +0,0 @@ -from __future__ import annotations - -import os -import sys -import time -from pathlib import Path -from typing import Final - -import pytest -from idp import Keycloak -from lifecycle import ResourceManager -from memory_readings import read_rss_everywhere -from proxy_client import ProxyClient -from setup.datastore import read_db_counters, read_redis_counters -from setup.observe import ( - Observation, - evaluate, - junit_properties, - read_mock_requests, - render, - reset_mock, - set_mock_latency, -) -from setup.plan import build_plan -from setup.podcpu import gateway_cpu_sampler -from setup.provision import MOCK_PROVIDERS, build_setup, mock_url -from setup.run import run_plan -from setup.upgrade import assert_finished, monitor_for, upgrade_case, upgrade_properties -from setup.upgrade_worker import connect, read_state - -pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] - - -@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only") -def test_database_upgrade( - proxy: ProxyClient, idp: Keycloak, resources: ResourceManager, request: pytest.FixtureRequest -) -> None: - case: Final = upgrade_case() - urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS) - for url in urls: - _ = reset_mock(url) - _ = set_mock_latency(url, case.mock_latency) - setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources) - metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip()) - with connect() as connection: - assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting" - rss_before: Final = read_rss_everywhere(proxy) - db_before: Final = read_db_counters() - redis_before: Final = read_redis_counters() - monitor: Final = monitor_for(connection, metrics) - with gateway_cpu_sampler() as cpu: - outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30) - assert outcome.locust_exit == 0, outcome.locust_output - remaining: Final = outcome.records[monitor.processed :] - _ = monitor.poll(remaining) - deadline: Final = time.monotonic() + 120 - while monitor.pending and time.monotonic() < deadline: - time.sleep(1) - _ = monitor.poll(()) - state: Final = read_state(connection) - _ = sys.stdout.write( - f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} " - f"unmatched_requests={len(monitor.pending)}\n" - ) - rss_after: Final = read_rss_everywhere(proxy) - observed: Final = Observation( - records=outcome.records, - mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)), - guardrail_calls=(), - rss_before=rss_before, - rss_after=rss_after, - wall_seconds=outcome.wall_seconds, - spend_logged_share=1, - db_before=db_before, - db_after=read_db_counters(), - redis_before=redis_before, - redis_after=read_redis_counters(), - gateway_cpu_seconds=cpu.seconds, - ) - rows, exceeded = evaluate(case, observed) - item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped - assert isinstance(item, pytest.Item) - item.user_properties[:] = [ - entry for entry in item.user_properties if not entry[0].startswith("perf_") - ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set - item.user_properties.extend(junit_properties(rows)) - item.user_properties.extend(upgrade_properties(state, monitor)) - _ = sys.stdout.write(render(case, observed, rows) + "\n") - _ = sys.stdout.write( - f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} " - f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} " - f"pending_indexes={state.pending_indexes}\n" - ) - artifact: Final = os.environ.get("E2E_ARTIFACT_DIR") - if artifact: - _ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8") - monitor.assert_complete(state) - assert_finished(connection, state) - assert {(r.replica, r.hostname) for r in rss_before.readings} == { - (r.replica, r.hostname) for r in rss_after.readings - }, "Gateway pods replaced during the upgrade" - assert not exceeded, "\n".join(exceeded) diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index db657a04b00..0a137e84572 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -48,7 +48,7 @@ from setup.ramp import step_properties from setup.run import run_plan from setup.telemetry import read_gateway_telemetry, telemetry_sampler -pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] +pytestmark = pytest.mark.perf _SPEND_SAMPLE: Final = 40 _WARMUP_SECONDS: Final = 60 diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/perf/test_reliability_memory_e2e.py similarity index 77% rename from tests/e2e/router/test_reliability_memory_e2e.py rename to tests/e2e/perf/test_reliability_memory_e2e.py index fb114092303..2066137bb72 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/perf/test_reliability_memory_e2e.py @@ -1,17 +1,19 @@ """Live e2e: a few hundred requests that fail before any provider answers must not -grow the proxy's resident memory past a loose ceiling once the proxy is warm. +grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads +the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack. The regression this guards shipped in v1.100.0: every retry breadcrumb copied the whole request and the copies nested into one router-global list, so a proxy under retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that shape: a model group whose deployments refuse at the socket (an unreachable base -URL) with cooldown_time 0 so the router keeps retrying them, per-request retries, -and a fallback group that refuses the same way, each request carrying a long chat +URL, with a literal key since the perf stack carries no provider keys) with +cooldown_time 0 so the router keeps retrying them, per-request retries, and a +fallback group that refuses the same way, each request carrying a long chat transcript so every whole-request copy costs hundreds of containers instead of a -handful. Under the stack's cooldown policy a -deployment that fails a handful of times in a row is benched (a bad-credential 401 -included), the router answers "No deployments available" without retrying, and the -retry loop that leaks stops running; cooldown_time 0 keeps it running. +handful. Under the stack's cooldown policy a deployment that fails a handful of +times in a row is benched (a bad-credential 401 included), the router answers "No +deployments available" without retrying, and the retry loop that leaks stops running; +cooldown_time 0 keeps it running. Four identical batches run back to back. The first two are the warmup that grows the proxy's caches and allocator arenas, the last two are the measured window, and @@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time, a worker that answered only one checkpoint would otherwise drop out of the comparison, which is where a leaking worker could hide. -RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, -json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS -by only about 15 MB per 300 failing requests, while every failing request's stored -request snapshot carried a copy of the request per failed attempt, over 100 KB on -the first call and a couple of MB once the copies nested, against tens of KB with -the fix. So the first check sends one failing request before the phases, reads its -spend log back through /spend/logs, and holds the stored request body to a fixed -size budget: the deterministic catch for a breadcrumb that copies the whole -request. It runs before the phases because the leaking writer drops its own rows -under the phases' traffic (a queue budget hit, a recursion limit on the nested -copies), which would turn the size check into a missing-row check. +The stored-request half of that regression (every failing request's spend log carrying a +copy of the request per failed attempt) is a deterministic size check, so it lives in +tests/integration/spend/test_failing_request_stored_snapshot.py. -A second, cheaper check holds the idle footprint: every worker's RSS as the harness +A cheaper check holds the idle footprint: every worker's RSS as the harness read it at collection time, before this pytest process sent any traffic (see conftest.pytest_collection_finish), must sit under a fixed budget. On the release gate that is a fresh stack right after its readiness gate, one worker per @@ -74,7 +68,6 @@ from typing import Final import pytest -from complexity_router_client import ComplexityRouterClient from e2e_config import ( MEMORY_CONCURRENCY, MEMORY_IDLE_RSS_BUDGET_MB, @@ -83,7 +76,6 @@ from e2e_config import ( MEMORY_RSS_CEILING_MB, MEMORY_RSS_SAMPLE_INTERVAL_SECONDS, MEMORY_RSS_SETTLE_SAMPLES, - MEMORY_STORED_REQUEST_BUDGET_KB, MEMORY_TRANSCRIPT_TURNS, PROXY_REPLICA_URLS, unique_marker, @@ -91,11 +83,12 @@ from e2e_config import ( from e2e_metadata import Domain, Mode, Provider, Subject, meta from lifecycle import ResourceManager from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere -from models import ChatMessage, RouterSettingsOverride, SpendLogRow +from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride from proxy_client import ProxyClient -from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = [pytest.mark.e2e, pytest.mark.perf] + +REFUSING_MODEL: Final = "openai/gpt-5.5" DEPLOYMENTS_PER_GROUP: Final = 2 RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES @@ -110,7 +103,6 @@ class FailedCall: status_code: int seconds: float body_head: str - call_id: str | None @dataclass(frozen=True, slots=True) @@ -123,8 +115,17 @@ class WorkerGrowth: return self.after.ram_usage_mb - self.warm.ram_usage_mb +def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str: + return proxy.create_model( + name, + LiteLLMParamsBody( + model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0 + ), + ) + + def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None: - for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): + for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): resources.defer(lambda model_id=model_id: proxy.delete_model(model_id)) @@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS) def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall: started: Final = time.perf_counter() - resp: Final = chat_override( - proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT + resp: Final = proxy.transport.send( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=model, + messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")], + max_tokens=512, + router_settings_override=override, + cache={"no-cache": True}, + ), ) - return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id) + return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300]) def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]: @@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall ) -def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float: - assert call.call_id, ( - f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}" - ) - rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id) - assert rows, ( - f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack " - "writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored " - "request outgrew the writer's queue budget" - ) - snapshot: Final = rows[0].proxy_server_request - assert snapshot, ( - f"spend log {call.call_id} stored no request body, so the stack is not running with " - "general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously" - ) - return len(json.dumps(snapshot).encode()) / 1024 - - class TestReliabilityMemory: @pytest.mark.covers("reliability.perf.idle_memory.under_slo") @meta( @@ -309,13 +300,13 @@ class TestReliabilityMemory: Subject( domain=Domain.ROUTING, providers=(Provider.OPENAI,), - models=(REAL_MODEL,), + models=(REFUSING_MODEL,), mode=Mode.NONSTREAM, ) ) - def test_failing_requests_do_not_grow_rss_or_stored_request( + def test_failing_requests_do_not_grow_rss( self, - client: ComplexityRouterClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, record_property: Callable[[str, object], None], @@ -323,35 +314,25 @@ class TestReliabilityMemory: marker: Final = unique_marker() primary: Final = f"reliability-memory-{marker}" fallback: Final = f"reliability-memory-fb-{marker}" - _register_refusing_group(client.proxy, resources, primary) - _register_refusing_group(client.proxy, resources, fallback) + _register_refusing_group(proxy, resources, primary) + _register_refusing_group(proxy, resources, fallback) override: Final = RouterSettingsOverride( num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}] ) - probe: Final = _fail_once(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback((probe,), fallback) - stored_kb: Final = _stored_request_kb(client.proxy, probe) - assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, ( - f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the " - f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with " - f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole " - f"request into the stored snapshot the way the v1.100.0 ones did" - ) - warm: Final = _settled_rss_after_batches( - client.proxy, + proxy, scoped_key, primary, override, fallback, record_property, - _settled_rss_per_worker(client.proxy), + _settled_rss_per_worker(proxy), 1, WARMUP_BATCHES, ) after: Final = _settled_rss_after_batches( - client.proxy, + proxy, scoped_key, primary, override, diff --git a/tests/e2e/perf/test_ui_journeys.py b/tests/e2e/perf/test_ui_journeys.py index 63eefedec68..11b2cd1051d 100644 --- a/tests/e2e/perf/test_ui_journeys.py +++ b/tests/e2e/perf/test_ui_journeys.py @@ -33,7 +33,6 @@ from transport import HttpTransport pytestmark = [ pytest.mark.perf, - pytest.mark.quiet_stack, pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"), ] diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 8cc7bf0c2b4..9b842a723d8 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -11,7 +11,7 @@ from __future__ import annotations import os import time import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime from functools import reduce @@ -155,6 +155,21 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +@dataclass(frozen=True, slots=True) +class ModelRegistered: + model_name: str + model_id: str + + +@dataclass(frozen=True, slots=True) +class ModelNotRegistered: + model_name: str + reason: str + + +ModelRegistration = ModelRegistered | ModelNotRegistered + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -461,6 +476,7 @@ class ProxyClient: model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT + settle: Callable[[float], None] = settle_propagation def with_caller(self, caller: Caller) -> ProxyClient: return replace(self, caller=caller) @@ -694,31 +710,59 @@ class ProxyClient: balancer address is configured (every request opens a fresh connection, so the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is what makes the model safe to use anywhere.""" - model_id = unwrap( - self.transport.post( - "/model/new", - headers=self.management_headers(), - json=body.model_copy(update={"litellm_params": route_cache_model( - body.litellm_params, provider_edge_base, - enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live, - mode=body.model_info.mode, - )}), - response_type=ModelNewResponse, - ) - ).model_id + model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id written_at = time.monotonic() try: self._await_model_servable(body.model_name, listed_for) except BaseException: self.delete_model(model_id) raise - settle_propagation(written_at) + self.settle(written_at) return model_id + @step("Add a batch of deployments, then wait once for all of them to propagate") + def register_models( + self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False + ) -> tuple[ModelRegistration, ...]: + posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies) + written_at: Final = time.monotonic() + registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted) + self.settle(written_at) + return registrations + + def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]: + return self.transport.post( + "/model/new", + headers=self.management_headers(), + json=body.model_copy(update={"litellm_params": route_cache_model( + body.litellm_params, provider_edge_base, + enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live, + mode=body.model_info.mode, + )}), + response_type=ModelNewResponse, + ) + + def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration: + if not isinstance(posted, Success): + return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}") + match self._servable_outcome(model_name, None): + case Servable(): + return ModelRegistered(model_name=model_name, model_id=posted.data.model_id) + case NotServableOn() as missing: + self.delete_model(posted.data.model_id) + return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing)) + def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" + match self._servable_outcome(model_name, listed_for): + case Servable(): + return + case NotServableOn() as missing: + raise AssertionError(self._not_servable_message(model_name, missing)) + + def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn: headers: Final = self.management_headers(listed_for) - outcome: Final = await_servable_everywhere( + return await_servable_everywhere( {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, timeout=self.model_servable_timeout, @@ -728,19 +772,15 @@ class ProxyClient: now=time.monotonic, sleep=time.sleep, ) - match outcome: - case Servable(): - return - case NotServableOn(replica=replica, last_result=last_result): - raise AssertionError( - servable_timeout_message( - model_name=model_name, - replica=replica, - timeout=self.model_servable_timeout, - db_sync_seconds=self.model_servable_db_sync_seconds, - last_result=last_result, - ) - ) + + def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str: + return servable_timeout_message( + model_name=model_name, + replica=missing.replica, + timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + last_result=missing.last_result, + ) @staticmethod def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller: diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index f15cfd4e98b..891c5fc282c 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -13,7 +13,6 @@ markers = prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set - quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set @@ -21,3 +20,4 @@ markers = secret_manager: needs a proxy booted from gateway/secret_manager__ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py) sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it + proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py similarity index 52% rename from tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py rename to tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py index 8f3bed3f89b..a77ca91ad3a 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py @@ -17,11 +17,6 @@ from models import ( ChatBody, ChatMessage, ChatResponse, - ConfigPatchResponse, - CostDiscountConfig, - CostDiscountConfigResponse, - CostMarginConfig, - CostMarginConfigResponse, GuardrailCreateBody, GuardrailCreateResponse, KeyDeleteBody, @@ -31,14 +26,11 @@ from models import ( ) from spend_e2e_client import SpendClient -pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] +pytestmark = pytest.mark.e2e BACKEND_MODEL: Final = "openai/gpt-6.1-sol" INPUT_RATE: Final = 0.00005 OUTPUT_RATE: Final = 0.0001 -DISCOUNT: Final = 0.25 -MARGIN_PERCENT: Final = 0.1 -MARGIN_FIXED: Final = 0.0005 def _register_bedrock_guardrail( @@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea ) -def _set_discount(client: SpendClient, values: dict[str, float]) -> float: - unwrap( - client.proxy.transport.patch( - "/config/cost_discount_config", - headers=client.proxy.transport.master, - json=CostDiscountConfig(values), - response_type=ConfigPatchResponse, - ) - ) - return time.monotonic() - - -def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float: - unwrap( - client.proxy.transport.patch( - "/config/cost_margin_config", - headers=client.proxy.transport.master, - json=CostMarginConfig(values), - response_type=ConfigPatchResponse, - ) - ) - return time.monotonic() - - -def _get_discount(client: SpendClient) -> dict[str, float]: - return unwrap( - client.proxy.transport.get( - "/config/cost_discount_config", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=CostDiscountConfigResponse, - ) - ).values - - -def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]: - return unwrap( - client.proxy.transport.get( - "/config/cost_margin_config", - headers=client.proxy.transport.master, - params=NoBody(), - response_type=CostMarginConfigResponse, - ) - ).values - - def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str: model: Final = f"{prefix}-{unique_marker()}" model_id: Final = client.proxy.create_model( @@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str: return key -@pytest.fixture -def restored_pricing_config(client: SpendClient) -> Iterator[None]: - discount: Final = _get_discount(client) - margin: Final = _get_margin(client) - try: - _set_discount(client, {}) - reset_at: Final = _set_margin(client, {}) - settle_propagation(reset_at) - yield - finally: - try: - _set_discount(client, discount) - finally: - restored_at: Final = _set_margin(client, margin) - settle_propagation(restored_at) - - -class TestPricingConfigSpend: - @pytest.mark.covers( - "quota_management.spend_tracking.discount_config.logs_cost", - exercised_on=["chat_completions"], - ) - @meta( - Subject( - domain=Domain.SPEND_BUDGETS, - route=Route.SPEND_REPORTING, - providers=(Provider.OPENAI,), - models=(BACKEND_MODEL,), - mode=Mode.NONSTREAM, - ) - ) - def test_configured_discount_reaches_persisted_spend_row( - self, - client: SpendClient, - strict_resources: ResourceManager, - scoped_key: str, - restored_pricing_config: None, - ) -> None: - discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT}) - settle_propagation(discount_set_at) - model: Final = _register_model(client, strict_resources, "discount-priced") - chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage - assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 - assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 - - row: Final = poll_cost_row(client.proxy, chat.id) - assert row is not None - base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) - breakdown: Final = row.breakdown - assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT) - assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT) - assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT)) - - @pytest.mark.covers( - "quota_management.spend_tracking.margin_config.logs_cost", - exercised_on=["chat_completions"], - ) - @meta( - Subject( - domain=Domain.SPEND_BUDGETS, - route=Route.SPEND_REPORTING, - providers=(Provider.OPENAI,), - models=(BACKEND_MODEL,), - mode=Mode.NONSTREAM, - ) - ) - def test_configured_margin_reaches_persisted_spend_row( - self, - client: SpendClient, - strict_resources: ResourceManager, - scoped_key: str, - restored_pricing_config: None, - ) -> None: - margin_set_at: Final = _set_margin( - client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}} - ) - settle_propagation(margin_set_at) - model: Final = _register_model(client, strict_resources, "margin-priced") - chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage - assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 - assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 - - row: Final = poll_cost_row(client.proxy, chat.id) - assert row is not None - base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) - breakdown: Final = row.breakdown - expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED - assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT) - assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED) - assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin) - assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin) - +class TestGuardrailCostSpend: @pytest.mark.covers( "quota_management.spend_tracking.guardrail_cost.logs_cost", exercised_on=["chat_completions"], @@ -311,7 +164,6 @@ class TestPricingConfigSpend: client: SpendClient, strict_resources: ResourceManager, scoped_key: str, - restored_pricing_config: None, ) -> None: identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"] diff --git a/tests/e2e/stack_lock.py b/tests/e2e/stack_lock.py deleted file mode 100644 index 06df7a20b6a..00000000000 --- a/tests/e2e/stack_lock.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Cross-process reader/writer lock over the proxy stack every xdist worker shares. -Every collected test holds it shared, marker or not, since the Claude Code cells and -other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive, -and the `gate` file makes a waiting exclusive holder win over readers that arrive -after it.""" - -from __future__ import annotations - -import fcntl -import hashlib -import tempfile -from collections.abc import Generator -from contextlib import ExitStack, contextmanager -from pathlib import Path -from typing import Final - -from e2e_config import PROXY_BASE_URL - -STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12] -LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}" -GATE_FILE: Final = LOCK_DIR / "gate" -STACK_FILE: Final = LOCK_DIR / "stack" - - -@contextmanager -def _flock(path: Path, operation: int) -> Generator[None]: - with path.open("a") as handle: - fcntl.flock(handle, operation) - try: - yield - finally: - fcntl.flock(handle, fcntl.LOCK_UN) - - -@contextmanager -def stack_lock(exclusive: bool) -> Generator[None]: - LOCK_DIR.mkdir(parents=True, exist_ok=True) - if exclusive: - with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX): - yield - return - with ExitStack() as held: - with _flock(GATE_FILE, fcntl.LOCK_SH): - held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH)) - yield diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 83e573ae249..4f507c2fe7b 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness tests -Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network +Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, `claude_code/` covers the driver, builder and probe, `cli_harness/` covers the version resolver, the CLI env lockdown and the rate limiter, and `codex/` covers how the Codex driver reads the CLI's JSONL events and the compactions in a thread's rollout file. Put a new harness test under the folder that mirrors the suite folder whose module it covers diff --git a/tests/e2e_harness/perf/test_run.py b/tests/e2e_harness/perf/test_run.py index 23f62c9c560..480af96a12a 100644 --- a/tests/e2e_harness/perf/test_run.py +++ b/tests/e2e_harness/perf/test_run.py @@ -8,7 +8,29 @@ from typing import Final import pytest from setup.run import run_until from setup.wire import WireRecord -from test_upgrade import record + + +def record(seq: int = 0, *, started_at: float = 100) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"request-{seq}", + caller="key", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=started_at, + status=200, + header_s=1, + first_byte_s=1, + first_token_s=None, + total_s=1, + text="answer", + ok=True, + mock_delay_s=1, + call_id=f"call-{seq}", + error=None, + response_id=f"response-{seq}", + ) class StopAfterRecord: diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py deleted file mode 100644 index 3414476ba4c..00000000000 --- a/tests/e2e_harness/perf/test_upgrade.py +++ /dev/null @@ -1,375 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest -from cases.small_bodies_5rps import CASE -from memory_readings import RssCapture, RssReading -from setup.observe import Observation, evaluate -from setup.upgrade import ( - MAX_TRAFFIC_SECONDS, - UpgradeMonitor, - matched_response_ids, - postgres_failures, - upgrade_case, - upgrade_properties, -) -from setup.upgrade_worker import UpgradeState -from setup.wire import PlanProgress, WireRecord, response_id - - -def record(seq: int = 0, *, started_at: float = 100) -> WireRecord: - return WireRecord( - seq=seq, - marker=f"request-{seq}", - caller="key", - endpoint="chat_completions", - stream=False, - blocked=False, - started_at=started_at, - status=200, - header_s=1, - first_byte_s=1, - first_token_s=None, - total_s=1, - text="answer", - ok=True, - mock_delay_s=1, - call_id=f"call-{seq}", - error=None, - response_id=f"response-{seq}", - ) - - -def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None: - records: Final = (record(0), record(1), record(2)) - rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call")) - assert matched_response_ids(records, rows) == frozenset({"response-0"}) - - -def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState: - return UpgradeState.model_validate( - { - "phase": phase, - "started_at": 0, - "required_indexes": ["index"], - "required_migrations": ["migration"], - "pending_indexes": ["index"] if pending else [], - "baseline_rows": {"seed-row": "hash"}, - } - ) - - -class Probe: - def __init__(self) -> None: - self.now = 0.0 - self.current = state() - self.builds: list[int] = [] - self.found: frozenset[str] = frozenset() - self.releases = 0 - self.metric_checks = 0 - self.errors = False - - def check_metrics(self) -> None: - self.metric_checks += 1 - assert not self.errors, "Failed PostgreSQL operation" - - def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]: - return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id) - - def release(self) -> None: - self.releases += 1 - - def monitor(self) -> UpgradeMonitor: - return UpgradeMonitor( - state=lambda: self.current, - check_metrics=self.check_metrics, - read_activity=lambda: (self.builds, 2), - read_persisted=self.persisted, - start_indexes=self.release, - clock=lambda: self.now, - wall_clock=lambda: self.now, - ) - - -def test_stop_drains_issued_requests_without_issuing_the_rest() -> None: - progress: Final = PlanProgress(100) - assert progress.take(stop=False) - assert progress.take(stop=False) - progress.completed = 1 - assert not progress.take(stop=True) - assert not progress.drained - assert not progress.take(stop=False) - progress.completed = 2 - assert progress.drained - assert progress.issued == 2 - - -def test_finite_plan_still_waits_for_its_last_response() -> None: - progress: Final = PlanProgress(1) - assert progress.take(stop=False) - assert not progress.take(stop=False) - assert not progress.drained - progress.completed = 1 - assert progress.drained - - -@pytest.mark.parametrize( - "payload,expected", - [ - ({"id": "chat-id"}, "chat-id"), - ({"type": "response.created", "response": {"id": "response-id"}}, "response-id"), - ({"type": "message_start", "message": {"id": "message-id"}}, "message-id"), - ({"type": "response.output_item.added", "item": {"id": "item-id"}}, None), - ({"type": "content_block_delta", "delta": {"text": "hello"}}, None), - ], -) -def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None: - assert response_id(payload) == expected - - -def metrics(*, failure: str = "0", batch: str = "0") -> str: - return f"""# TYPE litellm_postgres_failed_requests_total counter -litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure} -# TYPE litellm_batch_write_to_db_failed_requests_total counter -litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch} -""" - - -def test_failed_sql_and_async_batch_writes_are_both_counted() -> None: - assert postgres_failures(metrics(failure="2", batch="3")) == 5 - assert postgres_failures(metrics()) == 0 - - -def test_no_error_label_series_is_a_valid_zero() -> None: - assert ( - postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter -# TYPE litellm_batch_write_to_db_failed_requests_total counter -""") - == 0 - ) - - -@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"]) -def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None: - total: Final = """# TYPE litellm_postgres_total_requests_total counter -litellm_postgres_total_requests_total{postgres="postgres"} 1 -""" - failure: Final = ( - f"# TYPE litellm_{service}_failed_requests_total counter\n" - f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n' - if service is not None - else "" - ) - assert postgres_failures(total + failure) == (1 if service is not None else 0) - - -@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"]) -def test_missing_failure_families_require_observed_database_activity(total: str) -> None: - with pytest.raises(AssertionError): - _ = postgres_failures( - "# TYPE litellm_postgres_total_requests_total counter\n" - f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n' - ) - - -@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")]) -def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None: - with pytest.raises(AssertionError): - _ = postgres_failures(body) - - -def test_index_completion_stops_after_the_post_upgrade_window() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - assert not monitor.poll((record(),)) - assert probe.releases == 1 - probe.builds = [17] - probe.now = 101 - assert not monitor.poll(()) - probe.found = frozenset({"response-0", "response-1"}) - probe.now = 103 - assert not monitor.poll((record(1, started_at=102),)) - assert monitor.overlap_writes == 1 - assert monitor.persisted_count == 2 - assert not monitor.pending - assert monitor.max_lock_waiters == 2 - probe.current = state("complete") - assert not monitor.poll(()) - probe.now = 162 - assert not monitor.poll(()) - probe.now = 163 - assert not monitor.poll(()) - probe.found = frozenset(f"response-{i}" for i in range(300)) - assert monitor.poll(tuple(record(i) for i in range(2, 300))) - assert probe.releases == 1 - assert probe.metric_checks == 7 - - -def test_http_200_without_spend_row_cannot_complete() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - records: Final = tuple(record(i) for i in range(300)) - _ = monitor.poll(records) - with pytest.raises(AssertionError, match="no matching spend log"): - monitor.assert_complete(state("complete", pending=False)) - - -def test_no_pending_index_work_still_requires_300_persisted_requests() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - records: Final = tuple(record(i) for i in range(300)) - probe.found = frozenset(f"response-{i}" for i in range(300)) - _ = monitor.poll(records) - monitor.assert_complete(state("complete", pending=False)) - assert monitor.successes == monitor.persisted_count == 300 - with pytest.raises(AssertionError, match="during index work"): - monitor.assert_complete(state("complete", pending=True)) - - -def test_postgres_failure_fails_even_when_inference_succeeded() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - probe.errors = True - with pytest.raises(AssertionError, match="PostgreSQL"): - _ = monitor.poll((record(),)) - - -def test_failed_migration_and_operational_timeout_are_not_success() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - probe.current = state("failed") - with pytest.raises(AssertionError, match="worker failed"): - _ = monitor.poll(()) - probe.current = state() - probe.now = MAX_TRAFFIC_SECONDS - with pytest.raises(AssertionError, match="operational timeout"): - _ = monitor.poll(()) - - -def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None: - case: Final = upgrade_case() - assert case.traffic.rps == 5 - assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection( - limit.metric for limit in case.limits - ) - assert case.limits == tuple( - limit - for limit in CASE.limits - if limit.metric - in { - "unexpected_status_count", - "gateway_added_seconds", - "time_to_first_token_seconds", - "memory_growth_megabytes", - } - ) - - -def failed_record(seq: int) -> WireRecord: - return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None}) - - -def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None: - probe: Final = Probe() - monitor: Final = probe.monitor() - assert not monitor.poll(tuple(failed_record(i) for i in range(3))) - assert probe.releases == 0 - assert monitor.processed == 3 - assert monitor.successes == 0 - assert not monitor.pending - probe.found = frozenset(f"response-{i}" for i in range(3, 303)) - assert not monitor.poll(tuple(record(i) for i in range(3, 303))) - monitor.assert_complete(state("complete", pending=False)) - assert probe.releases == 1 - assert monitor.processed == 303 - assert monitor.successes == monitor.persisted_count == 300 - - -def test_fourth_non_200_fails_across_poll_batches() -> None: - monitor: Final = Probe().monitor() - assert not monitor.poll(tuple(failed_record(i) for i in range(3))) - with pytest.raises(AssertionError, match="Too many unexpected non-200"): - _ = monitor.poll((failed_record(3),)) - - -@pytest.mark.parametrize( - "failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})] -) -def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: - with pytest.raises(AssertionError, match="Failed or empty inference"): - _ = Probe().monitor().poll((failed,)) - - -def _ceiling(metric: str, percentile: int | None = None) -> float: - return next( - limit.max - for limit in upgrade_case().limits - if limit.metric == metric and limit.percentile == percentile and limit.max is not None - ) - - -@pytest.mark.parametrize( - "over,expected", - [ - ("", ()), - ("unexpected_status_count", ("unexpected_status_count",)), - ("gateway_added_seconds", ("gateway_added_seconds percentile 50",)), - ("time_to_first_token_seconds", ("time_to_first_token_seconds percentile 99",)), - ("memory_growth_megabytes", ("memory_growth_megabytes",)), - ], -) -def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples( - over: str, expected: tuple[str, ...] -) -> None: - """Each sample sits exactly on the case's ceiling, so only the one nudged past it goes over.""" - failures: Final = int(_ceiling("unexpected_status_count")) + (over == "unexpected_status_count") - overhead: Final = _ceiling("gateway_added_seconds", 50) * (1.01 if over == "gateway_added_seconds" else 1) - ttft: Final = _ceiling("time_to_first_token_seconds", 99) * (1.01 if over == "time_to_first_token_seconds" else 1) - growth: Final = _ceiling("memory_growth_megabytes") + (over == "memory_growth_megabytes") - nonstream: Final = tuple(record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100)) - stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft}) - observed: Final = Observation( - records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))), - mock_requests=(), - guardrail_calls=(), - rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()), - rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()), - wall_seconds=60, - ) - rows, exceeded = evaluate(upgrade_case(), observed) - assert tuple(row.label for row in rows if row.over or row.missing) == expected - assert len(exceeded) == len(expected) - - -def test_upgrade_properties_carry_the_upgrade_timings_and_what_the_monitor_saw() -> None: - state: Final = UpgradeState( - phase="complete", - started_at=0, - schema_seconds=12.5, - index_seconds=240.25, - required_indexes=("idx_a", "idx_b"), - required_migrations=("20260101_a",), - pending_indexes=("idx_b",), - baseline_rows={}, - ) - monitor: Final = UpgradeMonitor( - state=lambda: state, - check_metrics=lambda: None, - read_activity=lambda: ([], 0), - read_persisted=lambda records: frozenset(), - start_indexes=lambda: None, - ) - monitor.overlap_writes = 37 - monitor.max_lock_waiters = 3 - monitor.pending = (record(0), record(1)) - - assert dict(upgrade_properties(state, monitor)) == { - "perf_upgrade_schema_seconds": 12.5, - "perf_upgrade_index_seconds": 240.25, - "perf_upgrade_overlap_writes": 37, - "perf_upgrade_max_lock_waiters": 3, - "perf_upgrade_pending_indexes": 1, - "perf_upgrade_unmatched_requests": 2, - } diff --git a/tests/e2e_harness/test_proxy_client.py b/tests/e2e_harness/test_proxy_client.py index e1615ec5f65..d352e5c8d36 100644 --- a/tests/e2e_harness/test_proxy_client.py +++ b/tests/e2e_harness/test_proxy_client.py @@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything. from __future__ import annotations import json +import time from builtins import ExceptionGroup from collections.abc import Callable, Generator, Iterable, Mapping from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain, repeat from queue import SimpleQueue @@ -40,7 +41,10 @@ from models import ( LiteLLMParamsBody, McpServerCreateBody, McpServerUpdateBody, + ModelDeleteBody, + ModelInfoBody, ModelListEntry, + ModelNewBody, ModelsListResponse, OrgNewBody, OrgUpdateBody, @@ -59,6 +63,8 @@ from proxy_client import ( ConvergeOutcome, CredentialKind, EverywhereConverged, + ModelNotRegistered, + ModelRegistered, ModelsPoller, NeverConvergedOn, NotConverged, @@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: ) assert json.loads(bodies.get_nowait()) == expected assert bodies.empty() + + +@contextmanager +def model_registry( + listed: frozenset[str], refused: frozenset[str] +) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]: + writes: Final[SimpleQueue[float]] = SimpleQueue() + deleted: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def reply(self, status: int, body: object) -> None: + self.send_response(status) + self.end_headers() + self.wfile.write(json.dumps(body).encode()) + + def do_GET(self) -> None: + self.reply(200, {"data": [{"id": name} for name in sorted(listed)]}) + + def do_POST(self) -> None: + raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if self.path == "/model/delete": + deleted.put(ModelDeleteBody.model_validate_json(raw).id) + self.reply(200, {}) + return + writes.put(time.monotonic()) + name: Final = ModelNewBody.model_validate_json(raw).model_name + if name in refused: + self.reply(500, {"error": "refused"}) + return + self.reply(200, {"model_id": f"id-{name}"}) + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + proxy: Final = build_proxy_client( + base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap" + ) + try: + yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestRegisterModels: + def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None: + names: Final = ("first", "refused", "ghost", "last") + settled: Final[SimpleQueue[float]] = SimpleQueue() + with ( + model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as ( + proxy, + writes, + deleted, + ), + without_retries(), + ): + registrations: Final = replace(proxy, settle=settled.put).register_models( + tuple( + ModelNewBody( + model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody() + ) + for name in names + ) + ) + last_write: Final = max(writes.get_nowait() for _ in names) + assert writes.empty() + assert settled.get_nowait() >= last_write + assert settled.empty(), "a batch must wait out propagation once, not once per model" + assert tuple((type(r), r.model_name) for r in registrations) == ( + (ModelRegistered, "first"), + (ModelNotRegistered, "refused"), + (ModelNotRegistered, "ghost"), + (ModelRegistered, "last"), + ) + assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last") + assert deleted.get_nowait() == "id-ghost" + assert deleted.empty() diff --git a/tests/e2e_harness/test_stack_lock.py b/tests/e2e_harness/test_stack_lock.py deleted file mode 100644 index b275ca8c755..00000000000 --- a/tests/e2e_harness/test_stack_lock.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for -every reader and keeps them out, and a reader arriving behind a waiting exclusive holder -queues behind it instead of starving it.""" - -from __future__ import annotations - -import fcntl -import inspect -import os -import subprocess -import sys -import time -from contextlib import ExitStack -from pathlib import Path -from typing import Final - -import pytest -import stack_lock - -HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent -DEADLINE_SECONDS: Final = 30.0 -SETTLE_SECONDS: Final = 0.5 -HOLDER_SCRIPT: Final = """ -import sys, time -from pathlib import Path -from stack_lock import stack_lock -name, mode, release_path, log_path = sys.argv[1:] - - -def record(event): - with Path(log_path).open("a") as log: - log.write(f"{name} {event}\\n") - - -record("waiting") -with stack_lock(exclusive=mode == "exclusive"): - record("enter") - while not Path(release_path).exists(): - time.sleep(0.02) - record("exit") -""" - - -def _events(log_path: Path) -> tuple[str, ...]: - return tuple(log_path.read_text().splitlines()) if log_path.exists() else () - - -def _wait_for_event(log_path: Path, event: str) -> None: - deadline: Final = time.monotonic() + DEADLINE_SECONDS - while event not in _events(log_path): - if time.monotonic() > deadline: - pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}") - time.sleep(0.02) - - -def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None: - deadline: Final = time.monotonic() + DEADLINE_SECONDS - with gate_path.open("a") as handle: - while True: - try: - fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB) - except BlockingIOError: - return - fcntl.flock(handle, fcntl.LOCK_UN) - if time.monotonic() > deadline: - pytest.fail("no exclusive holder ever took the gate") - time.sleep(0.02) - - -def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]: - holder: Final = held.enter_context( - subprocess.Popen( - ( - sys.executable, - "-P", - "-c", - HOLDER_SCRIPT, - name, - mode, - str(tmp_path / f"release-{name}"), - str(tmp_path / "events"), - ), - cwd=HARNESS_DIR, - env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)}, - ) - ) - held.callback(holder.kill) - return holder - - -def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None: - lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}" - lock_dir.mkdir() - log_path: Final = tmp_path / "events" - with ExitStack() as held: - first_reader: Final = _start_holder(held, tmp_path, "A", "shared") - _wait_for_event(log_path, "A enter") - second_reader: Final = _start_holder(held, tmp_path, "R", "shared") - _wait_for_event(log_path, "R enter") - (tmp_path / "release-R").touch() - _wait_for_event(log_path, "R exit") - writer: Final = _start_holder(held, tmp_path, "W", "exclusive") - _wait_until_gate_is_held_exclusively(lock_dir / "gate") - late_reader: Final = _start_holder(held, tmp_path, "B", "shared") - _wait_for_event(log_path, "B waiting") - time.sleep(SETTLE_SECONDS) - (tmp_path / "release-A").touch() - _wait_for_event(log_path, "W enter") - (tmp_path / "release-W").touch() - _wait_for_event(log_path, "B enter") - (tmp_path / "release-B").touch() - for holder in (first_reader, second_reader, writer, late_reader): - assert holder.wait(timeout=DEADLINE_SECONDS) == 0 - events: Final = _events(log_path) - assert events.index("R enter") < events.index("A exit") - assert events.index("W enter") > events.index("A exit") - assert events.index("B enter") > events.index("W exit") diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py index 8552bc2197d..438d877386b 100644 --- a/tests/integration/_support/riva.py +++ b/tests/integration/_support/riva.py @@ -39,6 +39,7 @@ class RivaCall: rpc: Rpc audio: bytes cancelled: bool + authorization: tuple[str, ...] = () @dataclass(frozen=True, slots=True) @@ -68,6 +69,10 @@ def pcm_bytes(seconds: float) -> int: return int(seconds * SAMPLE_RATE_HZ) * 2 +def _authorization(context: grpc.ServicerContext) -> tuple[str, ...]: + return tuple(str(value) for key, value in context.invocation_metadata() if key == "authorization") + + def carries_clip(audio: bytes, seconds: float) -> bool: sent: Final = _clip_samples(seconds) if len(audio) != len(sent) * 2: @@ -100,9 +105,9 @@ class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): released: Final = threading.Event() context.add_callback(released.set) released.wait(timeout=_STALL_LIMIT_SECONDS) - self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active())) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, not context.is_active(), _authorization(context))) return riva_asr_pb2.RecognizeResponse() - self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False)) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, False, _authorization(context))) return riva_asr_pb2.RecognizeResponse( results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])] ) @@ -111,13 +116,12 @@ class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: if self._offline_only: - self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False)) + self._calls.put(RivaCall(Rpc.STREAMING, b"", False, _authorization(context))) context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL) requests: Final = tuple(request_iterator) config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig() - self._calls.put( - RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False) - ) + audio: Final = b"".join(request.audio_content for request in requests) + self._calls.put(RivaCall(Rpc.STREAMING, audio, False, _authorization(context))) yield riva_asr_pb2.StreamingRecognizeResponse( results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)] ) diff --git a/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py b/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py index 42fb55a5c4b..c8fbcbb1c97 100644 --- a/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py +++ b/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py @@ -535,8 +535,10 @@ def test_link_only_links_after_a_cached_rejection_from_the_previous_mode(gateway denied: Final = _chat(rejecting, model, token) assert denied.status_code == 403, denied.text with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: - cached: Final = cache.get(f"jwt_key_mapping:sub:{user}") - assert cached is not None and b"__NO_MAPPING__" in cached, cached + eventually( + lambda: cache.get(f"jwt_key_mapping:sub:{user}"), + lambda value: value is not None and b"__NO_MAPPING__" in value, + ) with owned_proxy( gateway, tmp_path, diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py deleted file mode 100644 index 15cd5413149..00000000000 --- a/tests/integration/database/test_accounting_ledger.py +++ /dev/null @@ -1,610 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import os -import subprocess -import sys -from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence -from contextlib import asynccontextmanager -from datetime import datetime, timezone -from decimal import Decimal -from itertools import chain -from pathlib import Path -from typing import TYPE_CHECKING, Final, TypeAlias - -import psycopg -import pytest -import pytest_asyncio -from psycopg import sql -from pydantic import JsonValue - -from litellm.accounting import ( - CostEvent, - DeliveryState, - DeliveryTarget, - EventKind, - InboundUsageEvent, - Metric, - MetricName, - OutboundUsageEvent, - Outcome, - PricingSnapshot, -) -from litellm.proxy.db import accounting_ledger -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite -from litellm.types.utils import CallTypes -from tests.integration._support.database import scratch_database -from tests.integration._support.prisma_clients import connected_proxy_prisma - -if TYPE_CHECKING: - from prisma import Prisma - - from litellm.proxy.utils import PrismaClient - -_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger] - -_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" -_EVENT_INSERT: Final = ( - 'INSERT INTO "LiteLLM_AccountingEvent" ' - "(event_id, kind, schema_version, payload_hash, occurred_at, payload) " - "VALUES ('event', %s, %s, 'hash', '2026-01-01', '{}') RETURNING kind::text, schema_version" -) -_DELIVERY_INSERT: Final = ( - 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state, attempts) ' - "VALUES ('event', %s, %s, %s, %s) RETURNING target::text, version, state::text, attempts" -) - - -@pytest.fixture(scope="module") -def ledger_database() -> Iterator[str]: - with scratch_database("accounting") as database_url: - result: Final = subprocess.run( - [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(_SCHEMA)], - env={**os.environ, "DATABASE_URL": database_url}, - capture_output=True, - text=True, - timeout=300, - ) - assert result.returncode == 0, result.stdout + result.stderr - yield database_url - - -@pytest.fixture -def ledger(ledger_database: str) -> Iterator[psycopg.Connection[tuple[object, ...]]]: - with psycopg.connect(ledger_database) as connection: - try: - yield connection - finally: - connection.rollback() - - -@pytest.mark.parametrize("kind", tuple(EventKind)) -def test_ledger_accepts_every_public_event_kind( - ledger: psycopg.Connection[tuple[object, ...]], kind: EventKind -) -> None: - row: Final = ledger.execute(_EVENT_INSERT, (kind.value, 1)).fetchone() - assert row == (kind.value, 1) - - -@pytest.mark.parametrize("target", tuple(DeliveryTarget)) -@pytest.mark.parametrize("state", tuple(DeliveryState)) -def test_delivery_accepts_every_public_target_and_state( - ledger: psycopg.Connection[tuple[object, ...]], target: DeliveryTarget, state: DeliveryState -) -> None: - row: Final = ledger.execute(_DELIVERY_INSERT, (target.value, 1, state.value, 0)).fetchone() - assert row == (target.value, 1, state.value, 0) - - -@pytest.mark.parametrize( - ("enum_name", "values"), - [ - ("AccountingEventKind", tuple(value.value for value in EventKind)), - ("AccountingDeliveryState", tuple(value.value for value in DeliveryState)), - ("AccountingDeliveryTarget", tuple(value.value for value in DeliveryTarget)), - ], -) -def test_migrated_enums_match_the_public_contract( - ledger: psycopg.Connection[tuple[object, ...]], enum_name: str, values: tuple[str, ...] -) -> None: - rows: Final = ledger.execute( - "SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid WHERE t.typname = %s", - (enum_name,), - ).fetchall() - assert frozenset(row[0] for row in rows) == frozenset(values) - - -@pytest.mark.parametrize("version", [0, -1]) -def test_event_rejects_nonpositive_schema_versions( - ledger: psycopg.Connection[tuple[object, ...]], version: int -) -> None: - with pytest.raises(psycopg.errors.CheckViolation, match="AccountingEvent_schema_version_check"): - ledger.execute(_EVENT_INSERT, (EventKind.OUTBOUND_USAGE.value, version)) - - -@pytest.mark.parametrize( - ("version", "attempts", "constraint"), - [ - (0, 0, "AccountingDelivery_version_check"), - (-1, 0, "AccountingDelivery_version_check"), - (1, -1, "AccountingDelivery_attempts_check"), - ], -) -def test_delivery_rejects_invalid_versions_and_attempts( - ledger: psycopg.Connection[tuple[object, ...]], version: int, attempts: int, constraint: str -) -> None: - with pytest.raises(psycopg.errors.CheckViolation, match=constraint): - ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, version, "pending", attempts)) - - -@pytest.mark.parametrize("column", ["state", "target"]) -def test_delivery_rejects_unknown_enum_values(ledger: psycopg.Connection[tuple[object, ...]], column: str) -> None: - ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, 1, "pending", 0)) - with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): - ledger.execute( - sql.SQL('UPDATE "LiteLLM_AccountingDelivery" SET {} = %s').format(sql.Identifier(column)), ("invalid",) - ) - - -def test_event_rejects_unknown_kind(ledger: psycopg.Connection[tuple[object, ...]]) -> None: - with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): - ledger.execute(_EVENT_INSERT, ("invalid", 1)) - - -def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Connection[tuple[object, ...]]) -> None: - row: Final = ledger.execute( - 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) ' - "VALUES ('event', %s, 1) RETURNING state::text, attempts", - (DeliveryTarget.KEY_SPEND.value,), - ).fetchone() - assert row == ("pending", 0) - - -_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) -_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc) -_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND) - - -def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent: - fields: Final = dict( - event_id=event_id, - request_id="request", - call_type=CallTypes.completion, - started_at=_AT, - ended_at=_AT, - outcome=Outcome.SUCCEEDED, - dimensions=(), - ) - if inbound: - return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),)) - return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None) - - -def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent: - return CostEvent( - event_id=event_id, - usage_event_id=source, - occurred_at=_AT, - dimensions=(), - charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),), - snapshot=PricingSnapshot( - calculator_version="fixture", - model="fixture", - provider=None, - pricing_source="cost_map", - priced_model_key="fixture", - total_cost=Decimal(amount), - ), - ) - - -def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str: - return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}}) - - -async def _insert( - store: AccountingLedger, - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...] = _TARGETS, - *, - context: Mapping[str, JsonValue] | None = None, -) -> str: - return await store.insert_event(event, targets, envelope=_envelope(event, context=context)) - - -def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]: - with psycopg.connect(database_url) as connection: - cursor: Final = connection.execute(statement, parameters) - return cursor.fetchall() if cursor.description is not None else [] - - -def _complete_and_age(database_url: str) -> None: - with psycopg.connect(database_url) as connection: - connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,)) - connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,)) - - -@pytest_asyncio.fixture(loop_scope="function") -async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]: - monkeypatch.setenv("DATABASE_URL", ledger_database) - monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) - _rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"') - async with connected_proxy_prisma(ledger_database) as client: - yield client, AccountingLedger.for_client(client) - - -@pytest.mark.asyncio -async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None: - _, store = writer - event: Final = _usage() - assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted" - rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"') - assert rows == [ - ("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"}))) - ] - deliveries: Final = _rows( - ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"' - ) - assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} - - -@pytest.mark.asyncio -async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: - _, store = writer - event: Final = _usage(inbound=True) - results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8))) - assert results.count("inserted") == 1 and results.count("duplicate") == 7 - _complete_and_age(ledger_database) - before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') - reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2) - assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate" - assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"]) -async def test_conflicting_acceptance_cannot_overwrite_or_add_work( - writer: _Writer, ledger_database: str, change: str -) -> None: - _, store = writer - event: Final = _usage() - await _insert(store, event) - altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event - targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS - context: Final = {"override": "0.5"} if change == "context" else {} - options: Final = {change: 2} if change in ("schema_version", "version") else {} - with pytest.raises(AccountingConflictError, match="conflicts"): - await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options) - assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)] - assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == { - (t.value,) for t in _TARGETS - } - - -@pytest.mark.asyncio -async def test_competing_cost_identities_accept_only_one_authoritative_charge( - writer: _Writer, ledger_database: str -) -> None: - _, store = writer - await _insert(store, _usage()) - results: Final = await asyncio.gather( - _insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True - ) - assert sum(result == "inserted" for result in results) == 1 - assert sum(isinstance(result, AccountingConflictError) for result in results) == 1 - accepted: Final = _rows( - ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" - ) - assert len(accepted) == 1 - event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"])) - assert await _insert(store, event) == "duplicate" - assert event.snapshot.total_cost == Decimal("0.125") - with pytest.raises(AccountingConflictError): - await _insert(store, _cost(event_id=event.event_id, amount="0.5")) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("source", ["missing", "inbound", "cost"]) -async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None: - _, store = writer - if source == "inbound": - await _insert(store, _usage("source", inbound=True)) - if source == "cost": - await _insert(store, _usage()) - await _insert(store, _cost(event_id="source")) - with pytest.raises(AccountingConflictError, match="outbound usage"): - await _insert(store, _cost(source="source", event_id="invalid")) - assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [ - (0,) - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) -async def test_interruption_rolls_back_or_retry_recognizes_commit( - writer: _Writer, ledger_database: str, failure: str -) -> None: - client, store = writer - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - if failure == "before_commit": - raise RuntimeError("injected pre-commit failure") - if failure == "cancellation": - raise asyncio.CancelledError - raise RuntimeError("injected lost acknowledgement") - - error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError - with pytest.raises(error): - await _insert(AccountingLedger(interrupted), _usage()) - expected: Final = 1 if failure == "after_commit" else 0 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)] - assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted") - - -@pytest.mark.asyncio -async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None: - from prisma.errors import RawQueryError - - _, store = writer - with psycopg.connect(ledger_database) as connection: - connection.execute( - 'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)' - ) - try: - with pytest.raises(RawQueryError, match="reject_delivery"): - await _insert(store, _usage()) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] - finally: - with psycopg.connect(ledger_database) as connection: - connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery') - - -@pytest.mark.asyncio -@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"]) -async def test_retention_preserves_incomplete_or_recovery_needed_families( - writer: _Writer, ledger_database: str, blocked: str -) -> None: - _, store = writer - await _insert(store, _usage()) - if blocked != "unpriced": - await _insert(store, _cost()) - _complete_and_age(ledger_database) - with psycopg.connect(ledger_database) as connection: - if blocked in ("pending", "parked"): - connection.execute( - "UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,) - ) - if blocked == "recent_cost": - connection.execute( - "UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,) - ) - before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') - assert ( - await store.purge_completed( - redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF - ) - == 0 - ) - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before - - -@pytest.mark.asyncio -async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage()) - await _insert(store, _cost()) - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] - - -@pytest.mark.asyncio -async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None: - client, store = writer - for name in ("a", "b", "c"): - await _insert(store, _usage(name, inbound=True)) - _complete_and_age(ledger_database) - async with client.writer_db.tx() as tx: - await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE") - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("a",), - ("c",), - ] - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - - -@pytest.mark.asyncio -async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage("a", inbound=True)) - await _insert(store, _usage("b")) - await _insert(store, _usage("c", inbound=True)) - _complete_and_age(ledger_database) - _rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'") - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("a",), - ("b",), - ] - - -@pytest.mark.asyncio -async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage(), targets=()) - event: Final = _cost(amount="0") - assert await _insert(store, event, targets=()) == "inserted" - assert await _insert(store, event, targets=()) == "duplicate" - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - - -@pytest.mark.asyncio -async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted( - writer: _Writer, ledger_database: str -) -> None: - client, store = writer - await _insert(store, _usage()) - _complete_and_age(ledger_database) - inserted: Final = asyncio.Event() - release: Final = asyncio.Event() - - @asynccontextmanager - async def paused_commit() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - inserted.set() - await release.wait() - - task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost())) - try: - await asyncio.wait_for(inserted.wait(), timeout=5) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0 - finally: - release.set() - await task - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("cost",), - ("usage",), - ] - - -@pytest.mark.asyncio -async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None: - client, store = writer - await _insert(store, _usage(inbound=True)) - _complete_and_age(ledger_database) - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - raise RuntimeError("injected cleanup failure") - - with pytest.raises(RuntimeError, match="cleanup failure"): - await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1 - - -@pytest.mark.asyncio -async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured( - writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch -) -> None: - with scratch_database("replica") as replica_url: - monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url) - async with connected_proxy_prisma(ledger_database) as client: - store: Final = AccountingLedger.for_client(client) - await _insert(store, _usage()) - assert await _insert(store, _usage()) == "duplicate" - assert await _insert(store, _cost()) == "inserted" - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)] - - -def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]: - if isinstance(plan, dict): - own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else () - children: Final = (plan.get("Plan"), *plan.get("Plans", ())) - return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None)) - return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else () - - -@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES]) -def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None: - explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s") - with psycopg.connect(ledger_database) as connection: - connection.execute("SET enable_seqscan = off") - plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone() - assert plan is not None - conditions: Final = _index_conditions(plan[0]) - assert any( - name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) - for name, condition in conditions - ), conditions - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "conflict"]) -async def test_usage_batch_is_atomic_and_replay_preserves_receipts( - writer: _Writer, ledger_database: str, failure: str -) -> None: - client, store = writer - events: Final = ( - _usage("caller", inbound=True), - _usage("attempt").model_copy( - update={"metrics": (Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),)} - ), - ) - writes: Final = tuple(UsageWrite(event, _envelope(event), _TARGETS) for event in events) - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - if failure == "before_commit": - raise RuntimeError("injected pre-commit failure") - if failure == "after_commit": - raise RuntimeError("injected lost acknowledgement") - - if failure == "conflict": - await store.insert_usage_batch((writes[1],)) - _complete_and_age(ledger_database) - rejected: Final = (writes[0], UsageWrite(events[1], _envelope(events[1], context={"changed": True}), _TARGETS)) - with pytest.raises(AccountingConflictError): - await store.insert_usage_batch(rejected) - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("attempt",)] - else: - with pytest.raises(RuntimeError, match="injected"): - await AccountingLedger(interrupted).insert_usage_batch(writes) - count: Final = 2 if failure == "after_commit" else 0 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(count,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(count * len(_TARGETS),)] - await store.insert_usage_batch(writes) - recovered: Final = await store.load_usage_batch(tuple(event.event_id for event in events)) - assert len(recovered) == 1 - assert recovered[0].event == events[1] - assert json.loads(recovered[0].envelope) == json.loads(writes[1].envelope) - assert set(recovered[0].targets) == set(_TARGETS) - _complete_and_age(ledger_database) - before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') - await store.insert_usage_batch((*writes, *writes)) - assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(2,)] - - -@pytest.mark.asyncio -async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit( - writer: _Writer, ledger_database: str -) -> None: - from prisma.errors import DataError - - from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings - - _, store = writer - invalid: Final = _usage("invalid") - valid: Final = _usage("valid") - meter: Final = DurableUsageMeter(store.insert_usage_batch, UsageCaptureSettings(batch_size=2)) - rejected: Final = meter.submit( - (UsageWrite(invalid, _envelope(invalid, context={"parent": "bad\x00text"}), _TARGETS),) - ) - accepted: Final = meter.submit((UsageWrite(valid, _envelope(valid), _TARGETS),)) - try: - with pytest.raises(DataError): - await rejected - await accepted - finally: - await meter.close() - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] - assert (meter.pending_bytes, meter.pending_records) == (0, 0) diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py new file mode 100644 index 00000000000..19f01a6b902 --- /dev/null +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -0,0 +1,512 @@ +import asyncio +import os +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import psycopg +import pytest +import pytest_asyncio +from prisma import Prisma + +from litellm.repositories.accounting_delivery_repository import ( + AccountingDeliveryRepository, + DeliveryProjection, + LedgerTransaction, +) +from litellm.repositories.spend_log_insert import spend_log_insert_statement +from tests.integration._support.database import scratch_database + +_MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/migrations" + + +@pytest.fixture +def delivery_database(request: pytest.FixtureRequest) -> Iterator[psycopg.Connection[tuple[object, ...]]]: + with scratch_database("spendlogs_delivery") as url, psycopg.connect(url, autocommit=True) as connection: + if getattr(request, "param", "public") == "ledger_tenant": + connection.execute("CREATE SCHEMA ledger_tenant") + connection.execute("SET search_path TO ledger_tenant") + connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text()) + connection.execute("""CREATE TABLE "LiteLLM_SpendLogs" ( + request_id text PRIMARY KEY, call_type text NOT NULL, spend float8 NOT NULL DEFAULT 0, + status text, + metadata jsonb DEFAULT '{}', "startTime" timestamp(3) NOT NULL, + created_at timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP + )""") + connection.execute((_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text()) + yield connection + + +def _insert( + connection: psycopg.Connection[tuple[object, ...]], rows: Sequence[Mapping[str, object]], *, ledger: bool = True +) -> None: + from psycopg.types.json import Jsonb + + statement: Final = spend_log_insert_statement(rows) + with connection.transaction(): + inserted: Final = { + row[0] for row in connection.execute(statement[0].replace("$1", "%s"), (statement[1],)).fetchall() + } + for row in rows: + if ledger and row["request_id"] in inserted: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch) ' + "VALUES (%s, %s, %s, %s)", + (row["request_id"], row["request_id"], Jsonb(dict(row)), row.get("call_type") == "aretrieve_batch"), + ) + + +def _with_payload(row: Mapping[str, object]) -> Mapping[str, object]: + return {**row, "accounting_payload": dict(row)} + + +def _row(request_id: str) -> Mapping[str, object]: + return { + "request_id": request_id, + "call_type": "acompletion", + "spend": 0.5, + "startTime": "2026-11-02T12:00:00+00:00", + "metadata": '{"accounting_savings":{"compression":0.2}}', + } + + +@pytest.mark.parametrize("call_type", ["acompletion", "aretrieve_batch"]) +@pytest.mark.parametrize("mode", ["local", "url", "disabled"]) +@pytest.mark.parametrize("partitioned", [False, True]) +@pytest.mark.asyncio +async def test_delivery_first_acceptance_deduplicates_all_storage_modes( + delivery_database, delivery_client, mode, partitioned, call_type +): + from types import SimpleNamespace + + from litellm.repositories.table_repositories import SpendLogsRepository + + if partitioned: + delivery_database.execute('ALTER TABLE "LiteLLM_SpendLogs" RENAME TO spend_template') + delivery_database.execute( + 'CREATE TABLE "LiteLLM_SpendLogs" (LIKE spend_template INCLUDING DEFAULTS INCLUDING CONSTRAINTS, PRIMARY KEY(request_id, "startTime")) PARTITION BY RANGE ("startTime")' + ) + delivery_database.execute('CREATE TABLE "LiteLLM_SpendLogs_p20261102" PARTITION OF "LiteLLM_SpendLogs" DEFAULT') + delivery_database.execute("DROP TABLE spend_template") + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + row = {**_with_payload({**_row("shared"), "call_type": call_type}), "accounting_log_disabled": mode == "disabled"} + results = await asyncio.gather( + repository.insert_with_deliveries((row,), write_spend_logs=mode != "url"), + repository.insert_with_deliveries( + ({**row, "startTime": "2026-11-03T12:00:00"},), write_spend_logs=mode != "url" + ), + ) + assert sorted(results) == [0, 1] + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(mode == "local"),) + assert delivery_database.execute( + "SELECT accounting_payload->>'spend', state::text FROM \"LiteLLM_SpendLogDelivery\"" + ).fetchall() == [("0.5", "pending")] + if partitioned: + delivery_database.execute('DROP TABLE "LiteLLM_SpendLogs_p20261102"') + else: + delivery_database.execute('DELETE FROM "LiteLLM_SpendLogs"') + + async def apply(transaction, deliveries): + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(deliveries[0].row["spend"])) + + worker = AccountingDeliveryRepository(delivery_client) + assert (await worker.apply_pending(apply)).applied == 1 + assert (await worker.apply_pending(apply)).applied == 0 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.parametrize("mode", ["local", "url", "disabled"]) +@pytest.mark.asyncio +async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(delivery_database, delivery_client, mode): + from types import SimpleNamespace + + from litellm.repositories.table_repositories import SpendLogsRepository + + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + first = {**_with_payload(_row("repeated")), "accounting_log_disabled": mode == "disabled"} + later = { + **_with_payload({**_row("repeated"), "spend": 9.0, "startTime": "2026-11-03T12:00:00+00:00"}), + "accounting_log_disabled": mode == "disabled", + } + assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1 + assert delivery_database.execute( + 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"' + ).fetchone() == (first["accounting_payload"],) + assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( + [(first["spend"],)] if mode == "local" else [] + ) + + +@pytest.mark.parametrize(("legacy_spend", "accepted"), [(0.5, 0), (0.0, 1)]) +@pytest.mark.asyncio +async def test_batch_charged_before_the_ledger_is_not_charged_again( + delivery_database, delivery_client, legacy_spend, accepted +): + from types import SimpleNamespace + + from litellm.repositories.table_repositories import SpendLogsRepository + + batch: Final = {**_row("batch_legacy"), "call_type": "aretrieve_batch", "status": "success"} + _insert(delivery_database, ({**batch, "spend": legacy_spend},), ledger=False) + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + assert await repository.insert_with_deliveries((_with_payload({**batch, "spend": 1.0}),)) == accepted + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (accepted,) + assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [(legacy_spend,)] + + +@pytest.mark.parametrize("table", ["LiteLLM_SpendLogs", "LiteLLM_SpendLogDelivery"]) +@pytest.mark.asyncio +async def test_transaction_a_rolls_back_payload_and_log_on_failure(delivery_database, delivery_client, table): + from types import SimpleNamespace + + from prisma.errors import PrismaError + from psycopg import sql + + from litellm.repositories.table_repositories import SpendLogsRepository + + delivery_database.execute( + sql.SQL("ALTER TABLE {} ADD CONSTRAINT reject_poison CHECK (request_id <> 'poison')").format( + sql.Identifier(table) + ) + ) + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + with pytest.raises(PrismaError): + await repository.insert_with_deliveries(tuple(_with_payload(_row(name)) for name in ("good", "poison"))) + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,) + + +@pytest.mark.asyncio +async def test_worker_rolls_back_destinations_and_parks_only_isolated_failure( + delivery_database: psycopg.Connection[tuple[object, ...]], +) -> None: + url: Final = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}")) + client: Final = Prisma(datasource={"url": url}) + delivery_database.execute("CREATE TABLE totals (spend float8 NOT NULL)") + delivery_database.execute("INSERT INTO totals VALUES (0)") + _insert(delivery_database, (_row("bad"), _row("good"))) + + async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw( + "UPDATE totals SET spend = spend + $1::float8", sum(float(str(row.row["spend"])) for row in rows) + ) + if any(row.row["request_id"] == "bad" for row in rows): + raise ValueError("Invalid projected destination") + + await client.connect() + try: + repository: Final = AccountingDeliveryRepository(client) + failed: Final = await repository.apply_pending(apply, max_attempts=2) + assert failed.failed == 2 + assert failed.parked == () + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) + assert delivery_database.execute( + 'SELECT DISTINCT state::text, attempts FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("pending", 1)] + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP') + isolated: Final = await repository.apply_pending(apply, max_attempts=2) + assert isolated.parked == ("bad",) + applied: Final = await repository.apply_pending(apply, max_attempts=2) + assert applied.applied == 1 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + assert delivery_database.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("bad", "parked"), ("good", "applied")] + assert (await repository.apply_pending(apply)).applied == 0 + finally: + await client.disconnect() + + +@pytest.mark.asyncio +async def test_a_locked_database_never_parks_a_valid_delivery(delivery_database, delivery_client): + _insert(delivery_database, (_row("valid"),)) + + async def apply(transaction, deliveries): + await transaction.execute_raw("SET LOCAL lock_timeout = '50ms'") + await transaction.execute_raw( + "UPDATE totals SET spend = spend + $1::float8", sum(float(row.row["spend"]) for row in deliveries) + ) + + repository = AccountingDeliveryRepository(delivery_client) + with psycopg.connect( + urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}")) + ) as locker: + locker.execute("LOCK TABLE totals IN ACCESS EXCLUSIVE MODE") + for _ in range(4): + result = await repository.apply_pending(apply, max_attempts=2) + assert result.failed == 1 and result.parked == () + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP') + assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("pending", 4) + ] + assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_an_unreachable_database_reported_as_a_data_error_never_parks_a_delivery( + delivery_database, delivery_client +): + from prisma.errors import DataError + + _insert(delivery_database, (_row("valid"),)) + + async def unreachable(transaction, deliveries): + raise DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}) + + repository = AccountingDeliveryRepository(delivery_client) + for _ in range(3): + assert (await repository.apply_pending(unreachable, max_attempts=2)).parked == () + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP') + assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("pending", 3) + ] + + +@pytest_asyncio.fixture(loop_scope="function") +async def delivery_client(delivery_database: psycopg.Connection[tuple[object, ...]]) -> AsyncIterator[Prisma]: + url: Final = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}")) + client: Final = Prisma(datasource={"url": url}) + delivery_database.execute("CREATE TABLE totals (spend float8 NOT NULL)") + delivery_database.execute("INSERT INTO totals VALUES (0)") + await client.connect() + try: + yield client + finally: + await client.disconnect() + + +@pytest.mark.asyncio +async def test_two_workers_claim_disjoint_rows_and_apply_once( + delivery_database: psycopg.Connection[tuple[object, ...]], + delivery_client: Prisma, +) -> None: + _insert(delivery_database, (_row("first"), _row("second"))) + entered: Final[asyncio.Queue[str]] = asyncio.Queue() + release: Final = asyncio.Event() + + async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + assert len(rows) == 1 + await entered.put(str(rows[0].row["request_id"])) + await release.wait() + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + + repository: Final = AccountingDeliveryRepository(delivery_client) + first: Final = asyncio.create_task(repository.apply_pending(apply, batch_size=1)) + first_id: Final = await asyncio.wait_for(entered.get(), timeout=10) + second: Final = asyncio.create_task(repository.apply_pending(apply, batch_size=1)) + try: + second_id: Final = await asyncio.wait_for(entered.get(), timeout=10) + assert first_id != second_id + finally: + release.set() + results: Final = await asyncio.gather(first, second) + assert sum(result.applied for result in results) == 2 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (1.0,) + assert delivery_database.execute( + "SELECT COUNT(*) FROM \"LiteLLM_SpendLogDelivery\" WHERE state = 'applied'" + ).fetchone() == (2,) + + +@pytest.mark.asyncio +async def test_cancelled_worker_rolls_back_and_restart_applies_once( + delivery_database: psycopg.Connection[tuple[object, ...]], + delivery_client: Prisma, +) -> None: + _insert(delivery_database, (_row("restart"),)) + wrote: Final = asyncio.Event() + hold: Final = asyncio.Event() + + async def interrupted(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + wrote.set() + await hold.wait() + + repository: Final = AccountingDeliveryRepository(delivery_client) + task: Final = asyncio.create_task(repository.apply_pending(interrupted)) + await asyncio.wait_for(wrote.wait(), timeout=10) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) + assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + "pending", + 0, + ) + + async def complete(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + + assert (await repository.apply_pending(complete)).applied == 1 + assert (await repository.apply_pending(complete)).applied == 0 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_worker_applies_accounting_after_spend_logs_are_deleted( + delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma +) -> None: + _insert(delivery_database, (_row("removed"),)) + delivery_database.execute('DELETE FROM "LiteLLM_SpendLogs"') + + async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + + repository: Final = AccountingDeliveryRepository(delivery_client) + assert (await repository.apply_pending(apply)).applied == 1 + assert (await repository.apply_pending(apply)).applied == 0 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_receipt_cleanup_keeps_pending_and_parked_rows_and_reports_lag( + delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma +) -> None: + _insert(delivery_database, tuple(_row(state) for state in ("pending", "parked", "applied"))) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" ' + 'SET state = request_id::"AccountingDeliveryState", ' + "created_at = CURRENT_TIMESTAMP - INTERVAL '40 days', " + "applied_at = CURRENT_TIMESTAMP - INTERVAL '40 days'" + ) + repository = AccountingDeliveryRepository(delivery_client) + assert await repository.cleanup_applied() == 1 + metrics = await repository.metrics() + assert metrics.pending_count == 1 + assert metrics.parked_count == 1 + assert metrics.oldest_pending_seconds >= 40 * 86400 + assert delivery_database.execute( + 'SELECT request_id FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("parked",), ("pending",)] + + +@pytest.mark.parametrize("delivery_database", ["public", "ledger_tenant"], indirect=True) +@pytest.mark.parametrize("state", [None, "pending", "parked", "applied"]) +def test_delivery_migration_replay_preserves_ledger_receipts(delivery_database, state): + if state is not None: + _insert(delivery_database, (_row("replay"),)) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET state=%s, attempts=3, last_error=%s', + (state, "retained failure"), + ) + before = delivery_database.execute('SELECT * FROM "LiteLLM_SpendLogDelivery"').fetchall() + migration = (_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text() + delivery_database.execute(migration) + delivery_database.execute(migration) + assert delivery_database.execute('SELECT * FROM "LiteLLM_SpendLogDelivery"').fetchall() == before + assert delivery_database.execute('SELECT count(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(state is not None),) + + +@pytest.mark.parametrize("state", ["pending", "parked"]) +def test_delivery_migration_preserves_unused_old_receipts(state): + with scratch_database("spendlogs_old_delivery") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text()) + connection.execute('CREATE TABLE "LiteLLM_SpendLogs" (request_id text PRIMARY KEY)') + connection.execute( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) VALUES (%s, %s, 1, %s)', + ("old-event", "LiteLLM_UserTable", state), + ) + connection.execute((_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text()) + assert connection.execute('SELECT event_id, state::text FROM "LiteLLM_AccountingDelivery"').fetchall() == [ + ("old-event", state) + ] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + + +@pytest.mark.asyncio +async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client): + _insert(delivery_database, (_row("bad"), _row("good"))) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$' + ) + + async def apply(transaction, deliveries): + await transaction.execute_raw( + "UPDATE totals SET spend = spend + $1::float8", sum(float(row.row["spend"]) for row in deliveries) + ) + + repository = AccountingDeliveryRepository(delivery_client) + assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') + assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",) + assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_retention_keeps_batch_receipts_indefinitely(delivery_database, delivery_client): + _insert(delivery_database, (_row("normal"), {**_row("batch"), "call_type": "aretrieve_batch"})) + delivery_database.execute( + "UPDATE \"LiteLLM_SpendLogDelivery\" SET state='applied', applied_at=CURRENT_TIMESTAMP - INTERVAL '400 days'" + ) + repository = AccountingDeliveryRepository(delivery_client) + assert await repository.cleanup_applied() == 1 + assert await repository.cleanup_applied() == 0 + assert delivery_database.execute('SELECT request_id FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("batch",)] + metrics = await repository.metrics() + assert metrics.batch_count == 1 + assert metrics.table_bytes > 0 + + +@pytest.mark.parametrize("disabled", [False, True]) +@pytest.mark.asyncio +async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode( + delivery_database, delivery_client, monkeypatch, disabled +): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.utils import jsonify_object + + completed = { + **_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "metadata": '{"status":"success","private":"request detail"}', + "accounting_payload": { + "request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "metadata": {"accounting_savings": {"compression": 0.2}}, + }, + } + monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}}) + client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object) + writer = DBSpendUpdateWriter() + first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled) + again = await writer._claim_settle_once_cost_row( + payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled + ) + assert (first, again) == (True, False) + assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == ( + [] if disabled else [(0.5, {"status": "success"})] + ) + assert delivery_database.execute( + 'SELECT delivery_id, accounting_payload FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("batch_completed", completed["accounting_payload"])] + + +@pytest.mark.parametrize("setting", ["retention_days", "batch_size"]) +@pytest.mark.asyncio +async def test_invalid_cleanup_settings_preserve_receipts(delivery_database, delivery_client, setting): + _insert(delivery_database, (_row("protected"),)) + with pytest.raises(ValueError, match="positive"): + await AccountingDeliveryRepository(delivery_client).cleanup_applied(**{setting: 0}) + assert delivery_database.execute('SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("protected", "pending") + ] + + +@pytest.mark.parametrize("setting", ["max_attempts", "batch_size"]) +@pytest.mark.asyncio +async def test_invalid_worker_settings_preserve_pending_work(delivery_database, delivery_client, setting): + from unittest.mock import AsyncMock + + _insert(delivery_database, (_row("protected"),)) + apply = AsyncMock() + with pytest.raises(ValueError, match="positive"): + await AccountingDeliveryRepository(delivery_client).apply_pending(apply, **{setting: 0}) + apply.assert_not_awaited() + assert delivery_database.execute( + 'SELECT request_id, state::text, attempts FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("protected", "pending", 0)] diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py new file mode 100644 index 00000000000..e1894aa66bf --- /dev/null +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -0,0 +1,1437 @@ +import contextlib +import asyncio +import json +import os +import uuid +import subprocess +import sys +import time +from collections.abc import Iterator +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final + +import psycopg +from psycopg import sql +import pytest +import litellm + +from litellm.accounting.application import apply_deliveries +from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy.spend_tracking.spend_tracking_utils import spend_log_row_with_retained_metadata +from litellm.proxy import proxy_server +from litellm.proxy._types import SpendLogsMetadataFields +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.model_usage_rollup import flush_model_usage_transactions +from litellm.proxy.guardrails.usage_tracking import PendingRollups, process_spend_logs_guardrail_usage +from litellm.proxy.utils import PrismaClient +from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows +from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter +from litellm.repositories.accounting_delivery_repository import AccountingDeliveryRepository +from litellm.repositories.table_repositories import SpendLogsRepository +from litellm.responses.streaming_iterator import drain_logging_tasks +from litellm.types.utils import ModelResponse, Usage +from tests.integration._support.database import scratch_database + +_ROOT: Final = Path(__file__).resolve().parents[3] +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) +_ENTITY_TABLES: Final = ( + "LiteLLM_VerificationToken", + "LiteLLM_UserTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembership", + "LiteLLM_EndUserTable", + "LiteLLM_AgentsTable", + "LiteLLM_TagTable", + "LiteLLM_ProjectTable", + "LiteLLM_ModelAccessGroupBudgetTable", +) +_DAILY_TABLES: Final = ( + "LiteLLM_DailyUserSpend", + "LiteLLM_DailyTeamSpend", + "LiteLLM_DailyOrganizationSpend", + "LiteLLM_DailyEndUserSpend", + "LiteLLM_DailyAgentSpend", + "LiteLLM_DailyTagSpend", + "LiteLLM_DailyModelUsage", + "LiteLLM_DailyGuardrailUsageUnits", + "LiteLLM_BudgetWindowSpend", +) + + +@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)]) +@pytest.mark.asyncio +async def test_worker_applies_the_configured_backlog_in_one_transaction( + parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + if batch_size is None: + monkeypatch.delenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", raising=False) + else: + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", batch_size) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + writer = DBSpendUpdateWriter() + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + stored = await client.writer_db.query_raw('SELECT to_jsonb(s) AS row FROM "LiteLLM_SpendLogs" s') + delivery = (await client.writer_db.query_raw('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"'))[0]["accounting_payload"] + assert await SpendLogsRepository(client).insert_with_deliveries(tuple( + {**stored[0]["row"], "request_id": f"backlog-{index:05}", "model": f"backlog-model-{index}" if distinct_models else stored[0]["row"]["model"], "accounting_payload": {**delivery, "request_id": f"backlog-{index:05}", "model": f"backlog-model-{index}" if distinct_models else delivery["model"]}} for index in range(row_count - 1) + )) == row_count - 1 + started = time.monotonic() + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}") + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,) + assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,) + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("ledger", [False, True]) +@pytest.mark.parametrize("exclude_messages", [False, True]) +@pytest.mark.asyncio +async def test_registered_proxy_callbacks_charge_model_windows_before_persisting_spend( + parity_database: str, monkeypatch: pytest.MonkeyPatch, ledger: bool, exclude_messages: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.utils import Rules, function_setup + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false") + for name in ("callbacks", "input_callback", "success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["messages"] if exclude_messages else []) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server.proxy_logging_obj, "db_spend_update_writer", writer) + limiter = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(), clock=lambda: _NOW.timestamp(), + shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: client), + ) + monkeypatch.setattr(proxy_server, "model_max_budget_limiter", limiter) + litellm.logging_callback_manager.add_litellm_callback(limiter) + proxy_server.cost_tracking() + logging, _ = function_setup("acompletion", Rules(), _NOW, model="ledger-test", messages=[], litellm_call_id="registered-window") + logging.update_environment_variables( + optional_params={}, custom_llm_provider="openai", + litellm_params={"metadata": {"user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", "user_api_key_model_max_budget": { + "ledger-test": {"max_budget": 100, "budget_duration": "1d"} + }}}, + ) + logging.model_call_details["response_cost"] = 0.25 + response = ModelResponse(id="registered-window", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5)) + await client.db.connect() + try: + await logging.async_success_handler(response, _NOW, _NOW, standard_logging_object={ + "model": "ledger-test", "model_group": "ledger-test", "response_cost": 0.25, + "metadata": {"user_api_key_hash": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}, + "model_map_information": {}, "hidden_params": {}, "request_tags": [], + }) + await drain_logging_tasks() + await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions)) + if ledger: + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type=\'model_max_budget\'').fetchall() == [(0.25,)] + assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [(0.25,)] + finally: + await client.db.disconnect() + + +@pytest.fixture(scope="module") +def ledger_schema() -> str: + return subprocess.run( + [ + sys.executable, + "-I", + "-m", + "prisma", + "migrate", + "diff", + "--from-empty", + "--to-schema-datamodel", + str(_ROOT / "schema.prisma"), + "--script", + ], + check=True, + capture_output=True, + text=True, + cwd=_ROOT, + ).stdout + + +@pytest.fixture +def parity_database(ledger_schema: str) -> Iterator[str]: + with scratch_database("spendlogs_parity") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute(ledger_schema) + connection.execute( + "INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, created_by, updated_by) VALUES ('budget', 'test', 'test')" + ) + connection.execute("INSERT INTO \"LiteLLM_UserTable\" (user_id) VALUES ('user')") + connection.execute("INSERT INTO \"LiteLLM_TeamTable\" (team_id) VALUES ('team')") + connection.execute( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by) " + "VALUES ('org', 'org', 'budget', 'test', 'test')" + ) + connection.execute("INSERT INTO \"LiteLLM_TeamMembership\" (team_id, user_id) VALUES ('team', 'user')") + connection.execute( + "INSERT INTO \"LiteLLM_OrganizationMembership\" (organization_id, user_id) VALUES ('org', 'user')" + ) + connection.execute("INSERT INTO \"LiteLLM_VerificationToken\" (token) VALUES ('2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683')") + connection.execute("INSERT INTO \"LiteLLM_EndUserTable\" (user_id) VALUES ('end-user')") + connection.execute( + 'INSERT INTO "LiteLLM_AgentsTable" ' + "(agent_id, agent_name, agent_card_params, created_by, updated_by) " + "VALUES ('agent', 'agent', '{}', 'test', 'test')" + ) + connection.execute( + 'INSERT INTO "LiteLLM_ProjectTable" (project_id, created_by, updated_by) ' + "VALUES ('project', 'test', 'test')" + ) + connection.execute("INSERT INTO \"LiteLLM_TagTable\" (tag_name) VALUES ('tag')") + connection.execute("INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name) VALUES ('group')") + yield url + + +def _snapshot(url: str) -> dict[str, list[object]]: + with psycopg.connect(url) as connection: + return { + table: [ + row[0] + for row in connection.execute( + sql.SQL( + "SELECT to_jsonb(t) - 'id' - 'created_at' - 'updated_at' - 'request_id' - 'last_active' FROM {} t " + "ORDER BY (to_jsonb(t) - 'id' - 'created_at' - 'updated_at' - 'request_id' - 'last_active')::text" + ).format(sql.Identifier(table)) + ).fetchall() + ] + for table in (*_ENTITY_TABLES, *_DAILY_TABLES) + } + + +def _reset(url: str) -> None: + with psycopg.connect(url, autocommit=True) as connection: + for table in _ENTITY_TABLES: + connection.execute(sql.SQL("UPDATE {} SET spend=0").format(sql.Identifier(table))) + for table in (*_DAILY_TABLES, "LiteLLM_SpendLogs", "LiteLLM_SpendLogDelivery"): + connection.execute(sql.SQL("TRUNCATE {}").format(sql.Identifier(table))) + connection.execute('UPDATE "LiteLLM_VerificationToken" SET total_spend=0') + connection.execute('UPDATE "LiteLLM_TeamMembership" SET total_spend=0') + + +async def _assert_legacy_database_parity( + parity_database: str, + monkeypatch: pytest.MonkeyPatch, + status: str, + origin: str | None, + cost: float, + retention: str, + reset_kind: str | None = None, + correction_order: str | None = None, + toggle_after_capture: bool = False, + fallback_narrowing: bool = False, + failing_table: str | None = None, + organization_id: str | None = "org", + export_url: str | None = None, + export_failure: bool = False, +): + monkeypatch.setitem( + litellm.model_cost, + "ledger-test", + { + "input_cost_per_token": 0.002, + "output_cost_per_token": 0.004, + "cache_read_input_token_cost": 0.0005, + "cache_creation_input_token_cost": 0.003, + "litellm_provider": "openai", + "mode": "chat", + "max_tokens": 4096, + }, + ) + if fallback_narrowing: + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[{ + "model_name": "ledger-test", + "litellm_params": {"model": "openai/ledger-test", "api_key": "test"}, + "model_info": {"id": "deployment", "access_groups": ["group"]}, + }])) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + "INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name) VALUES ('unserved-group')" + ) + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + await client.db.connect() + results = [] + try: + for path in ("legacy", "ledger"): + _reset(parity_database) + if reset_kind is not None: + _prime_reset(parity_database, reset_kind) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": path == "ledger"}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", path == "ledger" and retention == "lean") + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server.proxy_logging_obj, "db_spend_update_writer", writer) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(proxy_server, "spend_counter_cache", DualCache()) + for cache_key in ("2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", "team_id:team"): + await proxy_server.user_api_key_cache.async_set_cache( + key=cache_key, + value={ + "budget_limits": [ + {"budget_duration": "1d", "max_budget": 100, "reset_at": "2026-11-04T00:00:00+00:00"} + ] + }, + ) + request_kwargs = { + "model": "ledger-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "request", + "standard_logging_object": { + "model": "ledger-test", + "model_group": "ledger-test", + "metadata": { + "user_api_key_hash": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + "usage_object": { + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_read_input_tokens": 4, + "cache_creation_input_tokens": 1, + }, + }, + "model_map_information": {}, + "hidden_params": {}, + "request_tags": ["tag"], + "model_id": "deployment", + "response_cost": cost, + "autorouter_savings": 0.125, + "guardrail_information": [ + { + "guardrail_id": "guard", + "guardrail_name": "guard", + "guardrail_status": "success", + "guardrail_usage": {"text": 2}, + "guardrail_cost_by_unit": {"text": 0.02}, + "guardrail_cost_in_spend": True, + } + ], + }, + "agent_id": "agent", + "user": "end-user", + "litellm_params": { + "metadata": { + "user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + "user_api_key_model_max_budget": {"ledger-test": {"max_budget": 100, "budget_duration": "1d"}}, + "user_api_key_user_id": "user", + "user_api_key_team_id": "team", + "user_api_key_org_id": organization_id, + "user_api_key_project_id": "project", + "user_api_key_end_user_id": "end-user", + "tags": ["tag"], + "user_api_key_matched_model_access_groups": ["group", "unserved-group"] if fallback_narrowing else ["group"], + "model_info": {"id": "deployment"}, + "litellm_gateway_injected_cache": "deployment", + "compression_savings": {"tokens_before": 20, "tokens_after": 10, "tokens_saved": 10}, + "status": status, + **({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}), + **({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin == "auto_router" else {}), + }, + "user": "end-user", + }, + } + limiter = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=proxy_server.spend_counter_cache, + clock=lambda: _NOW.timestamp(), + shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: client), + ) + if status == "success": + await limiter.async_log_success_event(request_kwargs, None, _NOW, _NOW) + charged = await writer.update_database( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + user_id="user", + end_user_id="end-user", + team_id="team", + org_id=organization_id, + project_id="project", + kwargs=request_kwargs, + completion_response=ModelResponse( + id="request", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + response_cost=cost, + ) + assert charged + await proxy_server.increment_spend_counters( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + team_id="team", + user_id="user", + response_cost=cost, + request_started_at=_NOW, + **({"windows_from_ledger": True} if path == "ledger" else {}), + ) + await drain_logging_tasks() + rows = tuple(client.spend_log_transactions) + client.spend_log_transactions = [] + assert len(rows) == 1 + if path == "ledger" and retention == "lean": + assert rows[0]["accounting_log_disabled"] is True + frozen = rows[0]["accounting_payload"] + if isinstance(frozen, str): + frozen = json.loads(frozen) + assert not {"messages", "response", "proxy_server_request"}.intersection(frozen) + stored_rows = ( + tuple(spend_log_row_with_retained_metadata(row, SpendLogsMetadataFields(include=())) for row in rows) + if retention == "stripped" + else rows + ) + if toggle_after_capture: + for enabled in (True, False, True): + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": enabled}) + await SpendLogsRepository(client).insert_with_deliveries(stored_rows) + elif path == "ledger" and export_url is not None: + import httpx + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy.utils import ProxyUpdateSpend + + monkeypatch.setenv("SPEND_LOGS_URL", export_url) + handler = AsyncHTTPHandler(timeout=10) + wire_rows = [ + {key: value.isoformat() if isinstance(value, datetime) else value for key, value in row.items()} + for row in stored_rows + ] + try: + if export_failure: + with pytest.raises(httpx.HTTPStatusError): + await ProxyUpdateSpend.update_spend_logs(0, client, handler, proxy_server.proxy_logging_obj, wire_rows) + else: + await ProxyUpdateSpend.update_spend_logs(0, client, handler, proxy_server.proxy_logging_obj, wire_rows) + finally: + await handler.close() + else: + await SpendLogsRepository(client).insert_with_deliveries(stored_rows) + if path == "ledger" and (retention == "lean" or export_url is not None): + assert await client.writer_db.query_raw('SELECT COUNT(*) AS count FROM "LiteLLM_SpendLogs"') == [{"count": 0}] + if reset_kind is not None: + await _apply_reset(reset_kind, client) + if correction_order == "before": + await _publish_savings_correction(client, stored_rows[0]) + if path != "ledger": + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + await writer.commit_daily_tag_spend_to_db(client, 0, proxy_server.proxy_logging_obj) + await flush_model_usage_transactions(prisma_client=client, transactions=client.model_usage_transactions) + client.model_usage_transactions = [] + await process_spend_logs_guardrail_usage(client, rows, pending=PendingRollups()) + else: + + async def apply(transaction, deliveries, writer: DBSpendUpdateWriter = writer): + await apply_deliveries( + transaction, + deliveries, + prisma_client=client, + writer=writer, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + + if failing_table is not None: + before_failure = _snapshot(parity_database) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute(sql.SQL('ALTER TABLE {} ADD CONSTRAINT reject_spend CHECK (spend = 0)').format(sql.Identifier(failing_table))) + failed = await AccountingDeliveryRepository(client.writer_db).apply_pending(apply) + assert (failed.applied, failed.failed) == (0, 1) + assert _snapshot(parity_database) == before_failure + with psycopg.connect(parity_database, autocommit=True) as connection: + assert connection.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("pending", 1)] + connection.execute(sql.SQL('ALTER TABLE {} DROP CONSTRAINT reject_spend').format(sql.Identifier(failing_table))) + connection.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') + result = await AccountingDeliveryRepository(client.writer_db).apply_pending(apply) + assert result.applied == 1, result + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + if correction_order == "after": + await _publish_savings_correction(client, stored_rows[0]) + results.append(_snapshot(parity_database)) + differences = { + (index, table): (results[0][table], result[table]) + for index, result in enumerate(results[1:], 1) + for table in results[0] + if results[0][table] != result[table] + } + assert not differences, differences + prior_spend = 5 if reset_kind is not None and reset_kind != "scheduled" else 0 + assert results[-1]["LiteLLM_UserTable"][0]["spend"] == cost + prior_spend + daily = results[-1]["LiteLLM_DailyUserSpend"][0] + assert daily["compression_savings_spend"] > 0, results[-1]["LiteLLM_DailyUserSpend"] + assert daily["prompt_caching_savings_spend"] != 0 + assert daily["gateway_injected_caching_savings_spend"] != 0 + assert daily["autorouter_savings_spend"] == pytest.approx( + (0 if origin == "auto_router" else 0.125) + (0.2 if correction_order else 0) + ) + group_spend = {row["access_group_name"]: row["spend"] for row in results[-1]["LiteLLM_ModelAccessGroupBudgetTable"]} + assert group_spend == {"group": cost + prior_spend, **({"unserved-group": prior_spend} if fallback_narrowing else {})} + assert results[-1]["LiteLLM_DailyGuardrailUsageUnits"] + windows = results[-1]["LiteLLM_BudgetWindowSpend"] + assert len(windows) == (2 if cost else 0) + (1 if status == "success" else 0) + assert all(row["spend"] == cost for row in windows) + if reset_kind == "window_rollover": + model_window = next(row for row in windows if row["entity_type"] == "model_max_budget") + assert model_window["window_start"].startswith("2026-11-04") + if reset_kind in ("key", "global", "scheduled"): + assert results[-1]["LiteLLM_VerificationToken"][0]["spend"] == cost + if reset_kind in ("global", "scheduled"): + assert results[-1]["LiteLLM_TeamTable"][0]["spend"] == cost + if reset_kind == "team_member": + assert results[-1]["LiteLLM_TeamMembership"][0]["spend"] == cost + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize( + "status,origin,cost", + [ + ("success", None, 0.25), + ("failure", None, 0.25), + ("success", None, 0.0), + ("success", "auto_router", 0.25), + ("failure", "billed_retry", 0.25), + ], +) +@pytest.mark.parametrize("retention", ["full", "stripped", "lean"]) +@pytest.mark.asyncio +async def test_stored_ledger_matches_real_legacy_database_updates( + parity_database: str, monkeypatch: pytest.MonkeyPatch, status: str, origin: str | None, cost: float, retention: str +): + await _assert_legacy_database_parity(parity_database, monkeypatch, status, origin, cost, retention) + + +@pytest.mark.parametrize("retention", ["full", "stripped", "lean"]) +@pytest.mark.asyncio +async def test_missing_organization_matches_legacy_daily_rollups(parity_database, monkeypatch, retention): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, retention, organization_id=None + ) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT count(*) FROM "LiteLLM_DailyOrganizationSpend"').fetchone() == (0,) + + +@pytest.mark.parametrize("reset_kind", ["scheduled", "key", "team_member", "global", "window_rollover"]) +@pytest.mark.asyncio +async def test_pending_delivery_after_reset_matches_queued_legacy_spend( + parity_database: str, monkeypatch: pytest.MonkeyPatch, reset_kind: str +): + await _assert_legacy_database_parity(parity_database, monkeypatch, "success", None, 0.25, "full", reset_kind) + + +def _prime_reset(url: str, reset_kind: str) -> None: + with psycopg.connect(url, autocommit=True) as connection: + for table in _ENTITY_TABLES: + connection.execute(sql.SQL("UPDATE {} SET spend=5").format(sql.Identifier(table))) + if reset_kind == "scheduled": + for table in ("LiteLLM_VerificationToken", "LiteLLM_UserTable", "LiteLLM_TeamTable", "LiteLLM_BudgetTable"): + connection.execute( + sql.SQL("UPDATE {} SET budget_duration='1d', budget_reset_at='2026-11-01'").format( + sql.Identifier(table) + ) + ) + connection.execute("UPDATE \"LiteLLM_ModelAccessGroupBudgetTable\" SET budget_id='budget'") + + +async def _apply_reset(reset_kind: str, client: PrismaClient) -> None: + from litellm.proxy._types import ResetSpendRequest, UserAPIKeyAuth, LitellmUserRoles + from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob + from litellm.proxy.management_endpoints.key_management_endpoints import reset_key_spend_fn + from litellm.proxy.management_endpoints.team_endpoints import reset_team_member_spend_fn + from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + if reset_kind == "window_rollover": + from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row + + windows = await client.writer_db.query_raw('SELECT * FROM "LiteLLM_BudgetWindowSpend"') + assert windows + for window in windows: + await roll_window_spend_row( + client, window["entity_type"], window["entity_id"], window["window_duration"], _NOW + timedelta(days=2) + ) + elif reset_kind == "global": + await global_spend_reset() + elif reset_kind == "key": + await reset_key_spend_fn( + key="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", data=ResetSpendRequest(reset_to=0), user_api_key_dict=admin, litellm_changed_by=None + ) + elif reset_kind == "team_member": + await reset_team_member_spend_fn( + team_id="team", user_id="user", data=ResetSpendRequest(reset_to=0), user_api_key_dict=admin + ) + else: + await ResetBudgetJob(proxy_server.proxy_logging_obj, client).reset_budget() + + +@pytest.mark.parametrize("correction_order", ["before", "after"]) +@pytest.mark.asyncio +async def test_baseline_correction_order_preserves_frozen_request_savings( + parity_database: str, monkeypatch: pytest.MonkeyPatch, correction_order: str +): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "full", correction_order=correction_order + ) + + +async def _publish_savings_correction(client: PrismaClient, row) -> None: + from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING + from litellm.proxy.db.baseline_accounting import ( + BaselinePublication, + DailyBaselineAttribution, + DailyBaselineTarget, + _Change, + _publish, + ) + + metadata_before = ( + await client.writer_db.query_raw( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', row["request_id"] + ) + )[0]["metadata"] + if isinstance(metadata_before, str): + metadata_before = json.loads(metadata_before) + change = _Change( + request_id=row["request_id"], + publication=BaselinePublication( + comparison_id="autorouter-baseline:v3:" + "a" * 64, + comparison_started_at=_NOW.timestamp(), + status="estimated", + reason="modeled", + provenance="modeled", + actual_spend=0.25, + baseline_spend=0.575, + ), + api_key=row["api_key"], + session_id="session", + router_name="router", + baseline_model="ledger-test", + covered_delta=0, + actual_delta=0.0, + savings_delta=0.2, + daily=DailyBaselineAttribution( + date=_NOW.date().isoformat(), + api_key=row["api_key"], + model=row["model"], + custom_llm_provider=row["custom_llm_provider"], + model_group=row["model_group"], + endpoint=ROUTE_ENDPOINT_MAPPING.get(row["call_type"]), + targets=tuple( + DailyBaselineTarget(entity=entity, entity_id=identity) + for entity, identity in ( + ("user", "user"), + ("team", "team"), + ("org", "org"), + ("end_user", "end-user"), + ("agent", "agent"), + ("tag", "tag"), + ) + ), + ), + ) + async with client.writer_db.tx() as transaction: + await _publish(transaction, (change,)) + metadata_after = ( + await client.writer_db.query_raw( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', row["request_id"] + ) + )[0]["metadata"] + if isinstance(metadata_after, str): + metadata_after = json.loads(metadata_after) + assert metadata_after["autorouter_savings"] == pytest.approx(0.325) + if "accounting_payload" in row: + delivery = (await client.writer_db.query_raw( + 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery" WHERE request_id=$1', row["request_id"] + ))[0]["accounting_payload"] + assert delivery["metadata"]["accounting_savings"]["autorouter"] == pytest.approx(0.125) + + +@pytest.mark.asyncio +async def test_flag_changes_do_not_reclassify_or_reapply_buffered_rows( + parity_database: str, + monkeypatch: pytest.MonkeyPatch, +): + await _assert_legacy_database_parity( + parity_database, + monkeypatch, + "success", + None, + 0.25, + "stripped", + toggle_after_capture=True, + ) + + +@pytest.mark.parametrize("entity_type", ["key", "team"]) +@pytest.mark.asyncio +async def test_window_seed_excludes_pending_and_parked_but_preserves_legacy_cutoff( + parity_database: str, + monkeypatch: pytest.MonkeyPatch, + entity_type: str, +): + from datetime import timedelta + from litellm.proxy.db.budget_window_spend_writer import spend_logs_seed_totals, commit_window_spend_updates + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + await client.db.connect() + try: + for request_id, cost, offset in ( + ("old-legacy", 2.0, -1), + ("current-legacy", 5.0, 0), + ("applied", 3.0, 1), + ("pending", 7.0, 0), + ("parked", 11.0, 0), + ): + await client.writer_db.execute_raw( + 'INSERT INTO "LiteLLM_SpendLogs" ' + '(request_id,call_type,api_key,team_id,spend,"startTime","endTime") ' + "VALUES ($1,'acompletion','key','team',$2::float8,$3::timestamp,$3::timestamp)", + request_id, + cost, + _NOW + timedelta(seconds=offset), + ) + for request_id, state in (("applied", "applied"), ("pending", "pending"), ("parked", "parked")): + await client.writer_db.execute_raw( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id,state,accounting_payload,is_batch) VALUES ($1,$2::"AccountingDeliveryState",$$ {} $$::jsonb,false)', + request_id, + state, + ) + start = _NOW - timedelta(hours=1) + totals = await spend_logs_seed_totals(client, entity_type, entity_type, start, _NOW) + assert totals is not None and totals.total == 10.0 and totals.before_batch == 5.0 + unbounded = await spend_logs_seed_totals(client, entity_type, entity_type, start, None) + assert unbounded is not None and unbounded.total == unbounded.before_batch == 10.0 + await commit_window_spend_updates( + client, + ( + { + "entity_type": entity_type, + "entity_id": entity_type, + "window_duration": "1d", + "window_start": start.isoformat(), + "started_at": _NOW.isoformat(), + "spend": 5.0, + }, + ), + ) + rows = await client.writer_db.query_raw('SELECT spend FROM "LiteLLM_BudgetWindowSpend"') + assert rows[0]["spend"] == 10.0 + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("disable_logs", [False, True]) +@pytest.mark.asyncio +async def test_batch_claim_failure_queues_an_atomic_retry_and_repeated_retrieve_does_not_charge( + parity_database: str, monkeypatch: pytest.MonkeyPatch, disable_logs: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", disable_logs) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + + + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'ALTER TABLE "LiteLLM_SpendLogDelivery" ADD CONSTRAINT reject_batch ' + "CHECK (request_id <> 'batch_batch_cost')" + ) + assert await _retrieve_batch(writer, 0.25) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + assert len(client.spend_log_transactions) == 1 + queued = tuple(client.spend_log_transactions) + assert "accounting_payload" in queued[0] + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery" DROP CONSTRAINT reject_batch') + assert await SpendLogsRepository(client).insert_with_deliveries(queued) == 1 + assert await _retrieve_batch(writer, 0.25) is False + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=\'2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683\'').fetchone() == (0.25,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.25,) + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(not disable_logs),) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_failed_claim_over_a_legacy_zero_batch_row_still_charges_it_once( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": False}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0) + await drain_logging_tasks() + await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions)) + client.spend_log_transactions = [] + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + with psycopg.connect(parity_database, autocommit=True) as connection: + assert connection.execute( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=\'batch_batch_cost\'' + ).fetchall() == [(0.0,)] + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery" RENAME TO "LiteLLM_SpendLogDelivery_unavailable"') + assert await _retrieve_batch(writer, 0.25) + queued = tuple(client.spend_log_transactions) + assert ["accounting_payload" in row for row in queued] == [True] + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery_unavailable" RENAME TO "LiteLLM_SpendLogDelivery"') + assert await SpendLogsRepository(client).insert_with_deliveries(queued) == 1 + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert await _retrieve_batch(writer, 0.25) is False + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=\'2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683\'').fetchone() == (0.25,) + assert connection.execute( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=\'batch_batch_cost\'' + ).fetchall() == [(0.0,)] + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_worker_exposes_invalid_accounting_as_delayed_then_parked_work( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from prometheus_client import REGISTRY + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + await client.db.connect() + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, created_at, accounting_payload, is_batch) ' + "VALUES ('invalid-accounting', CURRENT_TIMESTAMP - INTERVAL '1 hour', " + '$$ {"metadata": {"accounting_windows": [{}]}} $$::jsonb, false)' + ) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert REGISTRY.get_sample_value("litellm_accounting_pending") == 1 + assert REGISTRY.get_sample_value("litellm_accounting_parked") == 0 + with psycopg.connect(parity_database, autocommit=True) as connection: + assert connection.execute( + 'SELECT state::text, attempts, last_error FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("pending", 1, "KeyError")] + connection.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET attempts=4, next_attempt_at=CURRENT_TIMESTAMP' + ) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert REGISTRY.get_sample_value("litellm_accounting_pending") == 0 + assert REGISTRY.get_sample_value("litellm_accounting_parked") == 1 + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT state::text, attempts, last_error FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("parked", 5, "KeyError")] + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0,) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_fallback_charges_only_groups_declared_by_the_served_deployment( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "stripped", fallback_narrowing=True + ) + + +@pytest.mark.asyncio +async def test_worker_skips_another_leader_and_releases_its_redis_lock_after_retry( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + + monkeypatch.setenv("DATABASE_URL", parity_database) + cache = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace=uuid.uuid4().hex) + other_leader = PodLockManager(cache) + writer = DBSpendUpdateWriter(redis_cache=cache) + writer.pod_lock_manager.redis_cache = cache + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + await client.db.connect() + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch) VALUES ($$invalid-row$$, $$ {"metadata": {"accounting_windows": [{}]}} $$::jsonb, false)') + assert await other_leader.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(0,)] + await other_leader.release_lock(cronjob_id="accounting_delivery") + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(1,)] + assert await other_leader.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False) + finally: + await other_leader.release_lock(cronjob_id="accounting_delivery") + await client.db.disconnect() + await cache.disconnect() + + +async def _retrieve_batch(writer: DBSpendUpdateWriter, cost: float) -> bool: + return await writer.update_database( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user_id="user", end_user_id=None, team_id="team", org_id=None, + kwargs={ + "model": "ledger-test", "call_type": "aretrieve_batch", "litellm_call_id": "batch", + "litellm_params": {"metadata": {"status": "success", "user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}}, + }, + completion_response={"id": "batch", "status": "completed"}, + start_time=_NOW, end_time=_NOW, response_cost=cost, + ) + + +@pytest.mark.parametrize("ledger_after_zero_row,delivery_unavailable", [(False, False), (True, False), (True, True)]) +@pytest.mark.asyncio +async def test_zero_cost_legacy_batch_row_is_charged_once_across_a_flag_change( + parity_database: str, monkeypatch: pytest.MonkeyPatch, ledger_after_zero_row: bool, delivery_unavailable: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": False}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0) + await drain_logging_tasks() + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger_after_zero_row}) + assert await _retrieve_batch(writer, 0.25) + await drain_logging_tasks() + assert await _retrieve_batch(writer, 0.25) is False + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0,) + assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [ + (0.0 if ledger_after_zero_row else 0.25,) + ] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + int(ledger_after_zero_row), + ) + if delivery_unavailable: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery" RENAME TO temporarily_unavailable_delivery') + try: + await writer.db_update_spend_transaction_handler(client, 0, proxy_server.proxy_logging_obj) + finally: + if delivery_unavailable: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE temporarily_unavailable_delivery RENAME TO "LiteLLM_SpendLogDelivery"') + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.25,) + assert connection.execute('SELECT spend, api_requests FROM "LiteLLM_DailyUserSpend"').fetchall() == [(0.25, 2)] + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("failing_table", [ + "LiteLLM_UserTable", "LiteLLM_VerificationToken", "LiteLLM_TeamTable", "LiteLLM_TeamMembership", + "LiteLLM_OrganizationMembership", "LiteLLM_OrganizationTable", "LiteLLM_ProjectTable", "LiteLLM_AgentsTable", "LiteLLM_EndUserTable", +]) +@pytest.mark.asyncio +async def test_destination_failure_rolls_back_all_earlier_updates_before_retry( + parity_database: str, monkeypatch: pytest.MonkeyPatch, failing_table: str +): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "full", failing_table=failing_table + ) + + +@pytest.mark.asyncio +async def test_spend_logs_cleanup_does_not_remove_pending_accounting(parity_database, monkeypatch): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + cleanup = SpendLogCleanup(clock=lambda: 0) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + delete_sql = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id IN (SELECT request_id FROM "LiteLLM_SpendLogs" WHERE "startTime" < $1::timestamp LIMIT $2::integer)' + assert await cleanup._execute_delete_batch(client, delete_sql, _NOW + timedelta(days=1), "LiteLLM_SpendLogs", 60) == 1 + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert await _retrieve_batch(writer, 0.25) is False + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=$$user$$').fetchone() == (0.25,) + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("mode", ["local", "url", "disabled"]) +@pytest.mark.parametrize("receipt_state", ["pending", "parked", "applied", "retained"]) +@pytest.mark.asyncio +async def test_ledger_charged_batch_is_not_charged_again_after_flag_disable( + parity_database: str, monkeypatch: pytest.MonkeyPatch, receipt_state: str, mode: str +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", mode == "disabled") + if mode == "url": + monkeypatch.setenv("SPEND_LOGS_URL", "http://127.0.0.1:1") + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0) + assert await _retrieve_batch(writer, 0.25) + if receipt_state in ("applied", "retained"): + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database, autocommit=True) as connection: + if receipt_state == "retained": + connection.execute("UPDATE \"LiteLLM_SpendLogDelivery\" SET applied_at = CURRENT_TIMESTAMP - INTERVAL '400 days'") + if receipt_state == "parked": + connection.execute('UPDATE "LiteLLM_SpendLogDelivery" SET state=\'parked\'') + assert await AccountingDeliveryRepository(client.writer_db).cleanup_applied() == 0 + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": False}) + assert await _retrieve_batch(writer, 0.25) is False + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( + [(0.25,)] if mode != "disabled" else [] + ) + expected_spend = 0 if receipt_state == "parked" else 0.25 + expected_state = "parked" if receipt_state == "parked" else "applied" + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected_spend,) + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(expected_state,)] + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("disabled", [False, True]) +@pytest.mark.parametrize("export_status", [200, 500]) +@pytest.mark.asyncio +async def test_export_failure_does_not_change_committed_accounting(parity_database, monkeypatch, export_status, disabled): + requests = asyncio.Queue() + + async def receive(reader, writer): + header = await reader.readuntil(b"\r\n\r\n") + length = next(int(line.split(b":", 1)[1]) for line in header.split(b"\r\n") if line.lower().startswith(b"content-length:")) + body = await reader.readexactly(length) + await requests.put(json.loads(body)) + writer.write(f"HTTP/1.1 {export_status} Result\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{{}}".encode()) + await writer.drain() + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(receive, "127.0.0.1", 0) + try: + port = server.sockets[0].getsockname()[1] + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "lean" if disabled else "full", + export_url=f"http://127.0.0.1:{port}", export_failure=export_status == 500 and not disabled, + ) + assert requests.qsize() == int(not disabled) + if disabled: + return + exported = await asyncio.wait_for(requests.get(), timeout=5) + assert len(exported) == 1 + assert exported[0]["spend"] == 0.25 + assert "accounting_payload" not in exported[0] + assert "accounting_log_disabled" not in exported[0] + finally: + server.close() + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_spend_logs_url_without_an_export_client_keeps_local_logs_and_charges_once( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.utils import ProxyUpdateSpend + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("SPEND_LOGS_URL", "http://127.0.0.1:1") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + try: + assert await writer.update_database( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user_id="user", end_user_id=None, team_id="team", org_id=None, + kwargs={ + "model": "ledger-test", "call_type": "acompletion", "litellm_call_id": "url-request", + "litellm_params": {"metadata": {"status": "success", "user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}}, + }, + completion_response=ModelResponse( + id="url-request", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, end_time=_NOW, response_cost=0.25, + ) + await drain_logging_tasks() + rows = [ + {key: value.isoformat() if isinstance(value, datetime) else value for key, value in row.items()} + for row in client.spend_log_transactions + ] + client.spend_log_transactions = [] + await ProxyUpdateSpend.update_spend_logs(0, client, None, proxy_server.proxy_logging_obj, rows) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=\'url-request\'' + ).fetchall() == [(0.25,)] + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + assert connection.execute('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=\'2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683\'').fetchone() == (0.25,) + finally: + await client.db.disconnect() + + +@contextlib.asynccontextmanager +async def _ledger_client(parity_database: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + try: + yield client + finally: + await client.db.disconnect() + + +async def _record_call( + writer: DBSpendUpdateWriter, call_id: str, *, token: str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user: str = "user", tags: object = None +) -> None: + metadata = {"status": "success", "user_api_key": token, **({"tags": tags} if tags is not None else {})} + assert await writer.update_database( + token=token, user_id=user, end_user_id=None, team_id="team", org_id=None, + kwargs={ + "model": "ledger-test", "call_type": "aembedding", "litellm_call_id": call_id, + "litellm_params": {"metadata": metadata}, + }, + completion_response=None, start_time=_NOW, end_time=_NOW, response_cost=0.25, + ) + await drain_logging_tasks() + + +def _wire(rows) -> list[dict[str, object]]: + return [ + {key: value.isoformat() if isinstance(value, datetime) else value for key, value in row.items()} + for row in rows + ] + + +async def _flush(writer: DBSpendUpdateWriter, client: PrismaClient) -> None: + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.utils import ProxyUpdateSpend + + rows = _wire(client.spend_log_transactions) + client.spend_log_transactions = [] + await ProxyUpdateSpend.update_spend_logs(0, client, None, proxy_server.proxy_logging_obj, rows) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + + +def _key_spend(parity_database: str) -> object: + with psycopg.connect(parity_database) as connection: + return connection.execute( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', ("2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683",) + ).fetchone() + + +@pytest.mark.asyncio +async def test_requests_sharing_a_client_call_id_are_each_charged_but_a_replay_is_not( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "shared-call-id") + await _record_call(writer, "shared-call-id") + buffered = _wire(client.spend_log_transactions) + await _flush(writer, client) + assert await SpendLogsRepository(client).insert_with_deliveries(buffered) == 0 + await _flush(writer, client) + assert _key_spend(parity_database) == (0.5,) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("shared-call-id", "applied"), ("shared-call-id", "applied")] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (1,) + + +@pytest.mark.asyncio +async def test_an_opaque_bearer_token_is_never_stored_in_a_delivery( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "opaque-call", token="opaque-bearer-secret") + await _flush(writer, client) + with psycopg.connect(parity_database) as connection: + stored = connection.execute('SELECT accounting_payload::text, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() + assert [state for _, state in stored] == ["applied"] + assert "opaque-bearer-secret" not in stored[0][0] + + +@pytest.mark.asyncio +async def test_an_evicted_spend_log_still_charges_its_request(parity_database: str, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.utils import enqueue_spend_logs + + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "evicted-call") + await _record_call(writer, "kept-call") + queued = tuple(client.spend_log_transactions) + client.spend_log_transactions = [] + await enqueue_spend_logs(client, queued, max_bytes=1) + assert [row.get("accounting_log_disabled") for row in client.spend_log_transactions] == [True, False] + await _flush(writer, client) + assert _key_spend(parity_database) == (0.5,) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT request_id FROM "LiteLLM_SpendLogs"').fetchall() == [("kept-call",)] + + +@pytest.mark.parametrize("isolation_budget", [5, 0], ids=["isolated", "budget_exhausted"]) +@pytest.mark.asyncio +async def test_a_spend_log_postgres_rejects_still_charges_its_request_and_the_rest_of_the_batch( + parity_database: str, monkeypatch: pytest.MonkeyPatch, isolation_budget: int +): + from prisma.errors import DataError + + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + class RejectsOneLog: + def __init__(self, repository: SpendLogsRepository) -> None: + self.repository = repository + + async def insert_with_deliveries(self, rows, *, write_spend_logs: bool = True) -> int: + if any(row["request_id"] == "poison-call" and not row.get("accounting_log_disabled") for row in rows): + raise DataError({"user_facing_error": {"message": 'invalid byte sequence for encoding "UTF8": 0x00'}}) + return await self.repository.insert_with_deliveries(rows, write_spend_logs=write_spend_logs) + + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "poison-call") + await _record_call(writer, "clean-call") + rows = tuple(client.jsonify_object(row) for row in client.spend_log_transactions) + client.spend_log_transactions = [] + await _create_spend_logs_with_poison_isolation(RejectsOneLog(SpendLogsRepository(client)), rows, isolation_budget) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert _key_spend(parity_database) == (0.5,) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT request_id FROM "LiteLLM_SpendLogs"').fetchall() == ( + [("clean-call",)] if isolation_budget else [] + ) + assert connection.execute( + 'SELECT state::text, COUNT(*) FROM "LiteLLM_SpendLogDelivery" GROUP BY 1' + ).fetchall() == [("applied", 2)] + + +@pytest.mark.asyncio +async def test_a_request_with_a_non_text_tag_is_still_charged(parity_database: str, monkeypatch: pytest.MonkeyPatch): + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "tagged-call", tags=[None, "real-tag"]) + await _flush(writer, client) + assert _key_spend(parity_database) == (0.25,) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT tag, spend FROM "LiteLLM_DailyTagSpend" WHERE tag=\'real-tag\'' + ).fetchall() == [("real-tag", 0.25)] + + +@pytest.mark.parametrize("threshold", ["count", "bytes"]) +@pytest.mark.asyncio +async def test_retained_batch_storage_threshold_alerts_without_deleting_receipts( + parity_database: str, monkeypatch: pytest.MonkeyPatch, threshold: str +): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.types.integrations.slack_alerting import AlertType + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1" if threshold == "count" else "1000000") + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1" if threshold == "bytes" else "1073741824") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + alerts = SimpleNamespace(alerting_handler=AsyncMock()) + await client.db.connect() + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, state, applied_at, accounting_payload, is_batch) ' + "VALUES ('old-batch', 'applied', CURRENT_TIMESTAMP - INTERVAL '400 days', '{}', true)" + ) + await flush_accounting_deliveries(DBSpendUpdateWriter(), client, alerts) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("old-batch", "applied") + ] + table_bytes = connection.execute( + 'SELECT pg_total_relation_size(\'"LiteLLM_SpendLogDelivery"\')' + ).fetchone()[0] + alerts.alerting_handler.assert_awaited_once_with( + message=f"Accounting delivery storage requires review. Retained batches=1, table bytes={table_bytes}.", + level="High", alert_type=AlertType.failed_tracking_spend, + ) + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("backlog", ["retries", "expired"]) +@pytest.mark.parametrize("batch_size", [1, 1100]) +@pytest.mark.asyncio +async def test_worker_drains_recovered_and_expired_work_at_its_configured_capacity( + parity_database, monkeypatch, backlog, batch_size +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", str(batch_size)) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + with psycopg.connect(parity_database, autocommit=True) as connection: + if backlog == "retries": + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, attempts) ' + "SELECT 'retry-' || i, accounting_payload, false, 1 " + 'FROM "LiteLLM_SpendLogDelivery" CROSS JOIN generate_series(1, 3) i' + ) + else: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'expired-' || i, '{}', false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' " + 'FROM generate_series(1, 1100) i' + ) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == ( + 0.25 * (1 + min(batch_size, 3)) if backlog == "retries" else 0.25, + ) + assert connection.execute( + 'SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\' OR request_id LIKE \'expired-%\'' + ).fetchone() == ((3 - min(batch_size, 3)) if backlog == "retries" else max(0, 1100 - batch_size * 2),) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_analytics_overflow_across_requests_still_applies_every_charge( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + writer = DBSpendUpdateWriter() + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + stored = await client.writer_db.query_raw('SELECT to_jsonb(s) AS row FROM "LiteLLM_SpendLogs" s') + delivery = (await client.writer_db.query_raw('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"'))[0][ + "accounting_payload" + ] + oversized_savings = {**delivery["metadata"], "compression_savings": {"tokens_saved": 2**62}} + assert await SpendLogsRepository(client).insert_with_deliveries( + tuple( + { + **stored[0]["row"], + "request_id": f"overflow-{index}", + "accounting_payload": {**delivery, "request_id": f"overflow-{index}", "metadata": oversized_savings}, + } + for index in range(3) + ) + ) == 3 + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT state::text, COUNT(*) FROM "LiteLLM_SpendLogDelivery" GROUP BY state' + ).fetchall() == [("applied", 4)] + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (1.0,) + assert connection.execute( + 'SELECT api_requests, compression_saved_tokens FROM "LiteLLM_DailyUserSpend"' + ).fetchone() == (4, 3 * 2**40) + finally: + await client.db.disconnect() diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_overloaded_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_overloaded_wire.py index bae775164dd..9d0e31058e4 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_overloaded_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_overloaded_wire.py @@ -140,6 +140,16 @@ def _assert_internal_server_error(response: httpx.Response) -> None: assert "InternalServerError" in string_value(error["message"]), response.text +def _assert_openai_overloaded(response: httpx.Response) -> None: + assert response.status_code == 503, response.text + error: Final = _openai_error(response) + assert set(error) == {"message", "type", "param", "code"}, response.text + assert string_value(error["message"]), response.text + assert error["type"] == "service_unavailable_error", response.text + assert error["param"] is None, response.text + assert error["code"] == "server_is_overloaded", response.text + + def _one_post(wire: Wire, target: str) -> None: assert [(request.method, unquote(request.target)) for request in wire.drain()] == [("POST", target)] @@ -361,7 +371,7 @@ def test_messages_overloaded_frame_before_content_without_a_retry_budget_is_rela ] -def test_chat_completions_overloaded_frame_before_content_answers_500(gateway: Gateway) -> None: +def test_chat_completions_overloaded_frame_before_content_answers_503(gateway: Gateway) -> None: prompt: Final = _prompt() with wire_server(_always(_overloaded_before_content(prompt))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) @@ -370,24 +380,24 @@ def test_chat_completions_overloaded_frame_before_content_answers_500(gateway: G "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True, "num_retries": 0}, ) - _assert_internal_server_error(response) + _assert_openai_overloaded(response) _one_post(wire, "/v1/messages") -def test_chat_completions_529_answers_500_to_the_sync_sdk(gateway: Gateway) -> None: +def test_chat_completions_529_answers_503_to_the_sync_sdk(gateway: Gateway) -> None: with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) with pytest.raises(openai.APIStatusError) as caught: _openai_sdk(gateway).chat.completions.create( model=model, messages=[{"role": "user", "content": _prompt()}], extra_body={"num_retries": 0} ) - assert caught.value.status_code == 500, caught.value.response.text - _assert_internal_server_error(caught.value.response) + assert caught.value.status_code == 503, caught.value.response.text + _assert_openai_overloaded(caught.value.response) _one_post(wire, "/v1/messages") - _assert_logged_failure(model, rows=1, error_class="InternalServerError", code="500") + _assert_logged_failure(model, rows=1, error_class="InternalServerError", code="503") -async def test_chat_completions_stream_529_answers_500_to_the_async_sdk(gateway: Gateway) -> None: +async def test_chat_completions_stream_529_answers_503_to_the_async_sdk(gateway: Gateway) -> None: with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) with pytest.raises(openai.APIStatusError) as caught: @@ -397,19 +407,31 @@ async def test_chat_completions_stream_529_answers_500_to_the_async_sdk(gateway: stream=True, extra_body={"num_retries": 0}, ) - assert caught.value.status_code == 500, caught.value.response.text - _assert_internal_server_error(caught.value.response) + assert caught.value.status_code == 503, caught.value.response.text + _assert_openai_overloaded(caught.value.response) + _one_post(wire, "/v1/messages") + + +def test_completions_529_answers_503(gateway: Gateway) -> None: + with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/completions", + {"model": model, "prompt": _prompt(), "max_tokens": 16, "num_retries": 0}, + ) + _assert_openai_overloaded(response) _one_post(wire, "/v1/messages") @pytest.mark.parametrize("shape", ("non-stream", "stream")) -def test_responses_529_answers_500(gateway: Gateway, shape: Shape) -> None: +def test_responses_529_answers_503(gateway: Gateway, shape: Shape) -> None: with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( "POST", "/v1/responses", {"model": model, "input": _prompt(), "stream": shape == "stream", "num_retries": 0} ) - _assert_internal_server_error(response) + _assert_openai_overloaded(response) _one_post(wire, "/v1/messages") diff --git a/tests/integration/observability/test_anthropic_overloaded_failure_signals.py b/tests/integration/observability/test_anthropic_overloaded_failure_signals.py index 044c58e44cc..e86ed98fae7 100644 --- a/tests/integration/observability/test_anthropic_overloaded_failure_signals.py +++ b/tests/integration/observability/test_anthropic_overloaded_failure_signals.py @@ -215,7 +215,7 @@ def _until(rig: _Rig, name: str, model: str) -> tuple[Sample, ...]: ( pytest.param("/v1/messages", _MESSAGES_ANTHROPIC, 529, "500", "InternalServerError", id="messages-anthropic"), pytest.param("/v1/messages", _MESSAGES_BEDROCK, 503, "503", "ServiceUnavailableError", id="messages-bedrock"), - pytest.param("/v1/chat/completions", _CHAT_ANTHROPIC, 500, "500", "InternalServerError", id="chat-anthropic"), + pytest.param("/v1/chat/completions", _CHAT_ANTHROPIC, 503, "500", "InternalServerError", id="chat-anthropic"), ), ) def test_the_failed_request_metric_counts_the_answered_status_and_the_deployment_metric_the_deployment_status( @@ -245,7 +245,7 @@ def _server_spans(rig: _Rig, since: int, trace_id: str) -> tuple[Span, ...]: ("route", "model", "answered"), ( pytest.param("/v1/messages", _TRACED_MESSAGES, 529, id="messages"), - pytest.param("/v1/chat/completions", _TRACED_CHAT, 500, id="chat"), + pytest.param("/v1/chat/completions", _TRACED_CHAT, 503, id="chat"), ), ) def test_the_server_span_carries_the_status_the_route_answered( @@ -313,7 +313,7 @@ def _datadog_status_codes(rig: _Rig, alias: str, expected: int) -> tuple[tuple[s ( pytest.param("/v1/messages", _DATADOG_MESSAGES_ANTHROPIC, 529, "500", id="messages-anthropic"), pytest.param("/v1/messages", _DATADOG_MESSAGES_BEDROCK, 503, "503", id="messages-bedrock"), - pytest.param("/v1/chat/completions", _DATADOG_CHAT_ANTHROPIC, 500, "500", id="chat-anthropic"), + pytest.param("/v1/chat/completions", _DATADOG_CHAT_ANTHROPIC, 503, "500", id="chat-anthropic"), ), ) def test_the_datadog_proxy_failure_log_carries_the_answered_status_and_the_llm_failure_log_the_deployment_status( @@ -361,7 +361,7 @@ def test_a_fail_closed_guardrail_outage_is_answered_and_recorded_with_the_routes ( pytest.param("/v1/messages", _TRACED_MESSAGES, 529, id="messages-anthropic"), pytest.param("/v1/messages", _MESSAGES_BEDROCK, 503, id="messages-bedrock"), - pytest.param("/v1/chat/completions", _TRACED_CHAT, 500, id="chat-anthropic"), + pytest.param("/v1/chat/completions", _TRACED_CHAT, 503, id="chat-anthropic"), ), ) def test_the_otel_v2_server_span_carries_the_status_the_route_answered( diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py index 79b0d9084aa..0095080d944 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_chaos.py @@ -1,12 +1,13 @@ from __future__ import annotations +import socket import threading import uuid from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from pathlib import Path +from queue import SimpleQueue from typing import Final -from urllib.parse import urlsplit import httpx import pytest @@ -44,7 +45,7 @@ from _openinference_support import ( ) from integration._support.client import Gateway, eventually from integration._support.process import graceful_stop_seconds -from integration._support.wire import Reply, Request, wire_server +from integration._support.wire import Reply, Request, Wire, wire_server def _call( @@ -249,11 +250,12 @@ def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: P _assert_response(response, marker, surface, stream, model_name) with ExitStack() as servers: - initial_stack: Final = servers.enter_context(ExitStack()) - stopped_destination: Final = initial_stack.enter_context(wire_server(_owned_sink_handler(sink))) - sink_port: Final = urlsplit(stopped_destination.url).port - assert sink_port is not None, stopped_destination.url - initial_stack.close() + outage_socket: Final = servers.enter_context(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) + outage_socket.bind(("127.0.0.1", 0)) # test-quality-ok: socket stays bound throughout the outage + sink_port: Final = outage_socket.getsockname()[1] + stopped_destination: Final = Wire( + f"http://127.0.0.1:{sink_port}", SimpleQueue(), SimpleQueue(), SimpleQueue(), SimpleQueue() + ) with _rig( gateway, tmp_path, @@ -272,6 +274,7 @@ def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: P with httpx.Client(trust_env=False) as client, pytest.raises(httpx.ConnectError): client.get(stopped_destination.url + "/health", timeout=2) recovered_stack: Final = servers.enter_context(ExitStack()) + outage_socket.close() recovered_destination: Final = recovered_stack.enter_context( wire_server(_owned_sink_handler(sink), port=sink_port) ) diff --git a/tests/integration/pricing/test_cost_margin_discount.py b/tests/integration/pricing/test_cost_margin_discount.py index b6cda8f125f..613fe72ed70 100644 --- a/tests/integration/pricing/test_cost_margin_discount.py +++ b/tests/integration/pricing/test_cost_margin_discount.py @@ -7,7 +7,7 @@ import httpx import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue class CostEstimate(BaseModel): @@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None: assert actual == pytest.approx(dict(expectations)), rows +class LoggedBreakdown(BaseModel): + model_config = ConfigDict(frozen=True) + + original_cost: float + discount_percent: float | None = None + discount_amount: float | None = None + margin_percent: float | None = None + margin_fixed_amount: float | None = None + margin_total_amount: float | None = None + + +def _logged_breakdown(request_id: str) -> LoggedBreakdown: + rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)) + assert len(rows) == 1, rows + metadata: Final = rows[0]["metadata"] + parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata) + return LoggedBreakdown.model_validate(parsed["cost_breakdown"]) + + def _estimated(gateway: Gateway, model: str) -> CostEstimate: response: Final = gateway.request( "POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20} @@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga (streamed_global_request_id, 0.12), ) ) + discounted: Final = _logged_breakdown(discounted_request_id) + assert discounted.original_cost == pytest.approx(0.06), discounted + assert discounted.discount_percent == pytest.approx(0.05), discounted + assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted + margined: Final = _logged_breakdown(provider_margin_request_id) + assert margined.original_cost == pytest.approx(0.06), margined + assert margined.discount_percent == pytest.approx(0.05), margined + assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined + assert margined.margin_percent == pytest.approx(0.1), margined + assert margined.margin_fixed_amount == pytest.approx(0.001), margined + assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined rejections: Final = ( ( diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 4e091460667..bb4973d6f2f 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -139,6 +139,7 @@ _PROVIDERS: Final = ( False, "databricks/databricks-openjev-qwen35-4b", ), + _Provider("empiriolabs", "empiriolabs/aplomb-test", "/v1/decisions", "aplomb-test", _API_KEY, False, None), ) _LOCAL_PROVIDERS: Final = ( pytest.param( diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py index ef36c449761..56d54a87600 100644 --- a/tests/integration/providers/test_nvidia_riva_wire.py +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -135,3 +135,12 @@ def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gatew assert timed_out.status_code == 408, timed_out.text calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30) assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls + + +@pytest.mark.parametrize("riva_offline", (False, True), ids=("streaming", "offline")) +def test_the_deployment_api_key_reaches_riva_as_exactly_one_bearer_token(gateway: Gateway, riva_offline: bool) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, api_key="riva-deployment-key", riva_offline=riva_offline) + served: Final = _transcribe(gateway, scenario.key(), model) + assert served.status_code == 200, served.text + assert [call.authorization for call in server.drain()] == [("Bearer riva-deployment-key",)] diff --git a/tests/integration/routing/test_least_busy_concurrent_decisions.py b/tests/integration/routing/test_least_busy_concurrent_decisions.py new file mode 100644 index 00000000000..bb923534db4 --- /dev/null +++ b/tests/integration/routing/test_least_busy_concurrent_decisions.py @@ -0,0 +1,90 @@ +import json +import threading +from collections import Counter +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +pytestmark = pytest.mark.timeout(240) + +_GROUP: Final = "least-busy-decider" +_DEPLOYMENTS: Final = ("least-busy-decider-a", "least-busy-decider-b") +_BURST: Final = 4 +_REPLY: Final = json.dumps( + { + "model": "Qwen/Qwen3-0.6B", + "answers": {"billing": {"type": "noul", "noul": 0.91}}, + "usage": {"input_tokens": 367, "output_tokens": 3}, + } +).encode() +_DECISION: Final = { + "model": _GROUP, + "input": "I was charged twice for order 4417.", + "questions": [{"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}], +} + + +def _held_until_the_burst_arrives(arrived: threading.Barrier) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + arrived.wait() + return Reply(body=_REPLY) + + return respond + + +def _write_config(path: Path, provider_url: str) -> Path: + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _GROUP, + "litellm_params": {"model": "hosted_vllm/Qwen/Qwen3-0.6B", "api_base": provider_url}, + "model_info": {"id": deployment_id}, + } + for deployment_id in _DEPLOYMENTS + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": {"routing_strategy": "least-busy", "num_retries": 0, "disable_cooldowns": True}, + } + ) + ) + return path + + +def _decide(proxy: Gateway) -> httpx.Response: + return proxy.client.post( + "/v1/decisions", json=_DECISION, headers={"Authorization": f"Bearer {proxy.key}"}, timeout=120 + ) + + +def test_least_busy_spreads_a_burst_of_decisions_calls_across_the_group(tmp_path: Path) -> None: + arrived: Final = threading.Barrier(_BURST, timeout=60) + with ( + wire_server(_held_until_the_burst_arrives(arrived)) as upstream, + gateway_from_environment() as rig, + owned_proxy( + rig, + tmp_path, + {"STORE_MODEL_IN_DB": "False"}, + config=_write_config(tmp_path / "least-busy.yaml", upstream.url), + ) as proxy, + ThreadPoolExecutor(max_workers=_BURST) as pool, + ): + responses: Final = tuple(pool.map(lambda _: _decide(proxy), range(_BURST))) + assert [response.status_code for response in responses] == [200] * _BURST, [r.text for r in responses] + assert Counter(response.headers["x-litellm-model-id"] for response in responses) == Counter( + {deployment_id: _BURST // len(_DEPLOYMENTS) for deployment_id in _DEPLOYMENTS} + ) diff --git a/tests/integration/spend/test_failing_request_stored_snapshot.py b/tests/integration/spend/test_failing_request_stored_snapshot.py new file mode 100644 index 00000000000..4f016c33f37 --- /dev/null +++ b/tests/integration/spend/test_failing_request_stored_snapshot.py @@ -0,0 +1,84 @@ +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from integration._support.process import owned_proxy +from pydantic import JsonValue + +UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1" +TRANSCRIPT_TURNS: Final = 40 +RETRIES: Final = 2 +STORED_REQUEST_BUDGET_KB: Final = 64 + + +def _config_storing_prompts(directory: Path) -> Path: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general_settings: Final = object_value(config["general_settings"]) + path: Final = directory / "store-prompts.yaml" + path.write_text( + yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}}) + ) + return path + + +@pytest.fixture(scope="module") +def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with gateway_from_environment() as parent: + directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot") + with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned: + yield owned + + +def _transcript(marker: str) -> list[JsonValue]: + turns: Final[list[JsonValue]] = [ + {"role": role, "content": f"turn {turn // 2} {role}"} + for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS) + ] + return [*turns, {"role": "user", "content": marker}] + + +def _stored_request(gateway: Gateway, call_id: str) -> JsonValue: + rows: Final = eventually( + lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=40, + ) + stored: Final = rows[0]["proxy_server_request"] + return json.loads(stored) if isinstance(stored, str) else stored + + +def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + primary: Final = f"stored-snapshot-{uuid.uuid4().hex}" + fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}" + for name in (primary, primary, fallback, fallback): + scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0) + key: Final = scenario.key(models=[primary, fallback]) + marker: Final = f"stored snapshot {uuid.uuid4().hex}" + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": primary, + "max_tokens": 16, + "messages": _transcript(marker), + "router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]}, + }, + key=key, + ) + assert response.status_code != 200, response.text + assert fallback in response.text, response.text + + serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"])) + stored_kb: Final = len(serialized.encode()) / 1024 + assert marker in serialized, serialized[:400] + assert stored_kb <= STORED_REQUEST_BUDGET_KB, ( + f"one failing request stored a {stored_kb:.0f} KB request snapshot for a " + f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback" + ) diff --git a/tests/unit/accounting/test_delivery_projection.py b/tests/unit/accounting/test_delivery_projection.py new file mode 100644 index 00000000000..0b1f36d6b48 --- /dev/null +++ b/tests/unit/accounting/test_delivery_projection.py @@ -0,0 +1,85 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final +from urllib.parse import unquote + +from litellm.accounting.delivery_projection import entity_transactions, project_deliveries +from litellm.accounting.projection import Contribution +from litellm.repositories.accounting_delivery_repository import DeliveryProjection + +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) + + +def _payload(spend: float, started: datetime = _NOW, window_start: datetime = _NOW) -> Mapping[str, object]: + return { + "request_id": f"request-{spend}-{started.isoformat()}", + "startTime": started, + "api_key": "key", + "user": "user/1", + "team_id": "team", + "organization_id": "org::1", + "spend": spend, + "metadata": json.dumps( + { + "accounting_entities": { + "api_key": "key", + "user": "user/1", + "team_id": "team", + "organization_id": "org::1", + }, + "accounting_windows": [ + { + "entity_type": "key", + "entity_id": "key", + "window_duration": "1d", + "window_start": window_start.isoformat(), + } + ], + } + ), + } + + +def _spend(contributions: tuple[Contribution, ...], table: str) -> tuple[float, ...]: + return tuple(float(dict(item.increments)["spend"]) for item in contributions if item.table == table) + + +def test_deliveries_for_one_destination_sum_and_keep_the_earliest_request_start() -> None: + later: Final = _NOW + timedelta(hours=1) + + contributions: Final = project_deliveries( + (DeliveryProjection(_payload(0.5, started=later)), DeliveryProjection(_payload(0.25))) + ) + + assert _spend(contributions, "LiteLLM_VerificationToken") == (0.75,) + windows: Final = tuple(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + assert len(windows) == 1 + assert dict(windows[0].increments) == {"spend": 0.75} + assert dict(windows[0].values)["started_at"] == str(_NOW) + + +def test_the_same_window_in_a_new_period_is_charged_to_its_own_row() -> None: + next_day: Final = _NOW + timedelta(days=1) + + contributions: Final = project_deliveries( + (DeliveryProjection(_payload(0.25)), DeliveryProjection(_payload(0.5, window_start=next_day))) + ) + + windows: Final = sorted( + (dict(item.values)["window_start"], dict(item.increments)["spend"]) + for item in contributions + if item.table == "LiteLLM_BudgetWindowSpend" + ) + assert windows == [(_NOW.isoformat(), 0.25), (next_day.isoformat(), 0.5)] + + +def test_entity_transactions_key_memberships_the_way_the_writer_applies_them() -> None: + transactions: Final = entity_transactions(project_deliveries((DeliveryProjection(_payload(0.25)),))) + + assert transactions["key_list_transactions"] == {"key": 0.25} + assert transactions["org_list_transactions"] == {"org::1": 0.25} + assert transactions["team_member_list_transactions"] == {"team_id::team::user_id::user/1": 0.25} + ((org_member_key, org_member_spend),) = transactions["org_member_list_transactions"].items() + _, quoted_org_id, _, quoted_user_id = org_member_key.split("::") + assert (unquote(quoted_org_id), unquote(quoted_user_id), org_member_spend) == ("org::1", "user/1", 0.25) diff --git a/tests/unit/accounting/test_measurements.py b/tests/unit/accounting/test_measurements.py deleted file mode 100644 index b65c3076c5d..00000000000 --- a/tests/unit/accounting/test_measurements.py +++ /dev/null @@ -1,235 +0,0 @@ -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import ( - CostEvent, - Dimension, - DimensionName, - InboundUsageEvent, - Measurement, - Metric, - MetricName, - Outcome, - OutboundUsageEvent, - PricingSnapshot, - measurements_for, -) -from litellm.types.utils import CallTypes - -_START: Final = datetime(2026, 1, 1, 23, 59, 59, tzinfo=timezone.utc) -_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) -_DIMENSIONS: Final = ( - Dimension(name=DimensionName.TEAM, value="team"), - Dimension(name=DimensionName.TAG, value="alpha"), - Dimension(name=DimensionName.TAG, value="beta"), -) - - -def _snapshot(amount: Decimal) -> PricingSnapshot: - return PricingSnapshot( - calculator_version="test", - model="test-model", - provider=None, - pricing_source="upstream_header", - priced_model_key="test-model", - upstream_reported_cost=amount, - total_cost=amount, - ) - - -@pytest.mark.parametrize( - ("outcome", "outcome_metric"), - [ - (Outcome.SUCCEEDED, MetricName.SUCCESSFUL_REQUESTS), - (Outcome.FAILED, MetricName.FAILED_REQUESTS), - (Outcome.CANCELLED, MetricName.FAILED_REQUESTS), - ], -) -def test_caller_measurements_keep_both_request_and_outcome_counts(outcome, outcome_metric): - event = InboundUsageEvent( - event_id="caller-event", - request_id="caller", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=outcome, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)), Metric(name=outcome_metric, value=Decimal(1))), - ) - assert measurements_for(event) == ( - Measurement( - source_event_id="caller-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.REQUESTS, value=Decimal(1)), - ), - Measurement( - source_event_id="caller-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=outcome_metric, value=Decimal(1)), - ), - ) - - -def test_usage_measurements_keep_every_metric_in_the_start_time_bucket(): - event = OutboundUsageEvent( - event_id="usage-event", - attempt_id="attempt", - request_id="caller", - retry_of_attempt_id=None, - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=_DIMENSIONS, - metrics=( - Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), - Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), - ), - ) - assert measurements_for(event) == ( - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), - ), - ) - - -@pytest.mark.parametrize("metrics", [None, ()]) -def test_unknown_or_empty_usage_does_not_fabricate_zero_measurements(metrics): - event = OutboundUsageEvent( - event_id="usage-event", - attempt_id="attempt", - request_id=None, - retry_of_attempt_id=None, - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=metrics, - ) - assert measurements_for(event) == () - - -@pytest.mark.parametrize("amount", [Decimal(0), Decimal("0.1234567890123456789")]) -def test_cost_measurements_preserve_authoritative_amount_identity_and_occurrence(amount): - event = CostEvent( - event_id="cost-event", - usage_event_id="usage-event", - occurred_at=_END, - dimensions=_DIMENSIONS, - charges=(Metric(name=MetricName.SPEND, value=amount, unit="USD"),), - snapshot=_snapshot(amount), - ) - assert measurements_for(event) == ( - Measurement( - source_event_id="cost-event", - occurred_at=_END, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.SPEND, value=amount, unit="USD"), - ), - ) - - -def test_billed_retry_and_success_contribute_one_caller_request_and_both_charges(): - caller = InboundUsageEvent( - event_id="caller-event", - request_id="caller", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - retry = OutboundUsageEvent( - event_id="retry-usage", - attempt_id="retry", - request_id="caller", - retry_of_attempt_id=None, - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.FAILED, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(4)),), - ) - success = OutboundUsageEvent( - event_id="success-usage", - attempt_id="success", - request_id="caller", - retry_of_attempt_id="retry", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),), - ) - retry_cost = CostEvent( - event_id="retry-cost", - usage_event_id="retry-usage", - occurred_at=_START, - dimensions=_DIMENSIONS, - charges=(Metric(name=MetricName.SPEND, value=Decimal("0.025"), unit="USD"),), - snapshot=_snapshot(Decimal("0.025")), - ) - success_cost = CostEvent( - event_id="success-cost", - usage_event_id="success-usage", - occurred_at=_START, - dimensions=_DIMENSIONS, - charges=(Metric(name=MetricName.SPEND, value=Decimal("0.125"), unit="USD"),), - snapshot=_snapshot(Decimal("0.125")), - ) - measurements = ( - *measurements_for(caller), - *measurements_for(retry), - *measurements_for(success), - *measurements_for(retry_cost), - *measurements_for(success_cost), - ) - assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.REQUESTS) == Decimal(1) - assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.TOTAL_TOKENS) == Decimal(16) - assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.SPEND) == Decimal("0.150") - assert {m.source_event_id for m in measurements} == { - "caller-event", - "retry-usage", - "success-usage", - "retry-cost", - "success-cost", - } diff --git a/tests/unit/accounting/test_metrics.py b/tests/unit/accounting/test_metrics.py new file mode 100644 index 00000000000..49b46eb36c2 --- /dev/null +++ b/tests/unit/accounting/test_metrics.py @@ -0,0 +1,28 @@ +from prometheus_client.metrics import MetricWrapperBase + +from litellm.accounting import metrics +from litellm.accounting.metrics import record_deliveries +from litellm.repositories.accounting_delivery_repository import DeliveryMetrics + + +def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: + family = metric.collect()[0] + return next(sample.value for sample in family.samples if sample.name == family.name + suffix) + + +def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None: + batches_before = _sample(metrics._BATCH_TIME, "_count") + + record_deliveries( + DeliveryMetrics(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192), + 0.5, + ) + + assert ( + _sample(metrics._PENDING), + _sample(metrics._PARKED), + _sample(metrics._OLDEST), + _sample(metrics._BATCH_COUNT), + _sample(metrics._TABLE_BYTES), + ) == (4.0, 1.0, 0.0, 7.0, 8192.0) + assert _sample(metrics._BATCH_TIME, "_count") == batches_before + 1 diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py new file mode 100644 index 00000000000..17186a89922 --- /dev/null +++ b/tests/unit/accounting/test_projection.py @@ -0,0 +1,557 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.accounting.projection import project +from litellm.accounting.snapshot import build_accounting_payload +from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.proxy_server import proxy_logging_obj +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import PrismaClient +from litellm.types.utils import ModelResponse, Usage + +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.parametrize("status", ["success", "failure"]) +@pytest.mark.parametrize("origin", [None, "auto_router", "billed_retry"]) +@pytest.mark.parametrize("spend", [0.0, 0.25]) +@pytest.mark.asyncio +async def test_daily_projection_matches_legacy_writer(status: str, origin: str | None, spend: float) -> None: + payload: Final = get_logging_payload( + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "projection-request", + "litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user", "status": status}}, + }, + response_obj=ModelResponse( + id="projection-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ) + metadata: Final = { + **_FIELDS.validate_json(payload["metadata"]), + **({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}), + **({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin not in (None, "billed_retry") else {}), + } + row: Final = {**payload, "spend": spend, "metadata": json.dumps(metadata), "request_duration_ms": 25} + legacy: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction( + row, PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj), "user" + ) + assert legacy is not None + accounting_input: Final[SpendLogsPayload] = {**row} + accounting: Final = build_accounting_payload( + accounting_input, + entities={"api_key": payload["api_key"], "user": payload["user"]}, + model_access_groups=(), + windows=(), + router=lambda: None, + ) + restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) + stored_daily: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend") + assert dict(stored_daily.increments) == {name: legacy[name] for name, _ in stored_daily.increments} + assert dict(stored_daily.key)["date"] == legacy["date"] + assert dict(stored_daily.key)["api_key"] == legacy["api_key"] + assert not {"messages", "response", "proxy_server_request", "requester_ip_address"}.intersection(accounting) + + +def test_projection_uses_frozen_savings_and_keeps_duplicate_entity_tags() -> None: + row: Final[Mapping[str, object]] = { + "request_id": "saved", + "startTime": _NOW, + "api_key": "key", + "user": "user", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "request_tags": '["tag", "tag"]', + "spend": 0.5, + "metadata": json.dumps( + { + "accounting_savings": { + "compression": 1.0, + "prompt_caching": -2.0, + "gateway_injected_caching": 3.0, + "autorouter": 4.0, + }, + "autorouter_savings": 9000, + } + ), + } + contributions: Final = project(row) + tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_TagTable") + daily_tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_DailyTagSpend") + assert sum(dict(item.increments)["spend"] for item in tags) == 1.0 + assert len(daily_tags) == 1 + assert {key: value for key, value in daily_tags[0].increments if key.endswith("savings_spend")} == { + "compression_savings_spend": 1.0, + "prompt_caching_savings_spend": -2.0, + "gateway_injected_caching_savings_spend": 3.0, + "autorouter_savings_spend": 4.0, + } + + +def test_projection_carries_windows_and_does_not_reprice_guardrails() -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "api_key": "key", + "spend": 0.75, + "metadata": json.dumps( + { + "accounting_windows": [ + { + "entity_type": "key", + "entity_id": "key", + "window_duration": "1d", + "window_start": _NOW.isoformat(), + } + ], + "guardrail_information": [ + { + "guardrail_id": "guard", + "guardrail_usage": {"text": 3, "images": 2}, + "guardrail_cost_by_unit": {"text": 0.3, "images": None}, + } + ], + } + ), + } + contributions: Final = project(row) + window: Final = next(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + units: Final = { + dict(item.key)["usage_unit"]: dict(item.increments) + for item in contributions + if item.table == "LiteLLM_DailyGuardrailUsageUnits" + } + assert dict(window.increments) == {"spend": 0.75} + assert dict(window.values)["window_start"] == _NOW.isoformat() + assert units == { + "text": {"units": 3, "cost": 0.3, "untracked_units": 0}, + "images": {"units": 2, "cost": 0.0, "untracked_units": 2}, + } + + +@pytest.mark.asyncio +async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.responses.streaming_iterator import drain_logging_tasks + + client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(client, "spend_log_transactions", []) + monkeypatch.setattr(client, "model_usage_transactions", []) + writer: Final = DBSpendUpdateWriter() + charged: Final = await writer.update_database( + token="sk-projection", + user_id="user", + end_user_id="end-user", + team_id="team", + org_id="org", + project_id="project", + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "entity-request", + "agent_id": "agent", + "user": "end-user", + "litellm_params": { + "metadata": { + "user_api_key": "sk-projection", + "user_api_key_user_id": "user", + "user_api_key_team_id": "team", + "user_api_key_org_id": "org", + "user_api_key_project_id": "project", + "user_api_key_end_user_id": "end-user", + "tags": ["tag", "tag"], + }, + "user": "end-user", + }, + }, + completion_response=ModelResponse( + id="entity-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + response_cost=0.25, + ) + await drain_logging_tasks() + assert charged + assert len(client.spend_log_transactions) == 1 + queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + logged: Final[SpendLogsPayload] = client.spend_log_transactions[0] + projected: Final = project( + _FIELDS.validate_json( + json.dumps( + build_accounting_payload( + logged, + entities={ + "api_key": logged["api_key"], + "user": "user", + "end_user": "end-user", + "team_id": "team", + "organization_id": "org", + "project_id": "project", + "agent_id": "agent", + }, + model_access_groups=(), + windows=(), + router=lambda: None, + ), + default=str, + ) + ) + ) + legacy: Final = {(item["entity_type"].value, item["entity_id"]): item["response_cost"] for item in queued} + assert legacy[("key", proxy_server.hash_token("sk-projection"))] == next( + dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_VerificationToken" + ) + for table, entity, identity in ( + ("LiteLLM_UserTable", "user", "user"), + ("LiteLLM_TeamTable", "team", "team"), + ("LiteLLM_OrganizationTable", "organization", "org"), + ("LiteLLM_EndUserTable", "end_user", "end-user"), + ("LiteLLM_AgentsTable", "agent", "agent"), + ("LiteLLM_ProjectTable", "project", "project"), + ): + assert legacy[(entity, identity)] == next( + dict(item.increments)["spend"] for item in projected if item.table == table + ) + assert sum(dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_TagTable") == sum( + item["response_cost"] or 0 for item in queued if item["entity_type"].value == "tag" + ) + + daily_queues: Final = ( + ("LiteLLM_DailyUserSpend", writer.daily_spend_update_queue), + ("LiteLLM_DailyTeamSpend", writer.daily_team_spend_update_queue), + ("LiteLLM_DailyOrganizationSpend", writer.daily_org_spend_update_queue), + ("LiteLLM_DailyEndUserSpend", writer.daily_end_user_spend_update_queue), + ("LiteLLM_DailyAgentSpend", writer.daily_agent_spend_update_queue), + ("LiteLLM_DailyTagSpend", writer.daily_tag_spend_update_queue), + ) + for table, queue in daily_queues: + transactions = await queue.flush_and_get_aggregated_daily_spend_update_transactions() + matching = tuple(item for item in projected if item.table == table) + assert matching, table + contribution = matching[0] + assert len(transactions) == 1 + transaction = next(iter(transactions.values())) + assert dict(contribution.increments) == {name: transaction[name] for name, _ in contribution.increments} + assert len(client.model_usage_transactions) == 1 + model: Final = client.model_usage_transactions[0] + model_projection: Final = next(item for item in projected if item.table == "LiteLLM_DailyModelUsage") + assert dict(model_projection.increments) == { + "spend": model.spend, + "prompt_tokens": model.prompt_tokens, + "completion_tokens": model.completion_tokens, + "request_count": int(model.counts_as_request), + "successful_requests": int(model.successful), + "failed_requests": int(not model.successful), + } + + +@pytest.mark.parametrize("served_id", ["requested", "fallback", "missing"]) +@pytest.mark.asyncio +async def test_projection_preserves_fallback_narrowed_model_access_groups(served_id: str) -> None: + from litellm import Router + from litellm.proxy.db.db_spend_update_writer import debitable_model_access_groups + + router: Final = Router( + model_list=[ + { + "model_name": "first", + "litellm_params": {"model": "openai/projection-test", "api_key": "unused"}, + "model_info": {"id": "requested", "access_groups": ["first-group"]}, + }, + { + "model_name": "second", + "litellm_params": {"model": "openai/projection-test", "api_key": "unused"}, + "model_info": {"id": "fallback", "access_groups": ["second-group", "unattributed"]}, + }, + ] + ) + attributed: Final = ("first-group", "second-group", "first-group") + writer: Final = DBSpendUpdateWriter() + client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) + await writer._update_model_access_group_db(0.5, attributed, served_id, client, router) + queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + frozen_groups: Final = debitable_model_access_groups(attributed, served_id, router) + projected: Final = project( + {"spend": 0.5, "metadata": json.dumps({"accounting_model_access_groups": frozen_groups})} + ) + assert {dict(item.key)["access_group_name"]: dict(item.increments)["spend"] for item in projected} == { + item["entity_id"]: item["response_cost"] for item in queued + } + assert "unattributed" not in frozen_groups + + +def test_guardrail_projection_skips_unmeasured_counters() -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "metadata": json.dumps( + { + "guardrail_information": [ + {"guardrail_usage": {"text": 3}}, + { + "guardrail_id": "guard", + "guardrail_usage": {"zero": 0, "negative": -1, "bool": True, "unknown": None, "text": 2}, + }, + ] + } + ), + } + units: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyGuardrailUsageUnits") + assert len(units) == 1 + assert dict(units[0].key)["usage_unit"] == "text" + assert dict(units[0].increments) == {"units": 2, "cost": 0.0, "untracked_units": 2} + + +def test_lean_row_without_model_does_not_fabricate_daily_entity_usage() -> None: + projected: Final = project( + { + "request_id": "lean", + "api_key": "key", + "startTime": _NOW, + "spend": 0.5, + "metadata": {"accounting_entities": {"api_key": "key"}}, + } + ) + assert tuple(item.table for item in projected) == ("LiteLLM_VerificationToken", "LiteLLM_DailyModelUsage") + + +@pytest.mark.parametrize( + "cost_fields", + [ + {"guardrail_cost_in_spend": False, "guardrail_cost_by_unit": {"text": 1.0}}, + {"guardrail_cost_by_unit": {"text": -1.0}}, + ], +) +def test_report_only_or_invalid_guardrail_prices_remain_untracked(cost_fields: Mapping[str, object]) -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "request_tags": '["task:debugging"]', + "metadata": json.dumps( + { + "guardrail_information": [ + {"guardrail_id": "guard", "guardrail_usage": {"text": 2}, **cost_fields}, + ] + } + ), + } + contributions: Final = project(row) + guardrail: Final = next(item for item in contributions if item.table == "LiteLLM_DailyGuardrailUsageUnits") + model: Final = next(item for item in contributions if item.table == "LiteLLM_DailyModelUsage") + assert dict(guardrail.increments) == {"units": 2, "cost": 0.0, "untracked_units": 2} + assert dict(model.key)["task_type"] == "debugging" + + +@pytest.mark.parametrize("organization_id", [None, "", "charged-org"]) +def test_daily_organization_uses_frozen_charged_identity(organization_id: str | None) -> None: + row: Final[Mapping[str, object]] = { + "startTime": _NOW, + "api_key": "key", + "model": "projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "organization_id": "stored-org", + "spend": 0.25, + "metadata": {"accounting_entities": {"organization_id": organization_id}}, + } + daily: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyOrganizationSpend") + if organization_id is None: + assert daily == () + else: + assert len(daily) == 1 + assert dict(daily[0].key)["organization_id"] == organization_id + assert dict(daily[0].increments)["spend"] == 0.25 + + +@pytest.mark.parametrize("start", [_NOW, "2026-11-02T12:00:00+00:00", "2026-11-02 12:00:00+00:00"]) +def test_every_daily_row_uses_the_calendar_date_of_the_request(start: object) -> None: + row: Final[Mapping[str, object]] = { + "request_id": "dated", + "startTime": start, + "api_key": "key", + "user": "user", + "team_id": "team", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "spend": 0.5, + "metadata": "{}", + } + dates: Final = { + item.table: dict(item.key)["date"] for item in project(row) if item.table.startswith("LiteLLM_Daily") + } + assert {"LiteLLM_DailyUserSpend", "LiteLLM_DailyModelUsage"} <= dates.keys() + assert set(dates.values()) == {"2026-11-02"} + + +def test_unparseable_start_time_skips_only_guardrail_usage() -> None: + projected: Final = project( + { + "request_id": "bad-time", + "startTime": "yesterday at noon", + "spend": 0.5, + "metadata": { + "accounting_entities": {"api_key": "key"}, + "guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}], + }, + } + ) + tables: Final = {item.table for item in projected} + assert "LiteLLM_VerificationToken" in tables + assert "LiteLLM_DailyGuardrailUsageUnits" not in tables + + +def test_guardrail_usage_is_dated_in_utc() -> None: + projected: Final = project( + { + "request_id": "offset", + "startTime": "2026-11-02T22:30:00-05:00", + "metadata": {"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 1}}]}, + } + ) + assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [ + "2026-11-03" + ] + + +@pytest.mark.parametrize( + ("untrusted", "metadata"), + [ + ({"request_tags": json.dumps(["x" * 16384, "normal"])}, {}), + ({"end_user": "x" * 3072}, {}), + ({}, {"compression_savings": {"tokens_saved": 2**63}}), + ], +) +def test_untrusted_analytics_values_never_block_the_charge( + untrusted: Mapping[str, str], metadata: Mapping[str, object] +) -> None: + row: Final[Mapping[str, object]] = { + "request_id": "untrusted", + "startTime": _NOW, + "api_key": "key", + "user": "user", + "team_id": "team", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "spend": 0.5, + **untrusted, + "metadata": { + **metadata, + "accounting_entities": { + "api_key": "key", + "user": "user", + "team_id": "team", + "end_user": untrusted.get("end_user"), + }, + }, + } + contributions: Final = project(row) + texts: Final = tuple(value for item in contributions for _, value in (*item.key, *item.values) if value) + numbers: Final = tuple(value for item in contributions for _, value in item.increments) + assert {"LiteLLM_VerificationToken", "LiteLLM_TeamTable", "LiteLLM_UserTable"} <= { + item.table for item in contributions + } + assert "LiteLLM_DailyUserSpend" in {item.table for item in contributions} + assert all(len(text.encode()) <= 2000 for text in texts) + assert all(abs(number) <= 2**40 for number in numbers) + + +def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> None: + payload: Final = get_logging_payload( + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "guardrail-request", + "litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user"}}, + }, + response_obj=ModelResponse( + id="guardrail-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ) + guardrail: Final = { + "guardrail_name": "pii", + "guardrail_usage": {"requests": 1}, + "guardrail_cost_by_unit": {"requests": 0.25}, + "guardrail_response": "card 4111 1111 1111 1111", + "match_details": [{"matched": "4111 1111 1111 1111"}], + } + headroom: Final = { + "guardrail_name": "compress", + "guardrail_provider": "headroom", + "guardrail_response": {"tokens_saved": 7, "compressed_prompt": "card 4111 1111 1111 1111"}, + } + metadata: Final = {**_FIELDS.validate_json(payload["metadata"]), "guardrail_information": [guardrail, headroom]} + accounting_input: Final[SpendLogsPayload] = { + **payload, + "spend": 0.5, + "metadata": json.dumps(metadata), + "request_tags": '["nul\x00tag"]', + } + accounting: Final = build_accounting_payload( + accounting_input, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None + ) + stored: Final = json.dumps(accounting, default=str) + restored: Final = _FIELDS.validate_json(stored) + usage_rows: Final = tuple(item for item in project(restored) if item.table == "LiteLLM_DailyGuardrailUsageUnits") + assert "4111" not in stored + assert "\\u0000" not in stored + assert [dict(item.increments) for item in usage_rows] == [{"units": 1, "cost": 0.25, "untracked_units": 0}] + daily_user: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend") + assert dict(daily_user.increments)["compression_saved_tokens"] == 7 + + +def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + + monkeypatch.setitem( + litellm.model_cost, + "savings-overflow-test", + {"input_cost_per_token": 0.002, "output_cost_per_token": 0.004, "litellm_provider": "openai", "mode": "chat"}, + ) + payload: Final[SpendLogsPayload] = { + **get_logging_payload( + kwargs={ + "model": "savings-overflow-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "overflow-request", + "litellm_params": {"metadata": {"user_api_key": "key"}}, + }, + response_obj=ModelResponse( + id="overflow-request", model="savings-overflow-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ), + "spend": 0.5, + "metadata": json.dumps({"compression_savings": {"tokens_saved": 10**400}}), + } + accounting: Final = build_accounting_payload( + payload, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None + ) + restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) + assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0 + assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [ + 0.5 + ] diff --git a/tests/unit/accounting/test_types.py b/tests/unit/accounting/test_types.py deleted file mode 100644 index f645ee6c2cc..00000000000 --- a/tests/unit/accounting/test_types.py +++ /dev/null @@ -1,484 +0,0 @@ -import re -from collections.abc import Mapping -from datetime import datetime, timedelta, timezone -from decimal import Decimal -from enum import Enum -from pathlib import Path -from typing import Final - -import pytest -from pydantic import ValidationError - -from litellm.accounting import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - DeliveryState, - DeliveryTarget, - Dimension, - DimensionFilter, - DimensionName, - EventKind, - Granularity, - InboundUsageEvent, - Metric, - MetricName, - Multiplier, - Outcome, - OutboundUsageEvent, - PricedQuantity, - PricingSnapshot, - Rate, - validate_aggregate_rows, -) -from litellm.accounting.types import PricingUnit -from litellm.types.utils import CallTypes - -_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) -_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) -_REQUEST: Final = Metric(name=MetricName.REQUESTS, value=Decimal(1)) -_TOKEN: Final = Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)) -_SPEND: Final = Metric(name=MetricName.SPEND, value=Decimal("0.1234567890123456789"), unit="USD") -_TEAM: Final = Dimension(name=DimensionName.TEAM, value="team") - - -def _inbound(metrics=(_REQUEST,), outcome=Outcome.SUCCEEDED): - return InboundUsageEvent( - event_id="inbound:request", - request_id="request", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_START, - outcome=outcome, - dimensions=(_TEAM,), - metrics=metrics, - ) - - -def _outbound(metrics=(_TOKEN,), **changes): - return OutboundUsageEvent.model_validate( - { - "event_id": "outbound:attempt", - "attempt_id": "attempt", - "request_id": "request", - "retry_of_attempt_id": None, - "call_type": CallTypes.acompletion, - "started_at": _START, - "ended_at": _END, - "outcome": Outcome.SUCCEEDED, - "dimensions": (_TEAM,), - "metrics": metrics, - **changes, - } - ) - - -def _query(**changes): - return AggregateQuery.model_validate( - { - "metric": MetricName.SPEND, - "filters": (DimensionFilter(name=DimensionName.TEAM, values=("one", "two")),), - "group_by": (DimensionName.MODEL,), - "start": _START, - "end": _END, - "granularity": Granularity.DAY, - **changes, - } - ) - - -def test_wire_values_preserve_decimal_precision_and_normalize_aware_times(): - offset = timezone(timedelta(hours=-8)) - event = CostEvent( - event_id="cost:attempt", - usage_event_id="outbound:attempt", - occurred_at=datetime(2025, 12, 31, 16, tzinfo=offset), - dimensions=(_TEAM,), - charges=(_SPEND,), - snapshot=_snapshot(), - ) - wire = event.model_dump(mode="json") - assert wire["occurred_at"] == "2026-01-01T00:00:00Z" - assert wire["charges"] == [{"name": "spend", "value": "0.1234567890123456789", "unit": "USD"}] - assert event.occurred_at == _START - assert CostEvent.model_validate_json(event.model_dump_json()).charges[0].value == _SPEND.value - - -@pytest.mark.parametrize( - ("name", "value", "unit", "message"), - [ - (MetricName.SPEND, Decimal(1), "count", "SPEND requires USD"), - (MetricName.REQUESTS, Decimal(1), "USD", "SPEND requires USD"), - (MetricName.PROMPT_TOKENS, Decimal("1.5"), "count", "integral"), - (MetricName.SPEND, Decimal("NaN"), "USD", "finite"), - (MetricName.SPEND, Decimal("Infinity"), "USD", "finite"), - (MetricName.SPEND, Decimal(-1), "USD", "greater than or equal"), - (MetricName.SPEND, 0.1, "USD", "instance of Decimal"), - ], -) -def test_metrics_reject_wrong_units_nonfinite_values_and_binary_float_inputs(name, value, unit, message): - with pytest.raises(ValidationError, match=message): - Metric(name=name, value=value, unit=unit) - - -@pytest.mark.parametrize("metrics", [(_TOKEN,), (_SPEND,), ()]) -def test_inbound_activity_cannot_own_usage_or_cost_or_omit_request_count(metrics): - with pytest.raises(ValidationError, match="Inbound events"): - _inbound(metrics) - - -@pytest.mark.parametrize( - ("outcome", "metric"), - [ - (Outcome.FAILED, Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(1))), - (Outcome.SUCCEEDED, Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(1))), - ], -) -def test_inbound_outcome_counts_cannot_contradict_outcome(outcome, metric): - with pytest.raises(ValidationError, match="count must match outcome"): - _inbound((_REQUEST, metric), outcome) - - -@pytest.mark.parametrize("metrics", [(_REQUEST,), (_SPEND,)]) -def test_outbound_attempts_cannot_own_caller_counts_or_charges(metrics): - with pytest.raises(ValidationError, match="Outbound events own only"): - _outbound(metrics) - - -@pytest.mark.parametrize( - ("changes", "message"), - [ - ({"started_at": datetime(2026, 1, 1)}, "timezone"), - ({"ended_at": _START - timedelta(seconds=1)}, "End must follow start"), - ({"retry_of_attempt_id": "attempt"}, "cannot retry itself"), - ({"dimensions": (_TEAM, _TEAM)}, "Duplicate dimension values"), - ({"dimensions": (_TEAM, Dimension(name=DimensionName.TEAM, value="other"))}, "Only tag and model access group"), - ( - { - "metrics": ( - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(5)), - Metric(name=MetricName.REASONING_TOKENS, value=Decimal(6)), - ) - }, - "reasoning_tokens cannot exceed completion_tokens", - ), - ( - { - "metrics": ( - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(5)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(6)), - ) - }, - "cache_creation_input_tokens_above_1hr cannot exceed cache_creation_input_tokens", - ), - ({"event_id": " "}, "pattern"), - ({"metrics": (_TOKEN, _TOKEN)}, "Duplicate metric names"), - ], -) -def test_outbound_events_reject_invalid_identity_time_and_duplicate_attribution(changes, message): - with pytest.raises(ValidationError, match=message): - _outbound(**changes) - - -def test_unknown_and_known_zero_usage_are_distinct_and_values_are_immutable(): - unknown = _outbound(None) - zero = _outbound((Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(0)),)) - assert unknown.metrics is None - assert zero.metrics is not None and zero.metrics[0].value == 0 - assert _inbound().started_at == _inbound().ended_at - with pytest.raises(ValidationError, match="frozen"): - zero.attempt_id = "different" - with pytest.raises(ValidationError, match="frozen"): - zero.metrics[0].value = Decimal(10) - - -@pytest.mark.parametrize("charges", [(), (_TOKEN,), (_SPEND, _SPEND)]) -def test_cost_events_require_one_authoritative_spend(charges): - with pytest.raises(ValidationError, match=r"authoritative USD spend|Duplicate metric"): - CostEvent( - event_id="cost", - usage_event_id="usage", - occurred_at=_START, - dimensions=(), - charges=charges, - snapshot=_snapshot(), - ) - - -def test_cost_event_cannot_reuse_its_usage_identity(): - with pytest.raises(ValidationError, match="distinct identities"): - CostEvent( - event_id="usage", - usage_event_id="usage", - occurred_at=_START, - dimensions=(), - charges=(_SPEND,), - snapshot=_snapshot(), - ) - - -@pytest.mark.parametrize("granularity", [Granularity.DAY, None]) -def test_daily_and_whole_interval_queries_preserve_filter_and_group_semantics(granularity): - query = _query(granularity=granularity) - assert query.granularity is granularity - assert query.filters[0].values == ("one", "two") - assert query.group_by == (DimensionName.MODEL,) - assert query.start == _START and query.end == _END - - -@pytest.mark.parametrize( - ("changes", "message"), - [ - ({"end": _START}, "End must follow start"), - ({"start": _START + timedelta(hours=1)}, "UTC day boundaries"), - ({"end": _END + timedelta(microseconds=1)}, "UTC day boundaries"), - ({"granularity": "hour"}, "day"), - ({"group_by": (DimensionName.TEAM, DimensionName.TEAM)}, "Duplicate group-by"), - ({"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one",)),) * 2}, "one IN filter"), - ({"filters": ({"name": "team", "values": ()},)}, "at least 1"), - ], -) -def test_queries_reject_unsupported_or_ambiguous_ranges_and_filters(changes, message): - with pytest.raises(ValidationError, match=message): - _query(**changes) - - -def test_rows_allow_adjacent_buckets_and_overlapping_distinct_groups(): - first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM,), value=Decimal(1)) - second = AggregateRow(start=_END, end=_END + timedelta(days=1), dimensions=(_TEAM,), value=Decimal(2)) - other = AggregateRow(start=_START, end=_END, dimensions=(), value=Decimal(3)) - rows = (second, other, first) - assert validate_aggregate_rows(rows) is rows - assert validate_aggregate_rows(()) == () - - -def test_rows_reject_overlap_within_a_group_regardless_of_dimension_order(): - model = Dimension(name=DimensionName.MODEL, value="model") - first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM, model), value=Decimal(1)) - overlapping = AggregateRow(start=_START, end=_END, dimensions=(model, _TEAM), value=Decimal(2)) - with pytest.raises(ValueError, match="overlap within a dimension group"): - validate_aggregate_rows((first, overlapping)) - - -def test_rows_reject_empty_intervals(): - with pytest.raises(ValidationError, match="End must follow start"): - AggregateRow(start=_START, end=_START, dimensions=(), value=Decimal(0)) - - -def test_admission_decision_preserves_a_denial_reason(): - decision = AdmissionDecision(allowed=False, reason="Team budget exhausted") - assert not decision.allowed - assert decision.reason == "Team budget exhausted" - - -def _snapshot(**changes: object) -> PricingSnapshot: - return PricingSnapshot.model_validate( - { - "calculator_version": "1.80.0", - "model": "gpt-4o", - "provider": "openai", - "pricing_source": "cost_map", - "priced_model_key": "openai/gpt-4o", - "multipliers": (Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")),), - "rates": (Rate(measure="input", unit="token", usd_per_unit=Decimal("0.0000025")),), - "usage": (PricedQuantity(measure="input", unit="token", value=_TOKEN.value),), - "input_cost": Decimal("0.000025"), - "total_cost": _SPEND.value, - **changes, - } - ) - - -def test_cost_event_carries_a_snapshot_whose_total_matches_the_spend(): - event = CostEvent( - event_id="cost:attempt", - usage_event_id="outbound:attempt", - occurred_at=_START, - dimensions=(_TEAM,), - charges=(_SPEND,), - snapshot=_snapshot(), - ) - assert event.snapshot is not None - assert event.snapshot.total_cost == _SPEND.value - assert event.snapshot.multipliers[0].factor == Decimal("1.1") - with pytest.raises(ValidationError, match="Snapshot total must equal the authoritative spend"): - CostEvent( - event_id="cost:attempt", - usage_event_id="outbound:attempt", - occurred_at=_START, - dimensions=(_TEAM,), - charges=(_SPEND,), - snapshot=_snapshot(total_cost=Decimal("0.5")), - ) - - -@pytest.mark.parametrize( - ("changes", "message"), - ( - ( - { - "rates": ( - Rate(measure="input", unit="token", usd_per_unit=Decimal(1)), - Rate(measure="input", unit="token", usd_per_unit=Decimal(2)), - ) - }, - "Duplicate rate measures and units", - ), - ({"usage": (_SPEND,)}, "PricedQuantity"), - ({"usage": (_REQUEST,)}, "PricedQuantity"), - ({"pricing_source": "guess"}, "pricing_source"), - ({"multipliers": ({"name": "x", "value": "y", "factor": Decimal(0)},)}, "factor"), - ({"schema_version": 2}, "schema_version"), - ( - {"additional_costs": ({"name": "fee", "amount": Decimal(1)}, {"name": "fee", "amount": Decimal(2)})}, - "Duplicate additional cost names", - ), - ), -) -def test_snapshots_reject_ambiguous_rates_wrong_usage_and_unknown_sources(changes: Mapping[str, object], message: str): - with pytest.raises(ValidationError, match=message): - _snapshot(**changes) - - -@pytest.mark.parametrize( - ("unit", "value"), - [("token", "10"), ("second", "1.25"), ("character", "25"), ("image", "2"), ("request", "1")], -) -def test_snapshot_preserves_matching_rates_and_quantities_for_each_billing_unit(unit: PricingUnit, value: str) -> None: - quantity: Final = PricedQuantity(measure="input", unit=unit, value=Decimal(value)) - rate: Final = Rate(measure="input", unit=unit, usd_per_unit=Decimal("0.003")) - amount: Final = quantity.value * rate.usd_per_unit - snapshot: Final = _snapshot(rates=(rate,), usage=(quantity,), total_cost=amount) - restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) - assert restored == snapshot - assert restored.usage[0].value * restored.rates[0].usd_per_unit == amount - - -@pytest.mark.parametrize( - ("usage", "message"), - [ - ((), "Every rate requires"), - (({"measure": "output", "unit": "token", "value": Decimal(10)},), "Every rate requires"), - (({"measure": "input", "unit": "second", "value": Decimal(10)},), "Every rate requires"), - ( - ({"measure": "input", "unit": "token", "value": Decimal(10)},) * 2, - "Duplicate priced quantity", - ), - ], -) -def test_snapshot_rejects_missing_mismatched_or_duplicate_priced_quantities( - usage: tuple[Mapping[str, object], ...], message: str -) -> None: - with pytest.raises(ValidationError, match=message): - _snapshot(usage=usage) - - -@pytest.mark.parametrize("unit", ["token", "character", "image", "request"]) -def test_priced_counts_remain_integral(unit: PricingUnit) -> None: - with pytest.raises(ValidationError, match="integral"): - PricedQuantity(measure="input", unit=unit, value=Decimal("1.25")) - - -@pytest.mark.parametrize( - ("value", "message"), - [ - (Decimal(-1), "greater than or equal"), - (Decimal("NaN"), "finite"), - (Decimal("Infinity"), "finite"), - (1.25, "instance of Decimal"), - ], -) -def test_pricing_values_reject_invalid_amounts(value: Decimal | float, message: str) -> None: - with pytest.raises(ValidationError, match=message): - PricedQuantity.model_validate({"measure": "input", "unit": "second", "value": value}) - with pytest.raises(ValidationError, match=message): - Rate.model_validate({"measure": "input", "unit": "second", "usd_per_unit": value}) - - -@pytest.mark.parametrize("snapshot_fields", [{}, {"snapshot": None}]) -def test_cost_event_cannot_omit_its_derivation(snapshot_fields: Mapping[str, object]) -> None: - with pytest.raises(ValidationError, match="snapshot"): - CostEvent.model_validate( - { - "event_id": "cost", - "usage_event_id": "usage", - "occurred_at": _START, - "dimensions": (), - "charges": (_SPEND,), - **snapshot_fields, - } - ) - - -def test_upstream_charge_snapshot_preserves_reported_cost_without_inventing_rates() -> None: - snapshot: Final = _snapshot( - pricing_source="upstream_header", rates=(), usage=(), upstream_reported_cost=_SPEND.value - ) - restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) - assert restored.upstream_reported_cost == restored.total_cost == _SPEND.value - assert restored.rates == restored.usage == () - - -def test_subset_token_counts_round_trip_without_changing_totals() -> None: - metrics: Final = ( - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(10)), - Metric(name=MetricName.REASONING_TOKENS, value=Decimal(3)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(8)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(2)), - ) - event: Final = _outbound(metrics=metrics) - assert OutboundUsageEvent.model_validate_json(event.model_dump_json()).metrics == metrics - - -@pytest.mark.parametrize( - "schema_path", - [ - "schema.prisma", - "litellm/proxy/schema.prisma", - "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", - ], -) -@pytest.mark.parametrize( - ("enum_name", "python_enum"), - [ - ("AccountingEventKind", EventKind), - ("AccountingDeliveryState", DeliveryState), - ("AccountingDeliveryTarget", DeliveryTarget), - ], -) -def test_prisma_accepts_exactly_the_public_accounting_enum_values( - schema_path: str, enum_name: str, python_enum: type[Enum] -) -> None: - schema: Final = (Path(__file__).resolve().parents[3] / schema_path).read_text() - declaration: Final = re.search(r"enum " + enum_name + r"\s*\{([^}]+)\}", schema) - assert declaration is not None, f"Missing persisted enum {enum_name}" - assert frozenset(declaration[1].split()) == frozenset(member.value for member in python_enum) - - -def test_one_event_carries_every_matched_model_access_group(): - groups = ( - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="gpt-group"), - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="claude-group"), - ) - assert _outbound(dimensions=(_TEAM, *groups)).dimensions == (_TEAM, *groups) - with pytest.raises(ValidationError, match="Duplicate dimension values"): - _outbound(dimensions=(groups[0], groups[0])) - - -def test_snapshots_reject_duplicate_multipliers_and_contradictory_token_counts(): - multiplier = Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")) - with pytest.raises(ValidationError, match="Duplicate multiplier names"): - _snapshot(multipliers=(multiplier, multiplier)) - with pytest.raises(ValidationError, match="reasoning_tokens cannot exceed completion_tokens"): - _snapshot( - usage=( - PricedQuantity(measure="input", unit="token", value=_TOKEN.value), - PricedQuantity(measure="completion_tokens", unit="token", value=Decimal(5)), - PricedQuantity(measure="reasoning_tokens", unit="token", value=Decimal(6)), - ) - ) - - -def test_agent_live_spend_is_a_delivery_target(): - assert DeliveryTarget.AGENT_SPEND.value == "LiteLLM_AgentsTable" diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 16b8cf3f0e1..0f85f1e3993 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -22,7 +22,6 @@ from typing import Final, Literal import httpx import pytest import respx -from openai.types.batch import BatchRequestCounts import litellm @@ -2028,61 +2027,6 @@ def test_extract_credentials_keeps_openai_workload_identity_params(): } -# --------------------------------------------------------------------------- # -# batch_cost_is_final -# --------------------------------------------------------------------------- # - - -def _retrieved_batch( - status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None -) -> LiteLLMBatch: - return LiteLLMBatch( - id="batch_abc", - completion_window="24h", - created_at=1, - endpoint="/v1/chat/completions", - input_file_id="file-in", - object="batch", - status="validating", - output_file_id=output_file_id, - request_counts=counts, - ).model_copy(update={"status": status}) - - -class TestBatchCostIsFinal: - """Every retrieve of one batch writes the same spend row, so the first retrieve - that prices it decides the row for good. A poll before the output exists must - therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" - - @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) - def test_in_flight_batch_is_not_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status)) is False - - @pytest.mark.parametrize("status", ["completed", "complete"]) - def test_completed_with_output_is_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True - - def test_completed_without_output_and_unknown_counts_is_not_final(self): - assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False - - def test_completed_without_output_and_zero_counts_is_not_final(self): - counts = BatchRequestCounts(total=0, completed=0, failed=0) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - - def test_completed_without_output_but_successful_lines_is_not_final(self): - counts = BatchRequestCounts(total=2, completed=2, failed=0) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - - @pytest.mark.parametrize("status", ["completed", "complete"]) - def test_completed_without_output_and_every_line_failed_is_final(self, status): - counts = BatchRequestCounts(total=2, completed=0, failed=2) - assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True - - @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) - def test_other_terminal_statuses_are_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status)) is True - - def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: @@ -2965,7 +2909,7 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( @@ -3060,7 +3004,6 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( wrong token sum, or mispriced model fails this test. """ from litellm.batches.batch_utils import handle_completed_batch - from litellm.types.utils import LiteLLMBatch batch = LiteLLMBatch( id="batch-real-cost-123", @@ -3105,7 +3048,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( @@ -3187,7 +3130,7 @@ async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown card renders blank for every batch. Regression for the split being dropped.""" from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes mock_batch = LiteLLMBatch( id="batch-breakdown-1", @@ -3243,7 +3186,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch, SpecialEnums + from litellm.types.utils import CallTypes, SpecialEnums # Create a proper unified file ID by encoding the correct prefix unified_id_str = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}:test_file_789;unified_id:batch-789" @@ -3322,7 +3265,7 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 2d02468fd8c..19aefb430c7 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -484,55 +484,6 @@ async def test_bridged_text_completion_is_logged_as_the_text_completion_it_serve assert logged.payload["response"]["choices"][0]["text"] == "pong" -@pytest.mark.asyncio -async def test_a_bridged_chat_response_waits_on_the_usage_ack_like_any_chat_response( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - logging_obj: Final = Logging( - model="openai/gpt-5.3-codex", - messages=[], - stream=False, - call_type="acompletion", - start_time=datetime.datetime.now(), - litellm_call_id="bridged-chat", - function_id="bridged-chat", - ) - with respx.mock(assert_all_called=True) as router: - router.post(url__startswith="https://api.openai.com/v1/responses").mock( - return_value=httpx.Response(200, json=_searched_responses_body("openai/gpt-5.3-codex")) - ) - served: Final = await litellm.acompletion( - model="openai/gpt-5.3-codex", - messages=[{"role": "user", "content": "ping"}], - api_key="test-key", - litellm_logging_obj=logging_obj, - ) - response: Final = asyncio.create_task(capture_proxy_response(logging_obj, served)) - try: - await database.entered.wait() - assert not response.done() - assert database.committed == () - database.release.set() - await response - assert "bridged-chat:inbound" in database.committed - finally: - database.release.set() - await runtime.stop(1) - await response - - def _served_message_start_id(events: Sequence[bytes]) -> object: lines: Final = b"".join(events).decode().splitlines() starts: Final = tuple( @@ -1905,6 +1856,51 @@ def test_transform_request_system_only_message_maps_to_system_input_item(): assert not result.get("instructions") +_JSON_SYSTEM_PROMPT: Final = "You output JSON objects." +_PLAIN_USER_ITEM: Final = { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "List three colors."}], +} +_COLORS_SCHEMA: Final = { + "type": "json_schema", + "json_schema": {"name": "colors", "schema": {"type": "object", "properties": {}}}, +} + + +@pytest.mark.parametrize( + ("response_format", "expected_input", "expected_instructions"), + [ + ( + {"type": "json_object"}, + [ + {"type": "message", "role": "system", "content": [{"type": "input_text", "text": _JSON_SYSTEM_PROMPT}]}, + _PLAIN_USER_ITEM, + ], + None, + ), + (_COLORS_SCHEMA, [_PLAIN_USER_ITEM], _JSON_SYSTEM_PROMPT), + ], + ids=["json_object_keeps_system_in_input", "json_schema_keeps_system_in_instructions"], +) +def test_transform_request_json_object_keeps_system_prompt_where_the_json_word_check_reads_it( + response_format, expected_input, expected_instructions +): + result = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5.4-mini", + messages=[ + {"role": "system", "content": _JSON_SYSTEM_PROMPT}, + {"role": "user", "content": "List three colors."}, + ], + optional_params={"response_format": response_format}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert (result["input"], result.get("instructions")) == (expected_input, expected_instructions), result + + def test_transform_request_single_char_keys_not_matched(): """Test that single-character keys are not incorrectly matched to 'metadata' or 'previous_response_id' diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 789d59bd613..0e2813e0a0d 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -6,8 +6,10 @@ import copy import json import re import threading +import uuid from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import wait +from datetime import datetime from functools import partial from types import MappingProxyType from typing import Final, TypeAlias @@ -132,6 +134,7 @@ _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( "https://openrouter.ai/api/alpha/decisions", "typesafe/jev-1.13", ), + ("empiriolabs", "empiriolabs/aplomb-1", "https://api.empiriolabs.ai/v1/decisions", "aplomb-1"), ) _OPENAI_RESPONSE: Final[Mapping[str, object]] = { @@ -299,6 +302,43 @@ async def test_router_dispatches_typesafe_decisions_without_api_base( assert response.answers["sentiment"].choice == "positive" +@pytest.mark.asyncio +async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_finishes( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "jev", + "litellm_params": {"model": "typesafe/jev-latest", "api_key": "k"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("jev-a", "jev-b") + ], + routing_strategy="least-busy", + ) + + async def call_through_proxy_logging() -> object: + logging_obj, _ = litellm.utils.function_setup( + original_function="adecisions", + rules_obj=litellm.utils.Rules(), + start_time=datetime(2026, 1, 1), + model="jev", + litellm_call_id=str(uuid.uuid4()), + ) + response: Final = await router.adecisions( + model="jev", + input="Is this a defect?", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + litellm_logging_obj=logging_obj, + ) + await _drain_logging_worker() + return response._hidden_params["model_id"] + + assert (await call_through_proxy_logging(), await call_through_proxy_logging()) == ("jev-a", "jev-a") + + def test_decisions_uses_the_same_wire_contract_for_sync_calls(respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) @@ -1549,6 +1589,7 @@ _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset( LlmProviders.LM_STUDIO.value, LlmProviders.OLLAMA.value, LlmProviders.OLLAMA_CHAT.value, + LlmProviders.EMPIRIOLABS.value, } ) _DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple( @@ -3449,6 +3490,147 @@ async def test_fireworks_ai_bills_cached_input_tokens_at_the_cache_read_rate( assert recording_logger.standard_logging_object["response_cost"] < 300 * float(prices["input_cost_per_token"]) +_EMPIRIOLABS_MODEL: Final = "empiriolabs/aplomb-1" +_EMPIRIOLABS_DEFAULT_URL: Final = "https://api.empiriolabs.ai/v1/decisions" +_EMPIRIOLABS_STATE: Final = "I was charged twice for order 4411 and want a refund today." +_EMPIRIOLABS_QUESTIONS: Final[Mapping[str, object]] = MappingProxyType( + {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}} +) +_EMPIRIOLABS_RESPONSE: Final[Mapping[str, object]] = { + "model": "aplomb-1", + "answers": {"is_refund": {"type": "noul", "noul": 0.97, "confidence": 0.94}}, + "usage": {"input_tokens": 24, "output_tokens": 0}, +} + + +@pytest.mark.parametrize( + ("api_base", "environment_api_base", "url"), + ( + (None, None, _EMPIRIOLABS_DEFAULT_URL), + (None, "https://environment-gateway.example/v1", "https://environment-gateway.example/v1/decisions"), + ( + "https://gateway.example/v1/", + "https://environment-gateway.example/v1", + "https://gateway.example/v1/decisions", + ), + ("https://gateway.example", None, "https://gateway.example/v1/decisions"), + ), + ids=("default_base", "environment_base", "explicit_base_over_environment", "explicit_base_without_v1"), +) +def test_empiriolabs_posts_the_systemone_body_to_v1_decisions_on_the_resolved_base( + api_base: str | None, + environment_api_base: str | None, + url: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + if environment_api_base is not None: + monkeypatch.setenv("EMPIRIOLABS_API_BASE", environment_api_base) + route: Final = respx_mock.post(url).respond(json=_EMPIRIOLABS_RESPONSE) + + response: Final = litellm.systemone( + model=_EMPIRIOLABS_MODEL, + state=_EMPIRIOLABS_STATE, + questions=_EMPIRIOLABS_QUESTIONS, + api_base=api_base, + api_key="caller-key", + ) + + assert route.call_count == 1 + assert json.loads(route.calls[0].request.content) == { + "model": "aplomb-1", + "state": _EMPIRIOLABS_STATE, + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + } + assert response.answers["is_refund"] == NoulAnswer(type="noul", noul=0.97, confidence=0.94) + assert response._hidden_params["custom_llm_provider"] == "empiriolabs" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_key", "authorization"), + ((None, "Bearer environment-key"), ("caller-key", "Bearer caller-key")), + ids=("environment_key", "caller_key_over_environment"), +) +async def test_empiriolabs_sends_the_caller_key_or_else_the_environment_key( + api_key: str | None, + authorization: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + monkeypatch.setenv("EMPIRIOLABS_API_KEY", "environment-key") + route: Final = respx_mock.post(_EMPIRIOLABS_DEFAULT_URL).respond(json=_EMPIRIOLABS_RESPONSE) + + response: Final = await litellm.adecisions( + model=_EMPIRIOLABS_MODEL, + input=_EMPIRIOLABS_STATE, + questions=[{"type": "predicate", "name": "is_refund", "instructions": "Is the customer asking for a refund?"}], + api_key=api_key, + ) + + assert route.call_count == 1 + assert route.calls[0].request.headers["authorization"] == authorization + assert isinstance(response, OpenAIDecisionResponse) + assert [answer.model_dump(mode="json") for answer in response.answers] == [ + {"type": "predicate", "name": "is_refund", "probability": 0.97} + ] + + +def test_empiriolabs_without_a_key_is_refused_before_http( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_KEY", raising=False) + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + + with pytest.raises(litellm.AuthenticationError, match="Missing API key for Decisions provider 'empiriolabs'"): + litellm.systemone(model=_EMPIRIOLABS_MODEL, state=_EMPIRIOLABS_STATE, questions=_EMPIRIOLABS_QUESTIONS) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + ("status", "error_type", "code", "exception"), + ( + (401, "authentication_error", "invalid_api_key", litellm.AuthenticationError), + (402, "insufficient_credits", "insufficient_credits", litellm.PaymentRequiredError), + (404, "not_found_error", "model_not_found", litellm.NotFoundError), + (422, "invalid_request_error", "invalid_request", litellm.BadRequestError), + (429, "rate_limit_error", "rate_limit_exceeded", litellm.RateLimitError), + (503, "service_unavailable", "model_unavailable", litellm.ServiceUnavailableError), + ), +) +def test_empiriolabs_upstream_errors_keep_their_status_and_message( + status: int, + error_type: str, + code: str, + exception: type[Exception], + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + message: Final = f"upstream refused the decision with {code}" + route: Final = respx_mock.post(_EMPIRIOLABS_DEFAULT_URL).respond( + status, + json={"error": {"message": message, "type": error_type, "code": code}}, + headers={"retry-after": "1"}, + ) + + with pytest.raises(exception, match=message) as caught: + litellm.systemone( + model=_EMPIRIOLABS_MODEL, + state=_EMPIRIOLABS_STATE, + questions=_EMPIRIOLABS_QUESTIONS, + api_key="caller-key", + num_retries=0, + ) + + assert getattr(caught.value, "status_code") == status + assert route.call_count == 1 + + _BOX_QUESTION_REQUESTS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( ( "decisions", diff --git a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py index 453fd28636d..3f3e2e0d465 100644 --- a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py @@ -512,6 +512,20 @@ def _provider_auth_failure_kwargs() -> dict[str, object]: } +def _deployment_entra_rejection_kwargs() -> dict[str, object]: + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + try: + raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from ( + EntraRejectionError(status_code=401, message=reason) + ) + except litellm.AuthenticationError as raised: + return { + key: value for key, value in _provider_auth_failure_kwargs().items() if key not in CALLED_THE_PROVIDER + } | {"exception": raised} + + def _metric_values(name: str, model_id: str) -> list[float]: return [ sample.value @@ -541,6 +555,20 @@ class TestProviderAuthFailures: assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0] + @pytest.mark.asyncio + async def test_log_failure_event_records_entra_rejecting_the_deployment_credential_before_the_call( + self, prometheus_logger + ): + await prometheus_logger.async_log_failure_event( + kwargs=_deployment_entra_rejection_kwargs(), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert _metric_values("litellm_llm_api_failed_requests_metric_total", "deployment-1") == [1.0] + assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] + def test_deployment_failure_metrics_record_a_provider_401(self, prometheus_logger): prometheus_logger.set_llm_deployment_failure_metrics(_provider_auth_failure_kwargs()) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 609b28270d7..3964f472bd9 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4310,6 +4310,37 @@ def test_azure_gpt_6_foundry_price_sheet(_local_model_cost_map, model_base): assert azure_eu_info[field] == pytest.approx(1.2 * base) +@pytest.mark.parametrize("model", ["azure/gpt-6-luna", "azure/gpt-6-luna-2026-09-22"]) +@pytest.mark.parametrize( + ("prompt_tokens", "input_rate", "cache_read_rate", "output_rate"), + [ + (100_000, 5e-08, 5e-09, 2.5e-07), + (300_000, 1e-07, 1e-08, 3.75e-07), + ], +) +def test_generic_cost_per_token_azure_gpt_6_luna_flex_tiers( + _local_model_cost_map, model, prompt_tokens, input_rate, cache_read_rate, output_rate +): + """azure/gpt-6-luna bills Azure's Flex Global rates, half of Standard, with the long context step past 272K.""" + cached_tokens = 20_000 + completion_tokens = 1_000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="azure", + service_tier="flex", + ) + expected_prompt = (prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate + assert prompt_cost == pytest.approx(expected_prompt) + assert completion_cost == pytest.approx(completion_tokens * output_rate) + + @pytest.mark.parametrize("region_prefix", ["azure/", "azure/us/", "azure/eu/"]) def test_azure_gpt_5_6_alias_matches_sol_pricing(_local_model_cost_map, region_prefix): """The bare gpt-5.6 alias routes to GPT-5.6 Sol, so every Azure region must bill the diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 9da2632943c..06e673c3e68 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1230,3 +1230,27 @@ def test_vertex_gemini_deployment_priced_web_search_stays_per_prompt( def test_built_in_tool_lookups_ignore_malformed_request_values(kwargs: dict[str, object]) -> None: assert StandardBuiltInToolCostTracking.get_web_search_options(kwargs) is None assert StandardBuiltInToolCostTracking.get_file_search_tool_call(kwargs) is None + + +def test_per_call_web_search_path_bills_the_count_the_usage_reports(local_model_cost_map): + """ + A provider with no web search calculator of its own (an OpenAI-compatible JSON provider) that + reports server_tool_use.web_search_requests is billed that many per-call searches, not one. + """ + from litellm.types.utils import ServerToolUse, Usage + + model = "empiriolabs/qwen3-7-plus" + per_search = litellm.model_cost[model]["search_context_cost_per_query"]["search_context_size_medium"] + response = ModelResponse(model="qwen3-7-plus") + usage = Usage( + prompt_tokens=4, completion_tokens=5, total_tokens=9, server_tool_use=ServerToolUse(web_search_requests=3) + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=usage, + custom_llm_provider="empiriolabs", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(3 * per_search) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index aab57b87a90..d7696c00dee 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -327,7 +327,7 @@ EMPTY_ASSISTANT = pytest.mark.parametrize( ) -def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_linear_time(): +def test_flagged_placement_turns_reminders_before_a_long_run_of_unanswered_tool_search_uses_into_user_turns(): search_uses: Final = tuple( { "role": "assistant", @@ -349,12 +349,9 @@ def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_ {"role": "user", "content": "end"}, ] - started: Final = time.perf_counter() placed: Final = place_mid_conversation_system(messages, supports_mid_conversation_system=True) - elapsed: Final = time.perf_counter() - started assert set(_roles(placed)[:21]) == {"user"} - assert elapsed < 1.0 def test_flagged_placement_merges_every_system_run_that_shares_a_user_turn(): diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index cd37d635c43..45461b50761 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -5,14 +5,22 @@ from io import BytesIO from typing import Dict, Final, List import pytest +from openai.types.batch import BatchRequestCounts import litellm from litellm.litellm_core_utils.batch_utils import ( + POLLED_BATCH_COST_IS_FINAL_KWARG, + POLLED_BATCH_ID_KWARG, + ChargeOnceJob, InMemoryFile, + batch_cost_is_final, + defers_batch_cost, get_router_metadata_variable_name, parse_jsonl_with_embedded_newlines, replace_model_in_jsonl, ) +from litellm.types.utils import LiteLLMBatch +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG, encode_video_id_with_provider from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -382,30 +390,241 @@ def setup_and_teardown(): asyncio.set_event_loop(None) -@pytest.mark.parametrize( - ("call_type", "cost", "eligible"), - (("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False), - ("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)), -) -def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: +def _batch(status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None): from openai.types import Batch - from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once - - batch: Final = Batch( - id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", - input_file_id="file-budget", object="batch", status="completed", + return Batch( + id="job-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + input_file_id="file-budget", object="batch", status=status, output_file_id=output_file_id, + request_counts=counts, ) - assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None) - assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None -def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: +def _completed_batch(): + return _batch("completed", output_file_id="file-out") + + +def _completed_video(): + from litellm.types.videos.main import VideoObject + + return VideoObject(id="job-budget", object="video", status="completed") + + +@pytest.mark.parametrize( + ("call_type", "polled", "cost", "kind"), + ( + ("aretrieve_batch", _completed_batch, 0.25, "batch"), + ("retrieve_batch", _completed_batch, 0.25, "batch"), + ("avideo_status", _completed_video, 0.25, "video"), + ("video_status", _completed_video, 0.25, "video"), + ("acompletion", _completed_batch, 0.25, None), + ("avideo_status", _completed_batch, 0.25, None), + ("aretrieve_batch", _completed_video, 0.25, None), + ("avideo_generation", _completed_video, 0.25, None), + ("aretrieve_batch", _completed_batch, 0, None), + ("avideo_status", _completed_video, 0, None), + ("avideo_status", _completed_video, -0.25, None), + ), +) +def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled, cost: float, kind) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + expected: Final = None if kind is None else ChargeOnceJob(kind=kind, job_id="job-budget") + kwargs: Final = { + "call_type": call_type, + POLLED_VIDEO_STATUS_KWARG: "completed", + "standard_logging_object": {"response_cost": cost}, + } + assert job_to_charge_once(kwargs, polled(), cost) == expected + assert job_to_charge_once(kwargs, {"id": "job-budget"}, cost) is None + + +@pytest.mark.parametrize( + ("call_type", "polled", "guardrail_only_poll", "kind", "repeat_cost"), + ( + ("aretrieve_batch", _completed_batch, lambda: _batch("in_progress"), "batch", 0.0), + ("avideo_status", _completed_video, _completed_video, "video", 0.01), + ), +) +def test_only_the_jobs_own_cost_is_charged_once_since_a_guardrail_charges_every_poll( + call_type: str, polled, guardrail_only_poll, kind, repeat_cost +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + kwargs: Final = { + "call_type": call_type, + POLLED_VIDEO_STATUS_KWARG: "completed", + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + assert job_to_charge_once(kwargs, guardrail_only_poll(), 0.01) is None + assert job_to_charge_once(kwargs, polled(), 0.43) == ChargeOnceJob( + kind=kind, job_id="job-budget", repeat_poll_cost=repeat_cost + ) + + +@pytest.mark.parametrize("status", ("completed", "failed", "cancelled", "expired")) +def test_a_finished_batch_that_cost_nothing_itself_is_still_charged_once_for_its_retrieves_guardrail( + status: str, +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + kwargs: Final = { + "call_type": "aretrieve_batch", + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + finished_batch: Final = _batch(status, output_file_id="file-out") + assert job_to_charge_once(kwargs, finished_batch, 0.01) == ChargeOnceJob(kind="batch", job_id="job-budget") + + +def test_a_completed_batch_still_waiting_on_its_output_does_not_take_the_charge_its_real_cost_needs() -> None: + """Its first retrieve can report completed before the output file exists, and settling on that guardrail-only + retrieve would leave the batch's real cost nothing to charge once the output arrives.""" + from litellm.litellm_core_utils.batch_utils import job_to_charge_once + + kwargs: Final = { + "call_type": "aretrieve_batch", + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + assert job_to_charge_once(kwargs, _batch("completed"), 0.01) is None + assert defers_batch_cost(kwargs, _batch("completed")) + + +def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() -> None: + """A provider may price a running video too, and a custom provider's own id can repeat across deployments.""" + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + from litellm.types.videos.main import VideoObject + + def poll(status: str, polled_id: str) -> ChargeOnceJob | None: + return job_to_charge_once( + {"call_type": "avideo_status", POLLED_VIDEO_STATUS_KWARG: status, "litellm_params": {"video_id": polled_id}}, + VideoObject(id="job-1", object="video", status=status), + 0.42, + ) + + assert poll("in_progress", "video_deployment_a") is None + assert poll("completed", "video_deployment_a") == ChargeOnceJob(kind="video", job_id="video_deployment_a") + assert poll("completed", "video_deployment_b") == ChargeOnceJob(kind="video", job_id="video_deployment_b") + + +@pytest.mark.parametrize("status", ["in_progress", "completed"]) +def test_a_video_poll_is_judged_by_its_logged_status_when_redaction_replaced_the_video(status: str) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + kwargs: Final = { + "call_type": "avideo_status", + POLLED_VIDEO_STATUS_KWARG: status, + "litellm_params": {"video_id": "video_deployment_a"}, + } + expected: Final = ChargeOnceJob(kind="video", job_id="video_deployment_a") if status == "completed" else None + + assert job_to_charge_once(kwargs, {"text": "redacted-by-litellm"}, 0.42) == expected + + +@pytest.mark.parametrize( + ("cost_is_final", "deferred", "expected"), + [(False, True, None), (True, False, ChargeOnceJob(kind="batch", job_id="batch_abc"))], +) +def test_a_batch_retrieve_is_judged_by_its_logged_batch_when_redaction_replaced_it( + cost_is_final: bool, deferred: bool, expected: ChargeOnceJob | None +) -> None: + from litellm.litellm_core_utils.batch_utils import job_to_charge_once + + kwargs: Final = { + "call_type": "aretrieve_batch", + POLLED_BATCH_ID_KWARG: "batch_abc", + POLLED_BATCH_COST_IS_FINAL_KWARG: cost_is_final, + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + redacted: Final = {"text": "redacted-by-litellm"} + + assert defers_batch_cost(kwargs, redacted) is deferred + assert job_to_charge_once(kwargs, redacted, 0.01) == expected + + +def test_a_video_polled_with_or_without_its_base64_padding_is_charged_as_one_job() -> None: + from litellm.litellm_core_utils.batch_utils import job_to_charge_once + + padded: Final = encode_video_id_with_provider("job-1", "custom_video", "deployment-a") + assert padded.endswith("=") + + def charged_job(polled_id: str): + return job_to_charge_once( + {"call_type": "avideo_status", POLLED_VIDEO_STATUS_KWARG: "completed", "litellm_params": {"video_id": polled_id}}, + None, + 0.42, + ) + + assert charged_job(padded.rstrip("=")) == charged_job(padded) + assert charged_job(padded) is not None + + +def test_job_charge_markers_keep_counters_jobs_and_kinds_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache - from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key + first_batch: Final = ChargeOnceJob(kind="batch", job_id="first-job") cache: Final = InMemoryCache() - cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) - assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1 - assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None - assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None + cache.set_cache(job_charged_once_marker_key("provider-spend", first_batch), 1) + assert cache.get_cache(job_charged_once_marker_key("provider-spend", first_batch)) == 1 + assert cache.get_cache(job_charged_once_marker_key("deployment-spend", first_batch)) is None + assert cache.get_cache(job_charged_once_marker_key("provider-spend", ChargeOnceJob("batch", "second-job"))) is None + assert cache.get_cache(job_charged_once_marker_key("provider-spend", ChargeOnceJob("video", "first-job"))) is None + + +def test_batch_charge_marker_keeps_the_key_markers_already_written_use() -> None: + """Markers written before the video kind existed must still block a re-charge after an upgrade.""" + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key + + assert job_charged_once_marker_key("provider-spend", ChargeOnceJob("batch", "first-job")) == ( + "provider-spend:batch:first-job" + ) + + +def _retrieved_batch( + status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None +) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="validating", + output_file_id=output_file_id, + request_counts=counts, + ).model_copy(update={"status": status}) + + +class TestBatchCostIsFinal: + """Every retrieve of one batch writes the same spend row, so the first retrieve + that prices it decides the row for good. A poll before the output exists must + therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" + + @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) + def test_in_flight_batch_is_not_final(self, status): + assert batch_cost_is_final(_retrieved_batch(status)) is False + + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_with_output_is_final(self, status): + assert batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True + + def test_completed_without_output_and_unknown_counts_is_not_final(self): + assert batch_cost_is_final(_retrieved_batch("completed")) is False + + def test_completed_without_output_and_zero_counts_is_not_final(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_but_successful_lines_is_not_final(self): + counts = BatchRequestCounts(total=2, completed=2, failed=0) + assert batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_without_output_and_every_line_failed_is_final(self, status): + counts = BatchRequestCounts(total=2, completed=0, failed=2) + assert batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True + + @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) + def test_other_terminal_statuses_are_final(self, status): + assert batch_cost_is_final(_retrieved_batch(status)) is True diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index 99fa849879c..327031cfaf9 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -9,7 +9,8 @@ resolution (get_model_info) including the shipped rules in the bundled cost map. import json import logging -from collections.abc import Callable +import re +from collections.abc import Callable, Mapping from pathlib import Path from typing import Final @@ -915,6 +916,7 @@ def test_shipped_claude_thinking_rules_backfill_only_anthropic(shipped_cost_map) assert info.get("supports_legacy_thinking") is None assert info.get("max_input_tokens") is None assert match_fill_missing_generalizations("claude-sonnet-4-6", "anthropic") == { + **_claude_baseline()["model_info"], "supports_adaptive_thinking": True, "supports_legacy_thinking": True, "supports_tool_search": True, @@ -961,3 +963,63 @@ def test_shipped_tool_search_rule_fills_mapped_claude_entries_without_flag(shipp assert match_fill_missing_generalizations("claude-opus-5", "bedrock")["supports_tool_search"] is True assert match_fill_missing_generalizations("claude-opus-5", "azure_ai")["supports_tool_search"] is True assert match_fill_missing_generalizations("claude-opus-5", "perplexity") is None + + +CLAUDE_BASELINE_FILL_PROVIDERS: Final = ( + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models", +) + + +def _claude_baseline() -> Mapping[str, object]: + return next(rule for rule in get_fallback_generalization_rules() if rule["name"] == "claude-family-baseline") + + +def _mapped_claude_chat_entries() -> tuple[tuple[str, Mapping[str, object]], ...]: + pattern: Final = re.compile(_claude_baseline()["pattern"], re.IGNORECASE) + return tuple( + (key, entry) + for key, entry in litellm.model_cost.items() + if isinstance(entry, dict) and entry.get("mode") == "chat" and pattern.search(key) is not None + ) + + +def _fields_overriding_baseline(entry: Mapping[str, object], baseline: Mapping[str, object]) -> frozenset[str]: + return frozenset(field for field in entry.keys() & baseline.keys() if entry[field] != baseline[field]) + + +@pytest.mark.parametrize("provider", CLAUDE_BASELINE_FILL_PROVIDERS) +def test_shipped_claude_baseline_fills_only_missing_keys_on_mapped_entries(shipped_cost_map, provider): + """A mapped Claude entry from a first-party host gets each baseline key it leaves out, + while every baseline key the entry sets keeps the entry's own value.""" + baseline: Final = _claude_baseline()["model_info"] + entries: Final = tuple( + (key, entry) for key, entry in _mapped_claude_chat_entries() if entry["litellm_provider"] == provider + ) + assert any(baseline.keys() - entry.keys() for _, entry in entries) + assert any(_fields_overriding_baseline(entry, baseline) for _, entry in entries) + + for key, entry in entries: + info = litellm.get_model_info(key) + assert info["key"] == key + assert {field: info[field] for field in baseline} == { + field: entry.get(field, value) for field, value in baseline.items() + }, key + + +def test_shipped_claude_baseline_leaves_other_providers_entries_unfilled(shipped_cost_map): + baseline: Final = _claude_baseline()["model_info"] + entries: Final = tuple( + (key, entry) + for key, entry in _mapped_claude_chat_entries() + if entry["litellm_provider"] not in CLAUDE_BASELINE_FILL_PROVIDERS + ) + assert any("supports_parallel_function_calling" not in entry for _, entry in entries) + + for key, entry in entries: + info = litellm.get_model_info(key) + missing = baseline.keys() - entry.keys() + assert {field: info.get(field) for field in missing} == dict.fromkeys(missing), key diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index cf9ad4f91a3..e8ea0c04f6f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14453,6 +14453,107 @@ def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_call assert shared_metadata == {"tag": "x"} +@pytest.mark.asyncio +async def test_a_redacted_video_poll_still_tells_the_spend_callbacks_the_video_finished( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + from litellm.types.videos.main import VideoObject + + class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[tuple[dict[str, object], object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events = (*self.events, (dict(kwargs), response_obj)) + + recorder: Final = Recorder() + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + logging_obj: Final = LitellmLogging( + model="custom-video", + messages=[], + stream=False, + call_type="avideo_status", + start_time=datetime.datetime.now(), + litellm_call_id="poll-1", + function_id="poll-1", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}, "video_id": "video_abc"}, optional_params={} + ) + + await logging_obj.async_success_handler(result=VideoObject(id="video_abc", object="video", status="completed")) + + ((kwargs, response_obj),) = recorder.events + assert not isinstance(response_obj, VideoObject) + assert job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="video", job_id="video_abc") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status", "output_file_id", "deferred", "charged"), + [("in_progress", None, True, False), ("completed", "file-out", False, True)], +) +async def test_a_redacted_batch_retrieve_still_tells_the_spend_callbacks_whether_the_batch_cost_is_final( + monkeypatch: pytest.MonkeyPatch, status: str, output_file_id: str | None, deferred: bool, charged: bool +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, defers_batch_cost, job_to_charge_once + from litellm.types.utils import LiteLLMBatch, Usage + + class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[tuple[dict[str, object], object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events = (*self.events, (dict(kwargs), response_obj)) + + recorder: Final = Recorder() + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="aretrieve_batch", + start_time=datetime.datetime.now(), + litellm_call_id="retrieve-1", + function_id="retrieve-1", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await logging_obj.async_success_handler( + result=LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ), + batch_cost=0.42, + batch_usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38), + batch_models=["gpt-4o-mini"], + ) + + ((kwargs, response_obj),) = recorder.events + assert not isinstance(response_obj, LiteLLMBatch) + assert defers_batch_cost(kwargs, response_obj) is deferred + assert (job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="batch", job_id="batch_abc")) is charged + + _FAILED_ATTEMPT: Final = (Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), 0.5) @@ -14493,3 +14594,50 @@ async def test_a_recovered_attempt_bills_a_failed_agentic_follow_up_but_not_a_fa assert own_cost > 0 assert after_a_dropped_stream == pytest.approx(own_cost) assert after_a_failed_follow_up == pytest.approx(own_cost + _FAILED_ATTEMPT[1]) + + +def test_standard_logging_payload_reports_the_per_call_prices_the_request_was_billed_from() -> None: + """A direct call that sets its own prices is logged at those prices, not at the catalog row's rates.""" + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + from litellm.litellm_core_utils.per_call_pricing import per_call_model_info + + rates: Final = {"input_cost_per_token": 3e-06, "output_cost_per_token": 7e-06} + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + logging_obj: Final = LitellmLogging( + model="gpt-5.4-nano", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=now, + litellm_call_id="per-call-prices", + function_id="fn", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-nano", user="", optional_params={}, litellm_params=rates, custom_llm_provider="openai" + ) + logging_obj.set_per_call_pricing(per_call_model_info("gpt-5.4-nano", "openai", rates, "chat")) + response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=6, total_tokens=16), + ) + + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-5.4-nano", + "custom_llm_provider": "openai", + "litellm_params": rates, + "response_cost": logging_obj.response_cost_calculator(result=response), + }, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(10 * 3e-06 + 6 * 7e-06) + model_map_value: Final = payload["model_map_information"]["model_map_value"] + assert model_map_value is not None + assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (3e-06, 7e-06) diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 2fd3f16ab17..491ea5315ff 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -54,9 +54,25 @@ SIGNATURE: Final = "opaque-thought-signature" def _provider_fields() -> dict[str, object]: return { "reasoning_content": SECRET, + "reasoning": SECRET, + "reasoning_text": SECRET, + "reasoning_details": [{"type": "reasoning.text", "text": SECRET, "index": 0, "format": "unknown"}], "thinking_blocks": [{"type": "thinking", "thinking": SECRET}], "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], + "reasoningContent": {"reasoningText": {"text": SECRET}}, + "annotations": [{"type": "file", "file": {"name": "notes.pdf", "content": [{"type": "text", "text": SECRET}]}}], "citations": [{"cited_text": SECRET}], + "citationsContent": [{"citations": [{"location": {}}], "content": [{"text": SECRET}]}], + "citation": {"type": "char_location", "cited_text": SECRET}, + "web_search_calls": [{"type": "web_search_call", "action": {"type": "search", "query": SECRET}}], + "web_search_results": [{"type": "web_search_tool_result", "content": [{"title": SECRET}]}], + "tool_results": [{"type": "web_fetch_tool_result", "content": {"text": SECRET}}], + "code_interpreter_results": [{"stdout": SECRET}], + "search_results": [{"object": "vector_store.search_results.page", "data": [{"content": [{"text": SECRET}]}]}], + "rerank_results": {"results": [{"document": {"text": SECRET}}]}, + "executed_tools": [{"type": "search", "search_results": {"results": [{"content": SECRET}]}}], + "custom_outputs": {"answer": SECRET}, + "tool_search_items": [{"type": "tool_search_output", "tools": [{"description": SECRET}]}], "compaction_blocks": [{"type": "compaction", "content": SECRET, "signature": "compaction-signature"}], "compaction_start": {"type": "compaction", "content": SECRET}, "compaction_delta": {"type": "compaction_delta", "content": SECRET}, diff --git a/tests/unit/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py index 9b29eac208e..958960af27a 100644 --- a/tests/unit/litellm_core_utils/test_url_utils.py +++ b/tests/unit/litellm_core_utils/test_url_utils.py @@ -16,11 +16,25 @@ from litellm.litellm_core_utils.url_utils import ( encode_url_path_segment, encode_url_path_segments, ensure_api_route, + is_https_url_with_host, join_api_base, validate_url, ) +@pytest.mark.parametrize( + ("url", "expected"), + ( + ("https://example.com/path", True), + ("http://example.com/path", False), + ("https:///path", False), + ("https://[", False), + ), +) +def test_is_https_url_with_host(url: str, expected: bool) -> None: + assert is_https_url_with_host(url) is expected + + @pytest.fixture def mock_dns_public(monkeypatch): """Resolve any hostname to 93.184.216.34 (public).""" diff --git a/tests/unit/llms/azure/decisions/test_transformation.py b/tests/unit/llms/azure/decisions/test_transformation.py index b6de6084b9c..3f5e97ee000 100644 --- a/tests/unit/llms/azure/decisions/test_transformation.py +++ b/tests/unit/llms/azure/decisions/test_transformation.py @@ -125,6 +125,29 @@ async def test_default_credential_refreshes_without_opt_in( ) +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("systemone", (False, True)) +async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason( + monkeypatch: pytest.MonkeyPatch, is_async: bool, systemone: bool +) -> None: + from azure.core.exceptions import ClientAuthenticationError + + app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" + credential = Mock(spec=["get_token"]) + credential.get_token.side_effect = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'." + ) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential)) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + await call_decisions(is_async, systemone, {}) + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert app_client_id not in str(excinfo.value) + + @pytest.mark.parametrize("kind", ("service-principal", "oidc", "provider")) async def test_azure_decisions_reuses_existing_azure_credential_resolution( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, kind: str diff --git a/tests/unit/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py index 29006901877..05c0817e19b 100644 --- a/tests/unit/llms/azure/response/test_azure_transformation.py +++ b/tests/unit/llms/azure/response/test_azure_transformation.py @@ -7,15 +7,18 @@ from unittest.mock import MagicMock, patch import httpx import pytest import respx +from openai.types.responses import ResponseReasoningItem +from pydantic import JsonValue, TypeAdapter import litellm from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.azure.responses.o_series_transformation import ( AzureOpenAIOSeriesResponsesAPIConfig, ) from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.azure_ai.responses.transformation import AzureAIAgentResponsesAPIConfig, AzureAIResponsesAPIConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams from litellm.types.router import GenericLiteLLMParams @@ -1547,3 +1550,267 @@ def test_openai_gpt5_resolves_to_responses_config_with_temperature() -> None: config: Final = ProviderConfigManager.get_provider_responses_api_config(provider=LlmProviders.OPENAI, model="gpt-5") assert type(config) is OpenAIResponsesAPIConfig assert "temperature" in config.get_supported_openai_params("gpt-5") + + +@pytest.mark.parametrize( + ("config", "model"), + [ + (AzureOpenAIResponsesAPIConfig(), "gpt-6.1-sol"), + (AzureOpenAIOSeriesResponsesAPIConfig(), "gpt-6.1-sol"), + (AzureAIResponsesAPIConfig(), "gpt-6.1-sol"), + (AzureAIAgentResponsesAPIConfig(), "agents/replay-test"), + ], + ids=["azure", "o_series", "azure_ai", "azure_ai_agent"], +) +@pytest.mark.parametrize("as_model", [False, True], ids=["dict", "pydantic"]) +@pytest.mark.parametrize("drop_params", [True, False], ids=["drop-params", "no-drop-params"]) +def test_azure_reasoning_replay_removes_output_content_only_under_drop_params_without_changing_history( + config: AzureOpenAIResponsesAPIConfig, model: str, as_model: bool, drop_params: bool +) -> None: + replayable_reasoning: Final = { + "id": "rs_saved_reasoning", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Check the file before changing it."}], + "encrypted_content": "opaque-saved-reasoning", + } + saved_reasoning: Final = { + **replayable_reasoning, + "content": [{"type": "reasoning_text", "text": "I should inspect the file first."}], + } + tool_call: Final = { + "type": "function_call", + "id": "fc_read", + "call_id": "call_read", + "name": "read_file", + "arguments": '{"path":"main.py"}', + } + tool_result: Final = {"type": "function_call_output", "call_id": "call_read", "output": "print(1)"} + answer: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "The file prints 1.", "annotations": []}], + } + prompt: Final = {"role": "user", "content": "Continue."} + history: Final = [ + ResponseReasoningItem.model_validate(saved_reasoning) if as_model else saved_reasoning, + tool_call, + tool_result, + answer, + prompt, + ] + original: Final = deepcopy(history) + + request: Final = config.transform_responses_api_request( + model=model, + input=history, + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + + sent_reasoning: Final = replayable_reasoning if drop_params else saved_reasoning + assert request["input"] == [sent_reasoning, tool_call, tool_result, answer, prompt] + assert history == original + + +@pytest.mark.parametrize("as_model", [False, True], ids=["dict", "pydantic"]) +def test_openai_replay_without_drop_params_keeps_reasoning_content_and_message_status(as_model: bool) -> None: + saved_reasoning: Final = { + "id": "rs_saved_reasoning", + "type": "reasoning", + "summary": [], + "content": [{"type": "reasoning_text", "text": "Inspect the file."}], + "encrypted_content": "opaque-saved-reasoning", + } + saved_message: Final = { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "The file is ready.", "annotations": []}], + } + history: Final = [ + ResponseReasoningItem.model_validate(saved_reasoning) if as_model else saved_reasoning, + saved_message, + ] + original: Final = deepcopy(history) + + request: Final = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-6.1-sol", + input=history, + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["input"] == [saved_reasoning, saved_message] + assert history == original + + +@pytest.mark.parametrize( + "config", + [AzureOpenAIResponsesAPIConfig(), AzureOpenAIOSeriesResponsesAPIConfig(), AzureAIResponsesAPIConfig()], + ids=["azure", "o_series", "azure_ai"], +) +def test_azure_family_forwards_input_items_that_are_not_objects_for_the_provider_to_judge( + config: AzureOpenAIResponsesAPIConfig, +) -> None: + saved_reasoning: Final = {"type": "reasoning", "summary": [], "status": "completed", "encrypted_content": "opaque"} + + request: Final = config.transform_responses_api_request( + model="gpt-6.1-sol", + input=["Continue.", saved_reasoning, ["nested"]], + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["input"] == [ + "Continue.", + {"type": "reasoning", "summary": [], "encrypted_content": "opaque"}, + ["nested"], + ] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.parametrize("async_call", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize(("reasoning_index", "include_id"), [(51, True), (10, False)], ids=["with-id", "without-id"]) +async def test_azure_responses_replays_reasoning_history_without_fields_azure_rejects( + respx_mock: respx.MockRouter, stream: bool, async_call: bool, reasoning_index: int, include_id: bool +) -> None: + prefix: Final = [{"role": "user", "content": f"Earlier context {index}."} for index in range(reasoning_index)] + replayable_reasoning: Final = { + "type": "reasoning", + "summary": [], + "encrypted_content": "opaque-saved-reasoning", + **({"id": "rs_saved_reasoning"} if include_id else {}), + } + saved_reasoning: Final = { + **replayable_reasoning, + "status": "completed", + "content": [{"type": "reasoning_text", "text": "Inspect the file before editing."}], + } + answer: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Ready to continue.", "annotations": []}], + } + prompt: Final = {"role": "user", "content": "Continue."} + replayable_followup: Final = { + "type": "reasoning", + "summary": [], + "encrypted_content": "opaque-followup-reasoning", + **({"id": "rs_followup_reasoning"} if include_id else {}), + } + saved_followup: Final = { + **replayable_followup, + "status": "completed", + "content": [{"type": "reasoning_text", "text": "Check the next file."}], + } + history: Final = [*prefix, saved_reasoning, answer, saved_followup, prompt] + original: Final = deepcopy(history) + + def respond(request: httpx.Request) -> httpx.Response: + body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + items: Final = body["input"] + assert isinstance(items, list) + reasoning: Final = items[reasoning_index] + assert isinstance(reasoning, dict) + if "status" in reasoning: + return httpx.Response( + 400, + json={ + "error": { + "message": f"Unknown parameter: 'input[{reasoning_index}].status'.", + "type": "invalid_request_error", + "param": f"input[{reasoning_index}].status", + "code": "unknown_parameter", + } + }, + ) + if reasoning.get("content"): + return httpx.Response( + 400, + json={ + "error": { + "message": f"Invalid 'input[{reasoning_index}].content': array too long. " + "Expected an array with maximum length 0, " + "but got an array with length 1 instead.", + "type": "invalid_request_error", + "param": f"input[{reasoning_index}].content", + "code": "array_above_max_length", + } + }, + ) + return httpx.Response( + 200, + text=_COMPLETED_RESPONSE_SSE if stream else json.dumps(_COMPLETED_RESPONSE), + headers={"content-type": "text/event-stream" if stream else "application/json"}, + ) + + route: Final = respx_mock.post(host="example-resource.openai.azure.com").mock(side_effect=respond) + request_params: Final = { + "model": "azure/gpt-6.1-sol", + "input": history, + "api_base": _AZURE_RESOURCE, + "api_key": "azure-test-key", + "stream": stream, + "store": False, + "include": ["reasoning.encrypted_content"], + "drop_params": True, + } + response: Final = await litellm.aresponses(**request_params) if async_call else litellm.responses(**request_params) + deltas: Final = ( + ( + [event.delta async for event in response if event.type == "response.output_text.delta"] + if async_call + else [event.delta for event in response if event.type == "response.output_text.delta"] + ) + if stream + else [] + ) + text: Final = "".join(deltas) if stream else response.output[0].content[0].text + + assert text == "Done." + assert route.call_count == 1 + assert TypeAdapter(dict[str, JsonValue]).validate_json(route.calls[0].request.content)["input"] == [ + *prefix, + replayable_reasoning, + answer, + replayable_followup, + prompt, + ] + assert history == original + + +@pytest.mark.parametrize( + ("as_model", "include_id"), + [(False, True), (False, False), (True, True)], + ids=["dictionary-with-id", "dictionary-without-id", "pydantic"], +) +@pytest.mark.parametrize("drop_params", [True, False], ids=["drop-params", "no-drop-params"]) +def test_azure_reasoning_replay_drops_status_always_and_content_under_drop_params( + as_model: bool, include_id: bool, drop_params: bool +) -> None: + replayable: Final = { + "type": "reasoning", + "summary": [], + "encrypted_content": "opaque-saved-reasoning", + **({"id": "rs_saved_reasoning"} if include_id else {}), + } + saved_content: Final = [{"type": "reasoning_text", "text": "Inspect the file first."}] + saved_reasoning: Final = {**replayable, "status": "completed", "content": saved_content} + history: Final = [ResponseReasoningItem.model_validate(saved_reasoning) if as_model else saved_reasoning] + original: Final = deepcopy(history) + + request: Final = AzureOpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-6.1-sol", + input=history, + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + + assert request["input"] == [replayable if drop_params else {**replayable, "content": saved_content}] + assert history == original diff --git a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py index dfd1474211f..53f426d8505 100644 --- a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py +++ b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py @@ -466,6 +466,12 @@ def _mock_response(): response.content = json.dumps(MOCK_BING_GROUNDING_RESPONSE).encode() return response +async def _search_what_is_litellm(is_async: bool) -> object: + if is_async: + return await litellm.asearch(query="what is litellm", search_provider="bing_grounding") + return litellm.search(query="what is litellm", search_provider="bing_grounding") + + @pytest.mark.usefixtures("_vcr_outcome_gate") class TestBingGroundingSearchTransformation: """ @@ -546,6 +552,36 @@ class TestBingGroundingSearchTransformation: } ] + @pytest.mark.parametrize("is_async", (False, True)) + async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason( + self, monkeypatch: pytest.MonkeyPatch, is_async: bool + ): + from azure.core.exceptions import ClientAuthenticationError + + from litellm.llms.azure import common_utils + + app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" + credential = Mock(spec=["get_token"]) + credential.get_token.side_effect = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'." + ) + monkeypatch.delenv("BING_GROUNDING_TOKEN") + for name in ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET", "AZURE_USERNAME", "AZURE_PASSWORD"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential)) + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + try: + with pytest.raises(litellm.AuthenticationError) as excinfo: + await _search_what_is_litellm(is_async) + finally: + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert app_client_id not in str(excinfo.value) + @pytest.mark.asyncio async def test_bing_grounding_asearch(self): with patch( # test-quality-ok: litellm.asearch has no client injection seam diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 0a41861e088..d66ba3f324e 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1,16 +1,23 @@ +import asyncio +import io import json import os +import sys import threading import traceback import uuid +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from queue import SimpleQueue -from typing import Callable, Final, Optional +from typing import Awaitable, Callable, Final, Optional from urllib.parse import parse_qs from unittest.mock import ANY, MagicMock, patch import httpx +import openai import pytest +import requests +from azure.core.pipeline.transport import RequestsTransport import litellm from litellm.llms.azure.common_utils import ( @@ -2500,3 +2507,487 @@ def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(m closer.reap() assert wrapper.is_closed() is True + + +_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" +_ENTRA_REJECTIONS: Final = { + "wrong_secret": ( + "AADSTS7000215", + "Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent in the " + "request is the client secret value, not the client secret ID, for a secret added to app " + f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID} Correlation ID: ecfd2c73-ec0a-474c-bba7-205be903159d " + "Timestamp: 2026-10-29 07:19:40Z", + ), + "expired_secret": ( + "AADSTS7000222", + f"Authentication failed: AADSTS7000222: The provided client secret keys for app '{_ENTRA_APP_CLIENT_ID}' " + f"are expired. Visit the Azure portal to create new keys for your app. Trace ID: {_ENTRA_TRACE_ID} " + "Correlation ID: 354f9b57-5fed-47e4-a67d-89a7c56f6ddc Timestamp: 2026-10-29 07:19:50Z", + ), +} + + +def _raising_token_provider(error: Exception) -> Callable[[], str]: + def token_provider() -> str: + raise error + + return token_provider + + +def _azure_chat(token_provider: Callable[[], str]) -> object: + return litellm.completion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=token_provider, + max_retries=0, + ) + + +def _azure_async_chat(token_provider: Callable[[], str]) -> object: + return asyncio.run( + litellm.acompletion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=token_provider, + ) + ) + + +def _azure_responses(token_provider: Callable[[], str]) -> object: + return litellm.responses( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + input="hi", + azure_ad_token_provider=token_provider, + ) + + +def _azure_ai_decisions(token_provider: Callable[[], str]) -> object: + return litellm.systemone( + model="azure_ai/decision-1-qa", + api_base="https://litellm-qa.services.ai.azure.com", + state={"ticket": "The login page returns a 500 error after deploy"}, + questions={"is_bug": {"type": "noul", "instructions": "Is this a software bug report?"}}, + azure_ad_token_provider=token_provider, + ) + + +_AZURE_CALLS: Final = { + "chat": _azure_chat, + "async_chat": _azure_async_chat, + "responses": _azure_responses, + "decisions": _azure_ai_decisions, +} + + +@pytest.fixture +def no_azure_api_key(monkeypatch): + for env_var in ( + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "azure_key", None) + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize(("aadsts_code", "entra_message"), _ENTRA_REJECTIONS.values(), ids=_ENTRA_REJECTIONS.keys()) +def test_entra_rejecting_the_credential_answers_401_with_only_the_aadsts_reason( + no_azure_api_key, call, aadsts_code, entra_message +): + from azure.core.exceptions import ClientAuthenticationError + + with pytest.raises(litellm.AuthenticationError) as excinfo: + call(_raising_token_provider(ClientAuthenticationError(message=entra_message))) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +@pytest.mark.parametrize("api_version", ["2025-04-01-preview", "v1"]) +def test_entra_rejecting_an_async_token_provider_answers_401(no_azure_api_key, api_version): + from azure.core.exceptions import ClientAuthenticationError + + aadsts_code, entra_message = _ENTRA_REJECTIONS["wrong_secret"] + + async def token_provider() -> str: + raise ClientAuthenticationError(message=entra_message) + + async_token_provider: Final[Callable[[], Awaitable[str]]] = token_provider + + with pytest.raises(litellm.AuthenticationError) as excinfo: + asyncio.run( + litellm.acompletion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version=api_version, + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=async_token_provider, + ) + ) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +def _azure_identity_wrapped_token_provider(inner_error: Exception) -> Callable[[], str]: + from azure.core.exceptions import ClientAuthenticationError + + def token_provider() -> str: + try: + raise inner_error + except Exception as error: + raise ClientAuthenticationError(message=f"Authentication failed: {error}") from error + + return token_provider + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize("failed_leg", ["request", "response"]) +def test_unreachable_entra_endpoint_stays_a_retryable_server_error(no_azure_api_key, call, failed_leg): + from azure.core.exceptions import ServiceRequestError, ServiceResponseError + + transport_error: Final = {"request": ServiceRequestError, "response": ServiceResponseError}[failed_leg]( + "('Unable to connect to proxy', NewConnectionError(\"Failed to establish a new connection: " + "[Errno 111] Connection refused\"))" + ) + + with pytest.raises(openai.APIError) as excinfo: + call(_azure_identity_wrapped_token_provider(transport_error)) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +_ENTRA_TENANT_ID: Final = "5b3e1c7a-2f4d-4e8b-9a61-0c2d7f8e9b14" +_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +_ENTRA_OPENID_CONFIGURATION: Final = { + "token_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token", + "authorization_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/authorize", + "issuer": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/v2.0", +} + + +class _FakeEntraAdapter(requests.adapters.BaseAdapter): + def __init__(self, status_code: int | None, body: Mapping[str, object] | str) -> None: + super().__init__() + self.status_code: Final = status_code + self.body: Final = body + + def send(self, request: requests.PreparedRequest, *args: object, **kwargs: object) -> requests.Response: + if "openid-configuration" in str(request.url): + return self._answer(request, 200, _ENTRA_OPENID_CONFIGURATION) + if self.status_code is None: + raise requests.exceptions.ConnectionError("Failed to establish a new connection: Connection refused") + return self._answer(request, self.status_code, self.body) + + def close(self) -> None: + pass + + @staticmethod + def _answer( + request: requests.PreparedRequest, status_code: int, body: Mapping[str, object] | str + ) -> requests.Response: + response: Final = requests.Response() + response.status_code = status_code + response.raw = io.BytesIO((body if isinstance(body, str) else json.dumps(body)).encode()) + response.headers = requests.structures.CaseInsensitiveDict( + {"Content-Type": "application/json", "Retry-After": "0"} + ) + response.url = str(request.url) + response.request = request + return response + + +def _entra_answering(status_code: int | None, body: Mapping[str, object] | str = "") -> RequestsTransport: + session: Final = requests.Session() + session.mount("https://", _FakeEntraAdapter(status_code, body)) + return RequestsTransport(session=session, session_owner=False) + + +def _entra_error(aadsts_code: str, description: str) -> Mapping[str, object]: + return { + "error": "invalid_client", + "error_description": f"{aadsts_code}: {description} Trace ID: {_ENTRA_TRACE_ID}", + "error_codes": [int(aadsts_code.removeprefix("AADSTS"))], + } + + +def _real_credential_token_provider(credential_kind: str, entra: RequestsTransport) -> Callable[[], str]: + from azure.identity import ( + ChainedTokenCredential, + ClientAssertionCredential, + ClientSecretCredential, + get_bearer_token_provider, + ) + + client_secret: Final = ClientSecretCredential( + _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, "wrong-secret", transport=entra, disable_instance_discovery=True + ) + credentials: Final = { + "client_secret": client_secret, + "default_azure_credential_chain": ChainedTokenCredential(client_secret), + "workload_identity": ClientAssertionCredential( + _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, lambda: "federated-token", transport=entra, retry_total=0 + ), + } + return get_bearer_token_provider(credentials[credential_kind], _ENTRA_SCOPE) + + +_REAL_CREDENTIAL_REJECTIONS: Final = { + "wrong_secret": ("client_secret", "AADSTS7000215", "Invalid client secret provided."), + "wrong_secret_through_a_credential_chain": ( + "default_azure_credential_chain", + "AADSTS7000215", + "Invalid client secret provided.", + ), + "unknown_app_through_workload_identity": ( + "workload_identity", + "AADSTS700016", + f"Application with identifier '{_ENTRA_APP_CLIENT_ID}' was not found in the directory.", + ), + "unknown_federated_subject_through_workload_identity": ( + "workload_identity", + "AADSTS700213", + "No matching federated identity record found for presented assertion subject 'system:serviceaccount:qa:proxy'.", + ), +} + + +@pytest.mark.parametrize( + ("credential_kind", "aadsts_code", "description"), + _REAL_CREDENTIAL_REJECTIONS.values(), + ids=_REAL_CREDENTIAL_REJECTIONS.keys(), +) +def test_entra_rejection_through_a_real_azure_identity_credential_answers_401( + no_azure_api_key, credential_kind, aadsts_code, description +): + entra: Final = _entra_answering(401, _entra_error(aadsts_code, description)) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + _azure_chat(_real_credential_token_provider(credential_kind, entra)) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +_REAL_CREDENTIAL_NON_REJECTIONS: Final = { + "unreachable_through_a_credential_chain": ("default_azure_credential_chain", None, ""), + "html_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + "Service Unavailable", + ), + "json_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS50196", "The server terminated an operation because it encountered a client request loop."), + ), + "tenant_throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS90055", "There are too many incoming requests."), + ), + "request_budget_exceeded_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90024", "A transient error has occurred. Try again."), + ), + "retryable_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50033", "A transient error has occurred."), + ), + "external_server_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90006", "The service is temporarily unavailable."), + ), + "graph_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90090", "The service is temporarily unavailable."), + ), + "token_issuance_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50000", "There was an error issuing a token or an issue with our sign-in service."), + ), + "throttled_without_a_code_through_a_credential_chain": ( + "default_azure_credential_chain", + 429, + {"error": "temporarily_unavailable", "error_description": "Too many requests."}, + ), + "json_503_through_workload_identity": ( + "workload_identity", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_workload_identity_naming_a_rejection_code": ( + "workload_identity", + 429, + _entra_error("AADSTS7000215", "Invalid client secret provided."), + ), +} + + +@pytest.mark.parametrize( + ("credential_kind", "status_code", "body"), + _REAL_CREDENTIAL_NON_REJECTIONS.values(), + ids=_REAL_CREDENTIAL_NON_REJECTIONS.keys(), +) +def test_entra_outage_or_throttle_through_a_real_azure_identity_credential_stays_retryable( + no_azure_api_key, credential_kind, status_code, body +): + with pytest.raises(openai.APIError) as excinfo: + _azure_chat(_real_credential_token_provider(credential_kind, _entra_answering(status_code, body))) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +def _azure_chat_with_oidc_token( + respx_mock, + monkeypatch, + entra_status_code: int, + entra_body: Mapping[str, object], + entra_headers: Mapping[str, str] | None = None, +): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + respx_mock.post(f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token").mock( + return_value=httpx.Response(entra_status_code, json=entra_body, headers=entra_headers) + ) + return litellm.completion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token="oidc/env/LITELLM_QA_FEDERATED_TOKEN", + tenant_id=_ENTRA_TENANT_ID, + client_id=_ENTRA_APP_CLIENT_ID, + max_retries=0, + ) + + +def test_entra_rejecting_an_oidc_token_answers_401_with_only_the_aadsts_reason(no_azure_api_key, respx_mock, monkeypatch): + rejection: Final = _entra_error( + "AADSTS700213", f"No matching federated identity record found for app '{_ENTRA_APP_CLIENT_ID}'." + ) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + _azure_chat_with_oidc_token(respx_mock, monkeypatch, 400, rejection) + + assert excinfo.value.status_code == 401 + assert "AADSTS700213" in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +def test_entra_answering_an_oidc_exchange_with_a_transient_code_keeps_its_status( + no_azure_api_key, respx_mock, monkeypatch +): + with pytest.raises(openai.APIError) as excinfo: + _azure_chat_with_oidc_token( + respx_mock, monkeypatch, 400, _entra_error("AADSTS90024", "A transient error has occurred. Try again.") + ) + + assert excinfo.value.status_code == 400 + assert not isinstance(excinfo.value, litellm.AuthenticationError) + + +def test_entra_throttling_an_oidc_exchange_keeps_its_retry_after(no_azure_api_key, respx_mock, monkeypatch): + with pytest.raises(litellm.RateLimitError) as excinfo: + _azure_chat_with_oidc_token( + respx_mock, + monkeypatch, + 429, + _entra_error("AADSTS90055", "TenantThrottlingError: There are too many pending requests."), + {"Retry-After": "7"}, + ) + + assert excinfo.value.response.headers.get("retry-after") == "7" + + +class _EntraServiceUnavailableAnswer: + status_code: Final = 503 + reason: Final = "Service Unavailable" + + def text(self, encoding: str | None = None) -> str: + return "Service Unavailable" + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize("credential_kind", ["client_secret", "managed_identity"]) +def test_entra_answering_a_server_error_stays_a_retryable_server_error(no_azure_api_key, call, credential_kind): + from azure.core.exceptions import ClientAuthenticationError + from msal.exceptions import MsalServiceError + + token_providers: Final = { + "client_secret": _azure_identity_wrapped_token_provider( + MsalServiceError("HTTP Error: 503", error=None, error_description=None) + ), + "managed_identity": _raising_token_provider( + ClientAuthenticationError(message="No token received.", response=_EntraServiceUnavailableAnswer()) + ), + } + + with pytest.raises(openai.APIError) as excinfo: + call(token_providers[credential_kind]) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +def test_token_provider_error_without_azure_core_installed_keeps_the_generic_token_error(monkeypatch): + monkeypatch.setitem(sys.modules, "azure.core.exceptions", None) + + with pytest.raises(RuntimeError, match="Failed to get Azure AD token: token service down"): + get_azure_ad_token( + GenericLiteLLMParams(azure_ad_token_provider=_raising_token_provider(ValueError("token service down"))) + ) + + +def test_image_generation_with_an_api_key_ignores_a_non_callable_azure_ad_token_provider(): + import respx + + with respx.mock: + image_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGk="}]}) + ) + response: Final = litellm.image_generation( + model="azure/gpt-image-1", + prompt="a red dot", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + api_key="configured-key", + azure_ad_token_provider="not-a-callable", + ) + + assert image_route.calls.last.request.headers["api-key"] == "configured-key" + assert response.data[0].b64_json == "aGk=" diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index faa33e559ca..7ce5d8b6a02 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -171,6 +171,79 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ ) +@pytest.mark.asyncio +async def test_entra_rejecting_the_deployment_oidc_token_becomes_a_401_error_response_without_counting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + with respx.mock: + respx.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 401, + json={ + "error": "invalid_client", + "error_description": "AADSTS700211: No matching federated identity record found for presented " + "assertion issuer 'https://oidc.litellm-qa.invalid'. Trace ID: 0f5c2b1e", + "error_codes": [700211], + }, + ) + ) + count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "api_base": API_BASE, + "tenant_id": "qa-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + }, + request_model="azure_ai/claude-sonnet-4-5", + ) + + assert result is not None + assert (result.error, result.status_code) == (True, 401) + assert "AADSTS700211" in (result.error_message or "") + assert not count_route.called + + +@pytest.mark.asyncio +async def test_entra_throttling_the_deployment_oidc_token_exchange_keeps_its_429( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + with respx.mock: + respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response(429, json={"error": "temporarily_unavailable", "error_codes": [90055]}) + ) + count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "api_base": API_BASE, + "tenant_id": "throttled-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + }, + request_model="azure_ai/claude-sonnet-4-5", + ) + + assert result is not None + assert (result.error, result.status_code) == (True, 429) + assert not count_route.called + + @pytest.mark.asyncio async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py index d0c1c4a6fb2..a8958b301db 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py @@ -254,6 +254,57 @@ def test_agent_entra_values_resolve_os_environ_references(monkeypatch): assert get_azure_ai_agent_entra_token({"azure_ad_token": "os.environ/FOUNDRY_AGENT_AD_TOKEN"}) == "token-from-env" +_AGENT_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_AGENT_CREDENTIALS: Final = { + "service_principal": ( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", + {"tenant_id": "tenant", "client_id": _AGENT_APP_CLIENT_ID, "client_secret": "wrong-secret"}, + ), + "username_password": ( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password", + {"client_id": _AGENT_APP_CLIENT_ID, "azure_username": "agent@example.com", "azure_password": "wrong"}, + ), +} + + +def _raising(error: Exception): + def token() -> str: + raise error + + return token + + +@pytest.mark.parametrize(("token_fetch", "agent_params"), _AGENT_CREDENTIALS.values(), ids=_AGENT_CREDENTIALS.keys()) +def test_agent_entra_rejection_raises_401_with_only_the_aadsts_reason(token_fetch, agent_params): + from azure.core.exceptions import ClientAuthenticationError + + rejection: Final = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent " + f"in the request is the client secret value, for a secret added to app '{_AGENT_APP_CLIENT_ID}'." + ) + + with patch(token_fetch, return_value=_raising(rejection)): # test-quality-ok: stands in for the Entra token endpoint answering a wrong credential, the shape azure-identity raises + with pytest.raises(litellm.AuthenticationError) as excinfo: + get_azure_ai_agent_entra_token(agent_params) + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert _AGENT_APP_CLIENT_ID not in str(excinfo.value) + + +def test_agent_entra_endpoint_unreachable_is_not_an_auth_error(): + from azure.core.exceptions import ClientAuthenticationError + + unreachable: Final = ClientAuthenticationError(message="Authentication failed: Connection refused") + token_fetch, agent_params = _AGENT_CREDENTIALS["service_principal"] + + with patch(token_fetch, return_value=_raising(unreachable)): # test-quality-ok: stands in for an unreachable Entra token endpoint, the shape azure-identity raises + with pytest.raises(ClientAuthenticationError) as excinfo: + get_azure_ai_agent_entra_token(agent_params) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + + def test_agent_entra_token_failure_names_the_credential_fields(): with pytest.raises(ValueError, match="client_secret"): get_azure_ai_agent_entra_token({"azure_scope": "https://ai.azure.com/.default"}) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index aec1672b335..1c0ce5a334a 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -1394,6 +1394,7 @@ class _ScriptedGuardrail(CustomGuardrail): self.skip_system_message_in_guardrail = skip_system self.scan_only_tool_results = scan_only_tool_results self.received: tuple[GenericGuardrailAPIInputs, ...] = () + self.request_data: tuple[Mapping[str, object], ...] = () async def apply_guardrail( self, @@ -1403,6 +1404,7 @@ class _ScriptedGuardrail(CustomGuardrail): logging_obj: LiteLLMLoggingObj | None = None, ) -> Mapping[str, object]: self.received = (*self.received, copy.deepcopy(inputs)) + self.request_data = (*self.request_data, request_data) return self.reply(inputs) @@ -1429,16 +1431,30 @@ def _converse_request(body: dict) -> dict: async def _sent_to_bedrock(data: dict) -> dict: - """The Converse body Bedrock receives when the passthrough route sends ``data``.""" + """The body Bedrock receives when the passthrough route sends ``data``.""" received: list[httpx.Request] = [] + action: Final = data["endpoint"].rsplit("/", 1)[-1] + def respond(request: httpx.Request) -> httpx.Response: received.append(request) - reply = { - "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, - } + reply: Final = ( + { + "id": "msg_bdrk_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + if action == "invoke" + else { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + ) return httpx.Response(200, json=reply, request=request) router = litellm.Router( @@ -1457,7 +1473,7 @@ async def _sent_to_bedrock(data: dict) -> dict: await router.allm_passthrough_route( model="bedrock-model", method="POST", - endpoint="/model/bedrock-model/converse", + endpoint=f"/model/bedrock-model/{action}", data=data["data"], client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), ) @@ -3843,6 +3859,260 @@ class TestInvokeStreamGuardrailRelay: assert rewritten_key == received_key +def _anthropic_invoke_data( + action: str = "invoke", +) -> dict[str, object]: # mutable-ok: tests write into the request the handler gets + return { + "endpoint": f"model/anthropic.claude-sonnet/{action}", + "custom_llm_provider": "bedrock", + "model": "anthropic.claude-sonnet", + "litellm_call_id": "call-1", + "data": { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 512, + "system": "You are a travel assistant. Card on file 4111 1111 1111 1111", + "tools": [ + { + "name": "find_trips", + "description": "Search trips", + "input_schema": {"type": "object", "properties": {"port": {"type": "string"}}}, + } + ], + "messages": [ + {"role": "user", "content": "Pay with 4111 1111 1111 1111"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "toolu_1", "name": "find_trips", "input": {"port": "Miami"}} + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card 4111 1111 1111 1111 charged"} + ], + }, + ], + }, + } + + +def _masked_anthropic_invoke_body() -> Mapping[str, object]: + body: Final = _anthropic_invoke_data()["data"] + return { + **body, + "system": "You are a travel assistant. Card on file [CARD] 1111 1111 1111", + "messages": [ + {"role": "user", "content": "Pay with [CARD] 1111 1111 1111"}, + body["messages"][1], + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card [CARD] 1111 1111 1111 charged"} + ], + }, + ], + } + + +def _rows_masked(inputs: GenericGuardrailAPIInputs) -> Mapping[str, object]: + rows = [ + {**row, "content": row["content"].replace("4111", "[CARD]")} if isinstance(row.get("content"), str) else row + for row in inputs["structured_messages"] + ] + return {"structured_messages": rows} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_anthropic_invoke_masks_are_written_back_into_the_body(action: str) -> None: + guardrail = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}) + + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + _anthropic_invoke_data(action), guardrail + ) + + assert result["data"] == _masked_anthropic_invoke_body() + + +@pytest.mark.asyncio +async def test_anthropic_invoke_masks_reach_bedrock() -> None: + data: Final = _anthropic_invoke_data() + + await BedrockPassthroughGuardrailHandler().process_input_messages( + data, _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}) + ) + + assert await _sent_to_bedrock(data) == _masked_anthropic_invoke_body() + + +@pytest.mark.asyncio +async def test_anthropic_invoke_guardrail_gets_roles_tools_and_the_body_at_the_top_level_of_request_data() -> None: + data: Final = _anthropic_invoke_data() + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + inputs = guardrail.received[-1] + assert [row["role"] for row in inputs["structured_messages"]] == ["system", "user", "assistant", "tool"] + assert [tool["function"]["name"] for tool in inputs["tools"]] == ["find_trips"] + assert inputs["model"] == data["model"] + request_data: Final = guardrail.request_data[-1] + assert request_data["messages"] is data["data"]["messages"] + assert request_data["endpoint"] == data["endpoint"] + assert request_data["litellm_call_id"] == data["litellm_call_id"] + assert "data" not in request_data + + +@pytest.mark.asyncio +async def test_anthropic_invoke_structured_message_rewrites_are_written_back() -> None: + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + _anthropic_invoke_data(), _ScriptedGuardrail(_rows_masked) + ) + + assert result["data"]["messages"][0] == { + "role": "user", + "content": [{"type": "text", "text": "Pay with [CARD] 1111 1111 1111"}], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["anthropic.claude-sonnet", "prod-model", ""]) +async def test_anthropic_invoke_structured_rewrite_keeps_cache_control_whatever_the_model_is_called( + model: str, +) -> None: + cached: Final = {"type": "ephemeral"} + data: Final = _anthropic_invoke_data() + data["model"] = model + data["data"]["system"] = [{"type": "text", "text": "Card on file 4111 1111 1111 1111", "cache_control": cached}] + data["data"]["tools"] = [ + { + "name": "find_trips", + "description": "Search trips", + "input_schema": {"type": "object", "properties": {}}, + "cache_control": cached, + } + ] + data["data"]["messages"] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Pay with 4111 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}} + ], + } + ] + + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + data, + _ScriptedGuardrail( + lambda inputs: { + "structured_messages": json.loads(json.dumps(inputs["structured_messages"]).replace("4111", "[CARD]")), + "tools": json.loads(json.dumps(inputs["tools"])), + } + ), + ) + + assert result["data"]["system"] == [ + {"type": "text", "text": "Card on file [CARD] 1111 1111 1111", "cache_control": cached} + ] + assert result["data"]["messages"][0]["content"] == [ + {"type": "text", "text": "Pay with [CARD] 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}} + ] + assert result["data"]["tools"][0]["cache_control"] == cached + + +@pytest.mark.asyncio +async def test_invoke_body_without_anthropic_version_is_scanned_whole() -> None: + data: Final = {**_anthropic_invoke_data(), "data": {"messages": [{"role": "user", "content": "hello"}]}} + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert "structured_messages" not in guardrail.received[-1] + + +_PLAIN_TEXT_DOCUMENT: Final = { + "type": "document", + "source": {"type": "text", "media_type": "text/plain", "data": "Card 4111 1111 1111 1111"}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", + [ + pytest.param([{"type": "text", "text": "Summarize this"}, _PLAIN_TEXT_DOCUMENT], id="document"), + pytest.param( + [ + { + "type": "search_result", + "source": "https://example.com/cards", + "title": "Cards", + "content": [{"type": "text", "text": "Card 4111 1111 1111 1111"}], + } + ], + id="search_result", + ), + pytest.param( + [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [_PLAIN_TEXT_DOCUMENT]}], + id="document_in_tool_result", + ), + ], +) +async def test_anthropic_invoke_body_with_blocks_the_messages_handler_does_not_scan_is_scanned_whole(content: Sequence[object]) -> None: + invoke: Final = _anthropic_invoke_data() + data: Final = {**invoke, "data": {**invoke["data"], "messages": [{"role": "user", "content": content}]}} + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert "structured_messages" not in guardrail.received[-1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "messages", + [ + pytest.param([{"content": "hi"}], id="message_without_role"), + pytest.param([{"role": "user", "content": ["hi"]}], id="bare_string_in_content"), + pytest.param([{"role": "user", "content": None}], id="null_content"), + pytest.param(["hi"], id="message_is_a_string"), + ], +) +async def test_anthropic_invoke_malformed_messages_are_scanned_whole_and_left_for_bedrock_to_reject( + messages: Sequence[object], +) -> None: + invoke: Final = _anthropic_invoke_data() + data: Final = {**invoke, "data": {**invoke["data"], "messages": messages}} + guardrail = _replying({}) + + result: Final = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"] + assert result["data"]["messages"] == messages + + +@pytest.mark.asyncio +async def test_anthropic_invoke_response_masks_are_written_back() -> None: + response = { + "id": "msg_bdrk_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Charged 4111 1111 1111 1111"}], + "stop_reason": "end_turn", + } + + result: Final = await BedrockPassthroughGuardrailHandler().process_output_response( + response=response, + guardrail_to_apply=_ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}), + request_data=_anthropic_invoke_data(), + ) + + assert result == {**response, "content": [{"type": "text", "text": "Charged [CARD] 1111 1111 1111"}]} + + def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None: event: Final = { "type": "content_block_delta", diff --git a/tests/unit/llms/empiriolabs/__init__.py b/tests/unit/llms/empiriolabs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/empiriolabs/decisions/__init__.py b/tests/unit/llms/empiriolabs/decisions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py b/tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py new file mode 100644 index 00000000000..02e9eab00ae --- /dev/null +++ b/tests/unit/llms/empiriolabs/decisions/test_empiriolabs_decisions_transformation.py @@ -0,0 +1,40 @@ +from collections.abc import Mapping + +import pytest + +from litellm.llms.empiriolabs.decisions.transformation import EmpirioLabsDecisionsConfig + + +@pytest.mark.parametrize( + ("environment", "connection"), + ( + ({}, ("https://api.empiriolabs.ai", None)), + ( + {"EMPIRIOLABS_API_BASE": "https://gateway.example/v1", "EMPIRIOLABS_API_KEY": "environment-key"}, + ("https://gateway.example/v1", "environment-key"), + ), + ), + ids=("public_api_without_a_key", "environment"), +) +def test_connection_comes_from_the_environment_before_the_public_api( + environment: Mapping[str, str], + connection: tuple[str, str | None], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + monkeypatch.delenv("EMPIRIOLABS_API_KEY", raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + assert EmpirioLabsDecisionsConfig().resolve_connection(None, None) == connection + + +@pytest.mark.parametrize( + "api_base", + ("https://api.empiriolabs.ai", "https://api.empiriolabs.ai/v1", "https://api.empiriolabs.ai/v1/"), +) +def test_requests_go_to_v1_decisions_whether_or_not_the_base_ends_in_v1(api_base: str) -> None: + assert ( + EmpirioLabsDecisionsConfig().get_complete_url(api_base=api_base, model="aplomb-1") + == "https://api.empiriolabs.ai/v1/decisions" + ) diff --git a/tests/unit/llms/github_copilot/test_github_host.py b/tests/unit/llms/github_copilot/test_github_host.py index df2efe4669b..daf5bfaab80 100644 --- a/tests/unit/llms/github_copilot/test_github_host.py +++ b/tests/unit/llms/github_copilot/test_github_host.py @@ -54,6 +54,7 @@ def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.Monke assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code" assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token" assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token" + assert endpoints.copilot_user_url == "https://api.octo.ghe.com/copilot_internal/user" assert endpoints.user_url == "https://api.octo.ghe.com/user" assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com" assert endpoints.client_id == "custom-client" diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index cf6fb9fa71f..5274f0c1cb8 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -307,6 +307,71 @@ async def test_aexchange_401_raises_and_evicts(): ) +ACCOUNT_URL = "https://api.github.com/copilot_internal/user" +_CUSTOM_CLIENT_ID: Final = "Ov23liCustomOrgApp" +_ACCOUNT_API_BASE: Final = "https://api.business.githubcopilot.com" + + +def _account_payload(access_type_sku="copilot_enterprise_seat", chat_enabled=True, api_base=_ACCOUNT_API_BASE): + return {"access_type_sku": access_type_sku, "chat_enabled": chat_enabled, "endpoints": {"api": api_base}} + + +def _account_responder(payload=None, status=200, model_payload=None): + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == ACCOUNT_URL: + return httpx.Response(status, json=payload or _account_payload(), request=request) + if model_payload is not None and request.url.host.endswith("githubcopilot.com"): + return httpx.Response(200, json=model_payload, request=request) + pytest.fail(f"unexpected request to {request.url}") + + return respond + + +@pytest.fixture +def custom_client_id(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", _CUSTOM_CLIENT_ID) + + +@pytest.mark.asyncio +async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_swap(custom_client_id): + endpoints: Final = github_host_endpoints("github.com") + with github_http(_account_responder()) as requests: + first = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints) + second = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints) + + assert [str(r.url) for r in requests] == [ACCOUNT_URL] + assert requests[0].headers["Authorization"] == "token gho_org_app" + assert first.token == second.token == "gho_org_app" + assert first.api_base == _ACCOUNT_API_BASE + + +@pytest.mark.parametrize( + "payload", + [ + _account_payload(access_type_sku="no_access", chat_enabled=False), + _account_payload(chat_enabled=False), + _account_payload(access_type_sku="no_access"), + ], +) +def test_custom_client_id_without_a_copilot_seat_raises_reconnect(custom_client_id, payload): + with github_http(_account_responder(payload)): + with pytest.raises(CallerCredentialAuthenticationError): + exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + + +@pytest.mark.parametrize("status", [401, 403, 404]) +def test_custom_client_id_rejected_account_lookup_raises_reconnect(custom_client_id, status): + with github_http(_account_responder(status=status)): + with pytest.raises(CallerCredentialAuthenticationError): + exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + + +def test_custom_client_id_account_api_base_must_stay_on_copilot_hosts(custom_client_id): + with github_http(_account_responder(_account_payload(api_base="https://evil.example"))): + session = exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + assert session.api_base == "https://api.githubcopilot.com" + + def _kwargs_with_connection( token: str = "gho_user_token", github_host: str | None = "github.com", @@ -740,6 +805,25 @@ def test_e2e_embedding_per_user_caller_authorization_cannot_override_session(per _assert_session_token_on_the_wire(requests) +def test_e2e_custom_client_id_sends_the_github_token_to_the_accounts_copilot_host( + per_user_credential, custom_client_id +): + with ( + github_http(_account_responder(model_payload=_embedding_response_payload())) as requests, + _no_authenticator(), + ): + litellm.embedding( + model="github_copilot/text-embedding-3-small", + input=["hi"], + litellm_credential_name="copilot-cred", + secret_fields=_secret_fields("user-a", "gho_org_app"), + ) + assert EXCHANGE_URL not in [str(r.url) for r in requests] + upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")] + assert upstream[0].url.host == httpx.URL(_ACCOUNT_API_BASE).host + assert upstream[0].headers["authorization"] == "Bearer gho_org_app" + + def test_e2e_responses_per_user_caller_authorization_cannot_override_session(per_user_credential): with ( github_http(_copilot_surface_responder(_responses_payload())) as requests, diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py index dbb5cde6060..40ee4d61ac4 100644 --- a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py @@ -1,780 +1,641 @@ -""" -End-to-end-ish tests for NvidiaRivaAudioTranscription. - -We mock ``riva.client`` so the test does not need the real gRPC SDK or a -running Riva server. The mock also lets us assert how Auth metadata is -constructed (NVCF vs self-hosted) and how the streaming generator output -is aggregated. -""" - import asyncio import io -from types import SimpleNamespace -from unittest.mock import MagicMock +import ipaddress +import threading +import uuid +import wave +from array import array +from collections.abc import Iterable, Iterator +from concurrent import futures +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timezone +from queue import SimpleQueue +from typing import Final +import grpc import httpx -import numpy as np import pytest -import soundfile as sf +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID +from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc, riva_audio_pb2 import litellm -from litellm.llms.nvidia_riva.audio_transcription import handler as handler_mod -from litellm.llms.nvidia_riva.audio_transcription.handler import ( - NvidiaRivaAudioTranscription, -) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.nvidia_riva.audio_transcription.handler import NvidiaRivaAudioTranscription from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException from litellm.types.utils import TranscriptionResponse +_SAMPLE_RATE_HZ: Final = 16000 +_SERVER_WAIT_SECONDS: Final = 30 +_STALL_LIMIT_SECONDS: Final = 60 +_NO_DEADLINE_SECONDS: Final = 10**9 +_STALL_TIMEOUT_SECONDS: Final = 0.5 +_API_KEY: Final = "nvapi-test" -def _make_wav_bytes(seconds: float = 1.0, sample_rate: int = 16000) -> bytes: - n = int(sample_rate * seconds) - samples = (0.05 * np.sin(np.linspace(0, 2 * np.pi * 220 * seconds, n))).astype( - np.float32 + +@dataclass(frozen=True, slots=True) +class _Script: + streaming: tuple[riva_asr_pb2.StreamingRecognizeResponse, ...] = () + offline: riva_asr_pb2.RecognizeResponse = field(default_factory=riva_asr_pb2.RecognizeResponse) + abort: tuple[grpc.StatusCode, str] | None = None + stall: bool = False + + +@dataclass(frozen=True, slots=True) +class _Call: + rpc: str + metadata: tuple[tuple[str, str], ...] + config: riva_asr_pb2.RecognitionConfig + audio: bytes + deadline_seconds: float | None + cancelled: bool + + def header(self, key: str) -> tuple[str, ...]: + return tuple(value for name, value in self.metadata if name == key) + + +@dataclass(frozen=True, slots=True) +class _Riva: + address: str + calls: SimpleQueue[_Call] + + def only_call(self) -> _Call: + call: Final = self.calls.get(timeout=_SERVER_WAIT_SECONDS) + assert self.calls.empty(), "Riva received more than one call" + return call + + +class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): + def __init__(self, script: _Script, calls: SimpleQueue[_Call]) -> None: + self._script: Final = script + self._calls: Final = calls + + def _settle( + self, context: grpc.ServicerContext, rpc: str, config: riva_asr_pb2.RecognitionConfig, audio: bytes + ) -> None: + remaining: Final = context.time_remaining() + released: Final = threading.Event() + if self._script.stall and context.add_callback(released.set): + released.wait(timeout=_STALL_LIMIT_SECONDS) + self._calls.put( + _Call( + rpc=rpc, + metadata=tuple((item.key, item.value) for item in context.invocation_metadata()), + config=config, + audio=audio, + deadline_seconds=None if remaining is None or remaining > _NO_DEADLINE_SECONDS else remaining, + cancelled=not context.is_active(), + ) + ) + if self._script.abort is not None: + context.abort(*self._script.abort) + + def Recognize( + self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext + ) -> riva_asr_pb2.RecognizeResponse: + self._settle(context, "Recognize", request.config, request.audio) + return self._script.offline + + def StreamingRecognize( + self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext + ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: + first, *rest = tuple(request_iterator) + self._settle( + context, "StreamingRecognize", first.streaming_config.config, b"".join(r.audio_content for r in rest) + ) + yield from self._script.streaming + + +@dataclass(frozen=True, slots=True) +class _Tls: + key_pem: bytes + cert_pem: bytes + + +def _self_signed_loopback_tls() -> _Tls: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]) + cert: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime(2000, 1, 1, tzinfo=timezone.utc)) + .not_valid_after(datetime(2099, 1, 1, tzinfo=timezone.utc)) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost"), x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]), + critical=False, + ) + .sign(key, hashes.SHA256()) ) - buf = io.BytesIO() - sf.write(buf, samples, sample_rate, format="WAV", subtype="PCM_16") - return buf.getvalue() - - -def _fake_word(word: str, start_ms: int, end_ms: int): - return SimpleNamespace(word=word, start_time=start_ms, end_time=end_ms) - - -def _fake_alternative(transcript: str, words=None): - return SimpleNamespace(transcript=transcript, words=words or []) - - -def _fake_result(is_final: bool, alternatives): - return SimpleNamespace(is_final=is_final, alternatives=alternatives) - - -def _fake_response(results): - return SimpleNamespace(results=results) - - -@pytest.fixture -def mock_riva(monkeypatch): - """ - Stand-ins for the bits of ``riva.client`` the handler touches: - - ``Auth`` (constructor) - - ``ASRService`` with ``streaming_response_generator`` - - ``RecognitionConfig``, ``StreamingRecognitionConfig``, ``EndpointingConfig`` - - ``AudioEncoding`` namespace with ``LINEAR_PCM`` - """ - auth_calls = {} - - class FakeAuth: - def __init__(self, *args, **kwargs): - # Support both keyword and positional Auth constructors. - if kwargs: - auth_calls["uri"] = kwargs.get("uri") - auth_calls["use_ssl"] = kwargs.get("use_ssl") - auth_calls["metadata_args"] = kwargs.get("metadata_args") - else: - # positional: (None, use_ssl, uri, metadata) - auth_calls["use_ssl"] = args[1] if len(args) > 1 else None - auth_calls["uri"] = args[2] if len(args) > 2 else None - auth_calls["metadata_args"] = args[3] if len(args) > 3 else None - - class FakeRecognitionConfig: - def __init__(self, **kwargs): - self._kwargs = kwargs - self.endpointing_config = SimpleNamespace(CopyFrom=lambda _: None) - - class FakeStreamingRecognitionConfig: - def __init__(self, config, interim_results): - self.config = config - self.interim_results = interim_results - - class FakeEndpointingConfig: - def __init__(self, **kwargs): - self._kwargs = kwargs - - class FakeAudioEncoding: - LINEAR_PCM = "LINEAR_PCM" - - streaming_responses_holder = {"value": []} - offline_response_holder = {"value": None} - offline_error_holder = {"value": None} - offline_future_holder = {"value": None} - streaming_calls = [] - offline_calls = {} - - class FakeRecognizeFuture: - def result(self, timeout: float | None = None) -> object: - offline_calls["timeout"] = timeout - if offline_error_holder["value"] is not None: - raise offline_error_holder["value"] - return offline_response_holder["value"] - - def done(self) -> bool: - return True - - def cancel(self) -> bool: - return False - - class FakeASRService: - def __init__(self, auth): - self.auth = auth - - def streaming_response_generator(self, audio_chunks, streaming_config): - # Drain audio_chunks generator so we exercise the chunking path. - streaming_calls.append(list(audio_chunks)) - yield from streaming_responses_holder["value"] - - def offline_recognize(self, audio_bytes: bytes, config: object, future: bool = False) -> object: - offline_calls["audio_bytes"] = audio_bytes - offline_calls["config"] = config - pending = offline_future_holder["value"] or FakeRecognizeFuture() - return pending if future else pending.result() - - fake_riva_client = SimpleNamespace( - Auth=FakeAuth, - ASRService=FakeASRService, - RecognitionConfig=FakeRecognitionConfig, - StreamingRecognitionConfig=FakeStreamingRecognitionConfig, - EndpointingConfig=FakeEndpointingConfig, - AudioEncoding=FakeAudioEncoding, - ) - - def fake_import_riva(): - return fake_riva_client, fake_riva_client - - monkeypatch.setattr(handler_mod, "_import_riva", fake_import_riva) - - return SimpleNamespace( - auth_calls=auth_calls, - responses=streaming_responses_holder, - offline_response=offline_response_holder, - offline_error=offline_error_holder, - offline_future=offline_future_holder, - streaming_calls=streaming_calls, - offline_calls=offline_calls, - client=fake_riva_client, + return _Tls( + key_pem=key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ), + cert_pem=cert.public_bytes(serialization.Encoding.PEM), ) -@pytest.fixture -def logging_obj(): - return MagicMock() +@contextmanager +def _riva_server(script: _Script, max_receive_bytes: int | None = None, tls: _Tls | None = None) -> Iterator[_Riva]: + calls: Final[SimpleQueue[_Call]] = SimpleQueue() + options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),) + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=4), options=options) + riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(script, calls), server) + port: Final = ( + server.add_insecure_port("127.0.0.1:0") + if tls is None + else server.add_secure_port("127.0.0.1:0", grpc.ssl_server_credentials(((tls.key_pem, tls.cert_pem),))) + ) + server.start() + try: + yield _Riva(f"127.0.0.1:{port}", calls) + finally: + server.stop(grace=None).wait() -def test_sync_path_aggregates_only_final_results(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - # Empty heartbeat chunk: ignore. - _fake_response(results=[]), - # Interim chunk (not final): ignore. - _fake_response( - results=[ - _fake_result( - is_final=False, alternatives=[_fake_alternative("partial...")] - ) - ] - ), - # Two final chunks aggregated. - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[ - _fake_alternative( - "Hello,", - words=[_fake_word("Hello,", 0, 320)], - ) - ], - ) - ] - ), - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[ - _fake_alternative( - " world.", - words=[_fake_word("world.", 480, 870)], - ) - ], - ) - ] - ), - ] +def _samples(seconds: float) -> array[int]: + return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * _SAMPLE_RATE_HZ)))) - impl = NvidiaRivaAudioTranscription() - response: TranscriptionResponse = impl.audio_transcriptions( - model="nvidia/parakeet-ctc-1_1b-asr", - audio_file=_make_wav_bytes(), - optional_params={ - "language_code": "en-US", - "enable_word_time_offsets": True, - "response_format": "verbose_json", - "timestamp_granularities": ["word"], - }, - litellm_params={}, + +def _wav(seconds: float = 1.0) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(_SAMPLE_RATE_HZ) + writer.writeframes(_samples(seconds).tobytes()) + return buffer.getvalue() + + +def _carries_clip(audio: bytes, seconds: float) -> bool: + sent: Final = _samples(seconds) + if len(audio) != len(sent) * 2: + return False + return all(abs(left - right) <= 1 for left, right in zip(sent, array("h", audio), strict=True)) + + +def _alternative( + transcript: str, words: tuple[tuple[str, int, int], ...] = () +) -> riva_asr_pb2.SpeechRecognitionAlternative: + return riva_asr_pb2.SpeechRecognitionAlternative( + transcript=transcript, + words=[riva_asr_pb2.WordInfo(word=word, start_time=start, end_time=end) for word, start, end in words], + ) + + +def _streamed(transcript: str, is_final: bool = True) -> riva_asr_pb2.StreamingRecognizeResponse: + return riva_asr_pb2.StreamingRecognizeResponse( + results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[_alternative(transcript)], is_final=is_final)] + ) + + +def _offline(*transcripts: str) -> riva_asr_pb2.RecognizeResponse: + return riva_asr_pb2.RecognizeResponse( + results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative(text)]) for text in transcripts] + ) + + +def _answers(transcript: str) -> _Script: + return _Script(streaming=(_streamed(transcript),), offline=_offline(transcript)) + + +def _logging() -> Logging: + return Logging( + model="m", + messages=[], + stream=False, + call_type="transcription", + start_time=datetime(2026, 1, 1), + litellm_call_id=str(uuid.uuid4()), + function_id="riva-test", + ) + + +def _transcribe( + riva: _Riva, + optional_params: dict[str, object] | None = None, + *, + api_key: str | None = None, + timeout: float | str | httpx.Timeout | None = 60, + litellm_params: dict[str, object] | None = None, +) -> TranscriptionResponse: + return NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_wav(), + optional_params={"language_code": "en-US", **(optional_params or {})}, + litellm_params=litellm_params or {}, model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", - ) - - assert response.text == "Hello, world." - # duration is propagated from the resampler. - assert response._hidden_params["audio_transcription_duration"] == pytest.approx( - 1.0, abs=0.05 - ) - # word timestamps converted from ms to seconds. - words = response["words"] - assert words[0]["start"] == pytest.approx(0.0) - assert words[1]["end"] == pytest.approx(0.87) - assert ( - logging_obj.pre_call.call_args.kwargs["additional_args"]["atranscription"] - is False + timeout=timeout, + logging_obj=_logging(), + api_key=api_key, + api_base=riva.address, ) -def test_riva_offline_uses_unary_recognize_and_keeps_results_without_is_final( - mock_riva: SimpleNamespace, logging_obj: MagicMock -) -> None: - """ - Offline-only NIMs (Parakeet TDT, Whisper) reject StreamingRecognize. With - ``riva_offline`` the handler must call ``offline_recognize`` with the whole - resampled buffer and a bare ``RecognitionConfig``, and must keep every - result even though ``RecognizeResponse`` results carry no ``is_final``. - """ - mock_riva.offline_response["value"] = _fake_response( +def test_streaming_aggregates_only_final_results_and_their_word_timings() -> None: + heartbeat: Final = riva_asr_pb2.StreamingRecognizeResponse() + hello: Final = riva_asr_pb2.StreamingRecognizeResponse( results=[ - SimpleNamespace(alternatives=[_fake_alternative("Hello,", words=[_fake_word("Hello,", 0, 320)])]), - SimpleNamespace(alternatives=[_fake_alternative(" world.", words=[_fake_word("world.", 480, 870)])]), + riva_asr_pb2.StreamingRecognitionResult( + alternatives=[_alternative("Hello,", (("Hello,", 0, 320),))], is_final=True + ) ] ) - - impl = NvidiaRivaAudioTranscription() - response: TranscriptionResponse = impl.audio_transcriptions( - model="nvidia/parakeet-0.6b-tdt", - audio_file=_make_wav_bytes(), - optional_params={ - "language_code": "multi", - "riva_offline": True, - "enable_word_time_offsets": True, - "response_format": "verbose_json", - "timestamp_granularities": ["word"], - }, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + world: Final = riva_asr_pb2.StreamingRecognizeResponse( + results=[ + riva_asr_pb2.StreamingRecognitionResult( + alternatives=[_alternative(" world.", (("world.", 480, 870),))], is_final=True + ) + ] ) + with _riva_server(_Script(streaming=(heartbeat, _streamed("partial...", is_final=False), hello, world))) as riva: + response: Final = _transcribe( + riva, + {"enable_word_time_offsets": True, "response_format": "verbose_json", "timestamp_granularities": ["word"]}, + ) + call: Final = riva.only_call() assert response.text == "Hello, world." - assert response["words"][1]["end"] == pytest.approx(0.87) - assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0, abs=0.05) - assert mock_riva.streaming_calls == [] - assert len(mock_riva.offline_calls["audio_bytes"]) == pytest.approx(32000, abs=1600) - assert mock_riva.offline_calls["config"]._kwargs["language_code"] == "multi" - assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is True - - -def test_default_mode_streams_and_never_calls_offline_recognize( - mock_riva: SimpleNamespace, logging_obj: MagicMock -) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + assert response["words"] == [ + {"word": "Hello,", "start": 0.0, "end": 0.32}, + {"word": "world.", "start": 0.48, "end": 0.87}, ] + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + assert call.rpc == "StreamingRecognize" + assert _carries_clip(call.audio, 1.0), "streamed chunks do not add up to the resampled clip" - impl = NvidiaRivaAudioTranscription() - response = impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + +def test_riva_offline_uses_recognize_and_keeps_results_without_is_final() -> None: + offline: Final = riva_asr_pb2.RecognizeResponse( + results=[ + riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative("Hello,", (("Hello,", 0, 320),))]), + riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative(" world.", (("world.", 480, 870),))]), + ] ) - assert response.text == "streamed" - assert mock_riva.offline_calls == {} - assert len(mock_riva.streaming_calls) == 1 - assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is False - - - -@pytest.mark.parametrize(("riva_offline", "expect_offline"), [("false", False), ("False", False), ("true", True)]) -def test_riva_offline_string_from_multipart_form_is_parsed( - mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: str, expect_offline: bool -) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) - ] - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) - - response = NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": riva_offline}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - - assert response.text == ("offline" if expect_offline else "streamed") - assert (mock_riva.offline_calls != {}) is expect_offline - - -@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) -def test_riva_offline_malformed_value_is_rejected_before_any_rpc( - mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: object -) -> None: - with pytest.raises(NvidiaRivaException) as excinfo: - NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": riva_offline}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + with _riva_server(_Script(offline=offline)) as riva: + response: Final = _transcribe( + riva, + { + "language_code": "multi", + "riva_offline": True, + "enable_word_time_offsets": True, + "response_format": "verbose_json", + "timestamp_granularities": ["word"], + }, ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 400 - assert "riva_offline" in excinfo.value.message - assert mock_riva.offline_calls == {} - assert mock_riva.streaming_calls == [] + assert response.text == "Hello, world." + assert response["words"][1] == {"word": "world.", "start": 0.48, "end": 0.87} + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + assert call.rpc == "Recognize" + assert call.config.language_code == "multi" + assert _carries_clip(call.audio, 1.0), "offline request does not carry the whole resampled clip" -def test_riva_offline_malformed_value_is_dropped_under_drop_params( - mock_riva: SimpleNamespace, logging_obj: MagicMock -) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) - ] +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_recognition_config_reaches_riva_as_nvidia_decodes_it(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe( + riva, + { + "riva_offline": riva_offline, + "enable_word_time_offsets": True, + "riva_model_name": "parakeet", + "sample_rate_hertz": 8000, + "profanity_filter": True, + }, + ) + call: Final = riva.only_call() - response = NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": "yes"}, - litellm_params={"drop_params": True}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + assert call.config == riva_asr_pb2.RecognitionConfig( + encoding=riva_audio_pb2.AudioEncoding.LINEAR_PCM, + sample_rate_hertz=_SAMPLE_RATE_HZ, + language_code="en-US", + max_alternatives=1, + profanity_filter=True, + audio_channel_count=1, + enable_word_time_offsets=True, + enable_automatic_punctuation=True, + model="parakeet", ) + +def test_default_mode_streams_and_never_calls_recognize() -> None: + with _riva_server(_answers("streamed")) as riva: + response: Final = _transcribe(riva) + call: Final = riva.only_call() + assert response.text == "streamed" - assert mock_riva.offline_calls == {} - - -def test_deployment_riva_offline_reaches_the_handler_through_transcription(mock_riva: SimpleNamespace) -> None: - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) - - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, - ) - - assert response.text == "offline" - assert mock_riva.streaming_calls == [] - assert mock_riva.offline_calls != {} + assert call.rpc == "StreamingRecognize" @pytest.mark.parametrize( - ("timeout", "waited_for"), - [("7.5", 7.5), (httpx.Timeout(30.0, read=7.5), 7.5), (httpx.Timeout(5.0, read=None), 600.0), (None, None)], + ("riva_offline", "rpc"), [("false", "StreamingRecognize"), ("False", "StreamingRecognize"), ("true", "Recognize")] ) -def test_offline_recognition_waits_the_timeout_in_seconds_whatever_its_form( - mock_riva: SimpleNamespace, timeout: str | httpx.Timeout | None, waited_for: float | None -) -> None: - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) +def test_riva_offline_string_from_multipart_form_is_parsed(riva_offline: str, rpc: str) -> None: + with _riva_server(_Script(streaming=(_streamed("streamed"),), offline=_offline("offline"))) as riva: + response: Final = _transcribe(riva, {"riva_offline": riva_offline}) + call: Final = riva.only_call() - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, - timeout=timeout, - ) + assert call.rpc == rpc + assert response.text == ("offline" if rpc == "Recognize" else "streamed") + + +@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) +def test_riva_offline_malformed_value_is_rejected_before_any_rpc(riva_offline: object) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": riva_offline}) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 400 + assert "riva_offline" in excinfo.value.message + + +def test_riva_offline_malformed_value_is_dropped_under_drop_params() -> None: + with _riva_server(_Script(streaming=(_streamed("streamed"),), offline=_offline("offline"))) as riva: + response: Final = _transcribe(riva, {"riva_offline": "yes"}, litellm_params={"drop_params": True}) + call: Final = riva.only_call() + + assert response.text == "streamed" + assert call.rpc == "StreamingRecognize" + + +def test_deployment_riva_offline_reaches_the_handler_through_transcription() -> None: + with _riva_server(_Script(offline=_offline("offline"))) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=True, + ) + call: Final = riva.only_call() assert response.text == "offline" - assert mock_riva.offline_calls["timeout"] == waited_for + assert call.rpc == "Recognize" -@pytest.mark.parametrize("timeout", ["abc", "30s"]) -def test_malformed_timeout_on_offline_call_is_a_400_before_any_rpc( - mock_riva: SimpleNamespace, timeout: str +@pytest.mark.parametrize( + ("timeout", "deadline_seconds"), + [ + (12.5, 12.5), + ("12.5", 12.5), + (httpx.Timeout(60.0, read=12.5), 12.5), + (httpx.Timeout(5.0, read=None), 600.0), + (None, None), + ], +) +def test_streaming_deadline_is_the_timeout_in_seconds_whatever_its_form( + timeout: float | str | httpx.Timeout | None, deadline_seconds: float | None ) -> None: - with pytest.raises(NvidiaRivaException) as excinfo: - litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, + with _riva_server(_answers("streamed")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, timeout=timeout, ) + call: Final = riva.only_call() + + assert response.text == "streamed" + assert call.deadline_seconds == (None if deadline_seconds is None else pytest.approx(deadline_seconds, rel=0.25)) + + +@pytest.mark.parametrize("riva_offline", [False, True]) +@pytest.mark.parametrize("timeout", ["abc", "30s"]) +def test_malformed_timeout_is_a_400_before_any_rpc(riva_offline: bool, timeout: str) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=riva_offline, + timeout=timeout, + ) + + assert riva.calls.empty() assert excinfo.value.status_code == 400 assert "timeout" in excinfo.value.message - assert "audio_bytes" not in mock_riva.offline_calls -def test_malformed_timeout_on_offline_call_waits_the_default_under_drop_params( - mock_riva: SimpleNamespace, -) -> None: - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) - - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, - timeout="abc", - drop_params=True, - ) - - assert response.text == "offline" - assert mock_riva.offline_calls["timeout"] == 600.0 - - -def test_malformed_timeout_still_streams_when_the_client_takes_no_timeout(mock_riva: SimpleNamespace) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) - ] - - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-ctc-1.1b", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - timeout="abc", - ) - - assert response.text == "streamed" - assert len(mock_riva.streaming_calls) == 1 - - -def test_oversized_offline_audio_is_a_non_retryable_413(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: - class FakeGrpcError(Exception): - def code(self) -> SimpleNamespace: - return SimpleNamespace(name="RESOURCE_EXHAUSTED") - - def details(self) -> str: - return "SERVER: Received message larger than max (4348825 vs. 4194304)" - - mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") - - with pytest.raises(NvidiaRivaException) as excinfo: - NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": True}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", +def test_malformed_streaming_timeout_falls_back_to_the_default_deadline_under_drop_params() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + timeout="abc", + drop_params=True, ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 413 + assert response.text == "ok" + assert call.rpc == "StreamingRecognize" + assert call.deadline_seconds == pytest.approx(600.0, rel=0.25) + + +def test_malformed_offline_timeout_is_ignored_under_drop_params() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=True, + timeout="abc", + drop_params=True, + ) + call: Final = riva.only_call() + + assert response.text == "ok" + assert call.rpc == "Recognize" + + +@pytest.mark.parametrize( + "timeout", [_STALL_TIMEOUT_SECONDS, str(_STALL_TIMEOUT_SECONDS), httpx.Timeout(30.0, read=_STALL_TIMEOUT_SECONDS)] +) +def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408(timeout: float | str | httpx.Timeout) -> None: + with _riva_server(_Script(stall=True)) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": True}, timeout=timeout) + call: Final = riva.only_call() + + assert excinfo.value.status_code == 408 + assert f"within {_STALL_TIMEOUT_SECONDS} seconds" in excinfo.value.message + assert call.cancelled is True + + +def test_stalled_streaming_recognition_hits_its_deadline_as_504() -> None: + with _riva_server(_Script(stall=True)) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, timeout=_STALL_TIMEOUT_SECONDS) + call: Final = riva.only_call() + + assert excinfo.value.status_code == 504, excinfo.value.message + assert "DEADLINE_EXCEEDED" in excinfo.value.message + assert call.cancelled is True + + +def test_oversized_offline_audio_is_a_non_retryable_413() -> None: + with _riva_server(_answers("ok"), max_receive_bytes=4096) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": True}) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 413, excinfo.value.message assert litellm.utils.should_retry(excinfo.value.status_code) is False -def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408( - mock_riva: SimpleNamespace, logging_obj: MagicMock +@pytest.mark.parametrize( + ("riva_offline", "code", "status_code"), + [ + (False, grpc.StatusCode.UNAUTHENTICATED, 401), + (True, grpc.StatusCode.INVALID_ARGUMENT, 400), + (False, grpc.StatusCode.RESOURCE_EXHAUSTED, 429), + ], +) +def test_grpc_error_is_wrapped_as_nvidia_riva_exception( + riva_offline: bool, code: grpc.StatusCode, status_code: int ) -> None: - class FutureTimeoutError(Exception): - pass + with _riva_server(_Script(abort=(code, "Riva said no"))) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": riva_offline}, api_key=_API_KEY) - class StalledRecognizeFuture: - def __init__(self) -> None: - self.waited_for: float | None = None - self.cancelled = False + assert excinfo.value.status_code == status_code + assert code.name in excinfo.value.message + assert "Riva said no" in excinfo.value.message - def result(self, timeout: float | None = None) -> object: - self.waited_for = timeout - raise FutureTimeoutError() - def done(self) -> bool: - return self.cancelled +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_self_hosted_api_key_is_sent_as_one_bearer_token(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"riva_offline": riva_offline}, api_key=_API_KEY) + call: Final = riva.only_call() - def cancel(self) -> bool: - self.cancelled = True - return True + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + assert call.header("function-id") == () - stalled = StalledRecognizeFuture() - mock_riva.offline_future["value"] = stalled +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_nvcf_function_id_and_api_key_are_each_sent_once(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe( + riva, {"riva_offline": riva_offline, "nvcf_function_id": "abc-123", "use_ssl": False}, api_key=_API_KEY + ) + call: Final = riva.only_call() + + assert call.header("function-id") == ("abc-123",) + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + + +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_no_api_key_sends_no_authorization(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"riva_offline": riva_offline}) + call: Final = riva.only_call() + + assert call.header("authorization") == () + assert call.header("function-id") == () + + +@pytest.mark.parametrize("riva_offline", [False, True]) +@pytest.mark.parametrize( + ("optional_params", "function_ids"), + [({"use_ssl": True}, ()), ({"nvcf_function_id": "abc-123"}, ("abc-123",))], + ids=["use_ssl", "nvcf_default_tls"], +) +def test_tls_api_key_is_sent_as_one_bearer_token( + monkeypatch: pytest.MonkeyPatch, + optional_params: dict[str, object], + function_ids: tuple[str, ...], + riva_offline: bool, +) -> None: + tls: Final = _self_signed_loopback_tls() + system_roots: Final = grpc.ssl_channel_credentials + monkeypatch.setattr(grpc, "ssl_channel_credentials", lambda: system_roots(root_certificates=tls.cert_pem)) + with _riva_server(_answers("ok"), tls=tls) as riva: + response: Final = _transcribe(riva, {**optional_params, "riva_offline": riva_offline}, api_key=_API_KEY) + call: Final = riva.only_call() + + assert response.text == "ok" + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + assert call.header("function-id") == function_ids + + +@pytest.mark.parametrize("optional_params", [{"use_ssl": True}, {"nvcf_function_id": "abc-123"}]) +def test_tls_never_falls_back_to_plaintext(optional_params: dict[str, object]) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, optional_params, api_key=_API_KEY, timeout=_SERVER_WAIT_SECONDS) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 503, excinfo.value.message + + +def test_explicit_endpointing_config_keeps_zero_thresholds() -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"endpointing_config": {"stop_threshold": 0.0, "stop_history": 800}}) + call: Final = riva.only_call() + + assert call.config.HasField("endpointing_config") + assert call.config.endpointing_config.HasField("stop_threshold") + assert call.config.endpointing_config == riva_asr_pb2.EndpointingConfig(stop_threshold=0.0, stop_history=800) + + +def test_unknown_endpointing_field_falls_back_to_riva_defaults() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = _transcribe(riva, {"endpointing_config": {"stop_threshold": 0.5, "not_a_field": 1}}) + call: Final = riva.only_call() + + assert response.text == "ok" + assert not call.config.HasField("endpointing_config") + + +def test_missing_api_base_raises_clear_error() -> None: with pytest.raises(NvidiaRivaException) as excinfo: NvidiaRivaAudioTranscription().audio_transcriptions( model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": True}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=7.5, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - - assert excinfo.value.status_code == 408 - assert stalled.waited_for == 7.5 - assert stalled.cancelled is True - -def test_riva_offline_grpc_error_is_wrapped(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: - class FakeGrpcError(Exception): - def code(self) -> SimpleNamespace: - return SimpleNamespace(name="INVALID_ARGUMENT") - - def details(self) -> str: - return "Unavailable model requested" - - mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") - - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "multi", "riva_offline": True}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - assert excinfo.value.status_code == 400 - assert "INVALID_ARGUMENT" in excinfo.value.message - - -def test_auth_nvcf_defaults_use_ssl_and_attaches_function_id(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[_fake_alternative("ok")], - ) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={ - "nvcf_function_id": "abc-123", - "language_code": "en-US", - }, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", - ) - - assert mock_riva.auth_calls["uri"] == "grpc.nvcf.nvidia.com:443" - assert mock_riva.auth_calls["use_ssl"] is True - metadata = dict(mock_riva.auth_calls["metadata_args"]) - assert metadata["function-id"] == "abc-123" - assert metadata["authorization"] == "Bearer nvapi-xxx" - - -def test_auth_self_hosted_defaults_no_ssl_and_no_function_id(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - - assert mock_riva.auth_calls["uri"] == "localhost:50051" - assert mock_riva.auth_calls["use_ssl"] is False - metadata = dict(mock_riva.auth_calls["metadata_args"]) - # No function-id, no authorization metadata. - assert "function-id" not in metadata - assert "authorization" not in metadata - - -def test_explicit_use_ssl_override_wins(mock_riva, logging_obj): - """ - Self-hosted Riva behind an ingress with TLS termination is a real - deployment topology. ``use_ssl=True`` must be honored even without an - NVCF function id. - """ - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"use_ssl": True, "language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="riva.internal.company.com:443", - ) - - assert mock_riva.auth_calls["use_ssl"] is True - - -def test_missing_api_base_raises_clear_error(mock_riva, logging_obj): - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), + audio_file=_wav(), optional_params={}, litellm_params={}, model_response=TranscriptionResponse(), timeout=60, - logging_obj=logging_obj, + logging_obj=_logging(), api_key=None, api_base=None, ) + + assert excinfo.value.status_code == 400 assert "api_base" in excinfo.value.message -def test_async_path_uses_to_thread(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result( - is_final=True, alternatives=[_fake_alternative("async ok")] - ) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - response = asyncio.run( - impl.async_audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - ) - assert response.text == "async ok" - assert ( - logging_obj.pre_call.call_args.kwargs["additional_args"]["atranscription"] - is True - ) - - -@pytest.mark.parametrize("timeout", [12.5, "12.5", httpx.Timeout(60.0, read=12.5)]) -def test_timeout_is_forwarded_to_streaming_generator_when_supported( - mock_riva: SimpleNamespace, logging_obj: MagicMock, timeout: float | str | httpx.Timeout -) -> None: - """ - Without a deadline the gRPC stream can block forever on a stalled Riva - server. The handler must forward the call-level ``timeout`` to - ``streaming_response_generator`` whenever the installed riva-client - accepts a ``timeout`` kwarg. - """ - captured_kwargs = {} - - def streaming_with_timeout(self, audio_chunks, streaming_config, timeout=None): - captured_kwargs["timeout"] = timeout - list(audio_chunks) - yield from [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] +def test_async_path_transcribes() -> None: + with _riva_server(_answers("async ok")) as riva: + response: Final = asyncio.run( + NvidiaRivaAudioTranscription().async_audio_transcriptions( + model="m", + audio_file=_wav(), + optional_params={"language_code": "en-US"}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=_logging(), + api_key=_API_KEY, + api_base=riva.address, ) - ] - - mock_riva.client.ASRService.streaming_response_generator = streaming_with_timeout - - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=timeout, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - assert captured_kwargs["timeout"] == pytest.approx(12.5) - - -def test_grpc_error_is_wrapped_as_nvidia_riva_exception(mock_riva, logging_obj): - class FakeGrpcError(Exception): - def code(self): - return SimpleNamespace(name="UNAUTHENTICATED") - - def details(self): - return "bad token" - - def raising_streaming_response_generator(self, audio_chunks, streaming_config): - list(audio_chunks) - raise FakeGrpcError("rpc fail") - - mock_riva.client.ASRService.streaming_response_generator = ( - raising_streaming_response_generator - ) - - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 401 - assert "UNAUTHENTICATED" in excinfo.value.message + assert response.text == "async ok" + assert call.header("authorization") == (f"Bearer {_API_KEY}",) diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py new file mode 100644 index 00000000000..124e3c28291 --- /dev/null +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py @@ -0,0 +1,41 @@ +from collections.abc import Iterator +from typing import Final + +from google.protobuf.descriptor import FieldDescriptor +from riva.client.proto import riva_asr_pb2, riva_audio_pb2 + +from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema + + +def _declared_fields() -> Iterator[tuple[str, FieldDescriptor]]: + for message_name, _ in riva_asr_schema._MESSAGES: + message: Final = riva_asr_schema._POOL.FindMessageTypeByName(f"{riva_asr_schema._PACKAGE}.{message_name}") + yield from ((message_name, field) for field in message.fields) + + +def _wire_shape(field: FieldDescriptor | None) -> tuple[object, ...] | None: + if field is None: + return None + return ( + field.number, + field.type, + field.label, + field.has_presence, + field.message_type.full_name if field.message_type else None, + field.enum_type.name if field.enum_type else None, + field.containing_oneof.name if field.containing_oneof else None, + ) + + +def _nvidia_field(message_name: str, field_name: str) -> FieldDescriptor | None: + return riva_asr_pb2.DESCRIPTOR.message_types_by_name[message_name].fields_by_name.get(field_name) + + +def test_every_declared_field_matches_nvidias_riva_schema() -> None: + declared: Final = tuple(_declared_fields()) + + assert len(declared) == sum(len(fields) for _, fields in riva_asr_schema._MESSAGES) + assert {f"{message}.{field.name}": _wire_shape(field) for message, field in declared} == { + f"{message}.{field.name}": _wire_shape(_nvidia_field(message, field.name)) for message, field in declared + } + assert riva_asr_schema.LINEAR_PCM == riva_audio_pb2.AudioEncoding.LINEAR_PCM diff --git a/tests/unit/llms/openai_like/test_dynamic_config.py b/tests/unit/llms/openai_like/test_dynamic_config.py index de70f98c3f1..f26f0821c4b 100644 --- a/tests/unit/llms/openai_like/test_dynamic_config.py +++ b/tests/unit/llms/openai_like/test_dynamic_config.py @@ -139,3 +139,24 @@ class TestForceStoreFalse: headers={}, ) assert params["store"] is True + + +class TestWebSearchRequestsFromToolUsage: + def test_counts_only_paid_search_tools(self): + payload = {"usage": {"tool_usage": {"web_search": 2, "image_search": 1, "code_interpreter": 3, "web_extractor": 1}}} + assert dynamic_config.web_search_requests_from_tool_usage(payload) == 3 + + @pytest.mark.parametrize( + "payload", + [ + {"usage": {"prompt_tokens": 1}}, + {"usage": {"tool_usage": {"code_interpreter": 2}}}, + {"usage": {"tool_usage": {"web_search": "2"}}}, + {"usage": {"tool_usage": {"web_search": True}}}, + {"usage": {"tool_usage": []}}, + {"usage": None}, + "not a mapping", + ], + ) + def test_anything_but_a_positive_search_count_is_none(self, payload): + assert dynamic_config.web_search_requests_from_tool_usage(payload) is None diff --git a/tests/unit/llms/openai_like/test_empiriolabs_provider.py b/tests/unit/llms/openai_like/test_empiriolabs_provider.py index a8acce5cc9d..58d8a2371d8 100644 --- a/tests/unit/llms/openai_like/test_empiriolabs_provider.py +++ b/tests/unit/llms/openai_like/test_empiriolabs_provider.py @@ -2,13 +2,19 @@ Unit tests for the EmpirioLabs OpenAI-like provider. """ +import json import os import sys +from typing import Final + +import pytest +import respx sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) ) +import litellm from litellm.llms.openai_like.dynamic_config import create_config_class from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -61,3 +67,114 @@ def test_empiriolabs_complete_url_appends_endpoint(): stream=False, ) assert url == f"{EMPIRIOLABS_BASE_URL}/chat/completions" + + +def test_empiriolabs_chat_models_still_route_through_the_json_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("EMPIRIOLABS_API_KEY", "environment-key") + monkeypatch.delenv("EMPIRIOLABS_API_BASE", raising=False) + + with respx.mock() as upstream: + route: Final = upstream.post(f"{EMPIRIOLABS_BASE_URL}/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_empiriolabs", + "object": "chat.completion", + "created": 1_789_550_000, + "model": "qwen3-7-plus", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello from EmpirioLabs AI"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 5, "total_tokens": 9}, + }, + ) + response: Final = litellm.completion( + model="empiriolabs/qwen3-7-plus", + messages=[{"role": "user", "content": "Say hello"}], + max_completion_tokens=16, + ) + + request: Final = route.calls.last.request + assert litellm.get_llm_provider("empiriolabs/qwen3-7-plus") == ( + "qwen3-7-plus", + "empiriolabs", + "environment-key", + EMPIRIOLABS_BASE_URL, + ) + assert route.call_count == 1 + assert request.headers["authorization"] == "Bearer environment-key" + assert json.loads(request.content) == { + "model": "qwen3-7-plus", + "messages": [{"role": "user", "content": "Say hello"}], + "max_tokens": 16, + } + assert response.choices[0].message.content == "Hello from EmpirioLabs AI" + + +EMBEDDING_RESPONSE: Final = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-v4", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, +} + + +def test_empiriolabs_embeddings_post_to_v1_embeddings_with_the_bearer_key(): + with respx.mock() as upstream: + route: Final = upstream.post(f"{EMPIRIOLABS_BASE_URL}/embeddings").respond(200, json=EMBEDDING_RESPONSE) + response: Final = litellm.embedding( + model="empiriolabs/text-embedding-v4", + input=["Say hello"], + api_key="empiriolabs-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert request.headers["authorization"] == "Bearer empiriolabs-test-key" + assert body["model"] == "text-embedding-v4" + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +SEARCHED_CHAT_RESPONSE: Final = { + "id": "chatcmpl_empiriolabs_search", + "object": "chat.completion", + "created": 1_789_550_000, + "model": "qwen3-7-plus", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Two searched answers"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 4, + "completion_tokens": 5, + "total_tokens": 9, + "tool_usage": {"web_search": 2, "code_interpreter": 1}, + }, +} + + +def test_empiriolabs_web_search_calls_in_tool_usage_are_billed_per_call(local_model_cost_map: None) -> None: + with respx.mock() as upstream: + upstream.post(f"{EMPIRIOLABS_BASE_URL}/chat/completions").respond(200, json=SEARCHED_CHAT_RESPONSE) + response: Final = litellm.completion( + model="empiriolabs/qwen3-7-plus", + messages=[{"role": "user", "content": "Search twice"}], + api_key="empiriolabs-test-key", + ) + + row: Final = litellm.model_cost["empiriolabs/qwen3-7-plus"] + token_cost: Final = 4 * row["input_cost_per_token"] + 5 * row["output_cost_per_token"] + per_search: Final = row["search_context_cost_per_query"]["search_context_size_medium"] + assert response.usage.server_tool_use.web_search_requests == 2 + assert litellm.completion_cost( + completion_response=response, model="empiriolabs/qwen3-7-plus", custom_llm_provider="empiriolabs" + ) == pytest.approx(token_cost + 2 * per_search) diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 2eecb83e2df..03bc308a89d 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1785,6 +1785,289 @@ def test_vertex_ai_map_tools(): assert tools == new_tools +def test_vertex_ai_map_tools_declares_a_namespace_by_its_members() -> None: + namespace = { + "type": "namespace", + "name": "github", + "tools": [ + {"type": "function", "name": "create_issue", "parameters": {"type": "object", "properties": {}}}, + {"type": "function", "function": {"name": "close_issue"}}, + ], + } + + tools = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + + assert [declaration["name"] for declaration in tools[0]["function_declarations"]] == [ + "create_issue", + "close_issue", + ] + + +def test_vertex_ai_map_tools_declares_a_namespace_member_like_the_same_chat_tool() -> None: + parameters = { + "type": "object", + "properties": {"repo": {"$ref": "#/$defs/Repo"}, "label": {"type": ["string", "null"]}}, + "$defs": {"Repo": {"type": "string"}}, + } + function = {"name": "create_issue", "description": "Open an issue", "parameters": parameters} + namespace = {"type": "namespace", "name": "github", "tools": [{"type": "function", **function}]} + + from_namespace = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + from_chat_tool = VertexGeminiConfig().map_function( + value=[{"type": "function", "function": function}], optional_params={} + ) + + assert from_namespace == from_chat_tool + + +_GET_WEATHER = {"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}} + + +_WEATHER_NAMESPACE = {"type": "namespace", "name": "weather", "tools": [_GET_WEATHER]} +_TOP_LEVEL_GET_WEATHER = {"type": "function", "function": {k: v for k, v in _GET_WEATHER.items() if k != "type"}} + + +@pytest.mark.parametrize( + ("tools", "declared"), + [ + ( + [{"type": "namespace", "name": "forecast", "tools": [_GET_WEATHER]}, _WEATHER_NAMESPACE], + ["forecast__get_weather", "weather__get_weather"], + ), + ([_TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], ["get_weather", "weather__get_weather"]), + ([_WEATHER_NAMESPACE, _TOP_LEVEL_GET_WEATHER], ["weather__get_weather", "get_weather"]), + ], + ids=["namespaces_share_a_member", "top_level_first", "namespace_first"], +) +def test_vertex_ai_map_tools_declares_qualified_member_names_when_another_tool_declares_a_member_name( + tools: list[dict], declared: list[str] +) -> None: + mapped = VertexGeminiConfig().map_function(value=tools, optional_params={}) + + assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == declared + + +def test_vertex_ai_map_tools_declares_a_member_a_namespace_lists_twice_once_from_its_first_entry() -> None: + first = {**_GET_WEATHER, "description": "Current weather"} + repeat = {**_GET_WEATHER, "description": "Forecast"} + namespace = {"type": "namespace", "name": "weather", "tools": [first, repeat]} + + mapped = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + + declarations = mapped[0]["function_declarations"] + assert [(declaration["name"], declaration["description"]) for declaration in declarations] == [ + ("get_weather", "Current weather") + ] + + +@pytest.mark.parametrize( + ("tools", "declared"), + [ + ( + [ + _WEATHER_NAMESPACE, + _TOP_LEVEL_GET_WEATHER, + {"type": "function", "function": {"name": "weather__get_weather", "description": "Top level"}}, + ], + [("get_weather", None), ("weather__get_weather", "Top level")], + ), + ( + [ + {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]}, + {"type": "namespace", "name": "mcp", "tools": [{"type": "function", "name": "ops__exec"}]}, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "ops__exec"}}, + ], + [("mcp__ops__exec", None), ("exec", None), ("ops__exec", None)], + ), + ( + [_TOP_LEVEL_GET_WEATHER, _TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], + [("get_weather", None), ("get_weather", None), ("weather__get_weather", None)], + ), + ], + ids=["top_level_function_has_the_qualified_name", "namespaces_build_one_qualified_name", "client_repeats_a_function"], +) +def test_vertex_ai_map_tools_never_repeats_a_generated_name( + tools: list[dict], declared: list[tuple[str, str | None]] +) -> None: + mapped = VertexGeminiConfig().map_function(value=tools, optional_params={}) + + declarations = mapped[0]["function_declarations"] + assert [(declaration["name"], declaration.get("description")) for declaration in declarations] == declared + + +def test_vertex_ai_map_openai_params_qualifies_members_by_the_namespace_top_level_name() -> None: + namespace = { + "type": "namespace", + "name": "safe", + "function": {"name": "lookup"}, + "tools": [{"type": "function", "name": "lookup"}], + } + + mapped = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [namespace], + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + declared = [declaration["name"] for tool in mapped["tools"] for declaration in tool["function_declarations"]] + assert declared == ["safe__lookup"] + assert mapped["tool_choice"]["functionCallingConfig"]["allowed_function_names"] == ["safe__lookup"] + + +def test_vertex_ai_map_openai_params_declares_qualified_member_names_when_functions_repeat_a_member() -> None: + mapped = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [{"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]}], + "functions": [{"name": "lookup", "parameters": {"type": "object", "properties": {}}}], + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + declared = [declaration["name"] for tool in mapped["tools"] for declaration in tool["function_declarations"]] + assert declared == ["mcp__weather__lookup", "lookup"] + + +def test_vertex_ai_map_tools_keeps_repeated_chat_tools_when_no_namespace_is_sent() -> None: + function = {k: v for k, v in _GET_WEATHER.items() if k != "type"} + + mapped = VertexGeminiConfig().map_function( + value=[{"type": "function", "function": function}, {"type": "function", "function": function}], + optional_params={}, + ) + + assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == ["get_weather", "get_weather"] + + +@pytest.mark.parametrize( + ("forced_name", "other_tools", "allowed_function_names"), + [ + ("github", [], ["create_issue", "close_issue"]), + ("close_issue", [], ["close_issue"]), + ("github", [{"type": "function", "function": {"name": "github"}}], ["github"]), + ( + "github", + [{"type": "function", "function": {"name": "close_issue"}}], + ["github__create_issue", "github__close_issue"], + ), + ( + "close_issue", + [{"type": "namespace", "name": "gitlab", "tools": [{"type": "function", "name": "close_issue"}]}], + ["github__close_issue", "gitlab__close_issue"], + ), + ( + "github", + [ + {"type": "function", "function": {"name": "close_issue"}}, + {"type": "function", "function": {"name": "github__close_issue"}}, + ], + ["github__create_issue"], + ), + ( + "close_issue", + [ + {"type": "namespace", "name": "gitlab", "tools": [{"type": "function", "name": "close_issue"}]}, + {"type": "function", "function": {"name": "github__close_issue"}}, + ], + ["gitlab__close_issue"], + ), + ( + "mcp", + [ + {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]}, + { + "type": "namespace", + "name": "mcp", + "tools": [{"type": "function", "name": "ops__exec"}, {"type": "function", "name": "status"}], + }, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "ops__exec"}}, + ], + ["mcp__status"], + ), + ], + ids=[ + "names_the_namespace", + "names_a_member", + "names_a_function_sharing_the_namespace_name", + "names_a_namespace_whose_member_a_function_repeats", + "names_a_member_two_namespaces_share", + "names_a_namespace_whose_generated_name_a_function_holds", + "names_a_member_whose_generated_name_a_function_holds", + "names_a_namespace_whose_generated_name_another_namespace_holds", + ], +) +def test_vertex_ai_tool_choice_forces_only_functions_gemini_is_sent( + forced_name: str, other_tools: list[dict[str, object]], allowed_function_names: list[str] +) -> None: + namespace = { + "type": "namespace", + "name": "github", + "tools": [{"type": "function", "name": "create_issue"}, {"type": "function", "name": "close_issue"}], + } + + mapped = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [namespace, *other_tools], + "tool_choice": {"type": "function", "function": {"name": forced_name}}, + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + declared = [declaration["name"] for declaration in mapped["tools"][0]["function_declarations"]] + assert mapped["tool_choice"]["functionCallingConfig"]["allowed_function_names"] == allowed_function_names + assert set(allowed_function_names) <= set(declared) + + +@pytest.mark.parametrize( + ("forced_name", "tools"), + [ + ( + "weather", + [ + {"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "lookup"}]}, + {"type": "function", "function": {"name": "lookup"}}, + {"type": "function", "function": {"name": "weather__lookup"}}, + ], + ), + ( + "ops__exec", + [ + {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]}, + {"type": "namespace", "name": "ops__exec", "tools": [{"type": "function", "name": "status"}]}, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "status"}}, + {"type": "function", "function": {"name": "ops__exec__status"}}, + ], + ), + ], + ids=["its_members_lose_to_functions", "another_namespace_generates_its_name"], +) +def test_vertex_ai_tool_choice_naming_a_namespace_gemini_gets_no_member_of( + forced_name: str, tools: list[dict[str, object]] +) -> None: + params = {"tools": tools, "tool_choice": {"type": "function", "function": {"name": forced_name}}} + + with pytest.raises(litellm.UnsupportedParamsError, match=f"tool_choice '{forced_name}'"): + VertexGeminiConfig().map_openai_params( + non_default_params=params, optional_params={}, model="gemini-3.8-flash", drop_params=False + ) + dropped = VertexGeminiConfig().map_openai_params( + non_default_params=params, optional_params={}, model="gemini-3.8-flash", drop_params=True + ) + + assert "tool_choice" not in dropped + assert dropped["tools"][0]["function_declarations"] + + def test_vertex_ai_map_tool_with_anyof(): """ Related issue: https://github.com/BerriAI/litellm/issues/11164 diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 88da07c612b..0aacc1f8bda 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1268,6 +1268,153 @@ async def test_task_method_failure_hook_uses_enriched_request_data(): assert failure_data.get("agent_id") == "test-agent" +_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" + + +class _JsonRpcError(BaseModel): + model_config = ConfigDict(frozen=True) + + code: int + message: str + + +class _JsonRpcErrorBody(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _JsonRpcError + + +async def _send_message_to_entra_agent( + litellm_params: Mapping[str, str], *patches: AbstractContextManager[object] +) -> tuple[int, _JsonRpcError]: + agent = _make_agent_mock() + agent.litellm_params = dict(litellm_params) + mock_request = _make_request_mock( + "message/send", + {"message": {"role": "user", "parts": [{"kind": "text", "text": "Hello"}], "messageId": "msg-1"}}, + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in (*_base_patches(agent), *patches): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"), + ) + + return response.status_code, _JsonRpcErrorBody.model_validate_json(response.body).error + + +@pytest.mark.asyncio +async def test_entra_rejecting_the_agent_credential_returns_401_with_only_the_aadsts_reason(): + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{_ENTRA_APP_CLIENT_ID}'." + ) + + status_code, error = await _send_message_to_entra_agent( + {"tenant_id": "tenant", "client_id": _ENTRA_APP_CLIENT_ID, "client_secret": "wrong-secret"}, + patch( # test-quality-ok: stands in for the Entra token endpoint answering a wrong secret, the shape azure-identity raises + "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=rejected_token + ), + ) + + assert status_code == 401 + assert "AADSTS7000215" in error.message + assert _ENTRA_APP_CLIENT_ID not in error.message + + +@pytest.mark.asyncio +async def test_entra_rejecting_the_agent_oidc_token_returns_401_with_only_the_aadsts_reason(monkeypatch, respx_mock): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + respx_mock.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 400, + json={ + "error": "invalid_client", + "error_description": "AADSTS700213: No matching federated identity record found for app " + f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID}", + }, + ) + ) + + status_code, error = await _send_message_to_entra_agent( + { + "tenant_id": "qa-tenant", + "client_id": _ENTRA_APP_CLIENT_ID, + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + ) + + assert status_code == 401 + assert error.code == -32000 + assert "AADSTS700213" in error.message + assert _ENTRA_APP_CLIENT_ID not in error.message + assert _ENTRA_TRACE_ID not in error.message + + +@pytest.mark.asyncio +async def test_upstream_authentication_error_without_an_entra_rejection_stays_an_internal_error(): + agent = _make_agent_mock() + mock_handler = MagicMock() + mock_handler.post = AsyncMock( + side_effect=litellm.AuthenticationError( + message="Incorrect API key provided.", llm_provider="openai", model="gpt-4.1" + ) + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: stands in for the upstream agent rejecting the provider key it was sent + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=mock_handler + ) + ) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=_make_request_mock("tasks/get", {"id": "task-1"}), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"), + ) + + assert response.status_code == 500 + assert json.loads(response.body.decode())["error"]["code"] == -32603 + + @pytest.mark.asyncio async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400(): from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py b/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py index 531d97efe64..04b81605e29 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py +++ b/tests/unit/proxy/anthropic_endpoints/test_skill_upload_archive.py @@ -8,10 +8,8 @@ import pytest from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import ( MAX_SKILL_ARCHIVE_COMPRESSION_RATIO, - MAX_SKILL_UPLOAD_BYTES, build_uploaded_skill_archive, ) -from litellm.proxy.discovery_endpoints.agent_skills_archive import MAX_ARCHIVE_ENTRIES, MAX_ARCHIVE_UNPACKED_BYTES MANIFEST: Final = b"---\nname: skill\ndescription: Notes\n---\nRead guide.md" @@ -64,29 +62,3 @@ def test_installable_archive_preserves_files_and_safe_modes( assert {entry.compress_type for entry in archive.infolist()} == {compression} assert sum(entry.file_size for entry in archive.infolist()) <= len(result) * MAX_SKILL_ARCHIVE_COMPRESSION_RATIO assert build_uploaded_skill_archive(upload, "registered-name") == result - - -_VALID: Final = _zip((("SKILL.md", MANIFEST),)) -_ENCRYPTED: Final = _zip((("__MACOSX/encrypted", b"secret"), ("SKILL.md", MANIFEST))) -_CENTRAL: Final = _ENCRYPTED.index(b"PK\x01\x02") - - -@pytest.mark.parametrize("content, name", [ - (b"not a zip", "skill"), - (_zip((("SKILL.md", MANIFEST),), zipfile.ZIP_STORED).replace(b"Read guide.md", b"Dead guide.md"), "skill"), - (_ENCRYPTED[:_CENTRAL + 8] + b"\x01\x00" + _ENCRYPTED[_CENTRAL + 10:], "skill"), - *[(_VALID, name) for name in ("../escape", "", "Not a slug", "x" * 65536)], - *[(_zip(((path, MANIFEST),)), "skill") for path in ("readme.md", "wrapper/nested/SKILL.md", "skills/example/SKILL.md")], - *[(_zip((("SKILL.md", MANIFEST), (path, b"unsafe"))), "skill") for path in ( - "../escape", "/absolute", "C:/escape", "a\\escape", "__MACOSX/../escape", "bad/./file", - )], - *[(_zip((("SKILL.md", MANIFEST), (_entry(path, stat.S_IFLNK | 0o777), b"../../outside"))), "skill") - for path in ("__MACOSX/link", "scripts/link")], - (_zip(((".claude-plugin/plugin.json", b'{"name":"plugin"}'), ("skills/example/SKILL.md", MANIFEST))), "skill"), - (_zip((("SKILL.md", MANIFEST), ("__MACOSX/a", b"one"), ("__MACOSX//a", b"two"))), "skill"), - (_zip((("SKILL.md", MANIFEST), ("asset.bin", bytes(MAX_SKILL_UPLOAD_BYTES))), zipfile.ZIP_STORED), "skill"), - (_zip((("SKILL.md", MANIFEST), ("__MACOSX/large", bytes(MAX_ARCHIVE_UNPACKED_BYTES)))), "skill"), - (_zip((("SKILL.md", MANIFEST),) + tuple((f"__MACOSX/{i}", b"") for i in range(MAX_ARCHIVE_ENTRIES))), "skill"), -]) -def test_rejects_unsafe_invalid_or_oversized_archives(content: bytes, name: str) -> None: - assert build_uploaded_skill_archive(content, name) is None diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 5b37cee88c5..92114b00398 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -24,12 +24,12 @@ import httpx import pytest import respx from fastapi import HTTPException +from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER import litellm from litellm.constants import MANAGED_OBJECT_STALENESS_CUTOFF_DAYS from litellm.types.router import LiteLLM_Params from litellm.types.utils import LiteLLMBatch -from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER if TYPE_CHECKING: from litellm.batches.batch_utils import BatchCostUsageResult @@ -254,7 +254,6 @@ class TestCheckBatchCost: @pytest.mark.asyncio async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" - from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) # Simulate column already known absent from a previous cycle @@ -2387,8 +2386,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: @pytest.mark.asyncio async def test_target_model_names_comes_from_input_file_not_provider_model(self): from litellm.proxy.openai_files_endpoints.common_utils import ( - is_base64_encoded_unified_file_id, get_models_from_unified_file_id, + is_base64_encoded_unified_file_id, ) output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) @@ -2422,8 +2421,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: @pytest.mark.asyncio async def test_falls_back_to_deployment_model_group_without_managed_input_file(self): from litellm.proxy.openai_files_endpoints.common_utils import ( - is_base64_encoded_unified_file_id, get_models_from_unified_file_id, + is_base64_encoded_unified_file_id, ) output_file_id = await self._run(self._job("file-raw-provider-input")) @@ -3309,7 +3308,9 @@ class TestMultiPodBatchCostClaim: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ["gpt-4"]), + return_value=_batch_cost_result( + 0.01, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ["gpt-4"] + ), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -3329,15 +3330,15 @@ class TestMultiPodBatchCostClaim: from datetime import datetime import httpx - import litellm import respx + import litellm from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.utils import PrismaClient + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import SpecialEnums @@ -3383,6 +3384,14 @@ class TestMultiPodBatchCostClaim: database_url="postgresql://unused:unused@localhost:5432/unused", proxy_logging_obj=dependency ) create_spend_rows: Final = AsyncMock(return_value=1) + from litellm.proxy.db.prisma_client import PrismaWrapper + + spend_prisma.db = PrismaWrapper( + original_prisma=MagicMock( + query_raw=AsyncMock(return_value=[]), + litellm_spendlogs=MagicMock(create_many=create_spend_rows), + ) + ) input_rate: Final = 0.000002 litellm.register_model( { @@ -3434,11 +3443,6 @@ class TestMultiPodBatchCostClaim: respx.mock, patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency), patch("litellm.proxy.proxy_server.prisma_client", spend_prisma), - patch.object( - spend_prisma.writer_db._original_prisma, - "litellm_spendlogs", - MagicMock(create_many=create_spend_rows), - ), ): respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock( return_value=httpx.Response(200, content=content) @@ -3483,9 +3487,9 @@ class TestMultiPodBatchCostClaim: from datetime import datetime import httpx - import litellm import respx + import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import LiteLLMBatch, SpecialEnums @@ -3587,9 +3591,9 @@ class TestMultiPodBatchCostClaim: import base64 import httpx - import litellm import respx + import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import SpecialEnums @@ -4423,9 +4427,10 @@ class TestBatchLineItemsFromThePoller: @pytest.mark.asyncio async def test_a_batch_bigger_than_one_part_reaches_the_callbacks_as_numbered_parts(self, monkeypatch): - import litellm from litellm_enterprise.proxy.common_utils import check_batch_cost + import litellm + journal: list = [] _, instance = self._setup(journal) logger = self._recording_logger(journal) @@ -4933,9 +4938,10 @@ async def test_failed_retirement_keeps_batch_available() -> None: @pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) @pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2")) def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None: - from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.exceptions import NotFoundError + error: Final = NotFoundError( message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai" ) @@ -4958,9 +4964,10 @@ def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other ), ) def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None: - from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.exceptions import NotFoundError + error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai") assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True @@ -4970,9 +4977,10 @@ def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, mes ("missing_key", "gone"), (("out/job-123/input.jsonl.out", True), ("out/job-123/other.jsonl.out", False)) ) def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: str, gone: bool) -> None: - from litellm.llms.bedrock.common_utils import BedrockError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.llms.bedrock.common_utils import BedrockError + # S3 GetObject 404 body shape, checked live against a missing key on 2026-10-20 error: Final = BedrockError( status_code=404, @@ -4985,9 +4993,10 @@ def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: @pytest.mark.parametrize("output_file_id", (None, "")) def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None: - from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.exceptions import NotFoundError + error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai") assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False @@ -5003,11 +5012,12 @@ async def test_batch_cost_reaches_configured_success_callbacks( initialized: bool, claimed: bool, ) -> None: + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.utils import LiteLLMBatch - from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost events: Final[SimpleQueue[tuple[str, object, int, int]]] = SimpleQueue() diff --git a/tests/unit/proxy/common_utils/test_model_deprecation.py b/tests/unit/proxy/common_utils/test_model_deprecation.py index 7e96ad591b7..b3c135f1fad 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation.py @@ -467,6 +467,52 @@ class TestFormatDeprecationAlertMessage: class TestSuccessorModel: + def test_unmatched_successor_logs_warning_by_default(self, monkeypatch, caplog): + monkeypatch.setattr(litellm, "model_cost", {}) + router = Router( + model_list=[ + { + "model_name": "old", + "litellm_params": {"model": "openai/old-model"}, + "model_info": { + "id": "1", + "deprecation_date": "2026-06-10", + "successor_model": "missing-model", + }, + } + ] + ) + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + collect_model_deprecations(llm_router=router, today=date(2026, 6, 1)) + + successor_records: Final = tuple( + record for record in caplog.records if "names successor missing-model" in record.getMessage() + ) + assert len(successor_records) == 1 + assert successor_records[0].levelno == logging.WARNING + + def test_empty_team_public_model_name_is_normalized_to_none(self, monkeypatch): + monkeypatch.setattr(litellm, "model_cost", {}) + router = Router( + model_list=[ + { + "model_name": "team-model", + "litellm_params": {"model": "openai/team-model"}, + "model_info": { + "id": "1", + "team_id": "team-1", + "team_public_model_name": "", + "deprecation_date": "2026-06-10", + }, + } + ] + ) + + snapshot: Final = collect_model_deprecations(llm_router=router, today=date(2026, 6, 1)) + + assert snapshot.imminent[0].team_public_model_name is None + def test_should_surface_the_successor_set_on_the_deployment(self, monkeypatch): monkeypatch.setattr(litellm, "model_cost", {}) router = Router( @@ -572,6 +618,10 @@ class TestSuccessorModel: "typo-old": "gpt-nwe", "wildcard-old": "anthropic/x", } + team_deprecation: Final = next( + info for info in snapshot.imminent if info.model_name == "model_name_t1_old" + ) + assert team_deprecation.team_public_model_name == "team-old" warnings = [record.getMessage() for record in caplog.records if "names successor" in record.getMessage()] assert len(warnings) == 2 assert any("global-old names successor team-new" in warning for warning in warnings) diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index 0a015255f2e..814fdcb4eda 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -2,7 +2,6 @@ Unit tests for ToolDiscoveryQueue. """ - import pytest from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( @@ -63,3 +62,19 @@ def test_aggregate_sums_call_counts_across_workers_flushes(): {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, {"tool_name": "search", "call_count": 1}, ) + + +def test_requeued_rows_go_ahead_of_and_merge_with_sightings_since_the_flush(queue: ToolDiscoveryQueue) -> None: + queue.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + failed_write = queue.flush() + queue.add_update({"tool_name": "translate", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + + queue.requeue(failed_write) + + assert queue.flush() == [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 1}, + {"tool_name": "search", "origin": "user_defined", "call_count": 2}, + {"tool_name": "translate", "origin": "user_defined", "call_count": 1}, + ] diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py deleted file mode 100644 index 6b301146301..00000000000 --- a/tests/unit/proxy/db/test_accounting_ledger.py +++ /dev/null @@ -1,149 +0,0 @@ -import json -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final, Literal, NoReturn - -import pytest - -from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingLedger, UsageWrite -from litellm.types.utils import CallTypes - - -def _no_transaction() -> NoReturn: - raise AssertionError("Invalid accounting input must not open a database transaction") - - -@pytest.fixture -def event() -> InboundUsageEvent: - return InboundUsageEvent( - event_id="caller", - request_id="caller", - call_type=CallTypes.completion, - started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"]) -async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None: - with pytest.raises(ValueError, match=r"Accounting envelope|validation error"): - await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("version", [0, -1, True, 1.5]) -@pytest.mark.parametrize("field", ["version", "schema_version"]) -async def test_invalid_versions_fail_before_database_access( - event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"] -) -> None: - with pytest.raises(ValueError, match="positive integers"): - await AccountingLedger(_no_transaction).insert_event( - event, - (), - envelope="{}", - **{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions - ) - - -@pytest.mark.asyncio -async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None: - invalid: Final = event.model_copy(update={"event_id": ""}) - with pytest.raises(ValueError, match="event_id"): - await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5]) -async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None: - with pytest.raises(ValueError, match="batch size"): - await AccountingLedger(_no_transaction).purge_completed( - limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"]) -async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None: - cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - deleted: Final = await AccountingLedger(_no_transaction).purge_completed( - redelivery_before=cutoff if missing == "recovery" else None, - recovery_before=cutoff if missing == "redelivery" else None, - ) - assert deleted == 0 - - -@pytest.mark.asyncio -async def test_retention_rejects_ambiguous_naive_watermarks() -> None: - with pytest.raises(ValueError, match="timezone"): - await AccountingLedger(_no_transaction).purge_completed( - redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc) - ) - - -@pytest.mark.asyncio -async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}}) - with pytest.raises(ValueError, match="Out of range float"): - await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) - - -@pytest.mark.asyncio -async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - envelope: Final = json.dumps({"event": event.model_dump(mode="json")}) - with pytest.raises(ValueError, match="not a valid DeliveryTarget"): - await AccountingLedger(_no_transaction).insert_event( - event, - ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target - envelope=envelope, - ) - - -@pytest.mark.asyncio -async def test_empty_usage_batches_do_not_open_a_transaction() -> None: - ledger: Final = AccountingLedger(_no_transaction) - await ledger.insert_usage_batch(()) - assert await ledger.load_usage_batch(()) == () - - -@pytest.mark.asyncio -async def test_conflicting_usage_within_one_batch_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - original: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) - changed: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json"), "pricing": {"rate": 2}})) - with pytest.raises(AccountingConflictError, match="conflicting copies"): - await AccountingLedger(_no_transaction).insert_usage_batch((original, changed)) - - -@pytest.mark.asyncio -async def test_loaded_usage_preserves_decimal_metrics_from_jsonb(event: InboundUsageEvent) -> None: - from collections.abc import AsyncIterator - from contextlib import asynccontextmanager - from typing import cast - from unittest.mock import AsyncMock - - from prisma import Prisma - - from litellm.accounting import OutboundUsageEvent - - outbound: Final = OutboundUsageEvent( - **event.model_dump(exclude={"metrics"}), - attempt_id="attempt", - retry_of_attempt_id=None, - metrics=(Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),), - ) - payload: Final = {"event": outbound.model_dump(mode="json"), "usage": {"prompt_tokens": 120000}} - database: Final = AsyncMock(spec=Prisma) - database.query_raw.return_value = [{"payload": payload, "targets": []}] - - @asynccontextmanager - async def transaction() -> AsyncIterator[Prisma]: - yield cast(Prisma, database) - - recovered: Final = await AccountingLedger(transaction).load_usage_batch((outbound.event_id,)) - assert len(recovered) == 1 - assert recovered[0].event == outbound - assert json.loads(recovered[0].envelope) == payload diff --git a/tests/unit/proxy/db/test_budget_window_spend_writer.py b/tests/unit/proxy/db/test_budget_window_spend_writer.py index 27f1472c183..7315d51c5dd 100644 --- a/tests/unit/proxy/db/test_budget_window_spend_writer.py +++ b/tests/unit/proxy/db/test_budget_window_spend_writer.py @@ -407,9 +407,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio -@pytest.mark.parametrize( - ("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))] -) +@pytest.mark.parametrize(("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))]) async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(replaces_window_start, replaced_param): db = _FakeDB() @@ -425,7 +423,9 @@ async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(re ((query, params),) = db.execute_raw_calls normalized = " ".join(query.split()) assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE" in normalized - assert normalized.endswith("WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')") + assert normalized.endswith( + "WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')" + ) assert params == ("model_max_budget", "k1", "1d", datetime(2026, 8, 31), replaced_param) @@ -561,7 +561,12 @@ async def test_seed_aggregate_sql_splits_the_window_at_the_batch_start(entity_ty ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" in normalized + assert ( + 'FILTER (WHERE EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery ' + 'WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id) ' + "OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" + in normalized + ) assert 'FROM "LiteLLM_SpendLogs"' in normalized # startTime is TIMESTAMP(3): the bound is floored to the second so the # batch's own earliest row cannot round under it. @@ -635,3 +640,21 @@ async def test_rolling_a_window_row_renders_a_postgres_update_span( new_window_start=WINDOW_B, ) assert await postgres_span_names() == ("postgres.update LiteLLM_BudgetWindowSpend",) + + +@pytest.mark.asyncio +async def test_window_reads_inside_an_enclosing_transaction_use_its_connection(): + client_db = _FakeDB() + transaction = _FakeDB(existing_rows=[]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(client_db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0, started_at=BATCH_STARTED_AT), + ), + outer_transaction=transaction, + ) + + assert client_db.query_raw_calls == [] + assert len(transaction.query_raw_calls) == 2 + assert len(transaction.batcher.calls) == 1 diff --git a/tests/unit/proxy/db/test_db_lookup_gate.py b/tests/unit/proxy/db/test_db_lookup_gate.py index 68903170840..3450d3e05b8 100644 --- a/tests/unit/proxy/db/test_db_lookup_gate.py +++ b/tests/unit/proxy/db/test_db_lookup_gate.py @@ -46,21 +46,37 @@ async def test_bounded_db_lookup_returns_a_prompt_answer_without_recording_a_sta assert tracker.stalled_within(30) is False +class _PendingLookups: + def __init__(self) -> None: + self.pending = 0 + self.peak = 0 + + async def never_answers(self) -> None: + self.pending += 1 + self.peak = max(self.peak, self.pending) + try: + await asyncio.Event().wait() + finally: + self.pending -= 1 + + @pytest.mark.asyncio async def test_bounded_db_lookup_fails_a_whole_stalled_burst_within_one_deadline(): tracker: Final = DBLookupStallTracker() burst: Final = 200 - started: Final = time.monotonic() + lookups: Final = _PendingLookups() results: Final = await asyncio.gather( *( - bounded_db_lookup(_never_answers(), name=f"key-{i}", deadline_seconds=0.1, tracker=tracker) + bounded_db_lookup(lookups.never_answers(), name=f"key-{i}", deadline_seconds=0.1, tracker=tracker) for i in range(burst) ), return_exceptions=True, ) - assert time.monotonic() - started < 2 + assert lookups.peak == burst, ( + f"only {lookups.peak} of {burst} stalled lookups were waiting on their deadline together" + ) assert len(results) == burst assert all(isinstance(result, DBLookupDeadlineExceeded) for result in results) assert tracker.stalled_within(30) is True diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 298b153f621..d90c7f6d21d 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1,10 +1,12 @@ import asyncio import copy +import functools +import itertools import json import logging import re -from collections.abc import AsyncIterator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace @@ -19,6 +21,7 @@ from redis.exceptions import DataError import litellm from litellm._logging import verbose_proxy_logger + from litellm._service_logger import ServiceTypes from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem @@ -4380,7 +4383,6 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bri assert prisma.spend_log_flush_requested.is_set() is expects_flush - @pytest.mark.asyncio @pytest.mark.timeout(30) async def test_a_chatgpt_responses_row_queues_without_a_default_chatgpt_login( @@ -4422,6 +4424,8 @@ def _batch_cost_payload() -> dict: def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock: prisma = _tool_usage_prisma() prisma.jsonify_object = lambda data: dict(data) + prisma.writer_db = prisma.db + prisma.db.query_raw = AsyncMock(return_value=[]) prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over) @@ -4456,7 +4460,7 @@ async def _update_database_with( end_user_id=None, team_id=None, org_id=None, - kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"}, + kwargs={"model": "gpt-5.6-luna", "call_type": payload["call_type"]}, completion_response=None, start_time=datetime.now(timezone.utc), end_time=datetime.now(timezone.utc), @@ -4707,6 +4711,392 @@ async def test_update_database_queues_every_other_spend_row_for_the_next_flush(p assert db_writer._batch_database_updates.await_count == 1 +def _video_status_payload(request_id: str = "video_abc_video_cost") -> dict: + return { + **_minimal_spend_payload(), + "request_id": request_id, + "call_type": "avideo_status", + "status": "success", + } + + +class _SpendLogsTable: + """LiteLLM_SpendLogs keyed by its request_id primary key, as the writer's three queries see it.""" + + def __init__(self) -> None: + self.rows: dict[str, dict] = {} + + async def create_many(self, data: list[dict], skip_duplicates: bool) -> int: + fresh = {row["request_id"]: dict(row) for row in data if row["request_id"] not in self.rows} + self.rows.update(fresh) + return len(fresh) + + async def find_unique(self, where: dict) -> SimpleNamespace | None: + row = self.rows.get(where["request_id"]) + return None if row is None else SimpleNamespace(**row) + + async def update_many(self, data: dict, where: dict) -> int: + row = self.rows.get(where["request_id"]) + if row is None or any(row.get(field) != value for field, value in where.items()): + return 0 + row.update(data) + return 1 + + +def _spend_logs_table_prisma() -> MagicMock: + prisma = _tool_usage_prisma() + prisma.jsonify_object = lambda data: dict(data) + prisma.db.litellm_spendlogs = _SpendLogsTable() + prisma.writer_db.query_raw = AsyncMock(return_value=[]) + return prisma + + +async def _flush_queued_spend_logs(prisma: MagicMock) -> None: + await prisma.db.litellm_spendlogs.create_many(data=prisma.spend_log_transactions, skip_duplicates=True) + prisma.spend_log_transactions.clear() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "flushes_before_completion", + [0, 1], + ids=["rendering_poll_still_queued", "rendering_poll_already_written"], +) +async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_completed( + flushes_before_completion: int, +): + """ + A custom video provider reports the video's cost on the status poll that sees it completed, + and clients poll a finished video again before they download it, so only that first completed + poll may charge, and its spend row has to carry the charge (LIT-10404). The $0 poll written + while the video was still rendering keeps the bare video id, so it lands on the create row. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + rendering = await _update_database_with(db_writer, prisma, _video_status_payload("video_abc"), response_cost=0.0) + for _ in range(flushes_before_completion): + await _flush_queued_spend_logs(prisma) + completed = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.42) + repolled = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.42) + await _flush_queued_spend_logs(prisma) + + assert (rendering, completed, repolled) == (True, True, False) + assert db_writer._batch_database_updates.await_count == 2 + rows = prisma.db.litellm_spendlogs.rows + assert {request_id: (row["call_type"], row["spend"]) for request_id, row in rows.items()} == { + "video_abc": ("avideo_status", 0.0), + "video_abc_video_cost": ("avideo_status", 0.42), + } + + +def _indexed_guardrail_request_ids(prisma: MagicMock, index_table: str) -> tuple[str, ...]: + calls: Final = getattr(prisma.db, index_table).create_many.await_args_list + batches: Final = tuple(call.kwargs["data"] for call in calls) + return tuple(row["request_id"] for row in itertools.chain.from_iterable(batches)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disable_spend_logs", [False, True], ids=["spend_logs_on", "spend_logs_off"]) +@pytest.mark.parametrize("left_uncharged", [False, True], ids=["claimed_row", "taken_over_row"]) +async def test_a_job_cost_row_counts_its_polls_guardrail_run_in_guardrail_usage_once( + disable_spend_logs: bool, left_uncharged: bool +): + """The job's cost row is written by its claim, not the spend log queue whose flush counts guardrail runs, so a + finished video's guardrail run must still reach guardrail usage, once, and not when spend logs are off.""" + from litellm.repositories.table_repositories import ( + DailyGuardrailMetricsRepository, + SpendLogGuardrailIndexRepository, + ) + + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + getattr(prisma.db, SpendLogGuardrailIndexRepository.table_name).create_many = AsyncMock(return_value=1) + getattr(prisma.db, DailyGuardrailMetricsRepository.table_name).upsert = AsyncMock() + if left_uncharged: + prisma.db.litellm_spendlogs.rows["video_abc_video_cost"] = {**_video_status_payload(), "spend": 0.0} + guarded_poll: Final = { + **_video_status_payload(), + "metadata": json.dumps( + {"guardrail_information": [{"guardrail_name": "priced", "guardrail_status": "success"}]} + ), + } + + charged = [ + await _update_database_with( + db_writer, prisma, guarded_poll, disable_spend_logs=disable_spend_logs, response_cost=0.43 + ) + for _ in range(2) + ] + + assert charged == [True, False] + assert _indexed_guardrail_request_ids(prisma, SpendLogGuardrailIndexRepository.table_name) == ( + () if disable_spend_logs else ("video_abc_video_cost",) + ) + + +def _job_poll( + litellm_call_id: str, + response_cost: float, + status: str = "completed", + polled_video_id: str = "video_abc", + guardrail_cost: float = 0.0, + call_type: str = "avideo_status", +) -> tuple[dict[str, object], object]: + """The logging kwargs a job poll reaches the spend writer with: litellm logs the provider's own cost as + response_cost and the guardrails' on top of it only in the standard logging payload.""" + from litellm.types.utils import LiteLLMBatch + from litellm.types.videos.main import VideoObject + + polled_job = ( + VideoObject(id="video_abc", object="video", status=status) + if call_type == "avideo_status" + else LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status=status, + output_file_id="file-out", + ) + ) + kwargs: Final[dict[str, object]] = { + "model": "sora-2", + "call_type": call_type, + "litellm_call_id": litellm_call_id, + "response_cost": response_cost - guardrail_cost, + "polled_video_status": status, + "litellm_params": {"metadata": {"user_api_key": "test-key"}, "video_id": polled_video_id}, + **( + { + "standard_logging_object": { + "response_cost": response_cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], + "metadata": {}, + "hidden_params": {}, + "request_tags": [], + "model_map_information": None, + } + } + if guardrail_cost + else {} + ), + } + return kwargs, polled_job + + +def _spend_writer_call(db_writer: DBSpendUpdateWriter, polled_job: object) -> Callable[..., Awaitable[bool]]: + return functools.partial( + db_writer.update_database, + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + completion_response=polled_job, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + ) + + +@contextmanager +def _proxy_globals(prisma: MagicMock) -> Iterator[None]: + with ( + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.prisma_client", prisma + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), + ): + yield + + +async def _update_database_for_video_poll( + db_writer: DBSpendUpdateWriter, + prisma: MagicMock, + litellm_call_id: str, + response_cost: float, + status: str = "completed", + polled_video_id: str = "video_abc", + guardrail_cost: float = 0.0, + call_type: str = "avideo_status", +) -> bool: + kwargs, polled_job = _job_poll(litellm_call_id, response_cost, status, polled_video_id, guardrail_cost, call_type) + with _proxy_globals(prisma): + charged = await _spend_writer_call(db_writer, polled_job)(kwargs=kwargs, response_cost=response_cost) + await asyncio.sleep(0) + return charged + + +@pytest.mark.asyncio +async def test_update_database_charges_the_completed_polls_cost_and_not_a_running_polls(): + """A provider that also prices a running video must still have the completed poll's cost billed, once.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + running = await _update_database_for_video_poll(db_writer, prisma, "poll-1", 0.1, status="in_progress") + completed = await _update_database_for_video_poll(db_writer, prisma, "poll-2", 0.42) + repolled = await _update_database_for_video_poll(db_writer, prisma, "poll-3", 0.42) + await _flush_queued_spend_logs(prisma) + + assert (running, completed, repolled) == (True, True, False) + assert prisma.db.litellm_spendlogs.rows["video_abc_video_cost"]["spend"] == 0.42 + + +@pytest.mark.asyncio +async def test_a_poll_paying_only_its_guardrail_cannot_take_the_jobs_charge_whatever_call_id_the_client_sent(): + """A client picks its own x-litellm-call-id, so a guardrail-only video poll naming the job's charge row, once or + twice, must still be charged each time and leave the finished job's charge to the poll that sees it done. An + unfinished batch retrieve never reaches the writer, and its row is named by the batch id the provider sent.""" + call_type: Final = "avideo_status" + charge_row_id: Final = "video_abc_video_cost" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll( + db_writer, prisma, call_id, cost, status=status, guardrail_cost=0.01, call_type=call_type + ) + for call_id, cost, status in ( + (charge_row_id, 0.01, "in_progress"), + (charge_row_id, 0.01, "in_progress"), + ("poll-3", 0.43, "completed"), + ) + ] + await _flush_queued_spend_logs(prisma) + + assert charged == [True, True, True] + assert prisma.db.litellm_spendlogs.rows[charge_row_id]["spend"] == 0.43 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "charge_row_id"), + [("avideo_status", "video_abc_video_cost"), ("aretrieve_batch", "batch_abc_batch_cost")], +) +async def test_a_finished_job_cheaper_than_its_polls_guardrail_is_still_charged_once(call_type, charge_row_id): + """A batch of a few cheap requests costs less than a priced guardrail on its retrieve, and the finished job's + cost must still settle under its charge row once instead of being billed again on every poll.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll( + db_writer, prisma, call_id, 0.01002775, guardrail_cost=0.01, call_type=call_type + ) + for call_id in ("poll-1", "poll-2") + ] + await _flush_queued_spend_logs(prisma) + + assert charged == [True, False] + assert {request_id: row["spend"] for request_id, row in prisma.db.litellm_spendlogs.rows.items()} == { + charge_row_id: 0.01002775 + } + + +@pytest.mark.asyncio +async def test_a_finished_batch_that_cost_nothing_itself_settles_its_retrieves_guardrail_once(): + """A batch whose every request failed costs nothing itself, and its finished retrieves must still settle under + the batch's charge row once, as they always have, instead of each paying the retrieve's guardrail again.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll( + db_writer, prisma, call_id, 0.01, status="failed", guardrail_cost=0.01, call_type="aretrieve_batch" + ) + for call_id in ("poll-1", "poll-2", "poll-3") + ] + await _flush_queued_spend_logs(prisma) + + assert charged == [True, False, False] + assert {request_id: row["spend"] for request_id, row in prisma.db.litellm_spendlogs.rows.items()} == { + "batch_abc_batch_cost": 0.01 + } + + +@pytest.mark.asyncio +async def test_a_repeat_poll_of_a_finished_video_charges_its_guardrail_on_a_row_of_its_own(): + from litellm.proxy.hooks.proxy_track_cost_callback import _charge_request_cost + + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + await _update_database_for_video_poll(db_writer, prisma, "poll-1", 0.43, guardrail_cost=0.01) + kwargs, polled_job = _job_poll("poll-2", 0.43, guardrail_cost=0.01) + + async def charge_spend( + *, kwargs: dict[str, object], response_cost: float, budget_reservation: dict[str, object] | None + ) -> bool: + return await _spend_writer_call(db_writer, polled_job)(kwargs=kwargs, response_cost=response_cost) + + with _proxy_globals(prisma): + repeat_charge = await _charge_request_cost(charge_spend, kwargs, 0.43, None) + await asyncio.sleep(0) + await _flush_queued_spend_logs(prisma) + + assert repeat_charge == 0.01 + assert {request_id: row["spend"] for request_id, row in prisma.db.litellm_spendlogs.rows.items()} == { + "video_abc_video_cost": 0.43, + "poll-2_guardrail_cost": 0.01, + } + + +@pytest.mark.asyncio +async def test_update_database_charges_each_deployments_video_though_their_provider_ids_match(): + """A custom provider's own ids can repeat across deployments, so the cost row keys off the id the client polled.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll(db_writer, prisma, call_id, 0.42, polled_video_id=polled_id) + for call_id, polled_id in (("poll-1", "video_dep_a"), ("poll-2", "video_dep_b"), ("poll-3", "video_dep_a")) + ] + + assert charged == [True, True, False] + assert { + request_id: row["spend"] + for request_id, row in prisma.db.litellm_spendlogs.rows.items() + if request_id.endswith("_video_cost") + } == {"video_dep_a_video_cost": 0.42, "video_dep_b_video_cost": 0.42} + + +@pytest.mark.asyncio +async def test_update_database_writes_no_row_of_its_own_for_a_video_poll_that_reports_no_cost(): + """ + Built-in video providers charge at create and report nothing on the status poll, so those + polls keep the ordinary queued write under the create row's id and still count as requests + in the daily rollups, leaving the video with its one create row. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + await prisma.db.litellm_spendlogs.create_many( + data=[{**_minimal_spend_payload(), "request_id": "video_abc", "call_type": "avideo_generation", "spend": 0.3}], + skip_duplicates=True, + ) + + assert await _update_database_for_video_poll(db_writer, prisma, "poll-1", response_cost=0.0) is True + assert await _update_database_for_video_poll(db_writer, prisma, "poll-2", response_cost=0.0) is True + await _flush_queued_spend_logs(prisma) + + assert db_writer._batch_database_updates.await_count == 2 + rows = prisma.db.litellm_spendlogs.rows + assert {request_id: (row["call_type"], row["spend"]) for request_id, row in rows.items()} == { + "video_abc": ("avideo_generation", 0.3) + } + + @pytest.mark.asyncio @pytest.mark.parametrize( "injected_deployment, attributed", @@ -5452,3 +5842,36 @@ async def test_completed_spend_survives_error_log_suppression( assert len(rows) == 1 assert rows[0]["spend"] == pytest.approx(0.23) assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (10, 5) + + +@pytest.mark.parametrize("ledger_enabled", [False, True]) +@pytest.mark.asyncio +async def test_request_error_counts_survive_the_ledger_path_switch(monkeypatch: pytest.MonkeyPatch, ledger_enabled: bool): + from litellm.proxy import proxy_server + from litellm.proxy.db import db_spend_update_writer as writer_module + from litellm.proxy.db.request_error_tracking import RequestErrorAccumulator + from litellm.types.proxy.request_errors import RequestErrorKey + + accumulator = RequestErrorAccumulator() + prisma = _tool_usage_prisma() + prisma.get_request_status.return_value = "failure" + monkeypatch.setattr(writer_module, "request_error_accumulator", accumulator) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger_enabled}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + now = datetime(2026, 11, 2, tzinfo=timezone.utc) + await DBSpendUpdateWriter().update_database( + token="key", user_id="user", team_id=None, org_id=None, end_user_id=None, + kwargs={ + "model": "ledger-test", "call_type": "acompletion", "litellm_call_id": "failed-request", + "litellm_params": {"metadata": { + "user_api_key": "key", "user_api_key_user_id": "user", "status": "failure", + "error_information": {"error_code": "503"}, + }}, + }, + completion_response=None, start_time=now, end_time=now, response_cost=0, + ) + await drain_logging_tasks() + assert accumulator.drain() == {RequestErrorKey( + date="2026-11-02", api_key=prisma.spend_log_transactions[0]["api_key"], team_id="", user_id="user", model_group="ledger-test", status_code=503, + ): 1} diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 46704bf52ef..abc0850194d 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -22,7 +22,7 @@ from prisma.errors import ( import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_never_sent # Test is_database_connection_error method @@ -996,3 +996,137 @@ def test_connection_closure_code_survives_prisma_data_error_wrapping() -> None: ) assert PrismaDBExceptionHandler.is_database_transport_error(error) is True assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is False + + +def _connector_error(message: str) -> DataError: + return DataError( + data={"user_facing_error": {"message": f"Error in connector: Error querying the database: FATAL: {message}"}} + ) + + +def _batched_postgres_error(sqlstate: str, message: str) -> DataError: + return DataError( + data={ + "user_facing_error": { + "message": "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: " + f'None, kind: QueryError(PostgresError {{ code: "{sqlstate}", message: "{message}", severity: "FATAL" }}) }})' + } + } + ) + + +def _localized_error(sqlstate: str) -> DataError: + return DataError(data={"user_facing_error": {"message": "FATAL: Datenbank", "meta": {"code": sqlstate}}}) + + +_NOT_ACCEPTING_CONNECTIONS: Final = ( + _connector_error("the database system is starting up"), + _connector_error("the database system is shutting down"), + _connector_error("the database system is not yet accepting connections"), + _connector_error("the database system is not accepting connections"), + _connector_error("the database system is in recovery mode"), + _localized_error("57P03"), + _batched_postgres_error("57P03", "the database system is starting up"), + EngineRequestError( + MagicMock(status=500), + '{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: the database system is ' + 'starting up","backtrace":null}', + ), + RawQueryError( + data={ + "user_facing_error": { + "message": "Raw query failed. Code: `57P03`. Message: `db error: FATAL: the database system is starting up`", + "meta": {"code": "57P03", "message": "FATAL: the database system is starting up"}, + "error_code": "P2010", + } + } + ), +) +_SERVER_TERMINATED: Final = ( + _connector_error("terminating connection due to administrator command"), + _connector_error("terminating connection due to immediate shutdown command"), + _connector_error("terminating connection because of unexpected SIGQUIT signal"), + _connector_error("terminating connection because of crash of another server process"), + _connector_error("terminating connection due to idle-session timeout"), + _localized_error("57P01"), + _localized_error("57P02"), + _localized_error("57P05"), + _batched_postgres_error("57P01", "terminating connection due to administrator command"), +) + + +@pytest.mark.parametrize( + ("error", "connection_error"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + ], +) +def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Exception, connection_error: bool) -> None: + """A server that is restarting, shutting down or cutting sessions is a transient outage: + auth answers 503 rather than 401, and the spend-log writer must not take the rows for poison. + Only a server refusing connections opens the ``allow_requests_on_db_unavailable`` fallback, + one terminated session does not prove the database is down.""" + assert { + "lifecycle": PrismaDBExceptionHandler.is_database_lifecycle_error(error), + "service_unavailable": PrismaDBExceptionHandler.is_database_service_unavailable_error(error), + "connection_error": PrismaDBExceptionHandler.is_database_connection_error(error), + "permanent": PrismaDBExceptionHandler.is_permanent_database_fault(error), + } == { + "lifecycle": True, + "service_unavailable": True, + "connection_error": connection_error, + "permanent": False, + } + + +@pytest.mark.parametrize( + ("error", "never_sent"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + (DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}), True), + (DataError(data={"user_facing_error": {"error_code": "P1001", "message": "unreachable"}}), True), + (_capacity_error("sorry, too many clients already"), True), + (_localized_error("53300"), True), + ( + DataError( + data={ + "user_facing_error": { + "error_code": "P2024", + "message": "Timed out fetching a new connection from the connection pool", + } + } + ), + True, + ), + (httpx.ReadTimeout("read timed out"), False), + (_localized_error("57014"), False), + ], +) +def test_db_write_never_sent_only_for_failures_before_any_statement_ran(error: Exception, never_sent: bool) -> None: + """A refused, unreachable or full server provably applied nothing, while a terminated session + or a timeout may have committed before the failure, so only the former lets a non-idempotent + write be replayed.""" + assert db_write_never_sent(error) is never_sent + + +@pytest.mark.parametrize( + "error", + [ + _localized_error("57P04"), + _localized_error("57014"), + _localized_error("23505"), + DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}), + RuntimeError("the database system is starting up"), + ], +) +def test_other_failures_are_not_postgres_lifecycle_errors(error: Exception) -> None: + assert PrismaDBExceptionHandler.is_database_lifecycle_error(error) is False + assert db_write_never_sent(error) is False + + +def test_lifecycle_error_is_seen_through_a_wrapping_exception() -> None: + wrapped: Final = RuntimeError("spend flush failed") + wrapped.__cause__ = _connector_error("the database system is starting up") + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(wrapped) is True diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index 76898b0eeab..a5ad897913e 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -69,9 +69,7 @@ def _make_prisma( prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) prisma.db.execute_raw = AsyncMock(return_value=1) - prisma.db.litellm_tooltable.find_many = AsyncMock( - return_value=find_many_rows if find_many_rows is not None else [] - ) + prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else []) prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row) prisma.db.litellm_verificationtoken = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows)) @@ -227,9 +225,7 @@ async def test_get_tool_found(): result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" - prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio @@ -266,9 +262,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): updated_by="admin", ) prisma = _make_prisma(find_unique_row=row) - result = await update_tool_policy( - prisma, "my_tool", updated_by="admin", input_policy="blocked" - ) + result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked") assert result is not None assert result.input_policy == "blocked" prisma.db.litellm_tooltable.upsert.assert_awaited_once() @@ -567,3 +561,42 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (_postgres_rejection("57P03"), True), + (_postgres_rejection("53300"), True), + (_postgres_rejection("57P01"), False), + (_postgres_rejection("54000"), False), + (_postgres_rejection("23502"), False), + ], + ids=[ + "restarting_server_keeps_the_rows", + "full_server_keeps_the_rows", + "terminated_session_may_have_counted_so_the_rows_are_dropped", + "oversized_value_is_dropped", + "rejected_row_is_dropped", + ], +) +async def test_failed_write_without_redis_keeps_the_rows_in_the_queue_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=_InMemoryRedis()), + pod_lock_manager=PodLockManager(redis_cache=_InMemoryRedis()), + use_redis_buffer=False, + ) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + expected = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2 if requeued else 1}] + assert queue.flush() == expected diff --git a/tests/unit/proxy/enterprise_billing/test_billing_metrics.py b/tests/unit/proxy/enterprise_billing/test_billing_metrics.py index 0446cfeeab0..14a8fea04f7 100644 --- a/tests/unit/proxy/enterprise_billing/test_billing_metrics.py +++ b/tests/unit/proxy/enterprise_billing/test_billing_metrics.py @@ -1,398 +1,379 @@ -""" -Tests for the enterprise billing-metrics recorder and its factory. - -These verify the license gate, the missing-config and missing-cert disable -paths, the OTLP/HTTP exporter wiring (client cert+key authenticate us to the -collector's mTLS-terminating front end; CA override optional for private -collectors), the metric attribute mapping, and that recording produces the -expected OTLP counter via an in-memory reader. -""" - -import os -import socket -import stat -from pathlib import Path -from typing import Dict, List, Optional +from collections.abc import Generator, Mapping +from typing import Final import pytest +import requests +import requests_mock from opentelemetry.sdk.metrics import MeterProvider -from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, Sum +from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.proxy._types import EnterpriseLicenseData from litellm.proxy.enterprise_billing import billing_metrics as bm from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory -_ENV_VARS = ( + +class _RecordedRequest(BaseModel): + model_config = ConfigDict(from_attributes=True, frozen=True) + + url: str | None + body: str | bytes | None + headers: dict[str, str] + + +def _request_history(http_mock: requests_mock.Mocker) -> tuple[_RecordedRequest, ...]: + return TypeAdapter(tuple[_RecordedRequest, ...]).validate_python(http_mock.request_history) + + +def _metric_category(attributes: Mapping[str, object] | None) -> str: + if attributes is None: + return "" + return str(attributes.get("litellm.endpoint.category")) + + +_ENV_VARS: Final = ( bm.ENDPOINT_ENV, - bm.CLIENT_CERT_ENV, - bm.CLIENT_KEY_ENV, - bm.CA_CERT_ENV, + bm.CLIENT_ID_ENV, + bm.CLIENT_SECRET_ENV, + bm.TOKEN_URL_ENV, bm.EXPORT_INTERVAL_ENV, ) @pytest.fixture(autouse=True) -def clear_env(monkeypatch): +def clear_env(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: for name in _ENV_VARS: monkeypatch.delenv(name, raising=False) yield bm.shutdown_billing_metrics_recorder() -def _write_certs(tmp_path: Path) -> Dict[str, str]: - files = { - bm.CA_CERT_ENV: ("ca.pem", b"ca-bytes"), - bm.CLIENT_CERT_ENV: ("client.pem", b"client-cert-bytes"), - bm.CLIENT_KEY_ENV: ("client.key", b"client-key-bytes"), - } - paths = {} - for env_name, (filename, content) in files.items(): - path = tmp_path / filename - path.write_bytes(content) - paths[env_name] = str(path) - return paths +def _set_full_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4318") + monkeypatch.setenv(bm.CLIENT_ID_ENV, "client-id") + monkeypatch.setenv(bm.CLIENT_SECRET_ENV, "client-secret") + monkeypatch.setenv(bm.TOKEN_URL_ENV, "https://auth.example/oauth2/token") -def _set_full_env(monkeypatch, tmp_path: Path) -> Dict[str, str]: - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) - return paths - - -def _config(tmp_path: Path, license_id: Optional[str] = "org-1") -> bm.BillingMetricsConfig: - paths = _write_certs(tmp_path) +def _config(license_id: str | None = "org-1") -> bm.BillingMetricsConfig: return bm.BillingMetricsConfig( - endpoint="https://collector.example:4317", - client_cert_path=paths[bm.CLIENT_CERT_ENV], - client_key_path=paths[bm.CLIENT_KEY_ENV], - ca_cert_path=paths[bm.CA_CERT_ENV], + endpoint="https://collector.example:4318", + token_url="https://auth.example/oauth2/token", + client_id="client-id", + client_secret="client-secret", export_interval_ms=60_000, litellm_version="1.2.3", license_id=license_id, ) -# ── Factory gating ──────────────────────────────────────────────────────────── +@pytest.mark.parametrize("missing_env", (bm.ENDPOINT_ENV, bm.CLIENT_ID_ENV, bm.CLIENT_SECRET_ENV, bm.TOKEN_URL_ENV)) +def test_missing_required_env_disables_metrics_without_logging_values( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + missing_env: str, +) -> None: + _set_full_env(monkeypatch) + monkeypatch.setenv(bm.CLIENT_SECRET_ENV, "secret-for-test") + monkeypatch.delenv(missing_env) + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + assert f"licensed deployment missing config ({missing_env})" in caplog.text + assert "secret-for-test" not in caplog.text -def test_not_premium_returns_none(tmp_path, monkeypatch): - _set_full_env(monkeypatch, tmp_path) +def test_empty_required_env_disables_metrics(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: + _set_full_env(monkeypatch) + monkeypatch.setenv(bm.ENDPOINT_ENV, "") + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + assert bm.ENDPOINT_ENV in caplog.text + + +@pytest.mark.parametrize( + ("env_name", "url"), + ( + (bm.ENDPOINT_ENV, "http://collector.example:4318"), + (bm.TOKEN_URL_ENV, "http://auth.example/oauth2/token"), + ), +) +def test_http_urls_disable_metrics( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + env_name: str, + url: str, +) -> None: + _set_full_env(monkeypatch) + monkeypatch.setenv(env_name, url) + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + assert f"{env_name} must use HTTPS" in caplog.text + + +def test_valid_config_contains_bearer_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + _set_full_env(monkeypatch) + + license_data: Final[EnterpriseLicenseData] = {"user_id": "org-1"} + config: Final = bm.load_billing_metrics_config(license_data=license_data, litellm_version="1.0") + assert config is not None + + assert config == bm.BillingMetricsConfig( + endpoint="https://collector.example:4318", + token_url="https://auth.example/oauth2/token", + client_id="client-id", + client_secret="client-secret", + export_interval_ms=bm.DEFAULT_EXPORT_INTERVAL_MS, + litellm_version="1.0", + license_id="org-1", + ) + assert "client-secret" not in repr(config) + assert "license-secret" not in repr(config) + + +def test_not_premium_returns_none_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + _set_full_env(monkeypatch) + assert bm.build_billing_metrics_recorder(premium=False, license_data=None, litellm_version="1.0") is None + assert "Enterprise billing metrics disabled" not in caplog.text -def test_premium_without_config_returns_none(monkeypatch): +def test_premium_without_config_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: assert bm.build_billing_metrics_recorder(premium=True, license_data={"user_id": "x"}, litellm_version="1.0") is None -def test_premium_with_missing_cert_files_returns_none(monkeypatch, tmp_path): - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CA_CERT_ENV, str(tmp_path / "missing-ca.pem")) - monkeypatch.setenv(bm.CLIENT_CERT_ENV, str(tmp_path / "missing-cert.pem")) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, str(tmp_path / "missing-key.pem")) - assert bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") is None +def test_premium_builds_recorder_and_logs_enabled( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + _set_full_env(monkeypatch) + caplog.set_level("INFO", logger=bm.verbose_proxy_logger.name) + token_url: Final = "https://auth.example/oauth2/token" + export_url: Final = "https://collector.example:4318/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, json={"access_token": "access-token", "expires_in": 3600}) + http_mock.post(export_url, status_code=200) + recorder: Final = bm.build_billing_metrics_recorder( + premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0" + ) + assert recorder is not None + assert "Enterprise billing metrics enabled: exporting to https://collector.example:4318 every 60000 ms" in caplog.text + bm.shutdown_billing_metrics_recorder() -def test_premium_with_full_config_builds_recorder(monkeypatch, tmp_path): - """Builds a real MeterProvider, so the exporter is stubbed: the live one - resolves the collector and opens a TLS connection during the shutdown flush. - The getaddrinfo spy keeps that stub from being quietly dropped later.""" - _set_full_env(monkeypatch, tmp_path) - monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class({})) +class _ShutdownCountingMeterProvider(MeterProvider): + def __init__(self) -> None: + super().__init__() + self.shutdown_calls = 0 - resolved: List[str] = [] - real_getaddrinfo = socket.getaddrinfo + def shutdown(self, timeout_millis: float = 30_000) -> None: + self.shutdown_calls += 1 + super().shutdown(timeout_millis=timeout_millis) - def _spy_getaddrinfo(host, port, *args, **kwargs): - resolved.append(str(host)) - return real_getaddrinfo(host, port, *args, **kwargs) - monkeypatch.setattr(socket, "getaddrinfo", _spy_getaddrinfo) +def test_shutdown_flushes_active_recorder_once(monkeypatch: pytest.MonkeyPatch) -> None: + provider: Final = _ShutdownCountingMeterProvider() + _set_full_env(monkeypatch) - recorder = bm.build_billing_metrics_recorder( - premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0" + def create_provider(config: bm.BillingMetricsConfig) -> MeterProvider: + return provider + + monkeypatch.setattr(bm, "build_bearer_meter_provider", create_provider) + recorder: Final = bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") + assert recorder is not None + + bm.shutdown_billing_metrics_recorder() + bm.shutdown_billing_metrics_recorder() + + assert provider.shutdown_calls == 1 + + +class _FakeClock: + def __init__(self, now: float) -> None: + self._now = now + + def __call__(self) -> float: + return self._now + + def set(self, now: float) -> None: + self._now = now + + +def test_bearer_auth_fetches_and_caches_client_credentials_token() -> None: + token_url: Final = "https://auth.example/oauth2/token" + export_url: Final = "https://collector.example/v1/metrics" + auth: Final = bm.BillingMetricsBearerAuth( + client_id="client-id", + client_secret="client-secret", + token_url=token_url, + clock=lambda: 100.0, ) - assert isinstance(recorder, bm.BillingMetricsRecorder) - recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id=None) - bm.shutdown_billing_metrics_recorder() + session: Final = requests.Session() + session.auth = auth + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, json={"access_token": "access-token", "expires_in": 3600}) + http_mock.post(export_url, status_code=200) + session.post(export_url, data=b"metrics") + session.post(export_url, data=b"metrics") - assert [host for host in resolved if "collector.example" in host] == [] + request_records: Final = _request_history(http_mock) + token_requests: Final = tuple(request for request in request_records if request.url == token_url) + export_requests: Final = tuple(request for request in request_records if request.url == export_url) + assert len(token_requests) == 1 + assert token_requests[0].headers["Authorization"].startswith("Basic ") + assert token_requests[0].body == "grant_type=client_credentials&scope=telemetry%2Fwrite" + assert tuple(request.headers["Authorization"] for request in export_requests) == ( + "Bearer access-token", + "Bearer access-token", + ) -def test_building_the_recorder_logs_an_affirmative_line(monkeypatch, tmp_path): - """ - Every disable path logs; a successful build must log too. Otherwise an - operator cannot tell a metering component from one that silently returned - None, which is how an unlicensed component looks healthy while exporting - nothing. - """ - _set_full_env(monkeypatch, tmp_path) +def test_bearer_auth_refreshes_before_token_expiry() -> None: + token_url: Final = "https://auth.example/oauth2/token" + export_url: Final = "https://collector.example/v1/metrics" + clock: Final = _FakeClock(1000.0) + auth: Final = bm.BillingMetricsBearerAuth( + client_id="client-id", + client_secret="client-secret", + token_url=token_url, + clock=clock, + ) + session: Final = requests.Session() + session.auth = auth + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post( + token_url, + [ + {"json": {"access_token": "first-token", "expires_in": 120}}, + {"json": {"access_token": "second-token", "expires_in": 120}}, + ], + ) + http_mock.post(export_url, status_code=200) + session.post(export_url, data=b"metrics") + clock.set(1059.0) + session.post(export_url, data=b"metrics") + clock.set(1060.0) + session.post(export_url, data=b"metrics") + + request_records: Final = _request_history(http_mock) + token_requests: Final = tuple(request for request in request_records if request.url == token_url) + export_requests: Final = tuple(request for request in request_records if request.url == export_url) + assert len(token_requests) == 2 + assert tuple(request.headers["Authorization"] for request in export_requests) == ( + "Bearer first-token", + "Bearer first-token", + "Bearer second-token", + ) + + +def test_bearer_export_sets_authorization_header() -> None: + config: Final = _config() + token_url: Final = config.token_url + export_url: Final = f"{config.endpoint}/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, json={"access_token": "export-token", "expires_in": 3600}) + http_mock.post(export_url, status_code=200) + provider: Final = bm.build_bearer_meter_provider(config) + try: + provider.get_meter("test").create_counter("requests").add(1) + assert provider.force_flush(timeout_millis=1000) + finally: + provider.shutdown(timeout_millis=1000) + + export_requests: Final = tuple(request for request in _request_history(http_mock) if request.url == export_url) + assert export_requests + assert all(request.headers["Authorization"] == "Bearer export-token" for request in export_requests) + + +def test_unauthorized_export_refetches_token_before_next_export() -> None: + config: Final = _config() + token_url: Final = config.token_url + export_url: Final = f"{config.endpoint}/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post( + token_url, + [ + {"json": {"access_token": "first-token", "expires_in": 3600}}, + {"json": {"access_token": "second-token", "expires_in": 3600}}, + ], + ) + http_mock.post(export_url, [{"status_code": 401}, {"status_code": 200}]) + provider: Final = bm.build_bearer_meter_provider(config) + try: + counter: Final = provider.get_meter("test").create_counter("requests") + counter.add(1) + provider.force_flush(timeout_millis=1000) + counter.add(1) + provider.force_flush(timeout_millis=1000) + finally: + provider.shutdown(timeout_millis=1000) + + requests: Final = _request_history(http_mock) + token_requests: Final = tuple(request for request in requests if request.url == token_url) + export_requests: Final = tuple(request for request in requests if request.url == export_url) + assert len(token_requests) == 2 + assert tuple(request.headers["Authorization"] for request in export_requests[:2]) == ( + "Bearer first-token", + "Bearer second-token", + ) + + +def test_failed_bearer_token_fetch_does_not_raise_from_export() -> None: + config: Final = _config() + token_url: Final = config.token_url + export_url: Final = f"{config.endpoint}/v1/metrics" + http_mock: Final = requests_mock.Mocker() + with http_mock: + http_mock.post(token_url, status_code=400) + provider: Final = bm.build_bearer_meter_provider(config) + try: + provider.get_meter("test").create_counter("requests").add(1) + assert provider.force_flush(timeout_millis=1000) + finally: + provider.shutdown(timeout_millis=1000) + + request_urls: Final = tuple(request.url for request in _request_history(http_mock)) + assert token_url in request_urls + assert export_url not in request_urls + + +def test_export_interval_default_and_override(monkeypatch: pytest.MonkeyPatch) -> None: + _set_full_env(monkeypatch) + default_config: Final = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert default_config is not None + assert default_config.export_interval_ms == bm.DEFAULT_EXPORT_INTERVAL_MS monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000") - monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class({})) - - infos: List[str] = [] - monkeypatch.setattr(bm.verbose_proxy_logger, "info", lambda msg, *args: infos.append(msg % args if args else msg)) - - recorder = bm.build_billing_metrics_recorder(premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0") - - assert recorder is not None - joined = "\n".join(infos) - assert "https://collector.example:4317" in joined - assert "5000" in joined - - -def test_unlicensed_build_does_not_warn(monkeypatch, tmp_path): - """Unlicensed is the common OSS case; warning there would be pure noise.""" - _set_full_env(monkeypatch, tmp_path) - - warnings: List[str] = [] - monkeypatch.setattr(bm.verbose_proxy_logger, "warning", lambda msg, *args: warnings.append(str(msg))) - - assert bm.build_billing_metrics_recorder(premium=False, license_data=None, litellm_version="1.0") is None - assert warnings == [] - - -def test_shutdown_flushes_active_recorder_once(monkeypatch, tmp_path): - """The shutdown hook must flush the recorder the factory built (buffered - counts are lost on restart otherwise) and be idempotent for repeat calls.""" - _set_full_env(monkeypatch, tmp_path) - shutdowns = [] - - class _SpyProvider: - def get_meter(self, name): - return MeterProvider().get_meter(name) - - def shutdown(self, timeout_millis=None): - shutdowns.append(timeout_millis) - - monkeypatch.setattr(bm, "build_mtls_meter_provider", lambda config: _SpyProvider()) - recorder = bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") - assert recorder is not None - - bm.shutdown_billing_metrics_recorder() - bm.shutdown_billing_metrics_recorder() - assert shutdowns == [bm.SHUTDOWN_FLUSH_TIMEOUT_MS] - - -def test_shutdown_without_active_recorder_is_noop(): - bm.shutdown_billing_metrics_recorder() - - -# ── Config loading ──────────────────────────────────────────────────────────── - - -def test_load_config_carries_license_id(monkeypatch, tmp_path): - _set_full_env(monkeypatch, tmp_path) - config = bm.load_billing_metrics_config(license_data={"user_id": "org-42"}, litellm_version="9.9") - assert config is not None and config.license_id == "org-42" and config.litellm_version == "9.9" - - -def test_load_config_with_empty_string_env_is_disabled(monkeypatch, tmp_path): - """An env var set to the empty string is as unusable as an unset one and - must disable metering rather than produce a config with a blank endpoint.""" - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - -_CLIENT_CERT_PEM = "-----BEGIN CERTIFICATE-----\nclient-cert-body\n-----END CERTIFICATE-----" -_CLIENT_KEY_PEM = "-----BEGIN PRIVATE KEY-----\nclient-key-body\n-----END PRIVATE KEY-----" -_CA_CERT_PEM = "-----BEGIN CERTIFICATE-----\nca-body\n-----END CERTIFICATE-----" - - -def test_load_config_materializes_inline_pem_content(monkeypatch): - """ - ECS and Cloud Run inject secrets as env content, not as mounted files, so the - cert env vars must accept PEM directly. The exporter takes paths, so the PEM - is written to disk and the config points at those files. - """ - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) - monkeypatch.setenv(bm.CA_CERT_ENV, _CA_CERT_PEM) - - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - - assert config is not None - assert config.ca_cert_path is not None - written = { - config.client_cert_path: _CLIENT_CERT_PEM, - config.client_key_path: _CLIENT_KEY_PEM, - config.ca_cert_path: _CA_CERT_PEM, - } - for path, pem in written.items(): - assert path != pem, "config must carry a file path, not the PEM itself" - assert os.path.isfile(path) - assert Path(path).read_text(encoding="utf-8") == f"{pem}\n" - - # The private key must not be world- or group-readable. - assert stat.S_IMODE(os.stat(config.client_key_path).st_mode) == 0o600 - - -def test_load_config_accepts_a_mix_of_pem_content_and_file_paths(monkeypatch, tmp_path): - """A deployment may mount the CA but inject the client credentials inline.""" - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) - monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV]) - - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - - assert config is not None - assert config.ca_cert_path == paths[bm.CA_CERT_ENV] - assert Path(config.client_cert_path).read_text(encoding="utf-8") == f"{_CLIENT_CERT_PEM}\n" - - -def test_load_config_leaves_file_paths_untouched(monkeypatch, tmp_path): - """Path-valued env vars keep working; nothing is copied or rewritten.""" - paths = _set_full_env(monkeypatch, tmp_path) - - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - - assert config is not None - assert config.client_cert_path == paths[bm.CLIENT_CERT_ENV] - assert config.client_key_path == paths[bm.CLIENT_KEY_ENV] - assert config.ca_cert_path == paths[bm.CA_CERT_ENV] - - -def test_load_config_with_inline_pem_disabled_when_unwritable(monkeypatch): - """A failure to materialize the PEM disables metering instead of raising.""" - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) - - def _explode(prefix=None): - raise OSError("read-only filesystem") - - monkeypatch.setattr(bm.tempfile, "mkdtemp", _explode) - - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - -def test_load_config_never_logs_credential_values(monkeypatch): - """ - A value that is neither a readable path nor `-----BEGIN`-prefixed PEM is - still secret material. The disable warning must name the env vars, never - echo their contents, or a malformed key lands in the proxy logs. - """ - secret_material = "MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQ-not-pem-prefixed" - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, secret_material) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, secret_material) - - logged: List[str] = [] - - def _capture(msg, *args): - logged.append(msg % args if args else msg) - - monkeypatch.setattr(bm.verbose_proxy_logger, "warning", _capture) - - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - joined = "\n".join(logged) - assert secret_material not in joined - assert bm.CLIENT_CERT_ENV in joined and bm.CLIENT_KEY_ENV in joined - - -def test_load_config_with_empty_pem_env_is_disabled(monkeypatch): - """Empty stays empty: an unset secret must not be mistaken for inline PEM.""" - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, "") - monkeypatch.setenv(bm.CLIENT_KEY_ENV, "") - - assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None - - -def test_export_interval_default_and_override(monkeypatch): - assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS - monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000") - assert bm._export_interval_ms() == 5000 + override_config: Final = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert override_config is not None + assert override_config.export_interval_ms == 5000 monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "not-a-number") - assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS + invalid_config: Final = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert invalid_config is not None + assert invalid_config.export_interval_ms == bm.DEFAULT_EXPORT_INTERVAL_MS -# ── OTLP/HTTP exporter wiring ───────────────────────────────────────────────── +def test_resource_attributes_include_license_id() -> None: + attributes: Final = bm.resource_attributes(_config(license_id="org-7")) + assert attributes == { + "service.name": "litellm-proxy", + "litellm.version": "1.2.3", + "litellm.license.id": "org-7", + } -def test_metrics_endpoint_appends_signal_path(): - assert bm._metrics_endpoint("https://telemetry.example.com") == "https://telemetry.example.com/v1/metrics" - assert bm._metrics_endpoint("https://telemetry.example.com/") == "https://telemetry.example.com/v1/metrics" - assert bm._metrics_endpoint("https://telemetry.example.com/v1/metrics") == "https://telemetry.example.com/v1/metrics" +def test_resource_attributes_omit_license_id_when_absent() -> None: + assert "litellm.license.id" not in bm.resource_attributes(_config(license_id=None)) -def _fake_exporter_class(captured: Dict[str, object]) -> type: - """A no-network stand-in for OTLPMetricExporter. Tests that build a real - MeterProvider must install this: the real exporter resolves the collector - host and opens a TLS connection on the reader's first export and on the - shutdown flush.""" - - class _FakeExporter: - # PeriodicExportingMetricReader probes these on the exporter it wraps. - _preferred_temporality: dict = {} - _preferred_aggregation: dict = {} - - def __init__(self, **kwargs): - captured.update(kwargs) - - def export(self, *args, **kwargs): - return None - - def shutdown(self, *args, **kwargs): - return None - - def force_flush(self, *args, **kwargs): - return True - - return _FakeExporter - - -def test_meter_provider_wires_client_cert_into_http_exporter(tmp_path, monkeypatch): - """Client cert+key authenticate us at the collector's mTLS front end; CA override rides certificate_file.""" - captured: Dict[str, object] = {} - - monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class(captured)) - config = _config(tmp_path) - provider = bm.build_mtls_meter_provider(config) - provider.shutdown() - - assert captured["endpoint"] == "https://collector.example:4317/v1/metrics" - assert captured["client_certificate_file"] == config.client_cert_path - assert captured["client_key_file"] == config.client_key_path - assert captured["certificate_file"] == config.ca_cert_path - - -def test_load_config_without_ca_is_valid(monkeypatch, tmp_path): - """The production collector presents a public web-PKI cert: no CA override required.""" - paths = _write_certs(tmp_path) - monkeypatch.setenv(bm.ENDPOINT_ENV, "https://telemetry.example.com") - monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) - monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) - config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") - assert config is not None and config.ca_cert_path is None - - -# ── Resource and metric attributes ──────────────────────────────────────────── - - -def test_resource_attributes_include_license_id(tmp_path): - attrs = bm._resource_attributes(_config(tmp_path, license_id="org-7")) - assert attrs["service.name"] == "litellm-proxy" - assert attrs["litellm.version"] == "1.2.3" - assert attrs["litellm.license.id"] == "org-7" - - -def test_resource_attributes_omit_license_id_when_absent(tmp_path): - attrs = bm._resource_attributes(_config(tmp_path, license_id=None)) - assert "litellm.license.id" not in attrs - - -def test_billable_attributes_with_model_id(): - attrs = bm._billable_attributes(BillableCategory.LLM, "/chat/completions", 200, "deploy-3") - assert attrs == { +def test_billable_attributes_include_model_id() -> None: + assert bm.billable_attributes( + BillableCategory.LLM, "/chat/completions", 200, "deploy-3" + ) == { "litellm.endpoint.category": "llm", "http.route": "/chat/completions", "http.response.status_code": 200, @@ -400,33 +381,35 @@ def test_billable_attributes_with_model_id(): } -def test_billable_attributes_omit_model_id_when_none(): - attrs = bm._billable_attributes(BillableCategory.MCP, "/mcp", 200, None) - assert "litellm.model_id" not in attrs +def test_billable_attributes_omit_model_id_when_none() -> None: + attributes: Final = bm.billable_attributes( + BillableCategory.MCP, "/mcp", 200, None + ) + assert "litellm.model_id" not in attributes -# ── End-to-end recording via in-memory reader ───────────────────────────────── +def _counter_points(reader: InMemoryMetricReader) -> tuple[tuple[str, int], ...]: + data: Final = reader.get_metrics_data() + if data is None: + return () + return tuple( + (_metric_category(point.attributes), int(point.value)) + for resource_metric in data.resource_metrics + for scope_metric in resource_metric.scope_metrics + for metric in scope_metric.metrics + if metric.name == bm.METRIC_NAME and isinstance(metric.data, Sum) + for point in metric.data.data_points + ) -def _counter_points(reader: InMemoryMetricReader): - data = reader.get_metrics_data() - for resource_metric in data.resource_metrics: - for scope_metric in resource_metric.scope_metrics: - for metric in scope_metric.metrics: - if metric.name == bm.METRIC_NAME: - return list(metric.data.data_points) - return [] - - -def test_record_increments_counter_with_attributes(): - reader = InMemoryMetricReader() - recorder = bm.BillingMetricsRecorder(MeterProvider(metric_readers=[reader])) +def test_record_increments_counter_with_attributes() -> None: + reader: Final = InMemoryMetricReader() + recorder: Final = bm.BillingMetricsRecorder(MeterProvider(metric_readers=[reader])) recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id="m1") recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id="m1") recorder.record(category=BillableCategory.MCP, route="/mcp", status_code=200, model_id=None) - points = _counter_points(reader) - by_category = {point.attributes["litellm.endpoint.category"]: point.value for point in points} - assert by_category["llm"] == 2 - assert by_category["mcp"] == 1 + points: Final = _counter_points(reader) + by_category: Final = dict(points) + assert by_category == {"llm": 2, "mcp": 1} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 91f00edfb39..46e766f753c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -2,11 +2,12 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ +import copy import json import logging import re from itertools import chain -from typing import Final +from typing import Final, Literal from unittest.mock import patch import pytest @@ -26,6 +27,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) from litellm.types.utils import ( + CallTypesLiteral, ChatCompletionCustomToolCallPayload, ChatCompletionDeltaToolCall, ChatCompletionMessageCustomToolCall, @@ -467,6 +469,346 @@ class TestToolPermissionGuardrail: ] assert kept == [[], ["Read"]] + @pytest.mark.asyncio + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_post_call_judges_a_returned_namespace_member_by_its_qualified_name(self, stream: bool) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + data = {"guardrails": ["post-mcp"], "tools": [namespace]} + function = Function(name="lookup", arguments='{"city": "Paris"}') + + if stream: + + async def _stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_lookup", index=0, type="function", function=function + ) + ] + ), + ) + ] + ) + + chunks = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data=data + ) + ] + deltas = [choice.delta for choice in chain.from_iterable(chunk.choices for chunk in chunks)] + called = [call.function.name for call in chain.from_iterable(delta.tool_calls or () for delta in deltas)] + content = "".join(delta.content or "" for delta in deltas) + else: + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_lookup", "type": "function", "function": function}]}, + ) + ] + ) + await guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + message = response.choices[0].message + called = [call.function.name for call in message.tool_calls or ()] + content = message.content or "" + + assert called == [] + assert "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" in content + + @pytest.mark.asyncio + async def test_post_call_judges_a_returned_namespace_member_by_the_rule_its_arguments_satisfy(self) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=[ + { + "id": "allow_ls", + "tool_name": r"^exec$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^ls$"}, + }, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_ops", "tool_name": r"^ops__.*$", "decision": "allow"}, + ], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + data = {"guardrails": ["post-exec"], "tools": [namespace]} + function = Function(name="exec", arguments='{"command": "rm -rf /data"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_exec", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=response) + + message = response.choices[0].message + assert [call.function.name for call in message.tool_calls or ()] == [] + assert "Tool 'exec'" in (message.content or "") + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rules", "command", "kept"), + [ + ( + [ + { + "id": "allow_ls", + "tool_name": r"^exec$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^ls$"}, + }, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "ls", + ["run"], + ), + ( + [ + { + "id": "allow_ls", + "tool_name": r"^exec$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^ls$"}, + }, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "pwd", + [], + ), + ( + [ + {"id": "allow_run", "tool_name": r"^run$", "decision": "allow"}, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "ls", + [], + ), + ], + ids=["namespace_rule_allows", "namespace_rule_denies", "namespace_deny_beats_member_allow"], + ) + async def test_post_call_judges_a_returned_namespace_member_by_a_rule_on_its_namespace_name( + self, rules: list[dict[str, object]], command: str, kept: list[str] + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=rules, + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "exec", "tools": [{"type": "function", "name": "run"}]} + function = Function(name="run", arguments=json.dumps({"command": command})) + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_run", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-exec"], "tools": [namespace]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + message = response.choices[0].message + assert [call.function.name for call in message.tool_calls or ()] == kept + assert ("Tool 'exec'" in (message.content or "")) == (not kept) + + @pytest.mark.asyncio + @pytest.mark.parametrize("namespace_first", [True, False], ids=["namespace_first", "top_level_first"]) + async def test_post_call_judges_a_returned_name_a_top_level_function_carries_as_that_function( + self, namespace_first: bool + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + top_level = {"type": "function", "function": {"name": "lookup"}} + tools = [namespace, top_level] if namespace_first else [top_level, namespace] + function = Function(name="lookup", arguments='{"country": "France"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_lookup", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-mcp"], "tools": tools}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == ["lookup"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "same_named_tool", + [ + {"type": "mcp", "server_label": "lookup", "server_url": "https://mcp.example.com/sse"}, + {"type": "mcp_toolset", "mcp_server_name": "lookup"}, + {"ignored": {"name": "lookup"}}, + {"type": "custom", "custom": {"name": "lookup"}}, + {"functionDeclarations": [{"name": "lookup"}]}, + ], + ids=["mcp", "mcp_toolset", "unknown_keyed_block", "chat_custom", "native_declaration"], + ) + async def test_post_call_judges_a_returned_namespace_member_as_that_member_beside_a_same_named_non_function( + self, same_named_tool: dict[str, object] + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + function = Function(name="lookup", arguments='{"country": "France"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_lookup", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-mcp"], "tools": [namespace, same_named_tool]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == [] + + @pytest.mark.asyncio + async def test_post_call_judges_a_call_tagged_with_a_namespace_as_that_member_beside_a_same_named_function( + self, + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + top_level = {"type": "function", "function": {"name": "lookup"}} + tagged_call = { + "id": "call_lookup", + "type": "function", + "function": Function(name="lookup", arguments='{"country": "France"}'), + "provider_specific_fields": {"namespace": "mcp__weather"}, + } + response = ModelResponse( + choices=[Choices(finish_reason="tool_calls", message={"tool_calls": [tagged_call]})] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-mcp"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + message = response.choices[0].message + assert [call.function.name for call in message.tool_calls or ()] == [] + assert "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" in (message.content or "") + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("decision", "default_action", "kept"), + [("allow", "deny", ["weather__get_weather"]), ("deny", "allow", [])], + ids=["bare_name_allowed", "bare_name_denied"], + ) + async def test_post_call_judges_a_returned_qualified_member_name_by_the_member_names_pre_call_judged( + self, decision: Literal["allow", "deny"], default_action: Literal["allow", "deny"], kept: list[str] + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-weather", + rules=[{"id": "get_weather", "tool_name": r"^get_weather$", "decision": decision}], + default_action=default_action, + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "get_weather"}]} + top_level = {"type": "function", "function": {"name": "get_weather"}} + function = Function(name="weather__get_weather", arguments='{"city": "Istanbul"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_weather", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-weather"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == kept + + @pytest.mark.asyncio + async def test_post_call_judges_a_top_level_function_by_its_own_name_beside_a_member_qualified_the_same( + self, + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=[{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + top_level = {"type": "function", "function": {"name": "ops__exec"}} + function = Function(name="ops__exec", arguments="{}") + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_ops", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-exec"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == ["ops__exec"] + @pytest.mark.asyncio async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule( self, @@ -1361,6 +1703,1723 @@ class TestToolPermissionGuardrailAnthropicMessages: def _tool_use(self, name, tool_id="tu_1"): return {"type": "tool_use", "id": tool_id, "name": name, "input": {"command": "ls"}} + def _always_on_pre_call(self, on_disallowed_action: Literal["block", "rewrite"]) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name=f"anthropic-pre-call-{on_disallowed_action}", + rules=self.rules, + default_action="deny", + on_disallowed_action=on_disallowed_action, + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "denied_tool", + [ + {"name": "Read", "input_schema": {"type": "object", "properties": {}}}, + {"type": "custom", "name": "Read", "input_schema": {"type": "object", "properties": {}}}, + {"type": "function", "name": "Read", "parameters": {"type": "object", "properties": {}}}, + ], + ids=["anthropic", "anthropic_custom_type", "responses_api_flat_function"], + ) + async def test_pre_call_blocks_denied_request_tool_in_flat_format(self, denied_tool: dict[str, object]) -> None: + data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [denied_tool]} + + with pytest.raises(HTTPException) as excinfo: + await self._always_on_pre_call("block").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="anthropic_messages", + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.detail["detection_message"] == "Tool 'Read' denied by rule 'deny_read'" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool_fields", "rule_tool_type", "call_type"), + [ + ({"format": {"type": "text"}}, r"^custom$", "responses"), + ({"input_schema": {"type": "object", "properties": {}}}, r"^function$", "anthropic_messages"), + ({"input_schema": {"type": "object", "properties": {}}}, r"^function$", "pass_through_endpoint"), + ({"custom": {"name": "apply_patch"}, "input_schema": {"type": "object"}}, r"^custom$", "acompletion"), + ], + ids=[ + "responses_custom_stays_custom", + "anthropic_custom_reads_as_function", + "anthropic_custom_reads_as_function_on_a_pass_through_route", + "chat_custom_with_a_stray_input_schema_stays_custom", + ], + ) + async def test_pre_call_tool_type_rules_follow_the_request_format( + self, tool_fields: dict[str, object], rule_tool_type: str, call_type: CallTypesLiteral + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name=f"tool-type-{call_type}", + rules=[{"id": "deny_type", "tool_type": rule_tool_type, "decision": "deny"}], + default_action="allow", + on_disallowed_action="block", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "custom", "name": "apply_patch", **tool_fields}], + } + + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type=call_type, + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.detail["detection_message"] == "Tool 'apply_patch' denied by rule 'deny_type'" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool_shape", "tool_choice", "call_type", "expected_tool_choice"), + [ + ( + {"input_schema": {"type": "object", "properties": {}}}, + {"type": "tool", "name": "Read"}, + "anthropic_messages", + {"type": "none"}, + ), + ( + {"type": "function", "parameters": {"type": "object", "properties": {}}}, + {"type": "function", "name": "Read"}, + "responses", + "none", + ), + ], + ids=["anthropic", "responses_api"], + ) + async def test_pre_call_rewrite_strips_denied_flat_tool_and_forced_choice( + self, + tool_shape: dict[str, object], + tool_choice: dict[str, str], + call_type: CallTypesLiteral, + expected_tool_choice: dict[str, str] | str, + ) -> None: + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "Bash", **tool_shape}, {"name": "Read", **tool_shape}], + "tool_choice": tool_choice, + } + + result = await self._always_on_pre_call("rewrite").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type=call_type, + ) + + assert [tool["name"] for tool in result["tools"]] == ["Bash"] + assert result["tool_choice"] == expected_tool_choice + + def _pre_call( + self, + rules: list[dict[str, str]], + on_disallowed_action: Literal["block", "rewrite"], + default_action: Literal["allow", "deny"] = "allow", + ) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name=f"request-shapes-{on_disallowed_action}", + rules=rules, + default_action=default_action, + on_disallowed_action=on_disallowed_action, + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + async def _run_pre_call( + self, guardrail: ToolPermissionGuardrail, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: + return await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type=call_type, + ) + + def _namespace(self, *member_names: str) -> dict[str, object]: + return { + "type": "namespace", + "name": "mcp__everything", + "tools": [{"type": "function", "name": name, "parameters": {"type": "object"}} for name in member_names], + } + + @pytest.mark.asyncio + async def test_pre_call_blocks_a_denied_tool_inside_a_responses_namespace(self) -> None: + guardrail = self._pre_call( + [{"id": "deny_sum", "tool_name": r"^mcp__everything__get_sum$", "decision": "deny"}], "block" + ) + data = {"model": "gpt-5", "input": "add 2 and 3", "tools": [self._namespace("echo", "get_sum")]} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + + assert excinfo.value.detail["detection_message"] == "Tool 'mcp__everything__get_sum' denied by rule 'deny_sum'" + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_denied_namespace_members_and_empty_namespaces(self) -> None: + guardrail = self._pre_call([{"id": "deny_sum", "tool_name": r"^.*get_sum$", "decision": "deny"}], "rewrite") + partly_denied = {**self._namespace("echo", "get_sum"), "name": "mcp__kept"} + data = {"model": "gpt-5", "input": "add", "tools": [partly_denied, self._namespace("get_sum")]} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{**partly_denied, "tools": partly_denied["tools"][:1]}] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "malformed"), + [ + ( + "agenerate_content", + {"toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": [{"a": 1}]}}}, + ), + ("responses", {"input": [{"type": {"x": 1}, "content": "hi"}]}), + ("acompletion", {"extra_body": {"input": [{"type": [1]}]}}), + ("acompletion", {"extra_body": {"function_call": {"name": 5}}}), + ], + ids=["unhashable_allowed_function_name", "unhashable_input_type", "extra_body_input_type", "int_function_call"], + ) + async def test_pre_call_rewrite_leaves_a_malformed_value_for_the_provider_to_reject( + self, call_type: CallTypesLiteral, malformed: dict[str, object] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], "rewrite", "deny" + ) + tools = [{"type": "function", "function": {"name": name}} for name in ("get_weather", "shell")] + + result = await self._run_pre_call(guardrail, {"model": "m", "tools": tools, **malformed}, call_type) + + assert result["tools"] == tools[:1] + assert {key: result[key] for key in malformed} == malformed + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rule", "tool", "call_type", "detection_message"), + [ + ( + {"id": "deny_mcp", "tool_type": r"^mcp$", "decision": "deny"}, + {"type": "mcp", "server_label": "weather", "server_url": "https://mcp.example.com"}, + "responses", + "Tool 'weather' denied by rule 'deny_mcp'", + ), + ( + {"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"}, + {"type": "mcp_toolset", "mcp_server_name": "weather"}, + "anthropic_messages", + "Tool 'weather' denied by default action", + ), + ( + {"id": "deny_weather", "tool_name": r"^weather$", "decision": "deny"}, + {"type": "mcp", "server_label": "weather", "server_url": "https://mcp.example.com"}, + "responses", + "Tool 'weather' denied by rule 'deny_weather'", + ), + ( + {"id": "deny_lookup", "tool_name": r"^mcp__weather__lookup$", "decision": "deny"}, + {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "custom", "name": "lookup"}]}, + "responses", + "Tool 'mcp__weather__lookup' denied by rule 'deny_lookup'", + ), + ( + {"id": "deny_lookup", "tool_name": r"^mcp__weather__lookup$", "decision": "deny"}, + {"functionDeclarations": [{"name": "mcp__weather__lookup"}, {"name": "get_weather"}]}, + "agenerate_content", + "Tool 'mcp__weather__lookup' denied by rule 'deny_lookup'", + ), + ( + {"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"}, + {"googleSearch": {}}, + "acompletion", + "Tool 'googleSearch' denied by default action", + ), + ( + {"id": "deny_custom", "tool_type": r"^custom$", "decision": "deny"}, + {"type": "custom", "custom": {"name": "grammar_tool"}}, + "acompletion", + "Tool 'grammar_tool' denied by rule 'deny_custom'", + ), + ( + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}, + {"type": "custom", "custom": {"name": "Read"}}, + "acompletion", + "Tool 'Read' denied by rule 'deny_read'", + ), + ], + ids=[ + "responses_mcp_by_type", + "anthropic_mcp_toolset_by_server_name", + "responses_mcp_by_server_label", + "responses_namespace_custom_member", + "gemini_function_declarations", + "gemini_type_only_tool_by_its_key", + "chat_custom_keeps_its_type", + "chat_custom_nested_name", + ], + ) + async def test_pre_call_checks_tools_that_carry_no_top_level_name( + self, rule: dict[str, str], tool: dict[str, object], call_type: CallTypesLiteral, detection_message: str + ) -> None: + guardrail = self._pre_call([rule], "block", default_action="deny") + data = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "input": "hi", "tools": [tool]} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, call_type) + + assert excinfo.value.detail["detection_message"] == detection_message + + @pytest.mark.asyncio + async def test_pre_call_lets_a_type_rule_allow_a_gemini_tool_by_its_key(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_search", "tool_type": r"^googleSearch$", "decision": "allow"}], "block", default_action="deny" + ) + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [{"googleSearch": {}}]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [{"googleSearch": {}}] + + @pytest.mark.asyncio + async def test_pre_call_blocks_a_denied_tool_loaded_by_responses_tool_search(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "block") + loaded = {"type": "tool_search_output", "tools": [{"type": "function", "name": "mcp__weather__lookup"}]} + data = {"model": "gpt-5", "input": [loaded], "tools": []} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + + assert excinfo.value.detail["detection_message"] == "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" + + @pytest.mark.asyncio + async def test_pre_call_rewrite_keeps_the_allowed_members_of_grouped_and_loaded_tools(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + get_weather = {"type": "function", "name": "get_weather"} + lookup = {"type": "function", "name": "mcp__weather__lookup"} + data = { + "model": "m", + "input": [{"role": "user", "content": "hi"}, {"type": "tool_search_output", "tools": [get_weather, lookup]}], + "tools": [{"functionDeclarations": [{"name": "get_weather"}, {"name": "mcp__weather__lookup"}]}], + } + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{"functionDeclarations": [{"name": "get_weather"}]}] + assert result["input"] == [ + {"role": "user", "content": "hi"}, + {"type": "tool_search_output", "tools": [get_weather]}, + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "message", "rewritten"), + [ + ( + "anthropic_messages", + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "tool_reference", "tool_name": "mcp__weather__lookup"}], + }, + {"type": "text", "text": "Tool loaded."}, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + { + "type": "text", + "text": "Permission denied: tool mcp__weather__lookup is not available", + } + ], + }, + {"type": "text", "text": "Tool loaded."}, + ], + }, + ), + ( + "anthropic_messages", + { + "role": "assistant", + "content": [ + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [ + {"type": "tool_reference", "tool_name": "mcp__weather__lookup"}, + {"type": "tool_reference", "tool_name": "get_weather"}, + ], + }, + } + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, + } + ], + }, + ), + ( + "acompletion", + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "tool_reference", "tool_name": "mcp__weather__lookup"}, + {"type": "tool_reference", "tool_name": "get_weather"}, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, + ), + ], + ) + async def test_pre_call_rewrite_drops_references_to_a_removed_tool(self, call_type, message, rewritten) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + data = { + "model": "m", + "messages": [{"role": "user", "content": "weather in Istanbul"}, message], + "tools": [ + {"name": "get_weather", "input_schema": {"type": "object"}, "defer_loading": True}, + {"name": "mcp__weather__lookup", "input_schema": {"type": "object"}, "defer_loading": True}, + ], + } + + result = await self._run_pre_call(guardrail, data, call_type) + + assert [tool["name"] for tool in result["tools"]] == ["get_weather"] + assert result["messages"] == [{"role": "user", "content": "weather in Istanbul"}, rewritten] + + @pytest.mark.asyncio + @pytest.mark.parametrize("tool_choice", ["required", {"type": "function", "name": "get_weather"}]) + async def test_pre_call_rewrite_keeps_tool_choice_when_only_loaded_tools_remain(self, tool_choice) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + get_weather = {"type": "function", "name": "get_weather"} + lookup = {"type": "function", "name": "mcp__weather__lookup"} + data = { + "model": "m", + "input": [{"role": "user", "content": "hi"}, {"type": "tool_search_output", "tools": [get_weather, lookup]}], + "tools": [lookup], + "tool_choice": tool_choice, + } + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [] + assert result["input"][1] == {"type": "tool_search_output", "tools": [get_weather]} + assert result["tool_choice"] == tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rule", "tools", "tool_choice"), + [ + ( + {"id": "deny_search", "tool_type": r"^web_search$", "decision": "deny"}, + [{"type": "web_search"}, {"type": "function", "name": "get_weather"}], + {"type": "web_search"}, + ), + ( + {"id": "deny_weather", "tool_name": r"^weather$", "decision": "deny"}, + [ + {"type": "mcp", "server_label": "weather", "server_url": "https://a.example.com"}, + {"type": "mcp", "server_label": "files", "server_url": "https://b.example.com"}, + ], + {"type": "mcp", "server_label": "weather"}, + ), + ], + ids=["hosted_tool_type", "hosted_mcp_server"], + ) + async def test_pre_call_rewrite_relaxes_a_hosted_tool_choice_it_removed( + self, rule: dict[str, str], tools: list[dict[str, str]], tool_choice: dict[str, str] + ) -> None: + guardrail = self._pre_call([rule], "rewrite") + data = {"model": "gpt-5", "input": "hi", "tools": tools, "tool_choice": tool_choice} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == tools[1:] + assert result["tool_choice"] == "none" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tools", "tool_choice", "call_type", "expected_tool_choice"), + [ + ( + [{"type": "custom", "name": "Bash"}, {"type": "custom", "name": "Read"}], + {"type": "custom", "name": "Read"}, + "responses", + "none", + ), + ( + [{"type": "function", "name": "Bash"}, {"type": "function", "name": "Read"}], + { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "Read"}, {"type": "function", "name": "Bash"}], + }, + "responses", + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "Bash"}]}, + ), + ( + [{"type": "function", "name": "Bash"}, {"type": "function", "name": "Read"}], + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "Read"}]}, + "responses", + "none", + ), + ( + [ + {"type": "function", "function": {"name": "Bash"}}, + {"type": "function", "function": {"name": "Read"}}, + ], + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "Read"}}]}, + }, + "acompletion", + "none", + ), + ( + [{"name": "Read", "input_schema": {"type": "object"}}], + {"type": "any"}, + "anthropic_messages", + {"type": "none"}, + ), + ( + [ + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": "Read", "input_schema": {"type": "object"}}, + ], + {"type": "any"}, + "anthropic_messages", + {"type": "any"}, + ), + ( + [{"type": "custom", "custom": {"name": "Bash"}}, {"type": "custom", "custom": {"name": "Read"}}], + {"type": "custom", "custom": {"name": "Read"}}, + "acompletion", + "none", + ), + ], + ids=[ + "responses_custom_choice", + "responses_allowed_tools_keeps_the_rest", + "responses_allowed_tools_all_denied", + "chat_allowed_tools_all_denied", + "anthropic_any_with_no_tools_left", + "anthropic_any_with_a_tool_left", + "chat_custom_choice", + ], + ) + async def test_pre_call_rewrite_leaves_no_tool_choice_that_needs_a_removed_tool( + self, + tools: list[dict[str, object]], + tool_choice: dict[str, object], + call_type: CallTypesLiteral, + expected_tool_choice: dict[str, object] | str, + ) -> None: + guardrail = self._pre_call(self.rules, "rewrite") + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": tools, + "tool_choice": tool_choice, + } + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tool_choice"] == expected_tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool", "tool_choice", "call_type", "allowed_tool_type"), + [ + ({"type": "custom", "name": "apply_patch"}, {"type": "custom", "name": "apply_patch"}, "responses", r"^custom$"), + ( + {"type": "web_search_20250305", "name": "web_search"}, + {"type": "tool", "name": "web_search"}, + "anthropic_messages", + r"^web_search_\d+$", + ), + ], + ids=["responses_custom_tool", "anthropic_server_tool"], + ) + async def test_pre_call_checks_a_forced_choice_with_its_tool_type( + self, + tool: dict[str, str], + tool_choice: dict[str, str], + call_type: CallTypesLiteral, + allowed_tool_type: str, + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_type", "tool_type": allowed_tool_type, "decision": "allow"}], "block", default_action="deny" + ) + data = {"model": "m", "input": "hi", "tools": [tool], "tool_choice": tool_choice} + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tools"] == [tool] + assert result["tool_choice"] == tool_choice + + @pytest.mark.asyncio + async def test_pre_call_rewrite_relaxes_a_choice_for_a_tool_a_type_rule_removed(self) -> None: + guardrail = self._pre_call([{"id": "deny_custom", "tool_type": r"^custom$", "decision": "deny"}], "rewrite") + data = { + "model": "gpt-5", + "input": "patch it", + "tools": [{"type": "function", "name": "Bash"}, {"type": "custom", "name": "apply_patch"}], + "tool_choice": {"type": "custom", "name": "apply_patch"}, + } + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{"type": "function", "name": "Bash"}] + assert result["tool_choice"] == "none" + + async def _blocked_and_rewritten( + self, rules: list[dict[str, str]], request: dict[str, object], call_type: CallTypesLiteral + ) -> tuple[str, dict[str, object]]: + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(self._pre_call(rules, "block"), copy.deepcopy(request), call_type) + rewritten = await self._run_pre_call(self._pre_call(rules, "rewrite"), request, call_type) + return excinfo.value.detail["detection_message"], rewritten + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("denied_tool", "detection_message"), + [ + ( + {"type": "function", "name": "mcp__weather__lookup"}, + "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'", + ), + ( + {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]}, + "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'", + ), + ({"type": "web_search"}, "Tool 'web_search' denied by rule 'deny_search'"), + ], + ids=["flat", "namespace_member", "hosted_type"], + ) + async def test_pre_call_checks_and_strips_tools_in_responses_additional_tools_items( + self, denied_tool: dict[str, object], detection_message: str + ) -> None: + rules = [ + {"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}, + {"id": "deny_search", "tool_type": r"^web_search$", "decision": "deny"}, + ] + get_weather = {"type": "function", "name": "get_weather"} + request = { + "model": "gpt-5", + "input": [ + {"role": "user", "content": "hi"}, + {"type": "additional_tools", "tools": [get_weather, denied_tool]}, + {"type": "additional_tools", "tools": [denied_tool]}, + ], + "tools": [], + "tool_choice": "required", + } + + message, rewritten = await self._blocked_and_rewritten(rules, request, "responses") + + assert message == detection_message + assert rewritten["input"] == [ + {"role": "user", "content": "hi"}, + {"type": "additional_tools", "tools": [get_weather]}, + ] + assert rewritten["tool_choice"] == "required" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rules", "default_action", "kept"), + [ + ([{"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}], "deny", ["run"]), + ([{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], "allow", []), + ( + [ + {"id": "allow_run", "tool_name": r"^run$", "decision": "allow"}, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "allow", + [], + ), + ( + [ + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_all", "tool_name": r"^.*$", "decision": "allow"}, + ], + "allow", + [], + ), + ( + [ + {"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}, + {"id": "deny_all", "tool_name": r"^.*$", "decision": "deny"}, + ], + "deny", + ["run"], + ), + ], + ids=[ + "namespace_allowed", + "namespace_denied", + "namespace_deny_beats_member_allow", + "namespace_deny_beats_a_catch_all_allow", + "namespace_allow_beats_a_catch_all_deny", + ], + ) + async def test_pre_call_judges_a_namespace_member_by_a_rule_on_its_namespace_name( + self, rules: list[dict[str, str]], default_action: Literal["allow", "deny"], kept: list[str] + ) -> None: + namespace = {"type": "namespace", "name": "exec", "tools": [{"type": "function", "name": "run"}]} + request = {"model": "gpt-5", "input": "Run pwd", "tools": [namespace]} + + rewritten = await self._run_pre_call(self._pre_call(rules, "rewrite", default_action), request, "responses") + + tools = rewritten.get("tools") or [] + assert [member["name"] for member in chain.from_iterable(tool["tools"] for tool in tools)] == kept + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "bash", "read"), + [ + ( + "acompletion", + {"type": "function", "function": {"name": "Bash"}}, + {"type": "function", "function": {"name": "Read"}}, + ), + ("responses", {"type": "function", "name": "Bash"}, {"type": "function", "name": "Read"}), + ( + "anthropic_messages", + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": "Read", "input_schema": {"type": "object"}}, + ), + ], + ids=["chat", "responses", "anthropic"], + ) + async def test_pre_call_checks_and_strips_tools_in_extra_body( + self, call_type: CallTypesLiteral, bash: dict[str, object], read: dict[str, object] + ) -> None: + request = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "input": "hi", + "extra_body": {"tools": [bash, read], "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(self.rules, request, call_type) + + assert message == "Tool 'Read' denied by rule 'deny_read'" + assert rewritten["extra_body"] == {"tools": [bash], "seed": 7} + + @pytest.mark.asyncio + async def test_pre_call_checks_and_strips_tools_loaded_in_extra_body_input(self) -> None: + rules = [{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}] + get_weather = {"type": "function", "name": "get_weather"} + lookup = {"type": "function", "name": "mcp__weather__lookup"} + user = {"role": "user", "content": "hi"} + request = { + "model": "gpt-5", + "input": "hi", + "tools": [], + "extra_body": {"input": [user, {"type": "tool_search_output", "tools": [get_weather, lookup]}], "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(rules, request, "responses") + + assert message == "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" + assert rewritten["extra_body"] == { + "input": [user, {"type": "tool_search_output", "tools": [get_weather]}], + "seed": 7, + } + + @pytest.mark.asyncio + async def test_rewrite_relaxes_an_extra_body_tool_choice_that_names_a_dropped_tool(self) -> None: + guardrail = self._pre_call([{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "function": {"name": "get_weather"}} + delete_db = {"type": "function", "function": {"name": "delete_db"}} + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "extra_body": { + "tools": [weather, delete_db], + "tool_choice": {"type": "function", "function": {"name": "delete_db"}}, + "seed": 7, + }, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["extra_body"] == {"tools": [weather], "tool_choice": "none", "seed": 7} + + @pytest.mark.asyncio + async def test_rewrite_relaxes_a_required_tool_choice_when_the_extra_body_tools_sent_lose_every_tool( + self, + ) -> None: + guardrail = self._pre_call([{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "function": {"name": "get_weather"}} + delete_db = {"type": "function", "function": {"name": "delete_db"}} + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": [weather], + "extra_body": {"tools": [delete_db], "tool_choice": "required", "seed": 7}, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + assert result["extra_body"] == {"tools": [], "tool_choice": "none", "seed": 7} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "tool_choice"), + [ + ("acompletion", {"type": "function", "function": {"name": "Read"}}), + ("responses", {"type": "function", "name": "Read"}), + ], + ids=["chat", "responses"], + ) + async def test_pre_call_checks_an_extra_body_tool_choice_that_names_an_unlisted_tool( + self, call_type: CallTypesLiteral, tool_choice: dict[str, object] + ) -> None: + request = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "input": "hi", + "tools": [{"type": "function", "name": "Bash"}], + "extra_body": {"tool_choice": tool_choice, "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(self.rules, request, call_type) + + assert message == "Tool 'Read' denied by rule 'deny_read'" + assert rewritten["extra_body"] == {"tool_choice": "none", "seed": 7} + + @pytest.mark.asyncio + async def test_pre_call_checks_and_strips_legacy_functions_in_extra_body(self) -> None: + rules = [{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}] + weather = {"name": "get_weather", "parameters": {"type": "object"}} + delete_db = {"name": "delete_db", "parameters": {"type": "object"}} + request = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "extra_body": {"functions": [weather, delete_db], "function_call": {"name": "delete_db"}, "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(rules, request, "acompletion") + + assert message == "Tool 'delete_db' denied by rule 'deny_delete'" + assert rewritten["extra_body"] == {"functions": [weather], "function_call": "none", "seed": 7} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "fields", + [ + {"functions": [{"name": "get_weather"}, {"name": "delete_db"}], "function_call": "delete_db"}, + { + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + "mcp_servers": [{"type": "url", "url": "https://mcp.example.com/mcp", "name": "delete_db"}], + "tool_choice": "required", + }, + { + "tools": [{"functionDeclarations": [{"name": "get_weather"}, {"name": "delete_db"}]}], + "toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["delete_db"]}}, + }, + ], + ids=["legacy_functions", "mcp_servers", "gemini_tool_config"], + ) + async def test_extra_body_fields_are_judged_and_rewritten_like_the_same_top_level_fields( + self, fields: dict[str, object] + ) -> None: + rules = [{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}] + base = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + top_request = {**base, **copy.deepcopy(fields)} + extra_request = {**base, "extra_body": {**copy.deepcopy(fields), "seed": 7}} + + top_message, top = await self._blocked_and_rewritten(rules, top_request, "acompletion") + extra_message, extra = await self._blocked_and_rewritten(rules, extra_request, "acompletion") + + top_fields = {key: top[key] for key in fields} + assert top_fields != fields + assert extra_message == top_message + assert extra["extra_body"] == {**top_fields, "seed": 7} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rule", + [ + {"id": "deny_deepwiki", "tool_name": r"^deepwiki$", "decision": "deny"}, + {"id": "deny_deepwiki", "tool_name": r"^deepwiki$", "tool_type": r"^mcp_toolset$", "decision": "deny"}, + ], + ids=["by_server_name", "by_server_name_and_mcp_toolset_type"], + ) + async def test_pre_call_checks_and_strips_anthropic_mcp_servers_like_an_mcp_toolset( + self, rule: dict[str, str] + ) -> None: + deepwiki = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + files = {"type": "url", "url": "https://files.example.com/mcp", "name": "files"} + request = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "mcp_servers": [deepwiki, files], + "tool_choice": {"type": "any"}, + } + + message, rewritten = await self._blocked_and_rewritten([rule], request, "anthropic_messages") + + assert message == "Tool 'deepwiki' denied by rule 'deny_deepwiki'" + assert rewritten["mcp_servers"] == [files] + assert rewritten["tool_choice"] == {"type": "any"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rule", "tool", "detection_message", "kept_tools"), + [ + ( + {"id": "deny_code", "tool_type": r"^codeExecution$", "decision": "deny"}, + {"googleSearch": {}, "codeExecution": {}}, + "Tool 'codeExecution' denied by rule 'deny_code'", + [{"googleSearch": {}}], + ), + ( + {"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}, + {"functionDeclarations": [{"name": "mcp__weather__lookup"}], "googleSearch": {}}, + "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'", + [{"googleSearch": {}}], + ), + ( + {"id": "deny_search", "tool_type": r"^googleSearch$", "decision": "deny"}, + {"functionDeclarations": [{"name": "get_weather"}], "googleSearch": {}}, + "Tool 'googleSearch' denied by rule 'deny_search'", + [{"functionDeclarations": [{"name": "get_weather"}]}], + ), + ( + {"id": "deny_builtin", "tool_type": r"^(googleSearch|codeExecution)$", "decision": "deny"}, + {"googleSearch": {}, "codeExecution": {}}, + "Tool 'googleSearch' denied by rule 'deny_builtin'", + [], + ), + ], + ids=["denied_key_beside_a_key", "every_declaration_denied", "denied_key_beside_declarations", "all_denied"], + ) + async def test_pre_call_checks_and_strips_each_part_of_a_multi_key_gemini_tool( + self, + rule: dict[str, str], + tool: dict[str, object], + detection_message: str, + kept_tools: list[dict[str, object]], + ) -> None: + request = {"model": "gemini-3", "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [tool]} + + message, rewritten = await self._blocked_and_rewritten([rule], request, "agenerate_content") + + assert message == detection_message + assert rewritten["tools"] == kept_tools + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("tool", "call_type"), + [({"type": "web_search"}, "responses"), ({"googleSearch": {}}, "acompletion")], + ids=["responses_hosted_type", "gemini_tool_key"], + ) + async def test_pre_call_matches_a_name_rule_against_a_nameless_tool_by_its_type( + self, + tool: dict[str, object], + call_type: CallTypesLiteral, + on_disallowed_action: Literal["block", "rewrite"], + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_all", "tool_name": r".*", "decision": "allow"}], on_disallowed_action, default_action="deny" + ) + data = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "input": "hi", "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rules", "default_action", "detection_message"), + [ + ( + [{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], + "allow", + "Tool 'exec' denied by rule 'deny_exec'", + ), + ([{"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}], "deny", None), + ( + [ + {"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}, + {"id": "deny_functions", "tool_name": r"^functions__.*$", "decision": "deny"}, + ], + "allow", + "Tool 'functions__exec' denied by rule 'deny_functions'", + ), + ( + [ + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_functions", "tool_type": r"^function$", "decision": "allow"}, + ], + "allow", + "Tool 'exec' denied by rule 'deny_exec'", + ), + ( + [ + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_all", "tool_name": r"^.*$", "decision": "allow"}, + ], + "deny", + "Tool 'exec' denied by rule 'deny_exec'", + ), + ( + [ + {"id": "allow_functions", "tool_name": r"^functions__.*$", "decision": "allow"}, + {"id": "deny_all", "tool_name": r"^.*$", "decision": "deny"}, + ], + "allow", + None, + ), + ], + ids=[ + "bare_name_deny_rule", + "bare_name_allow_rule", + "qualified_rule_wins_over_an_earlier_bare_rule", + "type_rule_keeps_an_earlier_bare_name_rule_in_charge", + "bare_name_deny_beats_a_later_catch_all_allow", + "qualified_allow_beats_a_later_catch_all_deny", + ], + ) + async def test_pre_call_judges_a_namespace_member_by_its_qualified_and_bare_names( + self, rules: list[dict[str, str]], default_action: Literal["allow", "deny"], detection_message: str | None + ) -> None: + guardrail = self._pre_call(rules, "block", default_action=default_action) + namespace = {"type": "namespace", "name": "functions", "tools": [{"type": "function", "name": "exec"}]} + data = {"model": "gpt-5", "input": "run date", "tools": [namespace]} + + if detection_message is None: + assert (await self._run_pre_call(guardrail, data, "responses"))["tools"] == [namespace] + return + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + assert excinfo.value.detail["detection_message"] == detection_message + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_a_namespace_member_a_bare_name_rule_denies(self) -> None: + guardrail = self._pre_call([{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], "rewrite") + members = [{"type": "function", "name": "exec"}, {"type": "function", "name": "read_file"}] + namespace = {"type": "namespace", "name": "functions", "tools": members} + data = {"model": "gpt-5", "input": "run date", "tools": [namespace]} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{"type": "namespace", "name": "functions", "tools": members[1:]}] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("choice", "relaxed"), + [ + ({"tool_choice": {"type": "function", "function": {"name": "exec"}}}, {"tool_choice": "none"}), + ({"tool_choice": {"type": "function", "function": {"name": "ops"}}}, {"tool_choice": "none"}), + ( + { + "tool_choice": { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [ + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "get_weather"}}, + ], + }, + } + }, + { + "tool_choice": { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + }, + } + }, + ), + ( + {"toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["exec"]}}}, + {"toolConfig": {"functionCallingConfig": {"mode": "NONE"}}}, + ), + ], + ids=["bare_member_choice", "namespace_choice", "allowed_tools_entry", "allowed_function_names"], + ) + async def test_pre_call_rewrite_relaxes_a_bare_name_a_dropped_namespace_member_was_sent_under( + self, choice: dict[str, object], relaxed: dict[str, object] + ) -> None: + guardrail = self._pre_call([{"id": "deny_exec", "tool_name": r"^ops__exec$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + data = { + "model": "gemini/gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + "tools": [weather, namespace], + **choice, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + assert {key: result[key] for key in relaxed} == relaxed + + @pytest.mark.asyncio + async def test_pre_call_rewrite_keeps_a_bare_name_another_kept_namespace_still_serves(self) -> None: + guardrail = self._pre_call([{"id": "deny_exec", "tool_name": r"^ops__exec$", "decision": "deny"}], "rewrite") + dropped_namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + kept_namespace = {"type": "namespace", "name": "dev", "tools": [{"type": "function", "name": "exec"}]} + tool_choice = {"type": "function", "function": {"name": "exec"}} + data = { + "model": "gemini/gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + "tools": [dropped_namespace, kept_namespace], + "tool_choice": tool_choice, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [kept_namespace] + assert result["tool_choice"] == tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize("exec_first", [True, False], ids=["exec_namespace_first", "exec_namespace_last"]) + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_denies_a_bare_name_another_namespace_shares_a_qualified_name_with( + self, on_disallowed_action: Literal["block", "rewrite"], exec_first: bool + ) -> None: + guardrail = self._pre_call( + [{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], on_disallowed_action + ) + exec_namespace = {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]} + twin_namespace = {"type": "namespace", "name": "mcp", "tools": [{"type": "function", "name": "ops__exec"}]} + tools = [exec_namespace, twin_namespace] if exec_first else [twin_namespace, exec_namespace] + data = {"model": "gpt-5", "input": "run date", "tools": tools} + + if on_disallowed_action == "rewrite": + assert (await self._run_pre_call(guardrail, data, "responses"))["tools"] == [] + return + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + assert excinfo.value.detail["detection_message"] == "Tool 'exec' denied by rule 'deny_exec'" + + @pytest.mark.asyncio + @pytest.mark.parametrize("placement", ["tools", "functions"]) + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("member", "detection_message"), + [ + ({"type": "function", "function": {"name": "delete_db"}}, "Tool 'delete_db' denied by rule 'deny_delete'"), + ( + {"type": "function", "name": "lookup", "function": {"name": "delete_db"}}, + "Tool 'delete_db' denied by rule 'deny_delete'", + ), + ({"googleSearch": {}}, "Tool 'googleSearch' denied by rule 'deny_search'"), + ( + { + "type": "namespace", + "name": "inner", + "tools": [{"type": "function", "name": "noop"}], + "function": {"name": "delete_db"}, + }, + "Tool 'delete_db' denied by rule 'deny_delete'", + ), + ], + ids=["chat_shaped_member", "function_object_behind_another_name", "hosted_key_member", "nested_namespace"], + ) + async def test_pre_call_judges_a_namespace_member_by_what_a_provider_reads_from_it( + self, + member: dict[str, object], + detection_message: str, + on_disallowed_action: Literal["block", "rewrite"], + placement: Literal["tools", "functions"], + ) -> None: + guardrail = self._pre_call( + [ + {"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}, + {"id": "deny_search", "tool_type": r"^googleSearch$", "decision": "deny"}, + ], + on_disallowed_action, + ) + weather = {"type": "function", "name": "get_weather"} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + placement: [{"type": "namespace", "name": "ops", "tools": [weather, member]}], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == detection_message + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result[placement] == [{"type": "namespace", "name": "ops", "tools": [weather]}] + + @pytest.mark.asyncio + async def test_pre_call_allows_a_chat_shaped_namespace_member_the_allowlist_names(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_create", "tool_name": r"^github__create_issue$", "decision": "allow"}], + "block", + default_action="deny", + ) + namespace = { + "type": "namespace", + "name": "github", + "tools": [{"type": "function", "function": {"name": "create_issue"}}], + } + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [namespace]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [namespace] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool_config", "rewritten_tool_config"), + [ + ( + { + "functionCallingConfig": { + "mode": "ANY", + "allowedFunctionNames": ["mcp__weather__lookup", "get_weather"], + } + }, + {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["get_weather"]}}, + ), + ( + {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["mcp__weather__lookup"]}}, + {"functionCallingConfig": {"mode": "NONE"}}, + ), + ( + {"function_calling_config": {"mode": "any", "allowed_function_names": ["mcp__weather__lookup"]}}, + {"function_calling_config": {"mode": "NONE"}}, + ), + ({"functionCallingConfig": {"mode": "ANY"}}, {"functionCallingConfig": {"mode": "ANY"}}), + ( + {"functionCallingConfig": {"mode": "AUTO", "allowedFunctionNames": []}}, + {"functionCallingConfig": {"mode": "AUTO", "allowedFunctionNames": []}}, + ), + ], + ids=[ + "some_allowed_names_left", + "every_allowed_name_removed", + "snake_case_fields", + "no_allowed_names", + "empty_allowed_names_list", + ], + ) + async def test_pre_call_rewrite_drops_removed_functions_from_a_gemini_tool_config( + self, tool_config: dict[str, object], rewritten_tool_config: dict[str, object] + ) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + config_key = "tool_config" if "function_calling_config" in tool_config else "toolConfig" + data = { + "model": "gemini-3", + "contents": [{"role": "user", "parts": [{"text": "weather in Paris?"}]}], + "tools": [{"functionDeclarations": [{"name": "mcp__weather__lookup"}, {"name": "get_weather"}]}], + config_key: tool_config, + } + + result = await self._run_pre_call(guardrail, data, "agenerate_content") + + assert result["tools"] == [{"functionDeclarations": [{"name": "get_weather"}]}] + assert result[config_key] == rewritten_tool_config + + @pytest.mark.asyncio + async def test_pre_call_rewrite_turns_off_forced_gemini_calling_when_no_function_is_left(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + data = { + "model": "gemini-3", + "contents": [{"role": "user", "parts": [{"text": "weather in Paris?"}]}], + "tools": [{"functionDeclarations": [{"name": "mcp__weather__lookup"}]}, {"googleSearch": {}}], + "toolConfig": {"functionCallingConfig": {"mode": "ANY"}}, + } + + result = await self._run_pre_call(guardrail, data, "agenerate_content") + + assert result["tools"] == [{"googleSearch": {}}] + assert result["toolConfig"] == {"functionCallingConfig": {"mode": "NONE"}} + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_judges_a_named_hosted_tool_by_its_type_as_well_as_its_name( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + "tools": [weather, {"type": "web_search", "name": "get_weather"}], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'web_search' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_judges_a_gemini_hosted_key_by_the_key_whatever_name_it_carries( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + computer_use = {"computerUse": {"name": "get_weather", "environment": "browser"}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [weather, computer_use]} + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'computerUse' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rules, tool", + [ + ( + [{"id": "allow_bash", "tool_name": r"^bash$", "decision": "allow"}], + {"type": "bash_20250124", "name": "bash"}, + ), + ( + [{"id": "allow_edit", "tool_name": r"^str_replace_based_edit_tool$", "decision": "allow"}], + {"type": "text_editor_20250728", "name": "str_replace_based_edit_tool"}, + ), + ( + [{"id": "allow_search", "tool_name": r"^web_search$", "decision": "allow"}], + {"type": "web_search_20250305", "name": "web_search"}, + ), + ( + [ + { + "id": "allow_search", + "tool_name": r"^web_search$", + "tool_type": r"^web_search_20250305$", + "decision": "allow", + } + ], + {"type": "web_search_20250305", "name": "web_search"}, + ), + ], + ids=["client_tool_by_name", "client_tool_with_its_own_name", "server_tool_by_name", "name_and_type_rule"], + ) + async def test_pre_call_allows_an_anthropic_defined_tool_that_an_allowlist_names( + self, rules: list[dict[str, object]], tool: dict[str, object] + ) -> None: + guardrail = self._pre_call(rules, "block", default_action="deny") + data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "tool, call_type", + [ + ({"type": "web_search_20250305", "name": "get_weather"}, "anthropic_messages"), + ( + { + "type": "computer_20251124", + "name": "get_weather", + "function": {"parameters": {"display_width_px": 1024, "display_height_px": 768}}, + }, + "acompletion", + ), + ({"type": "computer_use", "name": "get_weather", "environment": "browser"}, "acompletion"), + ({"type": "text_editor_20250728", "name": "get_weather"}, "anthropic_messages"), + ], + ids=["server_tool", "anthropic_computer_tool", "gemini_computer_use", "text_editor"], + ) + async def test_pre_call_blocks_a_provider_defined_tool_whose_name_is_not_its_own( + self, tool: dict[str, object], call_type: CallTypesLiteral + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + "block", + default_action="deny", + ) + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, call_type) + + assert excinfo.value.detail["detection_message"] == f"Tool '{tool['type']}' denied by default action" + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + "tool", + [ + {"type": "function", "function": {"name": "delete_db"}, "functionDeclarations": [{"name": "get_weather"}]}, + {"type": "function", "function": {"name": "delete_db"}, "function_declarations": [{"name": "get_weather"}]}, + {"type": "custom", "name": "get_weather", "function": {"name": "delete_db"}}, + {"type": "mcp", "server_label": "get_weather", "function": {"name": "delete_db"}}, + { + "type": "namespace", + "name": "github", + "tools": [{"type": "function", "name": "create_issue"}], + "function": {"name": "delete_db"}, + }, + { + "type": "computer_20251124", + "name": "computer", + "function": {"name": "delete_db", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + }, + ], + ids=[ + "declarations_list", + "snake_case_declarations_list", + "custom_top_level_name", + "mcp_server_label", + "namespace_with_members", + "computer_top_level_name", + ], + ) + async def test_pre_call_judges_the_function_a_tool_runs_behind_an_allowed_name( + self, tool: dict[str, object], on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_known", "tool_name": r"^(get_weather|computer|github__create_issue)$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [weather, tool]} + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'delete_db' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("tool", "hosted_key"), + [ + ({"type": "get_weather", "googleSearch": {}}, "googleSearch"), + ({"type": "function", "codeExecution": {}}, "codeExecution"), + ({"type": "get_weather", "url_context": {}}, "url_context"), + ], + ids=["named_type_google_search", "function_type_code_execution", "named_type_url_context"], + ) + async def test_pre_call_judges_a_gemini_hosted_tool_key_whatever_the_tool_type( + self, tool: dict[str, object], hosted_key: str, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [ + {"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}, + {"id": "allow_functions", "tool_type": r"^function$", "decision": "allow"}, + ], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [weather, tool]} + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == f"Tool '{hosted_key}' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_leaves_a_gemini_hosted_key_that_sits_beside_a_function( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + tool = {"type": "function", "function": {"name": "get_weather"}, "googleSearch": {}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("function", "denied"), + [ + ({"type": "function", "function": {"name": "delete_db"}}, "delete_db"), + ({"name": "get_weather", "function": {"name": "delete_db"}}, "delete_db"), + ({"googleSearch": {}}, "googleSearch"), + ], + ids=["function_object", "function_object_behind_an_allowed_name", "gemini_hosted_key"], + ) + async def test_pre_call_judges_a_legacy_functions_entry_like_a_tool( + self, function: dict[str, object], denied: str, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"name": "get_weather"} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + "functions": [weather, function], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == f"Tool '{denied}' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["functions"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_judges_a_legacy_functions_entry_as_a_function_whatever_its_type( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "deny_delete", "tool_name": r"^delete_db$", "tool_type": r"^function$", "decision": "deny"}], + on_disallowed_action, + ) + weather = {"name": "get_weather"} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + "functions": [weather, {"type": "custom", "name": "delete_db"}], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'delete_db' denied by rule 'deny_delete'" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["functions"] == [weather] + + @pytest.mark.asyncio + async def test_rewrite_keeps_an_allowed_function_whose_declaration_list_is_denied(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + "rewrite", + default_action="deny", + ) + tool = { + "type": "function", + "function": {"name": "get_weather"}, + "functionDeclarations": [{"name": "delete_db"}], + } + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [{"type": "function", "function": {"name": "get_weather"}}] + + @pytest.mark.asyncio + async def test_pre_call_allows_a_custom_tool_by_its_function_name(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + "block", + default_action="deny", + ) + tool = {"type": "custom", "function": {"name": "get_weather"}} + data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + async def test_pre_call_blocks_a_tool_whose_type_is_not_a_string_by_its_name(self) -> None: + guardrail = self._pre_call( + [ + {"id": "allow_search", "tool_type": r"^web_search$", "decision": "allow"}, + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}, + ], + "block", + ) + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": ["custom"], "name": "Read", "input_schema": {"type": "object"}}], + } + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "anthropic_messages") + + assert excinfo.value.status_code == 400 + assert excinfo.value.detail["detection_message"] == "Tool 'Read' denied by rule 'deny_read'" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("malformed_tool", "tool_choice", "call_type"), + [ + ({"type": ["custom"], "name": "Read", "input_schema": {"type": "object"}}, None, "anthropic_messages"), + ({"type": "function", "name": "Read"}, {"type": {"a": 1}}, "responses"), + ], + ids=["tool_type_is_a_list", "tool_choice_type_is_an_object"], + ) + async def test_pre_call_rewrite_strips_a_denied_tool_and_leaves_malformed_fields_for_the_provider( + self, malformed_tool: dict[str, object], tool_choice: dict[str, object] | None, call_type: CallTypesLiteral + ) -> None: + guardrail = self._pre_call([{"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}], "rewrite") + bash = {"type": "function", "name": "Bash"} + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": [bash, malformed_tool], + **({"tool_choice": tool_choice} if tool_choice is not None else {}), + } + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tools"] == [bash] + assert result.get("tool_choice") == tool_choice + + @pytest.mark.asyncio + async def test_pre_call_judges_a_tool_with_a_null_type_by_its_nested_function_name(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_read", "tool_name": r"^Read$", "decision": "allow"}], "block", default_action="deny" + ) + tool = {"type": None, "function": {"name": "Read", "parameters": {"type": "object"}}} + data = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_a_removed_hosted_tool_from_allowed_tools(self) -> None: + guardrail = self._pre_call([{"id": "deny_search", "tool_type": r"^web_search$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "name": "get_weather"} + tool_choice = {"type": "allowed_tools", "mode": "required", "tools": [weather, {"type": "web_search"}]} + data = {"model": "gpt-5", "input": "hi", "tools": [weather, {"type": "web_search"}], "tool_choice": tool_choice} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [weather] + assert result["tool_choice"] == {"type": "allowed_tools", "mode": "required", "tools": [weather]} + + @pytest.mark.asyncio + @pytest.mark.parametrize("tools", [{"not": "a list"}, '[{"type": "function"}]'], ids=["object", "json_string"]) + async def test_pre_call_rewrite_leaves_a_tools_field_that_is_not_a_list_for_the_provider_to_reject( + self, tools: object + ) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + data = { + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": tools, + "tool_choice": {"type": "function", "function": {"name": "mcp__weather__lookup"}}, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == tools + assert result["tool_choice"] == "none" + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_reads_bedrock_tool_blocks_by_name_and_skips_cache_points( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"toolSpec": {"name": "get_weather", "inputSchema": {"json": {}}}} + tools = [weather, {"cachePoint": {"type": "default"}}] + data = {"model": "bedrock/claude", "messages": [{"role": "user", "content": "hi"}], "tools": tools} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == tools + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_a_denied_bedrock_tool_block_by_its_name(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + kept = [{"toolSpec": {"name": "get_weather"}}, {"cachePoint": {"type": "default"}}] + data = { + "model": "bedrock/claude", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"toolSpec": {"name": "mcp__weather__lookup"}}, *kept], + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == kept + @pytest.mark.asyncio async def test_denied_anthropic_tool_use_is_blocked(self): response = self._response({"type": "text", "text": "reading"}, self._tool_use("Read")) diff --git a/tests/unit/proxy/guardrails/test_usage_tracking.py b/tests/unit/proxy/guardrails/test_usage_tracking.py index 69ec098b840..7e92c9be4e8 100644 --- a/tests/unit/proxy/guardrails/test_usage_tracking.py +++ b/tests/unit/proxy/guardrails/test_usage_tracking.py @@ -654,3 +654,15 @@ async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollu assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3 assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids) + + +@pytest.mark.asyncio +async def test_disabled_log_deliveries_do_not_enable_guardrail_analytics(): + prisma = _prisma() + row = _payload("lean", usage={"topicPolicyUnits": 2}) + row["accounting_payload"] = {"request_id": "lean"} + row["accounting_log_disabled"] = True + await process_spend_logs_guardrail_usage(prisma, [row], pending=PendingRollups()) + prisma.db.litellm_dailyguardrailmetrics.upsert.assert_not_awaited() + prisma.db.litellm_spendlogguardrailindex.create_many.assert_not_awaited() + prisma.db.litellm_dailyguardrailusageunits.upsert.assert_not_awaited() diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 45cc7eb8922..966e94fad80 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -8,6 +8,9 @@ Tests that session-scoped budget tracking works correctly: - Requests without agent_id pass through """ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final from unittest.mock import patch import logging @@ -22,6 +25,7 @@ from litellm.proxy.hooks.max_budget_per_session_limiter import ( ) from litellm.proxy.utils import InternalUsageCache from litellm.types.agents import AgentResponse +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG def _make_mock_agent(max_budget_per_session: float) -> AgentResponse: @@ -145,6 +149,255 @@ async def test_budget_per_session_independent_sessions(): assert result is None +@pytest.mark.asyncio +async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): + from openai.types import Batch + + from litellm.types.videos.main import VideoObject + + handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + video = VideoObject(id="job-1", object="video", status="completed") + batch = Batch( + id="job-1", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + input_file_id="file-1", object="batch", status="completed", output_file_id="file-out", + ) + polls = (("avideo_status", video, 0.4), ("aretrieve_batch", batch, 0.1), ("acompletion", None, 0.01)) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + for call_type, response_obj, cost in polls * 3: + await handler.async_log_success_event( + { + "call_type": call_type, + "response_cost": cost, + POLLED_VIDEO_STATUS_KWARG: "completed", + "litellm_params": {"metadata": {"session_id": "session-jobs", "agent_id": "agent-budget-123"}}, + }, + response_obj, + None, + None, + ) + + spend = await handler._get_current_spend(handler._make_cache_key("session-jobs")) + assert spend == pytest.approx(0.4 + 0.1 + 3 * 0.01) + + +@pytest.mark.asyncio +async def test_a_batch_retrieve_before_its_cost_is_final_charges_the_session_nothing(): + from openai.types import Batch + + handler: Final = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + + def polled(status: str, output_file_id: str | None = None) -> Batch: + return Batch( + id="batch-1", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + input_file_id="file-1", object="batch", status=status, output_file_id=output_file_id, + ) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + for response_obj, cost in ( + (polled("in_progress"), 0.01), + (polled("completed"), 0.01), + (polled("completed", "file-out"), 0.25), + (polled("completed", "file-out"), 0.25), + ): + await handler.async_log_success_event( + { + "call_type": "aretrieve_batch", + "response_cost": cost, + "litellm_params": {"metadata": {"session_id": "session-batch", "agent_id": "agent-budget-123"}}, + }, + response_obj, + None, + None, + ) + + spend: Final = await handler._get_current_spend(handler._make_cache_key("session-batch")) + assert spend == pytest.approx(0.25) + + +def _video_poll( + video_cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed" +) -> dict[str, object]: + return { + "call_type": "avideo_status", + "response_cost": video_cost, + POLLED_VIDEO_STATUS_KWARG: status, + "standard_logging_object": { + "response_cost": video_cost + guardrail_cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], + }, + "litellm_params": {"metadata": {"session_id": session_id, "agent_id": "agent-budget-123"}}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("video_cost", "guardrail_cost"), + [(0.40, 0.01), (0.05, 0.10)], + ids=["video_costs_more_than_its_guardrail", "video_costs_no_more_than_its_guardrail"], +) +async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video( + video_cost: float, guardrail_cost: float +): + from litellm.types.videos.main import VideoObject + + handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + completed_polls: Final = 3 + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + await handler.async_log_success_event( + _video_poll(0.0, guardrail_cost=guardrail_cost, status="in_progress"), + VideoObject(id="job-1", object="video", status="in_progress"), + None, + None, + ) + for _ in range(completed_polls): + await handler.async_log_success_event( + _video_poll(video_cost, guardrail_cost=guardrail_cost), + VideoObject(id="job-1", object="video", status="completed"), + None, + None, + ) + + spend = await handler._get_current_spend(handler._make_cache_key("session-jobs")) + assert spend == pytest.approx(video_cost + (1 + completed_polls) * guardrail_cost) + + +class _ClockedRedis: + """The Redis calls the session limiter makes, on a clock the test moves.""" + + def __init__(self) -> None: + self.now = 0.0 + self.entries: Mapping[str, tuple[float, float | None]] = MappingProxyType({}) + self.scripts_fail = False + + def _live(self, key: str) -> tuple[float, float | None] | None: + entry = self.entries.get(key) + if entry is None or (entry[1] is not None and entry[1] <= self.now): + return None + return entry + + def _store(self, key: str, value: float, expires_at: float | None) -> float: + self.entries = MappingProxyType({**self.entries, key: (value, expires_at)}) + return value + + async def async_get_cache(self, key: str, **kwargs: object) -> float | None: + entry = self._live(key) + return None if entry is None else entry[0] + + async def async_increment( + self, key: str, value: float, ttl: int | None = None, parent_otel_span: object = None, refresh_ttl: bool = False + ) -> float: + live = self._live(key) + kept_expiry = None if live is None else live[1] + expires_at = kept_expiry if ttl is None or (live is not None and not refresh_ttl) else self.now + ttl + return self._store(key, value + (0.0 if live is None else live[0]), expires_at) + + def _increment_expiring_on_first_write(self, key: str, amount: float, ttl: float) -> float: + live = self._live(key) + expires_at = self.now + ttl if live is None else live[1] + return self._store(key, amount + (0.0 if live is None else live[0]), expires_at) + + def async_register_script(self, script: str): + async def run(keys: list[str], args: list[object]) -> float: + if self.scripts_fail: + raise ConnectionError("Redis dropped the script call") + if len(keys) == 1: + return self._increment_expiring_on_first_write(keys[0], float(str(args[0])), float(str(args[1]))) + marker = self._live(keys[1]) + polls = self._store(keys[1], 1.0 + (0.0 if marker is None else marker[0]), self.now + float(str(args[2]))) + amount = float(str(args[0] if polls == 1 else args[1])) + if amount > 0: + self._increment_expiring_on_first_write(keys[0], amount, float(str(args[2]))) + return amount + + return run + + +@pytest.mark.asyncio +async def test_a_finished_video_polled_without_pause_is_not_charged_again_once_the_session_window_rolls_over(): + from litellm.types.videos.main import VideoObject + + redis = _ClockedRedis() + handler = PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis)) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + video = VideoObject(id="job-1", object="video", status="completed") + chat = { + "call_type": "acompletion", + "response_cost": 0.01, + "litellm_params": {"metadata": {"session_id": "session-jobs", "agent_id": "agent-budget-123"}}, + } + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + for _ in range(3): + await handler.async_log_success_event(_video_poll(0.4), video, None, None) + redis.now += 0.6 * handler.ttl + await handler.async_log_success_event(chat, None, None, None) + + assert await handler._get_current_spend(handler._make_cache_key("session-jobs")) == pytest.approx(0.01) + + +@pytest.mark.asyncio +async def test_a_finished_video_whose_charge_missed_redis_is_charged_on_the_next_poll_once_redis_answers(): + from litellm.types.videos.main import VideoObject + + redis: Final = _ClockedRedis() + handler: Final = PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis)) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + video: Final = VideoObject(id="job-1", object="video", status="completed") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + redis.scripts_fail = True + await handler.async_log_success_event(_video_poll(0.4), video, None, None) + redis.scripts_fail = False + for _ in range(2): + await handler.async_log_success_event(_video_poll(0.4), video, None, None) + + assert await handler._get_current_spend(handler._make_cache_key("session-jobs")) == pytest.approx(0.4) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), True]) +async def test_a_cost_that_is_not_a_finite_number_leaves_the_session_spend_alone(response_cost): + handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + cache_key = handler._make_cache_key("session-unpriced") + await handler._increment_spend(cache_key, 0.25) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + await handler.async_log_success_event( + { + "call_type": "acompletion", + "response_cost": response_cost, + "litellm_params": {"metadata": {"session_id": "session-unpriced", "agent_id": "agent-budget-123"}}, + }, + None, + None, + None, + ) + + assert await handler._get_current_spend(cache_key) == 0.25 + + @pytest.mark.asyncio async def test_no_agent_id_passes(): """ diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index f5a2ccf7b76..f1fd02981d6 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -27,6 +27,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.router import Router from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import LiteLLMBatch, Usage +from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG KEY_HASH: Final = "key-hash-batch" USER_ID: Final = "user-batch" @@ -47,6 +49,7 @@ def _batch(batch_id: str, status: str) -> LiteLLMBatch: input_file_id="file-batch", object="batch", status=status, + output_file_id="file-output", usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38), ) @@ -57,12 +60,14 @@ def _event( team_member_model_max_budget: Mapping[str, object] | None = None, user_id: str = USER_ID, team_id: str = TEAM_ID, + guardrail_information: tuple[Mapping[str, object], ...] = (), ) -> dict[str, object]: return { "call_type": call_type, "standard_logging_object": { "call_type": call_type, "response_cost": response_cost, + "guardrail_information": guardrail_information, "model": "openai/gpt-5.4-mini", "model_group": MODEL_GROUP, "metadata": { @@ -326,6 +331,57 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again( assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + +async def _poll_video( + limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, video_id: str, status: str, response_cost: float +) -> None: + await _log_success( + limiter, + {**_event("avideo_status", response_cost), POLLED_VIDEO_STATUS_KWARG: status}, + response_obj=VideoObject(id=video_id, object="video", status=status), + ) + + +@pytest.mark.asyncio +async def test_polls_of_a_finished_video_charge_each_per_model_budget_once(): + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + + await _poll_video(limiter, "video_first", "in_progress", response_cost=0) + for _ in range(3): + await _poll_video(limiter, "video_first", "completed", response_cost=BATCH_COST) + await _poll_video(limiter, "video_second", "completed", response_cost=BATCH_COST) + await _poll(limiter, _batch("video_first", "completed"), response_cost=BATCH_COST) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(3 * BATCH_COST) + assert await _spend(limiter, USER_SPEND_KEY) == pytest.approx(3 * BATCH_COST) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "polled", "guardrail_charges"), + ( + ("avideo_status", lambda status: VideoObject(id="job_first", object="video", status=status), 4), + ("aretrieve_batch", lambda status: _batch("job_first", status), 1), + ), +) +async def test_only_the_finished_poll_pays_the_job_and_only_a_video_repoll_pays_its_guardrail( + call_type, polled, guardrail_charges +): + """A running video poll pays its guardrail like any request, while a batch retrieve before its cost is final + charges nothing, as key spend defers it, so the retrieve that settles the batch pays the only guardrail.""" + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + guardrail_cost: Final = 0.01 + priced: Final = ({"guardrail_name": "priced", "guardrail_cost": guardrail_cost},) + + def poll(status: str, response_cost: float) -> dict[str, object]: + return {**_event(call_type, response_cost, guardrail_information=priced), POLLED_VIDEO_STATUS_KWARG: status} + + await _log_success(limiter, poll("in_progress", guardrail_cost), polled("in_progress")) + for _ in range(3): + await _log_success(limiter, poll("completed", BATCH_COST + guardrail_cost), polled("completed")) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST + guardrail_charges * guardrail_cost) + OPUS_GROUP: Final = "claude-opus" OPUS_GROUP_MEMBERS: Final = MappingProxyType( {"claude-opus-4-8": frozenset({OPUS_GROUP}), "claude-opus-5-5": frozenset({OPUS_GROUP})} diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index a6edd3db944..417ea5a9650 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -33,7 +33,7 @@ def test_managed_vector_stores_hook_registered(): assert hook_cls.__name__ == "_PROXY_LiteLLMManagedVectorStores" -def test_isolation_module_does_not_pull_in_proxy_utils(): +def test_isolation_module_does_not_pull_in_proxy_utils(monkeypatch): """Layering guard: litellm.llms.* must not transitively import litellm.proxy.utils, which would reintroduce the import cycle.""" import importlib @@ -44,7 +44,7 @@ def test_isolation_module_does_not_pull_in_proxy_utils(): "litellm.proxy.management_endpoints.common_utils", "litellm.llms.base_llm.managed_resources.isolation", ]: - sys.modules.pop(mod, None) + monkeypatch.delitem(sys.modules, mod, raising=False) importlib.import_module("litellm.llms.base_llm.managed_resources.isolation") assert "litellm.proxy.utils" not in sys.modules diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 21b5bd300c1..c6366a40504 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -4,7 +4,7 @@ import logging from collections.abc import Mapping from datetime import datetime, timedelta, timezone from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch @@ -24,6 +24,7 @@ from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( ProxyDBLogger, + _charge_request_cost, _get_budget_reservation_from_metadata, _should_track_cost_callback, _update_database_and_spend_counters, @@ -42,6 +43,7 @@ from litellm.proxy.utils import ( ) from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.videos.main import VideoObject @pytest.mark.asyncio @@ -1595,14 +1597,16 @@ def _retrieved_batch(status: str, output_file_id: str | None): @pytest.mark.asyncio @pytest.mark.parametrize( - ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"), + ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged", "redacted"), [ - ("aretrieve_batch", "in_progress", None, True, False, False), - ("aretrieve_batch", "completed", None, True, False, False), - ("aretrieve_batch", "completed", "file-out", False, True, False), - ("aretrieve_batch", "completed", "file-out", True, True, True), - ("aretrieve_batch", "failed", None, True, True, True), - ("acreate_batch", "validating", None, True, True, True), + ("aretrieve_batch", "in_progress", None, True, False, False, False), + ("aretrieve_batch", "completed", None, True, False, False, False), + ("aretrieve_batch", "completed", "file-out", False, True, False, False), + ("aretrieve_batch", "completed", "file-out", True, True, True, False), + ("aretrieve_batch", "failed", None, True, True, True, False), + ("acreate_batch", "validating", None, True, True, True, False), + ("aretrieve_batch", "in_progress", None, True, False, False, True), + ("aretrieve_batch", "completed", "file-out", True, True, True, True), ], ids=[ "retrieve_before_final", @@ -1611,20 +1615,31 @@ def _retrieved_batch(status: str, output_file_id: str | None): "retrieve_first_final", "retrieve_failed_batch", "create_before_final", + "redacted_retrieve_before_final", + "redacted_retrieve_first_final", ], ) async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate - call_type, status, output_file_id, row_claimed, spend_written, charged + call_type, status, output_file_id, row_claimed, spend_written, charged, redacted ): """ A poll before the batch is final used to pin its shared spend row at $0, and every completed retrieve after the first charged the key again (LIT-7048). Only retrieves are gated, since creating a batch is its own billable request, and a retrieve that - charges nothing hands its budget reservation back instead. + charges nothing hands its budget reservation back instead. Message redaction replaces + the logged batch, so the gate reads the batch logging stamped before redacting it. """ + from litellm.litellm_core_utils.batch_utils import POLLED_BATCH_COST_IS_FINAL_KWARG, POLLED_BATCH_ID_KWARG + logger = ProxyDBLogger() budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} - kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) + retrieved: Final = _retrieved_batch(status, output_file_id) + stamped: Final = ( + {POLLED_BATCH_ID_KWARG: retrieved.id, POLLED_BATCH_COST_IS_FINAL_KWARG: status != "in_progress"} + if redacted + else {} + ) + kwargs = {**_batch_retrieve_kwargs(call_type, reservation=budget_reservation), **stamped} with ( patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam @@ -1646,7 +1661,7 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # await logger._PROXY_track_cost_callback( kwargs=kwargs, - completion_response=_retrieved_batch(status, output_file_id), + completion_response={"text": "redacted-by-litellm"} if redacted else retrieved, start_time=datetime.now(), end_time=datetime.now(), ) @@ -3376,18 +3391,8 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski assert event.response_cost == 0.0125 -@pytest.mark.asyncio -async def test_async_log_success_event_keeps_batch_retrieves_in_process(): - producer = SpendEventProducer( - address=UnixAddress(path="/nonexistent/spend.sock"), - on_unavailable="drop", - buffer_size=10, - connect_timeout=1.0, - fallback=_no_fallback, - ) - logger = ProxyDBLogger(producer) - kwargs = {**_offload_kwargs(), "call_type": CallTypes.aretrieve_batch.value} - completed_batch = LiteLLMBatch( +def _completed_batch() -> LiteLLMBatch: + return LiteLLMBatch( id="batch_abc", completion_window="24h", created_at=1, @@ -3398,6 +3403,29 @@ async def test_async_log_success_event_keeps_batch_retrieves_in_process(): status="completed", ) + +def _completed_video() -> VideoObject: + return VideoObject(id="video_abc", object="video", status="completed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "polled_job"), + [(CallTypes.aretrieve_batch.value, _completed_batch), ("avideo_status", _completed_video)], +) +async def test_async_log_success_event_keeps_job_polls_in_process(call_type, polled_job): + """The spend event leaves out the polled video's status and id, which decide whether the poll charges the job.""" + producer = SpendEventProducer( + address=UnixAddress(path="/nonexistent/spend.sock"), + on_unavailable="drop", + buffer_size=10, + connect_timeout=1.0, + fallback=_no_fallback, + ) + logger = ProxyDBLogger(producer) + kwargs = {**_offload_kwargs(), "call_type": call_type} + completed_job = polled_job() + with ( patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam "litellm.proxy.proxy_server.proxy_logging_obj" @@ -3411,7 +3439,7 @@ async def test_async_log_success_event_keeps_batch_retrieves_in_process(): ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() - await logger.async_log_success_event(kwargs, completed_batch, datetime.now(), datetime.now()) + await logger.async_log_success_event(kwargs, completed_job, datetime.now(), datetime.now()) mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() assert producer.stats().queued == 0 @@ -3983,138 +4011,48 @@ async def test_async_post_call_failure_hook_disable_pre_routing_error_logs( assert len(writer.rows) == rows_written -@pytest.mark.asyncio -async def test_failure_capture_overload_does_not_skip_budget_release( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _write +class _RecordedCharges: + def __init__(self, outcomes: tuple[bool, ...]) -> None: + self.outcomes = outcomes + self.calls: tuple[tuple[object, float, dict[str, object] | None], ...] = () - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - meter: Final = runtime.meter - assert meter is not None - receipt: Final = meter.submit((_write("occupied-one"), _write("occupied-two"))) - auth: Final = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False}) - request: Final = { - "call_type": "acompletion", - "litellm_call_id": "failed", - "model": "capture-model", - "start_time": _CALL_START, - "litellm_params": {}, + async def __call__( + self, *, kwargs: dict, response_cost: float, budget_reservation: dict[str, object] | None + ) -> bool: + self.calls = (*self.calls, (kwargs.get("response_cost"), response_cost, budget_reservation)) + return self.outcomes[len(self.calls) - 1] + + +_GUARDED_POLL: Final = MappingProxyType( + { + "call_type": "avideo_status", + "response_cost": 0.43, + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, } - try: - await database.entered.wait() - await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), auth) - assert auth.budget_reservation is not None and auth.budget_reservation["finalized"] is True - assert "Usage capture failed on the request failure path" in caplog.text - assert "_accounting_capture" not in request - finally: - database.release.set() - await runtime.stop(1) - await receipt - assert database.committed == ("occupied-one", "occupied-two") +) @pytest.mark.asyncio -@pytest.mark.parametrize("failure", [False, True]) -@pytest.mark.parametrize("complete", [False, True]) -@pytest.mark.parametrize("deferred", [False, True]) -async def test_capture_and_legacy_share_resolved_key_attribution( - failure: bool, complete: bool, deferred: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - from pydantic import TypeAdapter +@pytest.mark.parametrize( + ("outcomes", "poll", "charged", "second_call"), + [ + ((True,), _GUARDED_POLL, 0.43, None), + ((False, True), _GUARDED_POLL, 0.01, (0.01, 0.01, None)), + ((False, False), _GUARDED_POLL, None, (0.01, 0.01, None)), + ((False,), {"response_cost": 0.43}, None, None), + ((False,), {**_GUARDED_POLL, "call_type": "aretrieve_batch"}, None, None), + ], + ids=[ + "first_poll_pays_the_job", + "repeat_poll_pays_its_guardrail", + "guardrail_charge_not_written", + "repeat_poll_free", + "repeat_batch_poll_free_whatever_its_guardrails_cost", + ], +) +async def test_a_poll_of_an_already_charged_job_still_charges_its_own_guardrails(outcomes, poll, charged, second_call): + charge_spend: Final = _RecordedCharges(outcomes) + reservation: Final = {"reserved": 0.43} - from litellm.accounting import DeliveryTarget, DimensionName - from litellm.proxy import proxy_server - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage - - key: Final = UserAPIKeyAuth( - api_key="a" * 64, - user_id="owner", - team_id="team", - org_id="org", - project_id="project", - key_alias="key", - team_alias="team-alias", - request_route="/v1/chat/completions", - ) - cache: Final = UserApiKeyCache() - await cache.async_set_cache(key="a" * 64, value=key) - monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - monkeypatch.setattr(proxy_server, "prisma_client", object()) - writer: Final = _RecordingSpendWriter() - monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(db_spend_update_writer=writer)) - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - rejected: Final = failure and deferred and not complete - owner: Final = "original-owner" if deferred and not rejected else "owner" - metadata: Final = { - "user_api_key": "b" * 64 if rejected else "a" * 64, - **( - { - "user_api_key_user_id": "forged-owner" if rejected else owner, - "user_api_key_team_id": "forged-team" if rejected else "team", - "user_api_key_org_id": "forged-org" if rejected else "org", - "user_api_key_project_id": None if deferred and not complete else "project", - "user_api_key_alias": "key", - "user_api_key_team_alias": "team-alias", - } - if complete or deferred - else {} - ), - } - kwargs: Final = { - "call_type": "aretrieve_batch" if deferred else "call_mcp_tool", - "litellm_call_id": "attributed", - "model": "MCP: echo", - "response_cost": 0.25, - "start_time": _CALL_START, - "litellm_params": {"metadata": metadata}, - } - if rejected: - from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY - - kwargs[REJECTED_BEFORE_REQUEST_SETUP_KEY] = REJECTED_BEFORE_REQUEST_SETUP - logger: Final = ProxyDBLogger(spend_writer=lambda: cast(DBSpendUpdateWriter, writer)) - try: - with track_batch_accounting("attributed" if deferred and not rejected else "other"): - if failure: - await logger.async_post_call_failure_hook( - kwargs, - HTTPException(401, "model not allowed"), - key if complete else UserAPIKeyAuth(api_key="a" * 64, request_route="/v1/chat/completions"), - ) - else: - await logger.async_log_success_event(kwargs, {"id": "tool"}, _CALL_START, _CALL_START) - finally: - await runtime.stop(1) - assert len(writer.rows) == 1 - assert len(database.rows) == 2 - dimensions: Final = {item.name: item.value for item in database.rows[-1].event.dimensions} - if rejected: - assert dimensions.get(DimensionName.KEY) == key.api_key - assert { - name: dimensions.get(name) - for name in (DimensionName.USER, DimensionName.TEAM, DimensionName.ORGANIZATION, DimensionName.PROJECT) - } == { - DimensionName.USER: owner, - DimensionName.TEAM: "team", - DimensionName.ORGANIZATION: "org", - DimensionName.PROJECT: None if deferred and not complete and not rejected else "project", - } - assert set(database.rows[-1].targets) >= {DeliveryTarget.DAILY_TEAM_SPEND, DeliveryTarget.DAILY_ORGANIZATION_SPEND} - legacy: Final = TypeAdapter(dict[str, object]).validate_python(writer.rows[0]["kwargs"]) - payload: Final = get_logging_payload(legacy, {"id": "tool"}, _CALL_START, _CALL_START) - assert (payload["user"], payload["team_id"], payload["organization_id"]) == (owner, "team", "org") + assert await _charge_request_cost(charge_spend, dict(poll), 0.43, reservation) == charged + assert charge_spend.calls == ((0.43, 0.43, reservation), *(() if second_call is None else (second_call,))) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4a76aa77f14..11b07472eaa 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7241,6 +7241,178 @@ class TestRouterModelRelayUpstreamContract: assert json.loads(result.body) == upstream_body assert result.headers["x-ms-request-id"] == "req-1" + @pytest.mark.asyncio + async def test_azure_relay_answers_401_when_entra_rejects_the_deployment_credential(self, monkeypatch): + from litellm.llms.azure.common_utils import EntraRejectionError + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + + class EntraRejectingRouter: + async def allm_passthrough_route(self, **kwargs): + try: + raise EntraRejectionError(status_code=401, message=reason) + except EntraRejectionError as rejection: + raise BaseLLMException(status_code=401, message=str(rejection)) from rejection + + self._install_router(monkeypatch, EntraRejectingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(ProxyException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert (excinfo.value.code, excinfo.value.message) == ("401", reason) + + @pytest.mark.asyncio + async def test_azure_relay_entra_rejection_hides_an_internal_issuer_host(self, monkeypatch): + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = ( + "Failed to get Azure AD token: AADSTS700211: No matching federated identity record found for " + "presented assertion issuer 'https://oidc.build.internal:8443/realms/ci'." + ) + + class EntraRejectingRouter: + async def allm_passthrough_route(self, **kwargs): + raise EntraRejectionError(status_code=401, message=reason) + + self._install_router(monkeypatch, EntraRejectingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(ProxyException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert excinfo.value.code == "401" + assert "AADSTS700211" in excinfo.value.message + assert "oidc.build.internal" not in excinfo.value.message + + @pytest.mark.asyncio + async def test_azure_relay_keeps_a_failure_entra_did_not_cause(self, monkeypatch): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + upstream_failure: Final = BaseLLMException(status_code=500, message="Connection reset by Azure") + + class FailingRouter: + async def allm_passthrough_route(self, **kwargs): + raise upstream_failure + + self._install_router(monkeypatch, FailingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(BaseLLMException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert excinfo.value is upstream_failure + + @pytest.mark.asyncio + async def test_azure_relay_keeps_entra_throttling_the_oidc_token_exchange_as_a_429(self, monkeypatch): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "throttled-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + }, + } + ], + num_retries=0, + ) + self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False}) + + with respx.mock: + respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 429, + headers={"Retry-After": "7"}, + json={"error": "temporarily_unavailable", "error_codes": [90055]}, + ) + ) + deployment_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert (response.status_code, response.headers.get("retry-after")) == (429, "7") + assert not deployment_route.called + + + @pytest.mark.asyncio + async def test_azure_relay_keeps_a_later_deployment_failure_after_an_entra_rejection( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from azure.core.exceptions import ClientAuthenticationError + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + "order": 1, + }, + }, + { + "model_name": "gpt-5", + "litellm_params": { + "model": "azure/gpt-5", + "api_base": "https://litellm-qa-unreachable.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "configured-key", + "order": 2, + }, + }, + ], + num_retries=1, + ) + self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False}) + + with respx.mock: + unreachable_route: Final = respx.post(url__startswith="https://litellm-qa-unreachable.openai.azure.com").mock( + side_effect=httpx.ConnectError("Connection refused") + ) + with pytest.raises(BaseLLMException) as excinfo: + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert unreachable_route.called + assert excinfo.value.status_code == 500 + NIM_INFER_BODY = { "input": [ diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 29413d2b189..c17a9ba9371 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -529,6 +529,7 @@ def _decisions_provider_entry(provider: str) -> ProviderCreateInfo: ("Nace", "Nace", LlmProviders.NACE), ("PARALLEL_AI", "Parallel AI", LlmProviders.PARALLEL_AI), ("EU_JEV", "eu/jev", LlmProviders.EU_JEV), + ("EMPIRIOLABS", "EmpirioLabs AI", LlmProviders.EMPIRIOLABS), ), ) def test_decisions_provider_fields(provider_key: str, display_name: str, provider: LlmProviders) -> None: diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index ef43fe56146..b7d0d9e62cd 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2556,6 +2556,66 @@ class TestResponsesWSFirstFrameModelAuth: "response": {"input": "hi [guarded]", "tools": tools[:1]}, } + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [True, False]) + async def test_frame_guard_forwards_the_tool_choice_a_guardrail_rewrote(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {**data, "tools": [], "tool_choice": "none"} + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + body = { + "input": "hi", + "tools": [{"type": "function", "name": "a"}], + "tool_choice": {"type": "function", "name": "a"}, + } + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + guarded = json.loads((await guard(json.dumps(frame))).frame) + + sent = guarded["response"] if nested else guarded + assert (sent["tools"], sent["tool_choice"]) == ([], "none") + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [True, False]) + async def test_frame_guard_forwards_the_extra_body_tools_a_rewrite_guardrail_stripped(self, nested: bool): + from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = ToolPermissionGuardrail( + guardrail_name="ws-rewrite", + rules=[{"id": "deny_shell", "tool_name": r"^shell$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=MagicMock(), data=data, call_type=call_type + ) + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + weather = {"type": "function", "name": "get_weather"} + body = {"input": "hi", "extra_body": {"tools": [weather, {"type": "function", "name": "shell"}]}} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + guarded = json.loads((await guard(json.dumps(frame))).frame) + + sent = guarded["response"] if nested else guarded + assert sent["extra_body"] == {"tools": [weather]} + @pytest.mark.asyncio async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard @@ -2991,7 +3051,7 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + @pytest.mark.parametrize("field", ["input", "instructions", "tools", "tool_choice", "extra_body"]) async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str): from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a87ac7fa7c0..9267a951459 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -588,3 +588,10 @@ async def test_group_reserves_on_its_current_deployments_after_another_vendors_d by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES) assert openai_only != by_own_name assert (await _input_tokens(router, "prod"))["prod"] == by_own_name + + +@pytest.mark.parametrize("window", [{}, {"budget_duration": "1d"}, {"reset_at": "2026-11-03T00:00:00Z"}]) +def test_fixed_budget_window_requires_duration_and_reset_boundary(window: object) -> None: + from litellm.proxy.spend_tracking.budget_reservation import fixed_budget_window + + assert fixed_budget_window(window) is None diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index a426d8f4e14..ef70a193280 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -92,12 +92,6 @@ def response(cost: float = 0.5) -> ModelResponse: async def test_header_and_logging_share_one_allowance_allocation( logging_obj: Logging, used: int, expected: float ) -> None: - import json - - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings - counter: Final = Counter(used) result: Final = response() costs: Final = await asyncio.gather(*(finalize_response_cost(logging_obj, result, counter, NOW) for _ in range(3))) @@ -116,24 +110,6 @@ async def test_header_and_logging_share_one_allowance_allocation( assert payload["cost_finalized"] is True assert payload["cost_breakdown"] == costs[0].logging_breakdown() - committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - committed.put_nowait(write) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings()) - assert logging_obj.prepared_cost is not None - prepared_usage: Final = logging_obj.prepared_cost.usage.model_dump() - result.usage = Usage(prompt_tokens=999) - capture_logging_usage(meter, logging_obj, result, NOW, NOW) - await meter.close() - committed.get_nowait() - event: Final = json.loads(committed.get_nowait().envelope) - assert event["finalized_cost"]["response_cost"] == pytest.approx(costs[0].amount) - assert event["allocation"]["discount"] == pytest.approx(costs[0].allowance_discount) - assert event["usage"] == prepared_usage - assert counter.calls == 1 @pytest.mark.asyncio diff --git a/tests/unit/proxy/spend_tracking/test_spend_event.py b/tests/unit/proxy/spend_tracking/test_spend_event.py index ff449235582..e1912d36d91 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_event.py +++ b/tests/unit/proxy/spend_tracking/test_spend_event.py @@ -194,9 +194,11 @@ def test_undecodable_line_is_an_error_value(): assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError) -def test_batch_retrieves_stay_in_process(): - assert is_offloadable_success(_response()) is True - assert is_offloadable_success(None) is True +def test_batch_retrieves_and_video_status_polls_stay_in_process(): + assert is_offloadable_success(_response(), "acompletion") is True + assert is_offloadable_success(None, None) is True + assert is_offloadable_success(None, "avideo_generation") is True + assert is_offloadable_success(None, "avideo_status") is False assert ( is_offloadable_success( LiteLLMBatch( @@ -207,7 +209,8 @@ def test_batch_retrieves_stay_in_process(): input_file_id="f", object="batch", status="completed", - ) + ), + "aretrieve_batch", ) is False ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 4e113ce03a4..45ded91b61e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -47,7 +47,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, is_master_key, - _redact_logged_api_key, + redact_logged_api_key, _redact_prompt_leaks_in_error_string, sanitize_error_information_for_spend_logs, _sanitize_guardrail_information_for_spend_logs, @@ -4038,7 +4038,7 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] -# ── _redact_logged_api_key unit tests ────────────────────────────────────── +# ── redact_logged_api_key unit tests ────────────────────────────────────── @pytest.mark.parametrize( @@ -4067,16 +4067,16 @@ def test_sanitize_error_information_persists_no_raw_model_for_an_unknown_model_r def test_redact_logged_api_key_none_returns_none(): - assert _redact_logged_api_key(None) is None + assert redact_logged_api_key(None) is None def test_redact_logged_api_key_empty_string_returns_none(): - assert _redact_logged_api_key("") is None + assert redact_logged_api_key("") is None def test_redact_logged_api_key_sk_key_is_hashed(): raw = "sk-9876secret" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result == hash_token(raw) assert result is not None assert not result.startswith("sk-") @@ -4085,15 +4085,15 @@ def test_redact_logged_api_key_sk_key_is_hashed(): def test_redact_logged_api_key_bearer_sk_equals_sk_hash(): raw = "sk-9876secret" - result_plain = _redact_logged_api_key(raw) - result_bearer = _redact_logged_api_key(f"Bearer {raw}") + result_plain = redact_logged_api_key(raw) + result_bearer = redact_logged_api_key(f"Bearer {raw}") assert result_bearer == result_plain def test_redact_logged_api_key_bearer_case_insensitive(): raw = "sk-9876secret" - result_lower = _redact_logged_api_key(f"bearer {raw}") - result_upper = _redact_logged_api_key(f"BEARER {raw}") + result_lower = redact_logged_api_key(f"bearer {raw}") + result_upper = redact_logged_api_key(f"BEARER {raw}") expected = hash_token(raw) assert result_lower == expected assert result_upper == expected @@ -4101,7 +4101,7 @@ def test_redact_logged_api_key_bearer_case_insensitive(): def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): raw = "anthropic-raw-key-xyz" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert result != raw assert len(result) == 64 @@ -4111,7 +4111,7 @@ def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(already_hashed, already_redacted=True) + result = redact_logged_api_key(already_hashed, already_redacted=True) assert result == already_hashed assert hash_token(already_hashed) != result # no double-hash @@ -4119,7 +4119,7 @@ def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): def test_redact_logged_api_key_sha256_without_flag_is_hashed(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(already_hashed) + result = redact_logged_api_key(already_hashed) assert result is not None assert result != already_hashed assert len(result) == 64 @@ -4129,7 +4129,7 @@ def test_redact_logged_api_key_sha256_without_flag_is_hashed(): def test_redact_logged_api_key_long_opaque_token_is_hashed(): raw = "x1" * 450 assert len(raw) == 900 - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert result != raw assert raw not in result @@ -4139,27 +4139,27 @@ def test_redact_logged_api_key_long_opaque_token_is_hashed(): def test_redact_logged_api_key_hashed_jwt_passes_through(): jwt_hash = "hashed-jwt-" + "a" * 64 - result = _redact_logged_api_key(jwt_hash, already_redacted=True) + result = redact_logged_api_key(jwt_hash, already_redacted=True) assert result == jwt_hash def test_redact_logged_api_key_hashed_jwt_shape_without_provenance_is_hashed(): lookalike = "hashed-jwt-" + "a" * 64 - result = _redact_logged_api_key(lookalike) + result = redact_logged_api_key(lookalike) assert result == hash_token(lookalike) assert result != lookalike def test_redact_logged_api_key_hashed_jwt_trailing_newline_is_hashed(): trailing = "hashed-jwt-" + "a" * 64 + "\n" - result = _redact_logged_api_key(trailing, already_redacted=True) + result = redact_logged_api_key(trailing, already_redacted=True) assert result == hash_token(trailing) assert result != trailing def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): short_jwt = "hashed-jwt-tooshort" - result = _redact_logged_api_key(short_jwt) + result = redact_logged_api_key(short_jwt) assert result is not None assert result != short_jwt assert len(result) == 64 @@ -4169,14 +4169,14 @@ def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): def test_redact_logged_api_key_master_key_alias_passes_through(): from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True) + result = redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True) assert result == LITELLM_PROXY_MASTER_KEY_ALIAS def test_redact_logged_api_key_master_key_alias_without_provenance_is_hashed(): from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS) + result = redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS) assert result == hash_token(LITELLM_PROXY_MASTER_KEY_ALIAS) assert result != LITELLM_PROXY_MASTER_KEY_ALIAS @@ -4282,9 +4282,9 @@ def test_get_logging_payload_reads_used_client_oauth_token_from_the_bucket_the_p def test_redact_logged_api_key_bearer_only_returns_none(): # "bearer " with nothing after stripping is equivalent to no key - assert _redact_logged_api_key("bearer ") is None - assert _redact_logged_api_key("Bearer ") is None - assert _redact_logged_api_key("BEARER ") is None + assert redact_logged_api_key("bearer ") is None + assert redact_logged_api_key("Bearer ") is None + assert redact_logged_api_key("BEARER ") is None # ── _get_spend_logs_metadata key-hash invariant tests ───────────────────── @@ -4645,7 +4645,7 @@ class TestSpendLogKeyRedaction: def test_bearer_prefixed_sk_key_is_hashed(self): raw = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert not result.startswith("Bearer") assert not result.startswith("sk-") @@ -4653,14 +4653,14 @@ class TestSpendLogKeyRedaction: def test_bare_sk_key_is_hashed(self): raw = "sk-WLi4iRn4JmbVlTaYw12IOA" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert not result.startswith("sk-") assert len(result) == 64 def test_bearer_lowercase_is_handled(self): raw = "bearer sk-WLi4iRn4JmbVlTaYw12IOA" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert not result.startswith("bearer") assert not result.startswith("sk-") @@ -4668,11 +4668,11 @@ class TestSpendLogKeyRedaction: def test_already_hashed_key_unchanged(self): hashed = "bcfe8173f5447f10be0e7fb37aaa8b97829d5c9e0498232152f9d123456789ab" - assert _redact_logged_api_key(hashed, already_redacted=True) == hashed + assert redact_logged_api_key(hashed, already_redacted=True) == hashed def test_bearer_prefixed_non_sk_key_is_hashed(self): raw = "Bearer some-other-token-format" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result == hash_token("some-other-token-format") assert result is not None assert not result.startswith("Bearer") @@ -4680,7 +4680,7 @@ class TestSpendLogKeyRedaction: def test_bearer_and_bare_produce_same_hash(self): bare = "sk-WLi4iRn4JmbVlTaYw12IOA" bearer = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA" - assert _redact_logged_api_key(bare) == _redact_logged_api_key(bearer) + assert redact_logged_api_key(bare) == redact_logged_api_key(bearer) @patch("litellm.proxy.proxy_server.master_key", None) @@ -4777,7 +4777,7 @@ def test_get_logging_payload_keeps_internal_service_account_key_readable(service def test_redact_logged_api_key_service_account_name_without_provenance_is_hashed(): - result = _redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME) + result = redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME) assert result == hash_token(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME) @@ -5130,6 +5130,105 @@ def test_get_spend_logs_id_prefers_the_response_id_for_batch_and_file_calls(call assert get_spend_logs_id(call_type, {"id": "batch_abc123"}, {"litellm_call_id": "call-id-1"}) == expected +@pytest.mark.parametrize("status_call_type", ["avideo_status", "video_status"]) +@pytest.mark.parametrize( + ("response_cost", "status_id"), + [ + (0.42, "video_abc_video_cost"), + (0.0, "video_abc"), + (None, "video_abc"), + (float("nan"), "video_abc"), + ], + ids=["charged", "free", "unpriced", "non_finite"], +) +def test_get_logging_payload_gives_only_a_charged_video_poll_a_row_of_its_own(status_call_type, response_cost, status_id): + """A custom provider charges the video on the poll that sees it completed, so that poll needs + a row apart from the create row, which answers with the same video id and would drop it on the + skip_duplicates insert. A poll that charges nothing keeps the bare id, so it still lands on the + create row and adds no row of its own.""" + from litellm.types.videos.main import VideoObject + + video = VideoObject(id="video_abc", object="video", status="completed") + create, status = ( + get_logging_payload( + kwargs={ + "call_type": call_type, + "litellm_call_id": call_id, + "response_cost": cost, + "polled_video_status": "completed", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=video, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + for call_type, call_id, cost in (("avideo_generation", "call-1", 0.42), (status_call_type, "call-2", response_cost)) + ) + + assert create["request_id"] == "video_abc" + assert status["request_id"] == status_id + + +@pytest.mark.parametrize(("response_cost", "pays_the_job"), [(0.01, False), (0.43, True)]) +def test_a_video_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own(response_cost, pays_the_job): + """A guardrail charges every poll, so a poll paying only that must neither land on the create row + nor take the row the video's own charge goes on, either of which the skip_duplicates insert drops.""" + kwargs = { + "litellm_call_id": "poll-call-id", + "response_cost": response_cost, + "polled_video_status": "completed", + "standard_logging_object": { + "response_cost": response_cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}], + }, + } + + assert get_spend_logs_id("avideo_status", {"id": "video_abc", "status": "completed"}, kwargs) == ( + "video_abc_video_cost" if pays_the_job else "poll-call-id_guardrail_cost" + ) + + +@pytest.mark.parametrize("status", ["failed", "completed"]) +def test_a_finished_batch_retrieve_paying_only_its_guardrail_charges_under_the_batch_row(status): + """The retrieve that settles a batch charges its whole cost once, guardrails included, so a batch whose own + cost is $0 must still land on the row that makes its repeat retrieves free.""" + from litellm.types.utils import LiteLLMBatch + + kwargs = { + "litellm_call_id": "poll-call-id", + "response_cost": 0.01, + "standard_logging_object": { + "response_cost": 0.01, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}], + }, + } + polled_batch = LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id="file-out", + ) + + assert get_spend_logs_id("aretrieve_batch", polled_batch.model_dump(), kwargs) == "batch_abc_batch_cost" + + +def test_a_video_polled_with_or_without_its_base64_padding_charges_one_row(): + from litellm.types.videos.utils import encode_video_id_with_provider + + padded = encode_video_id_with_provider("job-1", "custom_video", "deployment-a") + assert padded.endswith("=") + + def row_id(polled_id: str) -> str | None: + kwargs = {"response_cost": 0.42, "polled_video_status": "completed", "litellm_params": {"video_id": polled_id}} + return get_spend_logs_id("avideo_status", {"id": "job-1", "status": "completed"}, kwargs) + + assert row_id(padded.rstrip("=")) == row_id(padded) == f"{padded}_video_cost" + + def test_get_logging_payload_gives_redacted_batch_and_file_rows_distinct_request_ids(): """End to end at the payload level: a batch retrieve and a file create whose bodies were both flattened to the same redaction placeholder must still produce two @@ -5790,7 +5889,7 @@ def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal assert payload["api_key"] == "", ( f"Expected empty string but got {payload['api_key']!r}; " - "dropping _redact_logged_api_key's 'or \"\"' guard would yield 'None' here" + "dropping redact_logged_api_key's 'or \"\"' guard would yield 'None' here" ) @@ -5841,7 +5940,7 @@ def test_get_logging_payload_stores_user_email_in_spend_log_metadata(user_email, def test_redact_logged_api_key_partial_sha256_is_hashed(): partial_hex = "a" * 63 - result = _redact_logged_api_key(partial_hex) + result = redact_logged_api_key(partial_hex) assert result is not None assert result != partial_hex assert len(result) == 64 @@ -5851,7 +5950,7 @@ def test_redact_logged_api_key_partial_sha256_is_hashed(): def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True) + result = redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True) assert result == already_hashed assert hash_token(already_hashed) != result @@ -5859,7 +5958,7 @@ def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): def test_redact_logged_api_key_bearer_sha256_without_flag_is_hashed(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(f"Bearer {already_hashed}") + result = redact_logged_api_key(f"Bearer {already_hashed}") assert result is not None assert result != already_hashed assert result == hash_token(already_hashed) @@ -6471,12 +6570,12 @@ def test_get_logging_payload_never_lands_a_raw_cli_session_token(call_type: str) def test_redact_logged_api_key_cli_session_alias_needs_alias_provenance(): - assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == ( + assert redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == ( _CLI_SESSION_ALIAS ) - assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS) - assert _redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS) - assert _redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice") + assert redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS) + assert redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS) + assert redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice") def test_azure_spillover_stamped_from_response_headers(): @@ -7695,13 +7794,11 @@ def test_optional_mcp_metadata_does_not_change_legacy_spend_payload(tool_metadat assert actual == expected -@pytest.mark.parametrize("capture_enabled", [False, True]) @pytest.mark.parametrize("completion_start", [None, "absent", "reported"]) -def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, completion_start: str | None) -> None: +def test_spend_payload_preserves_unknown_completion_start(completion_start: str | None) -> None: start: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc) end: Final = start + datetime.timedelta(seconds=2) reported: Final = start + datetime.timedelta(seconds=1) - observations: Final[list[object]] = [] kwargs: Final = { "call_type": "acompletion", "litellm_call_id": "timestamp-call", @@ -7713,9 +7810,49 @@ def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, {"id": "timestamp-response"}, start, end, - capture_usage=(lambda row, usage: observations.append(row)) if capture_enabled else None, ) assert payload["completionStartTime"] == ( end if completion_start == "absent" else reported if completion_start else None ) - assert observations == ([payload] if capture_enabled else []) + + +def test_spend_log_metadata_retention_drops_accounting_inputs_it_was_not_told_to_keep() -> None: + metadata: Final = { + "status": "success", + "accounting_savings": { + "compression": 0.4, + "prompt_caching": 0.1, + "gateway_injected_caching": 0.2, + "autorouter": 0.3, + }, + "accounting_windows": [ + {"entity_type": "key", "entity_id": "key", "window_duration": "1d", "window_start": "2026-11-02T00:00:00"} + ], + "accounting_model_access_groups": ["group"], + "user_api_key_project_id": "project", + "usage_object": {"prompt_tokens_details": {"cached_tokens": 2}}, + "compression_savings": {"tokens_saved": 3}, + "guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}], + "internal_call_origin": "classifier", + "billed_retry_attempt": True, + "user_api_key_alias": "discard", + } + row: Final = { + "request_id": "retained-accounting", + "spend": 0.5, + "startTime": "2026-11-02T12:00:00", + "endTime": "2026-11-02T12:00:01", + "model": "model", + "model_group": "group", + "api_key": "key", + "custom_llm_provider": "openai", + "user": "user", + "team_id": "team", + "agent_id": "agent", + "request_tags": '["tag"]', + "metadata": json.dumps(metadata), + } + stored: Final = spend_log_row_with_retained_metadata(row, SpendLogsMetadataFields(include=("usage_object",))) + parsed: Final = json.loads(str(stored["metadata"])) + assert "user_api_key_alias" not in parsed + assert parsed == {"usage_object": metadata["usage_object"], "status": metadata["status"]} diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py deleted file mode 100644 index 0c59473114a..00000000000 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ /dev/null @@ -1,1738 +0,0 @@ -import asyncio -import json -from collections.abc import Awaitable, Mapping -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import Metric, MetricName, Outcome -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite -from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings -from litellm.types.utils import CallTypes -from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _manual_flush, _write - - -@pytest.mark.asyncio -@pytest.mark.parametrize("usage", [None, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) -async def test_capture_keeps_unknown_usage_distinct_from_measured_zero(usage: Mapping[str, int] | None) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - captured: Final = capture_callback( - meter, - {"call_type": "acompletion", "litellm_call_id": "caller", "model": "unknown-model", "litellm_params": {}}, - {"id": "provider", "usage": usage}, - timestamp, - timestamp, - ) - await meter.close() - outbound: Final = captured.writes[1].event - assert outbound.metrics == ( - None - if usage is None - else tuple(Metric(name=MetricName(key), value=Decimal(value)) for key, value in usage.items()) - ) - envelope: Final = json.loads(captured.writes[1].envelope) - assert envelope["usage"] == usage - assert envelope["pricing"]["reported_cost"] is None - assert database.committed == ("caller:inbound", "provider") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type", - [ - call.value - for call in ( - CallTypes.acompletion, - CallTypes.arealtime, - CallTypes.call_mcp_tool, - CallTypes.asend_message, - CallTypes.pass_through, - ) - ], -) -async def test_shared_callback_capture_retains_usage_and_attribution_without_bodies(call_type: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - captured: Final = capture_callback( - meter, - { - "call_type": call_type, - "litellm_call_id": "caller", - "model": "test-model", - "messages": "private prompt", - "litellm_params": {"api_key": "private-secret", "metadata": {"user_api_key_team_id": "team"}}, - }, - { - "id": "provider", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, - "output": "private answer", - }, - timestamp, - timestamp, - ) - await meter.close() - envelope: Final = captured.writes[1].envelope - assert all(secret not in envelope for secret in ("private prompt", "private-secret", "private answer")) - event: Final = captured.writes[1].event - assert event.call_type.value == call_type - assert event.metrics == ( - Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(3)), - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(5)), - ) - assert any(dimension.value == "team" for dimension in event.dimensions) - - -@pytest.mark.asyncio -async def test_cache_duplicate_uses_caller_identity_despite_timestamped_diagnostic_ids() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "test-model", - "cache_hit": True, - "litellm_params": {}, - } - response: Final = {"id": "cached-provider-id", "usage": {"total_tokens": 5}} - first: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) - second: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) - await meter.close() - assert first.writes == second.writes - assert first.writes[1].event.event_id == "caller_cache_hit" - assert json.loads(first.writes[1].envelope)["charge_eligible"] is False - - -NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - - -class StoredUsage: - def __init__(self) -> None: - self.rows: tuple[UsageWrite, ...] = () - - async def load(self, identities: tuple[str, ...]) -> tuple[UsageWrite, ...]: - return tuple(row for row in self.rows if row.event.event_id in identities) - - async def insert(self, writes: tuple[UsageWrite, ...]) -> None: - existing: Final = {row.event.event_id: row for row in self.rows} - if any(existing.get(row.event.event_id, row) != row for row in writes): - raise AccountingConflictError("Same identity has different content") - unique: Final = {row.event.event_id: row for row in writes} - if any(unique[row.event.event_id] != row for row in writes): - raise AccountingConflictError("Submission contains conflicting content") - self.rows += tuple(row for identity, row in unique.items() if identity not in existing) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("text", ["parent\x00value", "pärent-✓-日本", r"\u0000"]) -async def test_capture_normalizes_database_text_without_poisoning_neighbors(text: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import PricingCapture, capture_callback - - database: Final = StoredUsage() - batches: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - assert all("\x00" not in (json.loads(write.envelope).get("parent_request_id") or "") for write in writes) - batches.put_nowait(writes) - await database.insert(writes) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(), wait_for_batch=_manual_flush) - nested: Final = json.loads("[" * 100 + "3" + "]" * 100) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "search_context_size": text, - "_accounting_pricing_snapshot": PricingCapture( - "{}", "{}", json.dumps({"request": {"search_context_size": text}, "nested": nested}) - ), - "litellm_params": {"metadata": {"parent_request_id": text}}, - } - response: Final = {"id": "provider", "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}} - captured: Final = capture_callback(meter, kwargs, response, NOW, NOW) - neighbor: Final = capture_callback( - meter, - {**kwargs, "litellm_call_id": "neighbor", "litellm_params": {}}, - {**response, "id": "neighbor-provider"}, - NOW, - NOW, - ) - try: - await meter.flush() - await asyncio.gather(captured.receipt, neighbor.receipt) - assert batches.qsize() == 1 - assert len(batches.get_nowait()) == 4 - envelope: Final = json.loads(captured.writes[-1].envelope) - assert envelope["parent_request_id"] == text.replace("\x00", "") - assert envelope["pricing"]["request"]["search_context_size"] == text.replace("\x00", "") - assert envelope["pricing"]["nested"] == nested - assert envelope["event"] == captured.writes[-1].event.model_dump(mode="json") - assert envelope["usage"] == response["usage"] - assert kwargs["litellm_params"] == {"metadata": {"parent_request_id": text}} - repeated: Final = capture_callback(meter, kwargs, response, NOW, NOW) - await meter.flush() - await repeated.receipt - assert repeated.writes == captured.writes - assert len(database.rows) == 4 - finally: - await meter.stop(0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("invalid_field", ["identity", "context"]) -async def test_invalid_capture_text_is_rejected_before_admission(invalid_field: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "caller", "litellm_params": {}} - try: - accepted: Final = capture_callback(meter, kwargs, {"id": "provider"}, NOW, NOW) - invalid: Final = ( - {**kwargs, "litellm_call_id": "caller\x00"} - if invalid_field == "identity" - else {**kwargs, "litellm_params": {"metadata": {"parent_request_id": "parent\ud800"}}} - ) - with pytest.raises((UsageCaptureError, ValueError)): - capture_callback(meter, invalid, {"id": "other"}, NOW, NOW) - assert meter.pending_records == 2 - await meter.flush() - await accepted.receipt - assert tuple(write.event.event_id for write in database.rows) == ("caller:inbound", "provider") - finally: - await meter.stop(0) - - -def _logger(call_type: str = "acompletion", *, stream: bool = False) -> Logging: - from litellm.litellm_core_utils.litellm_logging import Logging - - logger: Final = Logging( - model="capture-model", - messages=[], - stream=stream, - call_type=call_type, - start_time=NOW, - litellm_call_id="caller", - function_id="function", - ) - logger.update_environment_variables( - model="capture-model", custom_llm_provider="openai", litellm_params={}, optional_params={} - ) - return logger - - -@pytest.mark.asyncio -async def test_terminal_retrieval_reuses_accepted_snapshot_but_rejects_changed_usage() -> None: - from datetime import timedelta - - from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = { - "call_type": "aget_responses", - "litellm_call_id": "poll-one", - "model": "capture-model", - "response_cost": 2.0, - "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, - } - response: Final = { - "id": "job", - "status": "completed", - "background": True, - "created_at": NOW.timestamp(), - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - } - first: Final = capture_callback(meter, kwargs, response, NOW, NOW) - await meter.flush() - later: Final = capture_callback( - meter, {**kwargs, "litellm_call_id": "poll-two", "response_cost": 9.0}, response, NOW, NOW + timedelta(hours=2) - ) - await meter.flush() - await later.receipt - assert database.rows == first.writes - assert len(database.rows) == 1 - assert database.rows[0].event.event_id == "job_background_cost" - changed: Final = capture_callback(meter, kwargs, {**response, "usage": {"total_tokens": 8}}, NOW, NOW) - await meter.close() - with pytest.raises(UsageCaptureError, match="identity"): - await changed.receipt - assert database.rows == first.writes - - -@pytest.mark.asyncio -@pytest.mark.parametrize("terminal", [False, True]) -async def test_competing_writers_reuse_only_terminal_pricing_snapshots(terminal: bool) -> None: - from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - database: Final = StoredUsage() - attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - competing: Final = asyncio.Event() - - async def insert(writes: tuple[UsageWrite, ...]) -> None: - attempts.put_nowait(writes) - if attempts.qsize() >= 2: - competing.set() - await competing.wait() - await database.insert(writes) - - meters: Final = tuple( - DurableUsageMeter(UsageBatchWriter(insert, database.load).write, UsageCaptureSettings()) for _ in range(2) - ) - kwargs: Final = { - "call_type": "aget_responses" if terminal else "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "litellm_params": { - "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} - if terminal - else {} - }, - } - response: Final = { - "id": "job", - "status": "completed", - "background": terminal, - "created_at": NOW.timestamp(), - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - } - captured: Final = tuple( - capture_callback(meter, {**kwargs, "response_cost": price}, response, NOW, NOW) - for meter, price in zip(meters, (2.0, 9.0), strict=True) - ) - try: - await asyncio.gather(*(meter.flush() for meter in meters)) - results: Final = await asyncio.gather(*(capture.receipt for capture in captured), return_exceptions=True) - assert sum(result is None for result in results) == (2 if terminal else 1) - assert sum(isinstance(result, UsageCaptureError) for result in results) == (0 if terminal else 1) - assert database.rows in tuple(capture.writes for capture in captured) - assert len(database.rows) == (1 if terminal else 2) - assert attempts.qsize() == (3 if terminal else 2) - finally: - await asyncio.gather(*(meter.stop(0) for meter in meters)) - - -@pytest.mark.asyncio -async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_chargeable() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = { - "call_type": "aretrieve_batch", - "litellm_call_id": "poll", - "model": "capture-model", - "litellm_params": {}, - } - pending: Final = capture_callback(meter, kwargs, {"id": "batch", "status": "in_progress"}, NOW, NOW) - completed: Final = capture_callback( - meter, kwargs, {"id": "batch", "status": "completed", "usage": {"total_tokens": 9}}, NOW, NOW - ) - await meter.close() - assert pending.writes[-1].event.event_id == "poll:status" - assert json.loads(pending.writes[-1].envelope)["charge_eligible"] is False - assert completed.writes[-1].event.event_id == "batch_batch_cost" - assert json.loads(completed.writes[-1].envelope)["charge_eligible"] is True - assert len(completed.writes) == 1 - - -@pytest.mark.asyncio -async def test_internal_capture_does_not_claim_reused_parent_logger(monkeypatch: pytest.MonkeyPatch) -> None: - import litellm - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.utils import ModelResponse, Usage - - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - { - "input_cost_per_token": 0.1, - "output_cost_per_token": 0.2, - "litellm_provider": "openai", - "mode": "chat", - }, - ) - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger() - response: Final = ModelResponse( - id="response", model="capture-model", usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3) - ) - with internal_sub_call("caller"): - capture_logging_usage(meter, logger, response, NOW, NOW) - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - { - "input_cost_per_token": 0.7, - "output_cost_per_token": 0.9, - "litellm_provider": "openai", - "mode": "chat", - }, - ) - litellm.register_model({"capture-model": litellm.model_cost["capture-model"]}) - capture_logging_usage(meter, logger, response, NOW, NOW) - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - assert tuple(row.event.event_id for row in database.rows) == ("response:included", "caller:inbound", "response") - child: Final = json.loads(database.rows[0].envelope) - parent: Final = json.loads(database.rows[-1].envelope) - assert (child["ownership"], child["charge_eligible"], child["parent_request_id"]) == ( - "included_in_parent", - False, - "caller", - ) - assert parent["charge_eligible"] is True - assert child["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.1 - assert parent["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.7 - - -@pytest.mark.asyncio -async def test_stream_capture_ignores_chunks_and_preserves_terminal_usage_extras() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.utils import ModelResponse, ModelResponseStream, Usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(stream=True) - capture_logging_usage(meter, logger, ModelResponseStream(id="chunk"), NOW, NOW) - assert meter.pending_records == 0 - capture_logging_usage( - meter, logger, ModelResponse(id="stream", usage=Usage(prompt_tokens=3, provider_units=4)), NOW, NOW - ) - await meter.close() - assert json.loads(database.rows[-1].envelope)["usage"]["provider_units"] == 4 - assert database.rows[-1].event.event_id == "stream" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["acompletion", "atext_completion", "aembedding", "aresponses"]) -async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy._types import ConfigGeneralSettings - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - deadlines: Final = asyncio.Queue[float | None]() - original_wait: Final = asyncio.wait_for - - async def wait_with_deadline(awaitable: Awaitable[object], timeout: float | None) -> object: - deadlines.put_nowait(timeout) - return await original_wait(awaitable, timeout) - - monkeypatch.setattr(asyncio, "wait_for", wait_with_deadline) - response: Final = asyncio.create_task( - capture_proxy_response( - _logger(call_type), - ModelResponse(id="provider", usage=Usage(prompt_tokens=3)), - ) - ) - try: - await database.entered.wait() - deadline: Final = deadlines.get_nowait() - assert deadline == ConfigGeneralSettings().accounting_sync_ack_timeout_seconds - assert deadline is not None and deadline > UsageCaptureSettings().flush_interval_seconds + 10 - assert not response.done() - assert database.committed == () - database.release.set() - await response - assert database.committed == ("caller:inbound", "provider") - finally: - database.release.set() - await runtime.stop(1) - await response - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure_policy", ["fail_closed", "log"]) -async def test_sync_ack_commit_failure_follows_response_policy( - failure_policy: str, monkeypatch: pytest.MonkeyPatch -) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier(failure=ValueError("invalid envelope")) - database.release.set() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - meter: Final = runtime.meter - assert meter is not None - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, - "general_settings", - {"accounting_sync_ack": True, "accounting_sync_ack_failure_policy": failure_policy}, - ) - logger: Final = _logger() - response: Final = ModelResponse(id="provider", usage=Usage(prompt_tokens=4)) - if failure_policy == "fail_closed": - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, response) - assert error.value.status_code == 503 - else: - await capture_proxy_response(logger, response) - await runtime.stop(1.0) - assert database.committed == () - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_success() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - for _ in range(2): - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4) - logger.carry_failed_attempt_usage() - logger.capture_usage_before_logging( - ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, total_tokens=3)), NOW, NOW - ) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ( - "caller-attempt-0", - "caller-attempt-1", - "caller:inbound", - "fallback", - ) - failed: Final = json.loads(database.rows[0].envelope) - assert failed["usage"]["total_tokens"] == 7 - assert failed["parent_request_id"] == "caller" - assert failed["request_count_eligible"] is False - assert failed["charge_eligible"] is True - assert failed["event"]["retry_of_attempt_id"] is None - assert failed["event"]["outcome"] == "failed" - assert failed["pricing"]["reported_cost"] == 0.4 - assert all(json.loads(row.envelope)["included_attempt_ids"] == [] for row in database.rows[:2]) - successful: Final = json.loads(database.rows[-1].envelope) - assert successful["related_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"] - assert successful["included_attempt_ids"] == [] - assert successful["usage"]["total_tokens"] == 3 - assert successful["charge_eligible"] is True - - -@pytest.mark.asyncio -async def test_success_that_bills_a_failed_follow_up_names_that_attempt_as_included() -> None: - from litellm.litellm_core_utils.hidden_params import set_hidden_param - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - fallback: Final = ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4)) - set_hidden_param(fallback, "response_cost", 0.1) - try: - logger.record_served_rounds_for_failure(((Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4),)) - logger.carry_failed_attempt_usage() - logger.prepare_response_cost(fallback, NOW) - logger.capture_usage_before_logging(fallback, NOW, NOW) - finally: - await runtime.stop(1) - envelopes: Final = {row.event.event_id: json.loads(row.envelope) for row in database.rows} - assert envelopes["caller-attempt-0"]["usage"]["total_tokens"] == 7 - assert envelopes["caller-attempt-0"]["charge_eligible"] is True - assert envelopes["fallback"]["usage"]["total_tokens"] == 4 + 7 - assert envelopes["fallback"]["included_attempt_ids"] == ["caller-attempt-0"] - assert database.rows[2].event.outcome is Outcome.SUCCEEDED - - -@pytest.mark.asyncio -@pytest.mark.parametrize("estimated", [False, True]) -async def test_cancelled_usage_keeps_measured_or_estimated_provenance(estimated: bool) -> None: - from fastapi import HTTPException - - from litellm.litellm_core_utils.litellm_logging import UsageObservation - from litellm.proxy.spend_tracking.usage_capture import capture_callback - from litellm.types.utils import Usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - capture_callback( - meter, - { - "litellm_call_id": "cancelled", - "call_type": "acompletion", - "model": "capture-model", - "litellm_params": {}, - "combined_usage_object": Usage(prompt_tokens=4, total_tokens=4), - "_accounting_usage_observed": UsageObservation.ESTIMATED if estimated else UsageObservation.KNOWN, - }, - HTTPException(499, "Client disconnected"), - NOW, - NOW, - ) - await meter.close() - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["event"]["outcome"] == "cancelled" - assert payload["usage"]["total_tokens"] == 4 - assert payload["usage_source"] == ("estimated" if estimated else "measured") - assert payload["pricing"]["reported_cost"] is None - - -@pytest.mark.asyncio -async def test_realtime_capture_keeps_metered_units_and_provider_usage_without_frames() -> None: - from litellm.constants import REALTIME_METERED_USAGE_KEY - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.llms.xai import RealtimeMeteredUsage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(CallTypes.arealtime.value) - logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 1.5, 2.5, 3) - capture_logging_usage( - meter, - logger, - [ - { - "type": "response.done", - "response": { - "id": "turn", - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - "output": "private realtime transcript", - }, - } - ], - NOW, - NOW, - ) - await meter.close() - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["usage"]["total_tokens"] == 7 - assert payload["usage"]["metered_usage"] == { - "push_to_talk": True, - "audio_seconds_sent": 1.5, - "audio_seconds_received": 2.5, - "text_messages": 3, - } - assert "private realtime transcript" not in database.rows[-1].envelope - - -@pytest.mark.asyncio -async def test_default_http_capture_returns_before_database_commit(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {}) - try: - await capture_proxy_response( - _logger(), - ModelResponse(id="asynchronous", usage=Usage(prompt_tokens=1)), - ) - await database.entered.wait() - assert database.committed == () - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("caller:inbound", "asynchronous") - - -@pytest.mark.asyncio -async def test_deferred_capture_freezes_pricing_before_configuration_changes(monkeypatch: pytest.MonkeyPatch) -> None: - import litellm - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_logging_usage - from litellm.types.utils import ModelResponse, Usage - from litellm.utils import dispatch_success_logging - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - logger.defer_async_logging = True - logger.custom_cost_per_second = 0.25 - monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1}) - response: Final = ModelResponse(id="deferred", usage=Usage(prompt_tokens=3)) - try: - dispatch_success_logging(logger, response, NOW, NOW, False, False) - assert database.rows == () - logger.custom_cost_per_second = 9 - monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.9}) - assert runtime.meter is not None - capture_logging_usage(runtime.meter, logger, response, NOW, NOW) - await runtime.meter.flush() - finally: - await runtime.stop(1) - pricing: Final = json.loads(database.rows[-1].envelope)["pricing"] - assert pricing["logging"]["custom_cost_per_second"] == 0.25 - assert pricing["adjustments"]["discount"] == 0.1 - - -@pytest.mark.asyncio -async def test_synchronous_internal_capture_reserves_capacity_before_scheduling() -> None: - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) - logger: Final = _logger() - - def capture_child() -> None: - with internal_sub_call("caller"): - logger.capture_usage_before_logging(ModelResponse(id="child", usage=Usage(prompt_tokens=2)), NOW, NOW) - - try: - await asyncio.to_thread(capture_child) - assert runtime.meter is not None - assert runtime.meter.pending_records == 1 - receipt: Final = await asyncio.to_thread(runtime.meter.submit, (_write("second"),)) - await database.entered.wait() - with pytest.raises(UsageCaptureError, match="capacity"): - await asyncio.to_thread(runtime.meter.submit, (_write("overflow"),)) - database.release.set() - await receipt - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("child:included", "second") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("has_metered_usage", [False, True]) -async def test_realtime_without_provider_token_usage_does_not_invent_zero(has_metered_usage: bool) -> None: - from litellm.constants import REALTIME_METERED_USAGE_KEY - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.llms.xai import RealtimeMeteredUsage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(CallTypes.arealtime.value) - if has_metered_usage: - logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 0, 0, 0) - capture_logging_usage(meter, logger, [], NOW, NOW) - await meter.close() - usage: Final = json.loads(database.rows[-1].envelope)["usage"] - assert usage == ( - { - "metered_usage": { - "push_to_talk": True, - "audio_seconds_sent": 0, - "audio_seconds_received": 0, - "text_messages": 0, - } - } - if has_metered_usage - else None - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("committed_before_failure", [False, True]) -async def test_database_recovery_retries_frozen_usage_without_callback_or_duplicate_rows( - committed_before_failure: bool, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - failures: Final = asyncio.Queue[None]() - failures.put_nowait(None) - retrying: Final = asyncio.Event() - recovered: Final = asyncio.Event() - delays: Final = asyncio.Queue[float]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - if not failures.empty(): - failures.get_nowait() - if committed_before_failure: - await database.insert(writes) - raise ConnectionError("database unavailable or acknowledgement lost") - await database.insert(writes) - - async def retry_delay(delay: float) -> None: - delays.put_nowait(delay) - retrying.set() - await recovered.wait() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(max_pending_records=2), wait_for_batch=_manual_flush, retry_delay=retry_delay - ) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "litellm_params": {}, - } - response: Final = {"id": "provider", "usage": {"total_tokens": 3}} - first: Final = capture_callback(meter, kwargs, response, NOW, NOW) - reserved: Final = meter.pending_bytes - flushing: Final = asyncio.create_task(meter.flush()) - await retrying.wait() - assert not first.receipt.done() - assert (meter.pending_records, meter.pending_bytes) == (2, reserved) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("overflow"),)) - recovered.set() - await flushing - await first.receipt - duplicate: Final = capture_callback(meter, {**kwargs, "_accounting_capture": first}, response, NOW, NOW) - await meter.close() - assert duplicate is first - assert database.rows == first.writes - assert delays.get_nowait() == 0.2 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -async def test_cut_stream_keeps_original_model_usage_as_part_of_the_parent_aggregate() -> None: - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.router import Router - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger(stream=True) - stream: Final = CustomStreamWrapper(completion_stream=iter(()), model="capture-model", logging_obj=logger) - partial: Final = ModelResponse( - id="cut", model="capture-model", usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) - ) - try: - Router._price_cut_partial(stream, partial) - finally: - await runtime.stop(1) - assert len(database.rows) == 1 - payload: Final = json.loads(database.rows[0].envelope) - assert payload["event"]["event_id"] == "cut:included" - assert payload["parent_request_id"] == "caller" - assert payload["charge_eligible"] is False - assert payload["usage"]["total_tokens"] == 5 - assert payload["pricing"]["model_map_information"]["model_map_key"] == "capture-model" - - -@pytest.mark.asyncio -async def test_direct_mcp_logging_captures_tool_pricing_and_identity_without_arguments() -> None: - from mcp.types import CallToolResult, TextContent - - from litellm.accounting import DimensionName - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger(CallTypes.call_mcp_tool.value) - logger.model_call_details["mcp_tool_call_metadata"] = { - "name": "lookup", - "namespaced_tool_name": "server.lookup", - "arguments": {"query": "private query"}, - "mcp_server_cost_info": {"default_cost_per_query": 0.2, "tool_name_to_cost_per_query": {"lookup": 0.3}}, - } - try: - await logger.async_success_handler( - CallToolResult(content=[TextContent(type="text", text="private result")]), NOW, NOW - ) - finally: - await runtime.stop(1) - outbound: Final = database.rows[-1] - assert any(d.name is DimensionName.TOOL and d.value == "server.lookup" for d in outbound.event.dimensions) - payload: Final = json.loads(outbound.envelope) - assert payload["pricing"]["tool"] == {"name": "lookup", "cost_per_query": 0.3} - assert payload["usage"] is None - assert all(secret not in outbound.envelope for secret in ("private query", "private result")) - - -@pytest.mark.asyncio -async def test_callback_overload_is_visible_without_suppressing_legacy_tracking( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, captured_callback_kwargs - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "one", "litellm_params": {}} - try: - for start, end in ((None, NOW), (NOW, None)): - missing_timestamp: Final = await captured_callback_kwargs(kwargs, {"id": "one"}, start, end) - assert missing_timestamp == kwargs - assert "Usage capture requires callback start and end timestamps" in caplog.text - admitted: Final = await captured_callback_kwargs(kwargs, {"id": "one"}, NOW, NOW) - await database.entered.wait() - rejected: Final = await captured_callback_kwargs({**kwargs, "litellm_call_id": "two"}, {"id": "two"}, NOW, NOW) - assert "_accounting_capture" in admitted - assert "_accounting_capture" not in rejected - assert "Usage capture failed before legacy cost tracking" in caplog.text - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("one:inbound", "one") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "streaming", "background"), - [("acompletion", True, False), ("aresponses", False, True), ("aocr", False, False)], -) -@pytest.mark.parametrize("overloaded", [False, True]) -async def test_async_capture_preserves_response_without_commit_ack( - call_type: str, streaming: bool, background: bool, overloaded: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, - "general_settings", - { - "accounting_sync_ack": True, - "accounting_sync_ack_timeout_seconds": 0, - }, - ) - meter: Final = runtime.meter - assert meter is not None - if overloaded: - meter.submit((_write("occupied-one"), _write("occupied-two"))) - try: - await capture_proxy_response( - _logger(call_type, stream=streaming), - ModelResponse(id="provider", usage=Usage(prompt_tokens=3)), - background=background, - ) - assert database.committed == () - assert meter.pending_records == 2 - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == (("occupied-one", "occupied-two") if overloaded else ("caller:inbound", "provider")) - - -@pytest.mark.asyncio -async def test_new_connection_turn_does_not_reuse_the_previous_turns_capture() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.capture_usage_before_logging(ModelResponse(id="first", usage=Usage(prompt_tokens=2)), NOW, NOW) - next_turn: Final = logger.copy_for_new_call("second-caller", NOW) - next_turn.capture_usage_before_logging(ModelResponse(id="second", usage=Usage(prompt_tokens=4)), NOW, NOW) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ( - "caller:inbound", - "first", - "second-caller:inbound", - "second", - ) - assert json.loads(database.rows[-1].envelope)["usage"]["prompt_tokens"] == 4 - - -@pytest.mark.asyncio -async def test_billed_failure_is_captured_before_legacy_persistence_gates(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - request: Final = { - "call_type": "acompletion", - "litellm_call_id": "failed", - "model": "capture-model", - "start_time": NOW, - "combined_usage_object": Usage(prompt_tokens=3, total_tokens=3), - "response_cost": 0.2, - "litellm_params": {}, - } - try: - await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), UserAPIKeyAuth()) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("failed:inbound", "failed") - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["event"]["outcome"] == "failed" - assert payload["usage"]["total_tokens"] == 3 - assert payload["charge_eligible"] is True - assert payload["pricing"]["reported_cost"] == 0.2 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("pages", [None, 0, 3]) -async def test_prepared_ocr_capture_preserves_reported_pages(pages: int | None) -> None: - from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger("aocr") - response: Final = OCRResponse(pages=[], model="capture-model", usage_info=OCRUsageInfo(pages_processed=pages)) - logger.prepare_response_cost(response, NOW) - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["pages_processed"]) == pages - assert envelope["usage_source"] == ("unknown" if pages is None else "measured") - assert response.usage_info is not None and response.usage_info.pages_processed == pages - - -@pytest.mark.asyncio -@pytest.mark.parametrize("measured", [None, {"total_tokens": 0}, {"total_tokens": 3}]) -async def test_batch_capture_provenance_tracks_observations_not_envelope_presence( - measured: Mapping[str, int] | None, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - usage: Final = { - "total_tokens": 0 if measured is None else measured["total_tokens"], - "batch_usage_by_model": {"missing-model": None, "other-model": measured}, - "batch_unknown_usage_requests": 2 if measured is None else 1, - } - capture_callback( - meter, - {"call_type": "aretrieve_batch", "litellm_call_id": "caller", "litellm_params": {}}, - {"id": "batch", "usage": usage}, - NOW, - NOW, - ) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert envelope["usage_source"] == ("unknown" if measured is None else "measured") - assert envelope["usage"]["batch_usage_by_model"] == usage["batch_usage_by_model"] - assert database.rows[-1].event.metrics == ( - None if measured is None else (Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(measured["total_tokens"])),) - ) - - -@pytest.mark.asyncio -async def test_ack_admission_failure_does_not_turn_provider_success_into_failed_usage( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - monkeypatch.setattr(proxy_server, "prisma_client", None) - meter: Final = runtime.meter - assert meter is not None - logger: Final = _logger() - try: - meter.submit((_write("occupied-one"), _write("occupied-two"))) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="caller", usage=Usage(prompt_tokens=3))) - assert error.value.status_code == 503 - await meter.flush() - request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} - await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("occupied-one", "occupied-two") - - -@pytest.mark.asyncio -async def test_response_ack_timeout_preserves_success_and_later_commits(monkeypatch: pytest.MonkeyPatch) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - logger: Final = _logger() - - async def expired_deadline(future: object, timeout: float) -> None: - raise TimeoutError("controlled response deadline") - - try: - with monkeypatch.context() as clock: - clock.setattr(asyncio, "wait_for", expired_deadline) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="provider", usage=Usage(prompt_tokens=3))) - assert error.value.status_code == 503 - await database.entered.wait() - captured: Final = logger.model_call_details["_accounting_capture"] - request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} - await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) - database.release.set() - await captured.receipt - assert tuple(row.event.outcome for row in captured.writes) == (Outcome.SUCCEEDED, Outcome.SUCCEEDED) - assert database.committed == ("caller:inbound", "provider") - assert database.transactions == 1 - finally: - await runtime.stop(1) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["generate_content", "agenerate_content"]) -@pytest.mark.parametrize("tokens", [None, 0, 7]) -async def test_prepared_native_google_capture_preserves_reported_usage( - call_type: str, tokens: int | None, monkeypatch: pytest.MonkeyPatch -) -> None: - import httpx - - import litellm - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(call_type) - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - {"input_cost_per_token": 0.1, "output_cost_per_token": 0.2, "litellm_provider": "openai", "mode": "chat"}, - ) - reported: Final = ( - {"trafficType": "ON_DEMAND"} - if tokens is None - else { - "promptTokenCount": tokens, - "candidatesTokenCount": tokens, - "totalTokenCount": 2 * tokens, - "cachedContentTokenCount": tokens, - "promptTokensDetails": [{"modality": "DOCUMENT", "tokenCount": tokens}], - "cacheTokensDetails": [{"modality": "document", "token_count": tokens}], - } - ) - response: Final = { - "responseId": "native-google", - "candidates": [{"content": {"parts": [{"text": "answer"}], "role": "model"}, "finishReason": "STOP"}], - "usageMetadata": reported, - } - native_response: Final = httpx.Response(200, json=response) - logger.model_call_details["httpx_response"] = native_response - prepared: Final = logger.prepare_response_cost(response, NOW) - assert prepared is not None - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == ( - None if tokens is None else 2 * tokens - ) - assert (None if envelope["usage"] is None else envelope["usage"]["prompt_tokens_details"]["text_tokens"]) == tokens - assert ( - None - if envelope["usage"] is None - else envelope["usage"]["prompt_tokens_details"]["cached_tokens_details"]["text_tokens"] - ) == tokens - assert envelope["usage_source"] == ("unknown" if tokens is None else "measured") - assert response == native_response.json() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("rejected", [False, True]) -async def test_retry_links_only_admitted_usage_and_preserves_later_measurement(rejected: bool) -> None: - from litellm.litellm_core_utils.litellm_logging import UsageObservation - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings(max_pending_records=4)) - logger: Final = _logger() - assert runtime.meter is not None - try: - if rejected: - runtime.meter.submit(tuple(_write(f"occupied-{index}") for index in range(4))) - logger.model_call_details["_accounting_usage_observed"] = UsageObservation.UNKNOWN - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - await runtime.meter.flush() - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - finally: - await runtime.stop(1) - successful: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "success")) - assert successful.get("related_attempt_ids", successful["included_attempt_ids"]) == ( - [] if rejected else ["caller-attempt-0"] - ) - assert successful["usage"]["prompt_tokens"] == 3 - if not rejected: - failed: Final = json.loads( - next(row.envelope for row in database.rows if row.event.event_id == "caller-attempt-0") - ) - assert failed["usage"]["prompt_tokens"] == 5 - assert failed["usage_source"] == "measured" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("current_tokens", [None, 0, 3]) -async def test_final_failure_captures_current_attempt_without_legacy_retry_aggregate( - current_tokens: int | None, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.utils import _failure_fields_to_lift - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - if current_tokens is not None: - logger.record_partial_usage_for_failure( - Usage(prompt_tokens=current_tokens, total_tokens=current_tokens), 0.2 - ) - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert lifted["combined_usage_object"].total_tokens == 5 + (current_tokens or 0) - assert runtime.meter is not None - capture_callback( - runtime.meter, - {**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"}, - None, - NOW, - NOW, - ) - finally: - await runtime.stop(1) - final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) - assert (None if final["usage"] is None else final["usage"]["total_tokens"]) == current_tokens - assert final["usage_source"] == ("unknown" if current_tokens is None else "measured") - assert final["pricing"]["reported_cost"] == (None if current_tokens is None else 0.2) - assert final["included_attempt_ids"] == [] - assert final["related_attempt_ids"] == ["caller-attempt-0"] - - -@pytest.mark.asyncio -async def test_estimated_final_attempt_keeps_its_provenance_after_measured_retry() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.utils import _failure_fields_to_lift - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - logger.messages = [{"role": "user", "content": "hello"}] - logger.model_call_details["messages"] = logger.messages - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - logger.model_call_details["current_attempt_api_call_start_time"] = NOW - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert runtime.meter is not None - capture_callback(runtime.meter, {**logger.model_call_details, **lifted}, None, NOW, NOW) - finally: - await runtime.stop(1) - final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) - assert final["usage_source"] == "estimated" - assert final["usage"]["total_tokens"] > 0 - assert lifted["combined_usage_object"].total_tokens == final["usage"]["total_tokens"] + 5 - assert final["pricing"]["reported_cost"] == 0 - - -@pytest.mark.asyncio -async def test_parent_receipt_cannot_acknowledge_when_a_linked_attempt_conflicts() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - assert runtime.meter is not None - await runtime.meter.flush() - other: Final = _logger() - other.record_partial_usage_for_failure(Usage(prompt_tokens=8), 0.8) - other.carry_failed_attempt_usage() - other.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - captured: Final = other.model_call_details["_accounting_capture"] - await runtime.meter.flush() - with pytest.raises(UsageCaptureError, match="dependency"): - await captured.receipt - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("caller-attempt-0",) - - -@pytest.mark.asyncio -async def test_unserializable_observation_does_not_interrupt_retry_or_legacy_usage() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - usage: Final = Usage(prompt_tokens=5, total_tokens=5, opaque_usage=object()) - try: - logger.record_partial_usage_for_failure(usage, 0.4) - assert logger.model_call_details["combined_usage_object"] is usage - logger.carry_failed_attempt_usage() - assert logger.model_call_details["failed_attempts_usage"] == ((usage, 0.4),) - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("caller:inbound", "success") - assert json.loads(database.rows[-1].envelope)["related_attempt_ids"] == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prior_committed", [False, True]) -@pytest.mark.parametrize("batch_size", [2, 3]) -async def test_retry_dependencies_do_not_resubmit_usage_or_expand_final_submission( - prior_committed: bool, batch_size: int -) -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - database.release.set() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=batch_size, max_pending_records=3)) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - assert runtime.meter is not None - if prior_committed: - await runtime.meter.flush() - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - captured: Final = logger.model_call_details.get("_accounting_capture") - assert captured is not None, "Final usage must be admitted independently of prior usage batch size" - await runtime.meter.close() - await captured.receipt - assert runtime.meter.pending_records == 0 - finally: - await runtime.stop(1) - assert database.committed == ("caller-attempt-0", "caller:inbound", "success") - final: Final = json.loads(captured.writes[-1].envelope) - assert final["related_attempt_ids"] == ["caller-attempt-0"] - assert final["usage"]["prompt_tokens"] == 3 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prepared", [False, True]) -@pytest.mark.parametrize( - ("call_type", "response", "field", "quantity"), - [ - ("arerank", {"id": "rerank", "results": [], "meta": {"tokens": {"input_tokens": 7}}}, "prompt_tokens", 7), - ("aocr", {"id": "ocr", "usage_info": {"pages_processed": 3}}, "pages_processed", 3), - ( - "acompletion", - {"id": "google", "usageMetadata": {"promptTokenCount": 7, "totalTokenCount": 7}}, - "prompt_tokens", - 7, - ), - ], -) -async def test_capture_preserves_native_usage_before_and_after_preparation( - prepared: bool, call_type: str, response: dict[str, object], field: str, quantity: int -) -> None: - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(call_type) - logger.model_call_details["response_cost"] = 0.25 - if prepared: - assert logger.prepare_response_cost(response, NOW) is not None - with internal_sub_call("parent"): - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - assert len(database.rows) == 1 - envelope: Final = json.loads(database.rows[0].envelope) - assert envelope["usage"][field] == quantity - assert envelope["usage_source"] == "measured" - assert envelope["ownership"] == "included_in_parent" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("tokens", [(120000, 120000), (210000, 30000)]) -@pytest.mark.parametrize("priced", [False, True]) -async def test_batch_capture_retains_request_basis_for_later_tiered_pricing( - tokens: tuple[int, int], priced: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - import litellm - from litellm.batches.batch_utils import batch_cost_and_usage - from litellm.cost_calculator import batch_cost_calculator - from litellm.proxy.spend_tracking.usage_capture import ( - UsageBatchWriter, - capture_batch_callback, - capture_logging_usage, - ) - from litellm.types.utils import Usage - - rates: Final = { - "input_cost_per_token_batches": 0.001, - "input_cost_per_token_above_200k_tokens_batches": 0.002, - "output_cost_per_token_batches": 0.0, - "litellm_provider": "openai", - "mode": "chat", - } - monkeypatch.setitem(litellm.model_cost, "capture-model", rates if priced else {"litellm_provider": "openai"}) - rows: Final = [ - { - "response": { - "status_code": 200, - "body": { - "model": "capture-model", - "usage": {"prompt_tokens": n, "completion_tokens": 0, "total_tokens": n}, - }, - } - } - for n in tokens - ] - result: Final = batch_cost_and_usage(rows, "openai", model_info=rates if priced else {}) - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger("aretrieve_batch") - logger.model_call_details["response_cost"] = result.cost - capture_logging_usage(meter, logger, {"id": "batch", "status": "completed", "usage": result.usage}, NOW, NOW) - await capture_batch_callback( - meter, - database.load, - logger.model_call_details, - {"id": "batch", "status": "completed", "usage": result.usage}, - NOW, - NOW, - ) - await meter.close() - retained: Final = tuple( - {"model": "capture-model", "usage": json.loads(write.envelope)["usage"]} for write in database.rows[1:] - ) - assert sorted(row["usage"]["prompt_tokens"] for row in retained) == sorted(tokens) - reconstructed: Final = sum( - sum(batch_cost_calculator(Usage(**row["usage"]), row["model"], "openai", model_info=rates)) for row in retained - ) - expected: Final = sum(n * (0.002 if n > 200000 else 0.001) for n in tokens) - assert reconstructed == pytest.approx(expected) - assert result.cost == pytest.approx(expected if priced else 0) - assert database.rows[0].event.metrics is None - assert json.loads(database.rows[0].envelope)["charge_eligible"] is False - assert all(json.loads(write.envelope)["finalized_cost"] is None for write in database.rows[1:]) - - -@pytest.mark.asyncio -async def test_large_batch_recovers_partial_commit_with_frozen_pricing_and_rejects_changed_usage() -> None: - from litellm.proxy.spend_tracking.usage_capture import ( - PricingCapture, - UsageBatchWriter, - capture_batch_callback, - capture_callback, - ) - - database: Final = StoredUsage() - transactions: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - failure: Final = asyncio.Queue[bool]() - failure.put_nowait(True) - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - assert sum(len(write.envelope.encode()) for write in writes) <= 65536 - assert len(writes) <= 32 - if len(database.rows) > 1 and not failure.empty(): - failure.get_nowait() - raise RuntimeError("Interrupted batch persistence") - await database.insert(writes) - transactions.put_nowait(writes) - - writer: Final = UsageBatchWriter(persist, database.load) - meter: Final = DurableUsageMeter( - writer.write, - UsageCaptureSettings(batch_size=32, max_pending_records=32, max_pending_bytes=65536), - wait_for_batch=_manual_flush, - ) - rows: Final = tuple( - { - "model": "capture-model" if index % 7 else None, - "usage": {"prompt_tokens": index % 5 + 1} if index % 7 else None, - } - for index in range(1000) - ) - response: Final = {"id": "large-batch", "status": "completed", "usage": {"batch_usage_by_request": rows}} - pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"capture-model": {"rate": 1}}})) - kwargs: Final = { - "call_type": "aretrieve_batch", - "litellm_call_id": "poll", - "litellm_params": {}, - "_accounting_pricing_snapshot": pricing, - } - failed_database: Final = CommitBarrier(failure=RuntimeError("Earlier capture failed")) - failed_database.release.set() - prior_meter: Final = DurableUsageMeter(failed_database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - prior: Final = capture_callback( - prior_meter, - {"call_type": "acompletion", "litellm_call_id": "earlier", "litellm_params": {}}, - {"id": "earlier"}, - NOW, - NOW, - ) - await prior_meter.close() - with pytest.raises(RuntimeError, match="Earlier capture failed"): - await capture_batch_callback( - meter, - database.load, - {**kwargs, "_accounting_prior_captures": (prior,)}, - response, - NOW, - NOW, - ) - assert database.rows == () - with pytest.raises(RuntimeError, match="Interrupted batch persistence"): - await capture_batch_callback(meter, database.load, kwargs, response, NOW, NOW) - partial: Final = database.rows - assert 1 < len(partial) < len(rows) + 1 - changed_pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"capture-model": {"rate": 2}}})) - retried: Final = await capture_batch_callback( - meter, - database.load, - {**kwargs, "_accounting_pricing_snapshot": changed_pricing}, - {**response, "usage": {"batch_usage_by_request": tuple(reversed(rows))}}, - NOW, - NOW, - ) - assert retried.receipt.done() and retried.receipt.exception() is None - assert len(database.rows) == len(rows) + 1 - assert database.rows[: len(partial)] == partial - envelopes: Final = tuple(json.loads(write.envelope) for write in database.rows[1:]) - assert all( - envelope["pricing"]["model_map_information"] == ({"rate": 1} if envelope["usage"] is not None else None) - for envelope in envelopes - ) - assert sum(envelope["usage"] is None for envelope in envelopes) == sum(row["usage"] is None for row in rows) - assert all(envelope["request_count_eligible"] is False for envelope in envelopes) - assert sum(envelope["usage"]["prompt_tokens"] for envelope in envelopes if envelope["usage"] is not None) == sum( - row["usage"]["prompt_tokens"] for row in rows if row["usage"] is not None - ) - before_conflict: Final = database.rows - with pytest.raises(UsageCaptureError, match="identity conflicts"): - await capture_batch_callback( - meter, - database.load, - kwargs, - {**response, "usage": {"batch_usage_by_request": rows[:-1]}}, - NOW, - NOW, - ) - await meter.close() - assert database.rows == before_conflict - assert transactions.qsize() > 2 - assert meter.pending_bytes == meter.pending_records == 0 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("outcome", ["drain", "cancel", "failure"]) -async def test_batch_callback_handoff_is_bounded_and_runtime_owns_persistence( - outcome: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, UsageCaptureRuntime, captured_callback_kwargs - - database: Final = StoredUsage() - entered: Final = asyncio.Event() - release: Final = asyncio.Event() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - entered.set() - await release.wait() - if outcome == "failure": - raise RuntimeError("Persistence unavailable") - await database.insert(writes) - - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), load=database.load) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - response: Final = { - "id": "batch", - "status": "completed", - "usage": {"batch_usage_by_request": ({"model": "model", "usage": {"prompt_tokens": 2}},) * 3}, - } - callback: Final = asyncio.create_task(captured_callback_kwargs(kwargs, response, NOW, NOW)) - try: - await entered.wait() - assert callback.done(), "Batch persistence must not consume the legacy callback deadline" - captured: Final = (await callback)["_accounting_capture"] - assert isinstance(captured, CapturedUsage) and not captured.receipt.done() - second: Final = await captured_callback_kwargs(kwargs, response, NOW, NOW) - assert isinstance(second.get("_accounting_capture"), CapturedUsage) - overflow: Final = await captured_callback_kwargs(kwargs, response, NOW, NOW) - assert "_accounting_capture" not in overflow - assert "Batch capture producer capacity exhausted" in caplog.text - if outcome != "cancel": - release.set() - await runtime.stop(0 if outcome == "cancel" else 1) - assert captured.receipt.done() - assert captured.receipt.cancelled() is (outcome == "cancel") - if outcome == "failure": - assert isinstance(captured.receipt.exception(), RuntimeError) - assert "Batch usage capture failed" in caplog.text - assert len(database.rows) == (4 if outcome == "drain" else 0) - with pytest.raises(UsageCaptureError, match="active ledger writer"): - runtime.submit_batch(kwargs, response, NOW, NOW) - finally: - release.set() - await asyncio.gather(callback, return_exceptions=True) - await runtime.stop(1) - - -@pytest.mark.asyncio -async def test_batch_producers_share_byte_capacity_and_release_it_after_drain() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - database: Final = StoredUsage() - release: Final = asyncio.Event() - runtime: Final = UsageCaptureRuntime() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - await release.wait() - assert runtime.meter is not None - assert runtime.meter.pending_bytes <= 65536 - await database.insert(writes) - - runtime.start(persist, UsageCaptureSettings(batch_size=32, max_pending_bytes=65536), load=database.load) - from litellm.batches.batch_utils import batch_cost_and_usage - - parsed: Final = batch_cost_and_usage( - [ - { - "response": { - "status_code": 200, - "body": { - "model": "capture-model", - "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, - }, - } - } - for _ in range(1000) - ], - "openai", - model_info={}, - ) - rows: Final = ({"model": "capture-model", "usage": {"prompt_tokens": 2}},) - response: Final = {"id": "batch", "status": "completed", "usage": parsed.usage} - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - try: - large: Final = asyncio.Queue[asyncio.Future[None]]() - for index in range(100): - try: - captured: Final = runtime.submit_batch(kwargs, {**response, "id": f"batch-{index}"}, NOW, NOW) - except UsageCaptureError: - break - large.put_nowait(captured.receipt) - large_count: Final = large.qsize() - assert 0 < large_count < 100 - assert runtime.meter is not None and 0 < runtime.meter.pending_bytes <= 65536 - release.set() - await asyncio.gather(*(large.get_nowait() for _ in range(large_count))) - assert len(database.rows) == 1001 * large_count - observations: Final = tuple(json.loads(write.envelope) for write in database.rows) - assert sum(row["usage"].get("prompt_tokens", 0) for row in observations) == 2000 * large_count - assert sum(row["usage"].get("completion_tokens", 0) for row in observations) == 1000 * large_count - assert runtime.meter.pending_bytes == 0 - admitted: Final = asyncio.Queue[asyncio.Future[None]]() - for index in range(100): - try: - small: Final = runtime.submit_batch( - kwargs, - {**response, "id": f"small-{index}", "usage": {"batch_usage_by_request": rows}}, - NOW, - NOW, - ) - except UsageCaptureError: - break - admitted.put_nowait(small.receipt) - count: Final = admitted.qsize() - assert 1 < count < 100 - await asyncio.gather(*(admitted.get_nowait() for _ in range(count))) - assert len(database.rows) == 1001 * large_count + 2 * count - assert runtime.meter.pending_bytes == 0 - finally: - release.set() - await runtime.stop(1) - - -@pytest.mark.asyncio -async def test_batch_replay_reserves_the_accepted_pricing_snapshot_size() -> None: - from litellm.proxy.spend_tracking.usage_capture import PricingCapture, UsageBatchWriter, UsageCaptureRuntime - - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - runtime: Final = UsageCaptureRuntime() - runtime.start(writer.write, UsageCaptureSettings(max_pending_bytes=16384), load=database.load) - response: Final = { - "id": "batch", - "status": "completed", - "usage": {"batch_usage_by_request": ({"model": "model", "usage": {"prompt_tokens": 2}},)}, - } - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"model": {"source": "x" * 1200}}})) - try: - initial: Final = runtime.submit_batch({**kwargs, "_accounting_pricing_snapshot": pricing}, response, NOW, NOW) - await asyncio.shield(initial.receipt) - accepted: Final = database.rows - replay: Final = runtime.submit_batch(kwargs, response, NOW, NOW) - await asyncio.shield(replay.receipt) - assert database.rows == accepted and len(accepted) == 2 - assert runtime.meter is not None and runtime.meter.pending_bytes == 0 - finally: - await runtime.stop(1) diff --git a/tests/unit/proxy/spend_tracking/test_usage_meter.py b/tests/unit/proxy/spend_tracking/test_usage_meter.py deleted file mode 100644 index bd3a0b9ed70..00000000000 --- a/tests/unit/proxy/spend_tracking/test_usage_meter.py +++ /dev/null @@ -1,473 +0,0 @@ -import asyncio -import json -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite -from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings -from litellm.types.utils import CallTypes - - -def _write(identity: str) -> UsageWrite: - event: Final = InboundUsageEvent( - event_id=identity, - request_id=identity, - call_type=CallTypes.acompletion, - started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - return UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) - - -class CommitBarrier: - def __init__(self, *, conflict: str | None = None, failure: Exception | None = None) -> None: - self.entered: Final = asyncio.Event() - self.release: Final = asyncio.Event() - self.committed: tuple[str, ...] = () - self.transactions = 0 - self.conflict: Final = conflict - self.failure: Final = failure - - async def write(self, writes: tuple[UsageWrite, ...]) -> None: - self.transactions += 1 - self.entered.set() - await self.release.wait() - if self.failure is not None: - raise self.failure - identities: Final = tuple(write.event.event_id for write in writes) - if self.conflict in identities: - raise AccountingConflictError("Conflicting test event") - self.committed += identities - - -async def _manual_flush(wakeup: asyncio.Event, timeout: float) -> None: - await wakeup.wait() - - -@pytest.mark.asyncio -async def test_receipts_wait_for_one_batched_commit_and_keep_inflight_capacity_reserved() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await database.entered.wait() - assert not first.done() and not second.done() - assert meter.pending_records == 2 - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("three"),)) - database.release.set() - await asyncio.gather(first, second) - await meter.close() - assert database.committed == ("one", "two") - assert database.transactions == 1 - assert (meter.pending_bytes, meter.pending_records) == (0, 0) - - -@pytest.mark.asyncio -async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> None: - database: Final = CommitBarrier() - write: Final = _write("large" * 300) - size: Final = len(write.envelope.encode()) - meter: Final = DurableUsageMeter( - database.write, UsageCaptureSettings(max_pending_bytes=size), wait_for_batch=_manual_flush - ) - reservation: Final = meter.reserve_producer(size // 2, size // 2) - with pytest.raises(UsageCaptureError, match="producer byte capacity"): - meter.reserve_producer(1, size) - with pytest.raises(UsageCaptureError, match="producer byte capacity"): - meter.resize_producer(reservation, size, 1) - assert meter.pending_bytes == size // 2 - meter.release_producer(reservation) - receipt: Final = meter.submit((write,)) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("extra"),)) - assert meter.pending_bytes == size - database.release.set() - await meter.close() - await receipt - assert database.committed == (write.event.event_id,) - with pytest.raises(UsageCaptureError, match="closed"): - meter.submit((_write("after-close"),)) - with pytest.raises(UsageCaptureError, match="producer byte capacity"): - meter.reserve_producer(1, 1) - - -@pytest.mark.asyncio -async def test_conflicting_group_rolls_back_and_valid_neighbor_still_commits() -> None: - database: Final = CommitBarrier(conflict="conflict") - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3)) - rejected: Final = meter.submit((_write("caller"), _write("conflict"))) - accepted: Final = meter.submit((_write("valid"),)) - await database.entered.wait() - database.release.set() - with pytest.raises(UsageCaptureError, match="identity"): - await rejected - await accepted - await meter.close() - assert database.committed == ("valid",) - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_permanent_persistence_failure_reaches_each_receipt_and_releases_capacity() -> None: - failure: Final = ValueError("invalid envelope") - database: Final = CommitBarrier(failure=failure) - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await database.entered.wait() - database.release.set() - results: Final = await asyncio.gather(first, second, return_exceptions=True) - await meter.close() - assert results == [failure, failure] - assert meter.pending_bytes == 0 - - -@pytest.mark.asyncio -async def test_cancelling_record_waiter_does_not_cancel_persistence() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - waiting: Final = asyncio.create_task(meter.record(_write("caller").event)) - await database.entered.wait() - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - database.release.set() - await meter.close() - assert database.committed == ("caller",) - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_shutdown_deadline_rejects_unacknowledged_work_and_releases_capacity() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("in-flight"), _write("in-flight-two"))) - await database.entered.wait() - queued: Final = meter.submit((_write("queued"),)) - await meter.stop(0) - results: Final = await asyncio.gather(first, queued, return_exceptions=True) - assert all(isinstance(result, UsageCaptureError) for result in results) - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - assert database.committed == () - - -@pytest.mark.asyncio -async def test_shutdown_during_database_backoff_releases_reserved_capacity() -> None: - retrying: Final = asyncio.Event() - blocked: Final = asyncio.Event() - - async def unavailable(writes: tuple[UsageWrite, ...]) -> None: - raise ConnectionError("database unavailable") - - async def backoff(delay: float) -> None: - retrying.set() - await blocked.wait() - - meter: Final = DurableUsageMeter(unavailable, UsageCaptureSettings(batch_size=2), retry_delay=backoff) - receipt: Final = meter.submit((_write("one"), _write("two"))) - await retrying.wait() - await meter.stop(0) - with pytest.raises(UsageCaptureError, match="interrupted"): - await receipt - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("wrapped", [False, True]) -async def test_closed_pooled_connection_retains_batch_until_automatic_recovery(wrapped: bool) -> None: - import httpx - from prisma.engine.errors import EngineRequestError - from prisma.http import Response - - database: Final = CommitBarrier() - attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - retrying: Final = asyncio.Event() - recovered: Final = asyncio.Event() - closure: Final = EngineRequestError( - Response(httpx.Response(500)), '{"error_code":"P1017","message":"Server has closed the connection."}' - ) - failure: Final = RuntimeError("ledger transaction failed") if wrapped else closure - if wrapped: - failure.__cause__ = closure - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - attempts.put_nowait(writes) - if attempts.qsize() == 1: - raise failure - await database.write(writes) - - async def backoff(delay: float) -> None: - retrying.set() - await recovered.wait() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), retry_delay=backoff - ) - writes: Final = (_write("caller"), _write("attempt")) - receipt: Final = meter.submit(writes) - retry_signal: Final = asyncio.create_task(retrying.wait()) - reserved: Final = meter.pending_bytes - try: - await asyncio.wait((receipt, retry_signal), return_when=asyncio.FIRST_COMPLETED) - assert not receipt.done(), "A dropped pooled connection must retain admitted usage for retry" - assert (meter.pending_records, meter.pending_bytes) == (2, reserved) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("overflow"),)) - database.release.set() - recovered.set() - await receipt - assert attempts.get_nowait() == attempts.get_nowait() == writes - assert database.committed == ("caller", "attempt") - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - retry_signal.cancel() - await asyncio.gather(retry_signal, return_exceptions=True) - await meter.stop(0) - await asyncio.gather(receipt, return_exceptions=True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("first_commits", [False, True]) -async def test_shutdown_during_conflict_isolation_completes_every_receipt(first_commits: bool) -> None: - isolating: Final = asyncio.Event() - blocked: Final = asyncio.Event() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - if len(writes) > 1: - raise AccountingConflictError("one producer conflicts") - if first_commits and writes[0].event.event_id == "one": - return - isolating.set() - await blocked.wait() - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await isolating.wait() - await meter.stop(0) - results: Final = await asyncio.gather(first, second, return_exceptions=True) - assert results[0] is None if first_commits else isinstance(results[0], UsageCaptureError) - assert isinstance(results[1], UsageCaptureError) - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -async def test_elapsed_flush_deadline_commits_a_partial_batch() -> None: - from litellm.proxy.spend_tracking.usage_meter import _wait_for_batch - - database: Final = CommitBarrier() - database.release.set() - - async def expired_flush(wakeup: asyncio.Event, timeout: float) -> None: - assert timeout == 5.0 - await _wait_for_batch(wakeup, 0) - - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=10), wait_for_batch=expired_flush) - receipt: Final = meter.submit((_write("partial-batch"),)) - try: - await receipt - assert database.committed == ("partial-batch",) - assert database.transactions == 1 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - await meter.stop(0) - - -@pytest.mark.asyncio -async def test_burst_drains_full_batches_without_waiting_for_partial_batch_deadline() -> None: - database: Final = CommitBarrier() - database.release.set() - waits: Final = asyncio.Queue[bool]() - - async def observe_wait(wakeup: asyncio.Event, timeout: float) -> None: - waits.put_nowait(wakeup.is_set()) - await wakeup.wait() - - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2), wait_for_batch=observe_wait) - receipts: Final = tuple(meter.submit((_write(str(index)),)) for index in range(5)) - try: - assert await waits.get() is True - await asyncio.gather(*receipts[:2]) - assert await waits.get() is True, "A full queued batch must not wait for the flush deadline" - await asyncio.gather(*receipts[2:4]) - assert await waits.get() is False - assert database.committed == ("0", "1", "2", "3") - assert meter.pending_records == 1 - assert not receipts[-1].done() - await meter.close() - await receipts[-1] - assert database.committed == ("0", "1", "2", "3", "4") - assert database.transactions == 3 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - await meter.stop(0) - await asyncio.gather(*receipts, return_exceptions=True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sqlstate", ["22P05", "22003", "42P01"]) -async def test_invalid_submission_isolated_without_retrying_systemic_faults(sqlstate: str) -> None: - from prisma.errors import DataError - - database: Final = CommitBarrier() - database.release.set() - attempts: Final = asyncio.Queue[tuple[str, ...]]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - identities: Final = tuple(write.event.event_id for write in writes) - attempts.put_nowait(identities) - if "invalid" in identities or sqlstate == "42P01": - raise DataError( - data={"user_facing_error": {"message": "database rejected input", "meta": {"code": sqlstate}}} - ) - await database.write(writes) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=3)) - rejected: Final = meter.submit((_write("caller"), _write("invalid"))) - neighbor: Final = meter.submit((_write("neighbor"),)) - try: - with pytest.raises(DataError): - await rejected - if sqlstate == "42P01": - with pytest.raises(DataError): - await neighbor - assert attempts.qsize() == 1 - else: - await neighbor - assert database.committed == ("neighbor",) - assert attempts.qsize() == 3 - finally: - await meter.close() - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("dependency_outcome", ["commit", "failure", "cancel"]) -async def test_commit_dependencies_preserve_capacity_and_isolate_failed_dependents(dependency_outcome: str) -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3, max_pending_records=4)) - prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) - neighbor: Final = meter.submit((_write("neighbor"),)) - assert meter.pending_records == 3 - assert not dependent.done() - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("excess-one"), _write("excess-two"))) - if dependency_outcome == "commit": - prerequisite.set_result(None) - elif dependency_outcome == "failure": - prerequisite.set_exception(ValueError("dependency failed")) - else: - prerequisite.cancel() - database.release.set() - await meter.close() - await neighbor - if dependency_outcome == "commit": - await dependent - assert database.committed == ("parent", "final", "neighbor") - else: - with pytest.raises(UsageCaptureError, match="dependency"): - await dependent - assert database.committed == ("neighbor",) - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - - -@pytest.mark.asyncio -async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capacity() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) - await meter.stop(0) - with pytest.raises(UsageCaptureError, match="acknowledgement"): - await dependent - assert not prerequisite.done() - assert database.committed == () - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - prerequisite.set_result(None) - - -@pytest.mark.asyncio -async def test_large_write_sequence_drains_within_byte_and_record_bounds() -> None: - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter( - database.write, - UsageCaptureSettings(batch_size=8, max_pending_records=2, max_pending_bytes=1024), - wait_for_batch=_manual_flush, - ) - writes: Final = tuple(_write(str(index)) for index in range(20)) - pending: Final = meter.submit((_write("pending-a"), _write("pending-b"))) - await meter.record_writes(writes) - await pending - with pytest.raises(UsageCaptureError, match="One usage observation"): - await meter.record_writes((_write("oversized" * 300),)) - await meter.close() - assert database.committed == ("pending-a", "pending-b", *(write.event.event_id for write in writes)) - assert database.transactions > 1 - assert (meter.pending_bytes, meter.pending_records) == (0, 0) - with pytest.raises(UsageCaptureError): - await meter.record_writes((_write("closed"),)) - - -@pytest.mark.asyncio -async def test_cancelling_bounded_capture_preserves_admitted_writes_without_admitting_remaining_rows() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter( - database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2), wait_for_batch=_manual_flush - ) - recording: Final = asyncio.create_task(meter.record_writes(tuple(_write(str(index)) for index in range(8)))) - await database.entered.wait() - recording.cancel() - with pytest.raises(asyncio.CancelledError): - await recording - database.release.set() - await meter.close() - assert database.committed == ("0", "1") - assert meter.pending_bytes == meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_capacity_wait_retries_admission_before_other_producers_finish() -> None: - entered: Final = asyncio.Queue[tuple[str, ...]]() - release: Final = asyncio.Queue[None]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - entered.put_nowait(tuple(write.event.event_id for write in writes)) - await release.get() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(batch_size=2, max_pending_records=3), wait_for_batch=_manual_flush - ) - initial: Final = meter.submit((_write("first"), _write("second"))) - remaining: Final = meter.submit((_write("remaining"),)) - assert await entered.get() == ("first", "second") - recording: Final = asyncio.create_task(meter.record_writes((_write("target"),))) - try: - await asyncio.sleep(0) - release.put_nowait(None) - assert await entered.get() == ("remaining",) - await asyncio.sleep(0) - assert meter.pending_records == 2 - release.put_nowait(None) - assert await entered.get() == ("target",) - release.put_nowait(None) - await recording - await asyncio.gather(initial, remaining) - finally: - recording.cancel() - await asyncio.gather(recording, return_exceptions=True) - release.put_nowait(None) - release.put_nowait(None) - await meter.close() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 13f6f840cb7..8143636a435 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -22,6 +22,7 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid +from litellm.exceptions import APIError, InternalServerError, MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame from litellm.caching.caching import DualCache @@ -5127,6 +5128,138 @@ class TestHandleLLMApiExceptionRetryAfter: assert "All deployments for selected model are in cooldown" not in body["message"] +class TestHandleLLMApiExceptionProviderOverload: + class _RecordingFailureHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.original_exception: Exception | None = None + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.original_exception = original_exception + + async def _invoke( + self, + exc: Exception, + route_type: str, + monkeypatch: pytest.MonkeyPatch, + recording_hook: _RecordingFailureHook | None = None, + ) -> ProxyException: + processor: Final = ProxyBaseLLMRequestProcessing(data={}) + processor.route_type = route_type + callback: Final = recording_hook if recording_hook is not None else self._RecordingFailureHook() + monkeypatch.setattr(litellm, "callbacks", [callback]) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + + with pytest.raises(ProxyException) as raised: + await processor.handle_llm_api_exception( + e=exc, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + return raised.value + + @pytest.mark.parametrize("route_type", ("acompletion", "atext_completion", "aresponses")) + async def test_openai_routes_answer_provider_overload_as_503( + self, route_type: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + exc: Final = InternalServerError( + message="Overloaded", + llm_provider="anthropic", + model="claude", + provider_status_code=529, + ) + + proxy_exc: Final = await self._invoke(exc, route_type, monkeypatch) + + assert proxy_exc.code == "503" + assert proxy_exc.type == "service_unavailable_error" + assert proxy_exc.to_dict()["code"] == "server_is_overloaded" + + @pytest.mark.parametrize( + ("route_type", "provider_status_code"), + (("acompletion", 500), ("anthropic_messages", 529)), + ) + async def test_non_overload_and_non_openai_routes_keep_their_status( + self, route_type: str, provider_status_code: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + exc: Final = InternalServerError( + message="Internal server error", + llm_provider="anthropic", + model="claude", + provider_status_code=provider_status_code, + ) + + proxy_exc: Final = await self._invoke(exc, route_type, monkeypatch) + + assert proxy_exc.code == "500" + assert proxy_exc.type == "internal_server_error" + assert proxy_exc.to_dict()["code"] == "500" + + async def test_openai_route_answers_wrapped_provider_overload_as_503( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + original: Final = InternalServerError( + message="Overloaded", + llm_provider="anthropic", + model="claude", + provider_status_code=529, + ) + exc: Final = MidStreamFallbackError( + message="Fallback failed", + llm_provider="anthropic", + model="claude", + original_exception=original, + ) + + proxy_exc: Final = await self._invoke(exc, "acompletion", monkeypatch) + + assert proxy_exc.code == "503" + assert proxy_exc.type == "service_unavailable_error" + assert proxy_exc.to_dict()["code"] == "server_is_overloaded" + + async def test_anthropic_text_api_overload_answers_503(self, monkeypatch: pytest.MonkeyPatch) -> None: + exc: Final = APIError( + status_code=529, + message="Overloaded", + llm_provider="anthropic_text", + model="claude", + ) + + proxy_exc: Final = await self._invoke(exc, "atext_completion", monkeypatch) + + assert proxy_exc.code == "503" + assert proxy_exc.type == "service_unavailable_error" + assert proxy_exc.to_dict()["code"] == "server_is_overloaded" + + @pytest.mark.parametrize( + ("route_type", "expected_status"), + (("acompletion", 503), ("anthropic_messages", 500)), + ) + async def test_failure_hook_receives_the_answered_overload_status( + self, route_type: str, expected_status: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + exc: Final = InternalServerError( + message="Overloaded", + llm_provider="anthropic", + model="claude", + provider_status_code=529, + ) + recording_hook: Final = self._RecordingFailureHook() + + await self._invoke(exc, route_type, monkeypatch, recording_hook) + + original_exception: Final = recording_hook.original_exception + assert original_exception is not None + assert original_exception.status_code == expected_status + + class TestHandleLLMApiExceptionFramingHeaders: """HTTP-framing headers on the provider exception must be stripped before the proxy builds its own response, or they conflict with the framing the proxy @@ -7755,7 +7888,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-synthetic"}, "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, }, - {"model_name": "chat-model", "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "sk-synthetic"}}, + { + "model_name": "chat-model", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "sk-synthetic"}, + }, ], fallbacks=[{"rate-limited": ["embed-alias", "chat-model"]}], ) @@ -9829,10 +9965,13 @@ def _hook_data(**kwargs: object) -> object: @pytest.mark.asyncio -@pytest.mark.parametrize("prompt, proxy_field, client_field", [ - ({"input": "hi"}, "litellm_metadata", "metadata"), - ({"messages": [{"role": "user", "content": "hi"}]}, "metadata", "litellm_metadata"), -]) +@pytest.mark.parametrize( + "prompt, proxy_field, client_field", + [ + ({"input": "hi"}, "litellm_metadata", "metadata"), + ({"messages": [{"role": "user", "content": "hi"}]}, "metadata", "litellm_metadata"), + ], +) async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape_picks( prompt: dict[str, str | list[dict[str, str]]], proxy_field: str, @@ -9842,7 +9981,9 @@ async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape logging: Final = MagicMock(spec=ProxyLogging) logging.pre_call_hook = AsyncMock(side_effect=_hook_data) - processor: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-x", "metadata": {"client": "cursor"}, **prompt}) + processor: Final = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-x", "metadata": {"client": "cursor"}, **prompt} + ) data, _ = await processor.common_processing_pre_call_logic( request=Request({"type": "http", "method": "POST", "path": "/cursor/chat/completions", "headers": []}), general_settings={}, @@ -9858,6 +9999,7 @@ async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape assert "queue_time_seconds" not in data.get(client_field, {}) assert data["metadata"]["client"] == "cursor" + class TestPerRequestModelGroupAlias: """``router_settings.model_group_alias`` on a key or team has to be resolved by the proxy: the Router resolves aliases from its own shared instance @@ -12693,7 +12835,9 @@ class TestResponseCostFromLoggingObjForCacheHits: proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), ) - cost_headers: Final = {name: value for name, value in headers.items() if name.startswith("x-litellm-response-cost")} + cost_headers: Final = { + name: value for name, value in headers.items() if name.startswith("x-litellm-response-cost") + } assert cost_headers == {"x-litellm-response-cost-output": "0.0012"} def test_a_fresh_zero_cost_v1_messages_call_still_sends_no_cost_header(self): diff --git a/tests/unit/proxy/test_model_deprecations_endpoint.py b/tests/unit/proxy/test_model_deprecations_endpoint.py index 7be90d0d1e4..705f5ec7b39 100644 --- a/tests/unit/proxy/test_model_deprecations_endpoint.py +++ b/tests/unit/proxy/test_model_deprecations_endpoint.py @@ -1,10 +1,12 @@ +import logging +from typing import Final from unittest.mock import MagicMock import pytest from fastapi.testclient import TestClient - import litellm +from litellm._logging import verbose_logger from litellm.proxy import proxy_server from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -140,3 +142,77 @@ def test_cost_map_successor_is_reported_unless_the_deployment_names_its_own(clie "dated-by-hand": "fresh-model", "orphan-alias": None, } + + +def test_unmatched_successor_logs_below_warning_through_endpoint(client_with_cost_map_successors, caplog, monkeypatch): + monkeypatch.setattr( + proxy_server.llm_router.get_model_list, + "side_effect", + lambda model_name=None, **kwargs: ( + proxy_server.llm_router.get_model_list.return_value if model_name is None else [] + ), + ) + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + response = client_with_cost_map_successors.get("/model/deprecations") + + assert response.status_code == 200, response.text + successor_records = tuple( + record + for record in caplog.records + if record.name == verbose_logger.name and "successor" in record.getMessage() + ) + assert successor_records + assert all(record.levelno < logging.WARNING for record in successor_records) + + +def test_mixed_successor_group_logs_below_warning_through_endpoint( + client_with_cost_map_successors: TestClient, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = proxy_server.llm_router + models: Final[tuple[dict[str, object], ...]] = ( + { + "model_name": "mixed-successor-group", + "litellm_params": {"model": "retiring-model"}, + "model_info": {"id": "10", "successor_model": "successor-a"}, + }, + { + "model_name": "mixed-successor-group", + "litellm_params": {"model": "retiring-model"}, + "model_info": {"id": "11", "successor_model": "successor-b"}, + }, + { + "model_name": "successor-a", + "litellm_params": {"model": "fresh-model"}, + "model_info": {"id": "12"}, + }, + { + "model_name": "successor-b", + "litellm_params": {"model": "fresh-model"}, + "model_info": {"id": "13"}, + }, + ) + + def filtered_model_list( + model_name: str | None = None, **_kwargs: object + ) -> tuple[dict[str, object], ...]: + return models if model_name is None else tuple(model for model in models if model["model_name"] == model_name) + + monkeypatch.setattr(router.get_model_list, "side_effect", filtered_model_list) + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + response = client_with_cost_map_successors.get("/model/deprecations") + + assert response.status_code == 200, response.text + successor_records: Final = tuple( + record + for record in caplog.records + if record.name == verbose_logger.name and "successor" in record.getMessage() + ) + assert successor_records + assert all( + "names several successors" in record.getMessage() and record.levelno < logging.WARNING + for record in successor_records + ) diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index a50db09128f..dae9ffcf7fd 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -268,76 +268,6 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm token = MASTER_KEY -@pytest.mark.asyncio -async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from datetime import datetime - - from litellm.litellm_core_utils.logging_worker import LoggingWorker - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from litellm.responses.streaming_iterator import spawn_logging_task - - worker: Final = LoggingWorker() - started: Final = asyncio.Event() - release: Final = asyncio.Event() - first_flush: Final = asyncio.Event() - pending: Final[asyncio.Queue[str]] = asyncio.Queue() - persisted: Final[asyncio.Queue[str]] = asyncio.Queue() - - captured: Final[asyncio.Queue[str]] = asyncio.Queue() - - async def persist_usage(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - captured.put_nowait(write.event.event_id) - - capture: Final = UsageCaptureRuntime() - capture.start(persist_usage, UsageCaptureSettings(flush_interval_seconds=10)) - monkeypatch.setattr(litellm.proxy.proxy_server, "usage_capture", capture) - - async def enqueue_spend() -> None: - assert capture.meter is not None - capture_callback( - capture.meter, - {"call_type": "acompletion", "litellm_call_id": "shutdown-usage", "litellm_params": {}}, - {"id": "shutdown-attempt", "usage": {"total_tokens": 3}}, - datetime(2026, 1, 1), - datetime(2026, 1, 1), - ) - pending.put_nowait("accounted request") - - async def callback() -> None: - started.set() - await release.wait() - spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting") - - async def flush_spend() -> None: - while not pending.empty(): - persisted.put_nowait(pending.get_nowait()) - first_flush.set() - - worker.ensure_initialized_and_enqueue(callback()) - await started.wait() - shutdown: Final = asyncio.create_task( - litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) - ) - try: - await first_flush.wait() - waited_for_callback: Final = not shutdown.done() - release.set() - await shutdown - assert waited_for_callback, "Shutdown returned while accounting callbacks were still running" - assert persisted.get_nowait() == "accounted request" - assert pending.empty() - assert (captured.get_nowait(), captured.get_nowait()) == ("shutdown-usage:inbound", "shutdown-attempt") - assert capture.meter is None - finally: - release.set() - await worker.flush() - await worker.stop() - await capture.stop(0) @pytest.mark.asyncio @@ -4377,276 +4307,11 @@ async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> await worker.stop(discard_pending=True) -@pytest.mark.asyncio -@pytest.mark.parametrize("unsupported", ["database", "collector", "remote", "free_tier"]) -async def test_capture_enablement_rejects_unrecoverable_writer_modes( - unsupported: str, monkeypatch: pytest.MonkeyPatch -) -> None: - import litellm - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - monkeypatch.setattr(proxy_server, "usage_capture", UsageCaptureRuntime()) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {}}) - monkeypatch.setattr(proxy_server, "prisma_client", None if unsupported == "database" else object()) - monkeypatch.setattr(litellm, "apply_search_free_tier", unsupported == "free_tier") - monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "true" if unsupported == "collector" else "false") - monkeypatch.delenv("SPEND_LOGS_URL", raising=False) - if unsupported == "remote": - monkeypatch.setenv("SPEND_LOGS_URL", "https://writer.invalid") - with pytest.raises(ValueError, match="requires local Prisma"): - proxy_server._configure_usage_capture() - assert proxy_server.usage_capture.meter is None -@pytest.mark.asyncio -async def test_capture_configuration_requires_restart_without_replacing_the_active_writer( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.utils import PrismaClient - - runtime: Final = UsageCaptureRuntime() - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - with pytest.raises(ValueError, match="requires accounting_capture"): - proxy_server._configure_usage_capture() - client: Final = PrismaClient("postgresql://localhost/unused", ProxyLogging(UserApiKeyCache())) - monkeypatch.setattr(proxy_server, "prisma_client", client) - monkeypatch.setattr(litellm, "apply_search_free_tier", False) - monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false") - monkeypatch.delenv("SPEND_LOGS_URL", raising=False) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 4}}) - try: - proxy_server._configure_usage_capture() - active: Final = runtime.meter - assert active is not None - assert active.settings.batch_size == 4 - proxy_server._configure_usage_capture() - assert runtime.meter is active - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 6}}) - with pytest.raises(ValueError, match="Changing accounting_capture requires a worker restart"): - proxy_server._configure_usage_capture() - monkeypatch.setattr(proxy_server, "general_settings", {}) - with pytest.raises(ValueError, match="Disabling accounting_capture requires a worker restart"): - proxy_server._configure_usage_capture() - assert runtime.meter is active - assert active.settings.batch_size == 4 - finally: - await runtime.stop(1) - assert runtime.meter is None -@pytest.mark.asyncio -@pytest.mark.parametrize("sync_ack", [False, True]) -@pytest.mark.parametrize( - ("path", "payload"), - [ - ( - "/v1/chat/completions", - {"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True}, - ), - ("/v1/responses", {"model": "openai/gpt-5-mini", "input": "hi", "background": True}), - ( - "/v1/messages", - {"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1}, - ), - ], -) -async def test_sync_ack_preserves_async_request_support( - sync_ack: bool, - path: str, - payload: dict[str, object], - monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, -) -> None: - from fastapi import Request - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - - persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - persisted.put_nowait(write) - - async def authorize(request: Request) -> UserAPIKeyAuth: - return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path) - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key") - upstream: Final = { - "id": "provider-result", - "object": "chat.completion", - "created": 1, - "model": "gpt-5-mini", - "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, - } - if payload.get("stream"): - chunk: Final = { - **upstream, - "object": "chat.completion.chunk", - "choices": [{"index": 0, "finish_reason": "stop", "delta": {"role": "assistant", "content": "hi"}}], - } - respx_mock.post().respond( - 200, text=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n", headers={"content-type": "text/event-stream"} - ) - elif path == "/v1/responses": - respx_mock.post().respond( - 200, - json={ - "id": "resp_background", - "object": "response", - "created_at": 1, - "model": "gpt-5-mini", - "status": "in_progress", - "background": True, - "output": [], - "usage": None, - }, - ) - else: - respx_mock.post(url__regex=r".*/responses$").respond( - 200, - json={ - "id": "resp_message", - "object": "response", - "created_at": 1, - "model": "gpt-5-mini", - "status": "completed", - "output": [ - { - "id": "msg_reply", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "hi", "annotations": []}], - } - ], - "usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, - }, - ) - respx_mock.post(url__regex=r".*/chat/completions$").respond(200, json=upstream) - runtime: Final = UsageCaptureRuntime() - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, "general_settings", {"accounting_sync_ack": sync_ack, "accounting_sync_ack_timeout_seconds": 0} - ) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr( - proxy_server, - "llm_router", - litellm.Router( - model_list=[ - { - "model_name": payload["model"], - "litellm_params": {"model": payload["model"], "api_key": "test-provider-key"}, - } - ], - num_retries=0, - ), - ) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache())) - monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()]) - monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize) - runtime.start(persist, UsageCaptureSettings()) - try: - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test" - ) as client: - response: Final = await client.post(path, json=payload) - finally: - await runtime.stop(1) - assert response.status_code == 200, response.text - assert len(respx_mock.calls) == 1 - if payload.get("background"): - assert response.json()["status"] == "in_progress" - else: - assert "hi" in response.text - assert not persisted.empty(), "Terminal usage must still be captured" - -@pytest.mark.asyncio -@pytest.mark.parametrize("sync_ack", [False, True]) -async def test_text_completion_provider_rejection_captures_unknown_usage( - sync_ack: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter -) -> None: - from fastapi import Request - - from litellm.accounting import Outcome - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - - persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - persisted.put_nowait(write) - - async def authorize(request: Request) -> UserAPIKeyAuth: - return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path) - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key") - respx_mock.post().respond( - 404, - json={ - "error": {"message": "The model is unavailable", "type": "invalid_request_error", "code": "model_not_found"} - }, - ) - runtime: Final = UsageCaptureRuntime() - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack}) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr( - proxy_server, - "llm_router", - litellm.Router( - model_list=[ - { - "model_name": "instruct", - "litellm_params": {"model": "openai/gpt-3.5-turbo-instruct", "api_key": "test-provider-key"}, - } - ], - num_retries=0, - ), - ) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache())) - monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()]) - monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize) - runtime.start(persist, UsageCaptureSettings()) - try: - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test" - ) as client: - response: Final = await client.post( - "/v1/completions", json={"model": "instruct", "prompt": "hi", "max_tokens": 1} - ) - finally: - await runtime.stop(1) - assert response.status_code == 404, response.text - assert len(respx_mock.calls) == 1 - assert persisted.qsize() == 2 - inbound: Final = persisted.get_nowait() - outbound: Final = persisted.get_nowait() - assert outbound.event.outcome is Outcome.FAILED - assert outbound.event.call_type.value == "atext_completion" - assert outbound.event.metrics is None - assert inbound.event.request_id == outbound.event.request_id - envelope: Final = json.loads(outbound.envelope) - assert envelope["usage"] is None - assert envelope["usage_source"] == "unknown" - assert envelope["pricing"]["reported_cost"] is None @pytest.mark.parametrize( diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 9dd15649d3e..62d8db6f810 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -16589,3 +16589,63 @@ async def test_initialize_prints_the_banner_as_one_json_record_in_json_mode( banner_records: Final = [record for record in records if record["message"] == LITELLM_BANNER] assert len(banner_records) == 1, records assert banner_records[0]["level"] == "INFO" + + +@contextlib.contextmanager +def _accounting_window_env(cached: Mapping[str, object], client: MagicMock) -> Iterator[None]: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache: Final = UserApiKeyCache() + for key, value in cached.items(): + cache.set_cache(key=key, value=value) + originals: Final = (proxy_server_module.user_api_key_cache, proxy_server_module.prisma_client) + proxy_server_module.user_api_key_cache = cache + proxy_server_module.prisma_client = client + try: + yield + finally: + proxy_server_module.user_api_key_cache, proxy_server_module.prisma_client = originals + + +@pytest.mark.asyncio +async def test_accounting_window_keys_cover_the_stored_key_and_the_cached_team(): + key_reset: Final = datetime(2026, 11, 13, tzinfo=timezone.utc) + team_reset: Final = datetime(2026, 11, 6, tzinfo=timezone.utc) + key_row: Final = _StoredRow( + budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": key_reset.isoformat()}] + ) + team_obj: Final = MagicMock() + team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": team_reset.isoformat()}] + + with _accounting_window_env({"team_id:team-1": team_obj}, _prisma_with_rows(keys={"hashed-key": key_row})): + windows = await proxy_server_module.accounting_window_keys("hashed-key", "team-1") + + assert windows == ( + { + "entity_type": "key", + "entity_id": "hashed-key", + "window_duration": "30d", + "window_start": (key_reset - timedelta(days=30)).isoformat(), + }, + { + "entity_type": "team", + "entity_id": "team-1", + "window_duration": "7d", + "window_start": (team_reset - timedelta(days=7)).isoformat(), + }, + ) + + +@pytest.mark.asyncio +async def test_accounting_window_keys_remember_a_key_without_windows_but_retry_a_failed_read(): + client: Final = _prisma_with_rows(keys={"no-windows": _StoredRow(budget_limits=[])}) + + with _accounting_window_env({}, client): + assert await proxy_server_module.accounting_window_keys("no-windows", None) == () + assert await proxy_server_module.accounting_window_keys("no-windows", None) == () + assert client.db.litellm_verificationtoken.find_unique.await_count == 1 + + client.db.litellm_verificationtoken.find_unique.side_effect = RuntimeError("database unavailable") + assert await proxy_server_module.accounting_window_keys("unreadable", None) == () + assert await proxy_server_module.accounting_window_keys("unreadable", None) == () + assert client.db.litellm_verificationtoken.find_unique.await_count == 3 diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index d28619c0938..3b000ba741f 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -3735,42 +3735,6 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) -@pytest.mark.asyncio -async def test_a_billed_responses_route_failure_captures_usage_from_the_newest_compaction_item() -> None: - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from litellm.proxy.utils import _failure_fields_to_lift - - async def persist(_: tuple[UsageWrite, ...]) -> None: - return None - - def captured_prompt_tokens(request_input: list[dict[str, str]]) -> int: - logging_obj: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params()) - logging_obj.model_call_details["call_type"] = "aresponses" - logging_obj.model_call_details["messages"] = request_input - lifted: Final = _failure_fields_to_lift( - {"litellm_logging_obj": logging_obj}, provider_billed=True, route_call_type="aresponses" - ) - snapshot: Final = lifted["_accounting_usage_snapshot"] - assert isinstance(snapshot, UsageCaptureObservation) and snapshot.usage_json is not None - return json.loads(snapshot.usage_json)["prompt_tokens"] - - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings()) - try: - sent: Final = captured_prompt_tokens( - [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] - ) - since_compaction: Final = captured_prompt_tokens( - [COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION] - ) - finally: - await runtime.stop(1) - - assert sent == since_compaction @pytest.mark.parametrize( @@ -4405,57 +4369,3 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None: ProxyException(message="x", type="x", param=None, code="429") ) is True assert _is_client_error_exception(ValueError("x")) is False - - -@pytest.mark.asyncio -@pytest.mark.parametrize("tokens", [None, 0, 3]) -async def test_failure_capture_preserves_prepared_usage_observation(tokens: int | None) -> None: - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - - timestamp: Final = real_datetime.datetime(2026, 1, 1) - committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - for write in writes: - committed.put_nowait(write) - - logger: Final = Logging( - model=FAILURE_USAGE_MODEL, - messages=[], - stream=False, - call_type="acompletion", - start_time=timestamp, - litellm_call_id="post-response-failure", - function_id="post-response-failure", - ) - logger.update_environment_variables( - model=FAILURE_USAGE_MODEL, custom_llm_provider="openai", litellm_params={}, optional_params={} - ) - logger.model_call_details["response_cost"] = 0.2 - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings()) - try: - prepared: Final = logger.prepare_response_cost( - {"id": "provider", "usage": None if tokens is None else {"prompt_tokens": tokens, "total_tokens": tokens}}, - timestamp, - ) - assert prepared is not None - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert lifted["combined_usage_object"] is logger.model_call_details["combined_usage_object"] - assert lifted["response_cost"] == 0.2 - assert runtime.meter is not None - capture_callback( - runtime.meter, - {**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"}, - None, - timestamp, - timestamp, - ) - finally: - await runtime.stop(1) - committed.get_nowait() - envelope: Final = json.loads(committed.get_nowait().envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == tokens - assert envelope["usage_source"] == ("unknown" if tokens is None else "measured") diff --git a/tests/unit/proxy/test_update_spend.py b/tests/unit/proxy/test_update_spend.py index 571341e85e7..e2c027408c1 100644 --- a/tests/unit/proxy/test_update_spend.py +++ b/tests/unit/proxy/test_update_spend.py @@ -351,3 +351,125 @@ async def test_update_spend_logs_drops_row_prisma_cannot_encode_and_keeps_the_re assert sorted(written) == ["good-1", "good-2"] assert prisma_client.spend_log_transactions == [] proxy_logging_obj.failure_handler.assert_not_called() + + +def _raw_query_error(sqlstate: str): + from prisma.errors import RawQueryError + + return RawQueryError( + { + "user_facing_error": { + "error_code": "P2010", + "meta": {"code": sqlstate, "message": f"ERROR: sqlstate {sqlstate}"}, + } + } + ) + + +class _RawInsertRepository: + def __init__(self, sqlstate: str): + self.sqlstate = sqlstate + self.written: list[str] = [] + + async def insert_with_deliveries(self, rows): + if any(row["request_id"] == "rejected" for row in rows): + raise _raw_query_error(self.sqlstate) + self.written.extend(str(row["request_id"]) for row in rows) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sqlstate", ["22P02", "23502", "54000"]) +async def test_raw_insert_rejecting_one_row_drops_only_that_row(sqlstate): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _RawInsertRepository(sqlstate) + rows = [{"request_id": "good-1"}, {"request_id": "rejected"}, {"request_id": "good-2"}] + + await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) + + assert sorted(repo.written) == ["good-1", "good-2"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sqlstate", ["42P01", "40P01"]) +async def test_systemic_raw_insert_failure_is_retried_not_dropped(sqlstate): + from prisma.errors import RawQueryError + + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _RawInsertRepository(sqlstate) + rows = [{"request_id": "good-1"}, {"request_id": "rejected"}] + + with pytest.raises(RawQueryError): + await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) + + assert repo.written == [] + + +class _RescueFailsRepository: + def __init__(self, rescue_error: Exception): + self.rescue_error = rescue_error + + async def insert_with_deliveries(self, rows): + if any(row.get("accounting_log_disabled") for row in rows): + raise self.rescue_error + if any(row["request_id"] == "rejected" for row in rows): + raise _raw_query_error("22P02") + + +@pytest.mark.asyncio +async def test_a_transient_failure_while_keeping_a_dropped_rows_charge_is_retried_not_lost(): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + outage = _raw_query_error("40P01") + repo = _RescueFailsRepository(outage) + rows = [{"request_id": "rejected", "accounting_payload": {"spend": 1.0}}] + + with pytest.raises(type(outage)): + await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) + + +@pytest.mark.asyncio +async def test_a_rejected_charge_rescue_is_logged_and_the_flush_continues(): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _RescueFailsRepository(_raw_query_error("22P02")) + rows = [{"request_id": "rejected", "accounting_payload": {"spend": 1.0}}] + + assert await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) == 9 + + +class _StatementValidatingRepository: + def __init__(self): + self.written = [] + self.charges_kept = [] + + async def insert_with_deliveries(self, rows): + from litellm.repositories.spend_log_insert import spend_log_insert_statement + + if all(row.get("accounting_log_disabled") for row in rows): + self.charges_kept.extend(row["request_id"] for row in rows) + return len(rows) + spend_log_insert_statement(rows) + self.written.extend(row["request_id"] for row in rows) + return len(rows) + + +@pytest.mark.asyncio +async def test_a_row_the_insert_statement_rejects_is_dropped_alone_and_keeps_its_charge(): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _StatementValidatingRepository() + rows = [ + {"request_id": "good", "call_type": "acompletion", "messages": "[]"}, + { + "request_id": "unparseable", + "call_type": "acompletion", + "messages": "not json", + "accounting_payload": {"spend": 1.0}, + }, + ] + + assert await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) < 10 + assert repo.written == ["good"] + assert repo.charges_kept == ["unparseable"] diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py index fc7ae531c34..293f2d49f43 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -23,9 +23,9 @@ import pytest from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, + drain_spend_logs_queue, monitor_spend_logs_queue, raise_failed_update_spend_exception, - drain_spend_logs_queue, recover_parked_spend_logs, update_daily_tag_spend, update_spend, @@ -985,6 +985,85 @@ async def test_update_spend_logs_job_requeues_whole_batch_when_postgres_is_out_o } == {"create_many_calls": 1, "queue_after": ["r1", "r2", "r3"], "backoff_sleeps": []} +def _postgres_restarting() -> Exception: + """prisma's shape for Postgres SQLSTATE 57P03 while the server runs crash recovery after + a restart: a base ``DataError`` carrying only the connector message.""" + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: the database system is starting up", + "backtrace": None, + } + } + ) + + +def _postgres_unreachable() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Can't reach database server at `db.internal`:`5432`", + "error_code": "P1001", + } + } + ) + + +def _postgres_terminated_session() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: terminating connection due to " + "administrator command", + "meta": {"code": "57P01", "message": "terminating connection due to administrator command"}, + } + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_terminated_session()], ids=["57P03", "57P01"]) +async def test_update_spend_logs_job_keeps_every_row_while_postgres_restarts( + error: Exception, + mock_prisma_client: MagicMock, + make_spend_log_row: Callable[..., dict[str, object]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A restarting server refuses or cuts every statement, so the batch is not bisected + into per-row drops: it is retried whole and then goes back to the queue head.""" + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + proxy_logging: Final = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [ + make_spend_log_row(request_id="r1"), + make_spend_log_row(request_id="r2"), + make_spend_log_row(request_id="r3"), + make_spend_log_row(request_id="r4"), + ] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=error) + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert { + "create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count, + "queue_after": [row["request_id"] for row in mock_prisma_client.spend_log_transactions], + } == {"create_many_calls": 4, "queue_after": ["r1", "r2", "r3", "r4"]} + + def _tool_usage_transaction(request_id: str) -> object: from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -1177,6 +1256,53 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out } == {"queue_after": [first, second], "autorouter_queue_after": 1} +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"]) +async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes( + error: Exception, + mock_prisma_client: MagicMock, +) -> None: + """A server that is restarting or unreachable never saw the rollup batch, so the tool + usage increments are kept for the next flush, and the job stops before draining the + auto-router queue into the same outage.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=error) + mock_prisma_client.spend_log_transactions = [] + first, second = _tool_usage_transaction("r1"), _tool_usage_transaction("r2") + mock_prisma_client.tool_usage_transactions = [first, second] + mock_prisma_client.autorouter_turn_transactions = [MagicMock(name="autorouter_turn")] + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert { + "queue_after": mock_prisma_client.tool_usage_transactions, + "autorouter_queue_after": len(mock_prisma_client.autorouter_turn_transactions), + } == {"queue_after": [first, second], "autorouter_queue_after": 1} + + +@pytest.mark.asyncio +async def test_rollups_drop_a_batch_whose_session_the_server_terminated( + mock_prisma_client: MagicMock, +) -> None: + """A terminated session (57P01) may have committed before it was cut, and the rollup + increments are not idempotent, so the batch is dropped rather than counted twice.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=_postgres_terminated_session()) + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [_tool_usage_transaction("r1")] + + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_model_usage_flush_still_drops_a_batch_that_may_have_applied( mock_prisma_client: MagicMock, diff --git a/tests/unit/repositories/test_spend_log_insert.py b/tests/unit/repositories/test_spend_log_insert.py new file mode 100644 index 00000000000..a099551086b --- /dev/null +++ b/tests/unit/repositories/test_spend_log_insert.py @@ -0,0 +1,31 @@ +from datetime import datetime, timezone +import json + +import pytest + +from litellm.repositories.spend_log_insert import spend_log_insert_statement + + +@pytest.mark.parametrize( + "rows", + [(), ({"request_id": "request", "unknown_column": 1},), ({"request_id": "first"}, {"spend": 1})], +) +def test_insert_rejects_empty_unknown_or_inconsistent_columns(rows): + with pytest.raises(ValueError, match="Spend log insert requires"): + spend_log_insert_statement(rows) + + +def test_insert_keeps_json_and_timestamps_in_bound_data(): + timestamp = datetime(2026, 11, 2, tzinfo=timezone.utc) + request_id = "request'); DROP TABLE unrelated; --" + statement, encoded = spend_log_insert_statement(({ + "request_id": request_id, "startTime": timestamp, "metadata": '{"project":"project"}', + "request_tags": '["tag"]', "messages": '[{"role":"user","content":"hello"}]', + },)) + row = json.loads(encoded)[0] + assert request_id not in statement + assert row["request_id"] == request_id + assert row["startTime"] == timestamp.isoformat() + assert row["metadata"] == {"project": "project"} + assert row["request_tags"] == ["tag"] + assert row["messages"] == [{"role": "user", "content": "hello"}] diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index c0583566531..d24d0598d87 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -7,6 +7,7 @@ the callback raise before any spend was recorded, so those budgets never moved. """ from datetime import datetime, timezone +from functools import partial from typing import Final import litellm @@ -17,6 +18,7 @@ from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.utils import BudgetConfig +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG import sys, os, asyncio, time, random @@ -38,6 +40,7 @@ def _success_kwargs( call_type: str = "aresponses", response_cost: float = 0.25, model_id: str = "deployment-1", + guardrail_information: tuple[dict[str, object], ...] = (), ) -> dict[str, object]: provider_params: Final[dict[str, str]] = ( {} if provider_in_litellm_params is None else {"custom_llm_provider": provider_in_litellm_params} @@ -51,6 +54,7 @@ def _success_kwargs( "response_cost": response_cost, "model_id": model_id, "custom_llm_provider": provider_in_payload, + "guardrail_information": guardrail_information, }, } @@ -462,6 +466,7 @@ def _completed_budget_batch(batch_id: str): input_file_id="file-budget", object="batch", status="completed", + output_file_id="file-budget-output", ) @@ -498,6 +503,130 @@ async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) + +@pytest.mark.asyncio +async def test_a_batch_retrieve_before_its_cost_is_final_charges_no_router_budget_even_for_its_guardrail(): + from litellm.types.utils import LiteLLMBatch + + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "batch-budget", + "litellm_params": {"model": "openai/batch-budget", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + retrieve: Final = partial( + _success_kwargs, + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type="aretrieve_batch", + guardrail_information=({"guardrail_name": "priced", "guardrail_cost": 0.01},), + ) + + def polled(status: str, output_file_id: str | None = None) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-running", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-budget", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + for response_cost, batch in ( + (0.01, polled("in_progress")), + (0.01, polled("completed")), + (0.26, polled("completed", "file-out")), + (0.26, polled("completed", "file-out")), + ): + await limiter.async_log_success_event(retrieve(response_cost=response_cost), batch, None, None) + + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") + assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.26, 0.26)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ("avideo_status", "video_status")) +async def test_video_polls_charge_each_router_budget_once(call_type): + from litellm.types.videos.main import VideoObject + + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "video-budget", + "litellm_params": {"model": "openai/sora-2", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + kwargs: Final = { + **_success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type), + POLLED_VIDEO_STATUS_KWARG: "completed", + } + video: Final = VideoObject(id="video-first", object="video", status="completed") + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") + + await asyncio.gather(*(limiter.async_log_success_event(kwargs, video, None, None) for _ in range(10))) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.25, 0.25) + + await limiter.async_log_success_event( + kwargs, VideoObject(id="video-second", object="video", status="completed"), None, None + ) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.5, 0.5) + + +@pytest.mark.asyncio +async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_router_for_the_video(): + from litellm.types.videos.main import VideoObject + + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "video-budget", + "litellm_params": {"model": "openai/sora-2", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + priced: Final = ({"guardrail_name": "priced", "guardrail_cost": 0.01},) + poll: Final = partial( + _success_kwargs, + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type="avideo_status", + guardrail_information=priced, + ) + + await limiter.async_log_success_event( + {**poll(response_cost=0.01), POLLED_VIDEO_STATUS_KWARG: "in_progress"}, + VideoObject(id="video-first", object="video", status="in_progress"), + None, + None, + ) + for _ in range(3): + await limiter.async_log_success_event( + {**poll(response_cost=0.26), POLLED_VIDEO_STATUS_KWARG: "completed"}, + VideoObject(id="video-first", object="video", status="completed"), + None, + None, + ) + + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") + assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.29, 0.29)) + @pytest.mark.asyncio async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): from fakeredis import FakeServer, FakeRedis as FakeSyncRedis diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 5b3af388ae5..f156a5ee405 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,9 +1,16 @@ import asyncio +import contextlib +import contextvars +import datetime import json import logging +import threading +from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx +import openai import pytest import respx @@ -11,11 +18,19 @@ import litellm from litellm import Router from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler -from litellm.types.utils import ModelResponse +from litellm.router_strategy.least_busy import ( + IN_FLIGHT_COUNT_TTL_SECONDS, + SELECTION_RESERVATION_TTL_SECONDS, + LeastBusyLoggingHandler, + owning_picks, + picks_reserve, + release_unstarted_picks, +) +from litellm.types.utils import CallTypes, ModelResponse GROUP: Final = "least-busy-group" DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} @@ -23,8 +38,11 @@ DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] -def _call_kwargs(deployment_id: str) -> dict[str, object]: - return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} +def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]: + return { + "litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}, + "litellm_call_id": call_id, + } class SharedRedisCounters: @@ -58,6 +76,20 @@ class SharedRedisCounters: return self.batch_get_counts(key_list) +class SlowRedisRead(SharedRedisCounters): + """Answers a count read with the counts from when the read started, after letting other work run.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + counts_when_asked: Final = self.batch_get_counts(key_list) + await asyncio.sleep(0) + self.while_reading() + return counts_when_asked + + def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double return LeastBusyLoggingHandler(router_cache=cache) @@ -81,6 +113,16 @@ async def test_worker_routes_around_a_request_another_worker_started() -> None: assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +@pytest.mark.asyncio +async def test_a_request_that_starts_while_redis_is_read_still_counts_on_its_deployment() -> None: + redis: Final = SlowRedisRead() + worker: Final = _worker(redis) + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + redis.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + def test_sync_pick_reads_the_shared_counts() -> None: shared: Final = SharedRedisCounters() streaming_worker: Final = _worker(shared) @@ -98,6 +140,44 @@ def test_sync_pick_reads_the_shared_counts() -> None: assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]: + return { + "litellm_params": { + "metadata": {"tags": ["caller-tag"]}, + "litellm_metadata": {"model_group": GROUP}, + "model_info": {"id": deployment_id}, + }, + "litellm_call_id": "call-1", + } + + +def test_a_request_routed_through_litellm_metadata_counts_while_in_flight() -> None: + shared: Final = SharedRedisCounters() + calling_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + assert calling_worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + calling_worker.log_pre_api_call(model="m", messages=[], kwargs=_routed_through_litellm_metadata("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + calling_worker.log_success_event(_routed_through_litellm_metadata("dep-a"), None, None, None) + + assert calling_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_closed_realtime_session_leaves_its_deployment_idle() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + session_event: Final = {**_routed_through_litellm_metadata("dep-a"), "call_type": CallTypes.arealtime.value} + + for _ in range(5): + worker.log_pre_api_call(model="m", messages=[], kwargs=session_event) + worker.log_success_event(session_event, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that stuck count an hour after the key was created, which only works while nothing writes the TTL @@ -259,12 +339,16 @@ ROUTER_GROUP: Final = "least-busy-router-group" ROUTER_DEPLOYMENT_IDS: Final = ("1", "2", "3") -def _least_busy_router() -> Router: +def _least_busy_router(group: str = ROUTER_GROUP, *, use_in_pass_through: bool = False) -> Router: return Router( model_list=[ { - "model_name": ROUTER_GROUP, - "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": f"key-{deployment_id}"}, + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": f"key-{deployment_id}", + "use_in_pass_through": use_in_pass_through, + }, "model_info": {"id": deployment_id}, } for deployment_id in ROUTER_DEPLOYMENT_IDS @@ -391,3 +475,575 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts assert response._hidden_params["model_id"] == "1" assert route.calls.last.request.headers["authorization"] == "Bearer key-1" assert _router_counts(router) == seeded + + +@pytest.mark.asyncio +async def test_router_spreads_requests_that_start_together( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + burst: Final = 2 * len(ROUTER_DEPLOYMENT_IDS) + all_started: Final = asyncio.Barrier(burst) + + async def completion_after_all_started(_: httpx.Request) -> httpx.Response: + await all_started.wait() + return httpx.Response( + 200, + json={ + "id": "chatcmpl-least-busy", + "object": "chat.completion", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=completion_after_all_started + ) + router: Final = _least_busy_router() + + await asyncio.gather( + *(router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(burst)) + ) + + one_per_deployment: Final = tuple(f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS) + assert sorted(call.request.headers["authorization"] for call in route.calls) == sorted(one_per_deployment * 2) + + +class _ManualClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +def test_a_reservation_that_never_reaches_the_api_call_expires() -> None: + clock: Final = _ManualClock() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=clock) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = 1.0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + clock.now = 2.0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.5 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +def test_a_fresh_reservation_outlives_the_ttl_of_the_first_one_on_its_deployment() -> None: + clock: Final = _ManualClock() + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + worker: Final = LeastBusyLoggingHandler(router_cache=cache, clock=clock) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(kwargs=_call_kwargs("dep-a"), response_obj=None, start_time=None, end_time=None) + clock.now = SELECTION_RESERVATION_TTL_SECONDS - 0.1 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.1 + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> None: + worker: Final = _worker(None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None: + cache: Final = DualCache() + reservations: Final = InMemoryCache() + picker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + registered: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + registered.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +class _StoreWhereThreadsReadBeforeWriting: + def __init__(self) -> None: + self._store: Final = InMemoryCache() + self._first_reads: Final = threading.Barrier(2, timeout=0.5) + self._thread_state: Final = threading.local() + self._meeting: Final = threading.Event() + + @contextlib.contextmanager + def first_reads_meet(self) -> Generator[None, None, None]: + self._meeting.set() + yield + self._meeting.clear() + + def get_cache(self, key: str) -> object: + value: Final = self._store.get_cache(key) + if self._meeting.is_set() and not getattr(self._thread_state, "has_read", False): + self._thread_state.has_read = True + with contextlib.suppress(threading.BrokenBarrierError): + self._first_reads.wait() + return value + + def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: + self._store.set_cache(key, value, ttl=ttl) + + def delete_cache(self, key: str) -> None: + self._store.delete_cache(key) + + +def test_two_threads_picking_at_once_reserve_different_deployments() -> None: + store: Final = _StoreWhereThreadsReadBeforeWriting() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) + + with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: + picks: Final = tuple(pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2))) + + assert sorted(pick["model_info"]["id"] for pick in picks if pick is not None) == ["dep-a", "dep-b"] + + +def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> None: + store: Final = _StoreWhereThreadsReadBeforeWriting() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b")) + + with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: + reserved: Final = pool.submit(worker.get_available_deployments, GROUP, HEALTHY, True) + started: Final = pool.submit(worker.log_pre_api_call, "m", [], _call_kwargs("dep-a")) + assert reserved.result() is DEPLOYMENT_A + started.result() + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + + +class _SlowLocalRead(InMemoryCache): + """Answers a count read with the counts from when the read started, once another thread has begun its request.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + self.other_thread: threading.Thread | None = None + + def batch_get_cache(self, keys: list[str], **kwargs: object) -> list[object]: + counts_when_asked: Final = super().batch_get_cache(keys, **kwargs) + began: Final = threading.Event() + + def begin_then_run() -> None: + began.set() + self.while_reading() + + self.other_thread = threading.Thread(target=begin_then_run) + self.other_thread.start() + began.wait() + self.other_thread.join(timeout=0.2) + return counts_when_asked + + +def test_a_request_that_starts_on_another_thread_while_counts_are_read_still_counts_on_its_deployment() -> None: + local: Final = _SlowLocalRead() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(in_memory_cache=local)) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + local.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + picked_while_the_other_thread_started: Final = worker.get_available_deployments(GROUP, HEALTHY, reserve=True) + if local.other_thread is not None: + local.other_thread.join() + + assert picked_while_the_other_thread_started is DEPLOYMENT_B + + +def _router_pick(router: Router, group: str = ROUTER_GROUP) -> str: + deployment: Final = router.get_available_deployment(group, request_kwargs={"metadata": {}}) + return deployment["model_info"]["id"] + + +def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None: + router: Final = _least_busy_router() + + with picks_reserve(True): + first: Final = _router_pick(router) + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): + router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60) + + assert _router_pick(router) != first + + +@pytest.mark.asyncio +async def test_looking_up_a_deployment_does_not_hold_it_for_later_picks() -> None: + router: Final = _least_busy_router() + + sync_picks: Final = [_router_pick(router) for _ in range(3)] + async_picks: Final = [ + (await router.async_get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}))["model_info"]["id"] + for _ in range(3) + ] + + assert sync_picks + async_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6 + + +@pytest.mark.asyncio +async def test_looking_up_a_pass_through_deployment_does_not_hold_it_for_later_picks() -> None: + router: Final = _least_busy_router(use_in_pass_through=True) + + selected: Final = [ + (await router.select(ROUTER_GROUP, {"metadata": {}}, kind="pass_through"))["model_info"]["id"] for _ in range(2) + ] + async_picks: Final = [ + (await router.async_get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}}))[ + "model_info" + ]["id"] + for _ in range(2) + ] + sync_picks: Final = [ + router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"]["id"] + for _ in range(2) + ] + + assert selected + async_picks + sync_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6 + + +@pytest.mark.asyncio +async def test_a_cache_hit_releases_the_pick_it_reserved() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + await worker.async_log_success_event({**_call_kwargs("dep-b"), "cache_hit": True}, None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + worker.log_success_event({**_call_kwargs("dep-a", "cached-call"), "cache_hit": True}, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.parametrize("outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event]) +def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick( + outcome: Callable[[LeastBusyLoggingHandler, dict[str, object], None, None, None], None], +) -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + outcome(worker, _call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behind() -> None: + worker: Final = _worker(None) + rejected: Final = {**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True} + + worker.log_pre_api_call(model="m", messages=[], kwargs=rejected) + worker.log_failure_event(rejected, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments("tier-group", HEALTHY) is DEPLOYMENT_A + + +def test_a_fallback_attempt_that_fails_before_reaching_its_provider_frees_its_pick() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "busy-call")) + first_attempt_started_at: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + first_attempt: Final = {**_call_kwargs("dep-a", "fallback-call"), "api_call_start_time": first_attempt_started_at} + worker.log_pre_api_call(model="m", messages=[], kwargs=first_attempt) + worker.log_failure_event(first_attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + second_attempt: Final = {**_call_kwargs("dep-b", "fallback-call"), "api_call_start_time": first_attempt_started_at} + worker.log_failure_event(second_attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_ending_one_request_frees_only_the_picks_it_never_started() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="ended-call") is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="running-call") is DEPLOYMENT_A + + release_unstarted_picks(reservations, "ended-call") + + assert [worker.get_available_deployments(GROUP, HEALTHY, reserve=True) for _ in range(2)] == [ + DEPLOYMENT_A, + DEPLOYMENT_B, + ] + + +def test_a_call_without_a_call_id_turns_its_own_pick_into_its_in_flight_count() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "attempt-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_another_call_starting_leaves_the_pick_of_a_call_without_a_call_id_reserved() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + other_call: Final = contextvars.copy_context() + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + other_call.run(worker.log_pre_api_call, model="m", messages=[], kwargs=_call_kwargs("dep-a", "other-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_a_call_started_without_a_pick_of_its_own_leaves_another_calls_pick_reserved() -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="picked-call") is DEPLOYMENT_A + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "looked-up-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def _fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": "unreachable-group", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "key-unreachable"}, + "model_info": {"id": "unreachable"}, + }, + *( + { + "model_name": ROUTER_GROUP, + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": f"key-{deployment_id}"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ROUTER_DEPLOYMENT_IDS + ), + ], + routing_strategy="least-busy", + num_retries=0, + fallbacks=[{"unreachable-group": [ROUTER_GROUP]}], + ) + + +@pytest.mark.asyncio +async def test_a_fallback_hop_rejected_before_its_provider_frees_its_pick_when_every_hop_shares_one_logging_object( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(500, json={"error": {"message": "upstream down"}}) + ) + router: Final = _fallback_router() + logging_obj, request = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.datetime(2026, 1, 1), + model="unreachable-group", + messages=[{"role": "user", "content": "hi"}], + logit_bias={"1": 1}, + litellm_call_id="fallback-call", + ) + + with pytest.raises(litellm.InternalServerError): + await router.acompletion(**request, litellm_logging_obj=logging_obj) + await _drain_logging() + + assert {call.request.headers["authorization"] for call in route.calls} == {"Bearer key-unreachable"} + assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + + +def test_a_retry_that_starts_before_the_failed_attempt_is_logged_counts_each_attempt_once() -> None: + worker: Final = _worker(None) + attempt: Final = _call_kwargs("dep-a", "retried-call") + + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_failure_event(attempt, None, None, None) + worker.log_success_event(attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +async def _drain_logging() -> None: + for _ in range(10): + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() + + +def _called_deployment(route: respx.Route) -> str: + return route.calls.last.request.headers["authorization"].removeprefix("Bearer key-") + + +_BATCH: Final = { + "id": "batch_least_busy", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-least-busy", + "completion_window": "24h", + "status": "validating", + "created_at": 1, +} +_REJECTED_BATCH: Final = {"error": {"message": "Invalid input file", "type": "invalid_request_error"}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)]) +async def test_a_batch_create_does_not_keep_its_deployment_busy( + provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/batches").mock(return_value=provider_response) + router: Final = _least_busy_router() + + with contextlib.suppress(openai.BadRequestError): + await router.acreate_batch( + model=ROUTER_GROUP, completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-least-busy" + ) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +_VECTOR_STORE: Final = { + "id": "vs_least_busy", + "object": "vector_store", + "created_at": 1, + "name": "least-busy", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", +} +_REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}} +_CALL_IDENTITIES: Final = ({"litellm_call_id": "least-busy-call"}, {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +@pytest.mark.parametrize( + "provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)] +) +async def test_a_vector_store_create_does_not_keep_its_deployment_busy( + provider_response: httpx.Response, + call_identity: dict[str, str], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response) + router: Final = _least_busy_router() + + with contextlib.suppress(openai.BadRequestError): + await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", **call_identity) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +async def test_listing_fine_tuning_jobs_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.get(url__startswith="https://api.openai.com/v1/fine_tuning/jobs").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "has_more": False}) + ) + router: Final = _least_busy_router() + + await router.alist_fine_tuning_jobs(model=ROUTER_GROUP, **call_identity) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +TEXT_GROUP: Final = "openai/least-busy-text-group" + + +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +def test_a_sync_text_completion_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "cmpl-least-busy", + "object": "text_completion", + "created": 1, + "model": "least-busy-text-group", + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + ) + ) + router: Final = _least_busy_router(TEXT_GROUP) + + router.text_completion(model=TEXT_GROUP, prompt="Say ok", **call_identity) + + assert _router_pick(router, TEXT_GROUP) == _called_deployment(route) + + + +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its_pick( + call_identity: dict[str, str], +) -> None: + router: Final = _least_busy_router() + + with pytest.raises(litellm.InternalServerError): + router.responses(model=ROUTER_GROUP, input="hi", api_base="not-a-url", **call_identity) + + assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + +@pytest.mark.asyncio +async def test_a_batch_cancel_does_not_keep_its_deployment_busy( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/batches/batch_least_busy/cancel").mock( + return_value=httpx.Response(200, json={**_BATCH, "status": "cancelling"}) + ) + router: Final = _least_busy_router() + + await router.acancel_batch(model=ROUTER_GROUP, batch_id="batch_least_busy") + + assert _router_pick(router) == _called_deployment(route) + + +@pytest.mark.asyncio +async def test_picking_a_deployment_without_calling_it_does_not_keep_it_busy() -> None: + router: Final = _least_busy_router() + + picked: Final = await router.select(ROUTER_GROUP, {"metadata": {}}) + + assert _router_pick(router) == picked["model_info"]["id"] diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 5eeb46c3518..951182051c4 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -21,6 +21,7 @@ from litellm import Router from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext +from litellm.router_strategy.least_busy import picks_reserve from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -482,6 +483,38 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 +def _filtered_model_pick(router: Router) -> str: + deployment: Final = router.get_available_deployment( + "filtered-model", request_kwargs=_proxy_shaped_request(model="filtered-model") + ) + return deployment["model_info"]["id"] + + +def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = _build_router( + routing_strategy="least-busy", + routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}], + ) + + request: Final = _proxy_shaped_request(model="filtered-model") + with picks_reserve(True): + first: Final = router.get_available_deployment("filtered-model", request_kwargs=request)["model_info"]["id"] + kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}, + "litellm_call_id": request["litellm_call_id"], + } + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert _filtered_model_pick(router) == first + + def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): """ Least-busy hangs its counting off litellm's global callback lists, and those lists keep one @@ -614,6 +647,7 @@ async def test_select_deployment_async_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -628,6 +662,7 @@ def test_select_deployment_sync_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -662,6 +697,7 @@ def test_select_deployment_sync_does_not_dispatch_cost_based_routing(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -708,6 +744,7 @@ async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1(): messages=[{"role": "user", "content": "hi"}], input=None, request_kwargs={}, + reserve=False, ) assert v1_spy.called, "async dispatch must route v1 strategy through sync method" @@ -1022,7 +1059,11 @@ async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_fo await _mock_completion(router, routing_strategy="usage-based-routing-v2") -def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): +def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(monkeypatch): + monkeypatch.setattr( + "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", + lambda: datetime.datetime(2026, 1, 1, 12, 0, 30, tzinfo=datetime.timezone.utc), + ) router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) messages = [{"role": "user", "content": "hi"}] diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 26dbb96c60a..e9dd91d4362 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -302,6 +302,63 @@ def _two_deployment_router() -> Router: ) +def _entra_fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) + + +def _raised_from_an_entra_rejection(failed_deployment_id: str) -> litellm.AuthenticationError: + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + try: + raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from ( + EntraRejectionError(status_code=401, message=reason) + ) + except litellm.AuthenticationError as raised: + raised.failed_deployment_id = failed_deployment_id + return raised + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_params", "cooled_down"), + [ + pytest.param({"client_secret": "not-the-credential"}, [], id="caller-secret"), + pytest.param({"client_secret": "configured-secret"}, ["fallback-dep-1"], id="deployment-secret"), + ], +) +async def test_fallback_hop_cools_down_only_a_deployment_entra_rejected_its_own_credential( + call_params: Mapping[str, object], cooled_down: list[str] +) -> None: + router: Final = _entra_fallback_router() + + _trigger_cooldown_for_failed_deployment( + ctx=RoutingContext(router), + kwargs={}, + exception=_raised_from_an_entra_rejection("fallback-dep-1"), + model_call_details={"litellm_params": call_params}, + ) + + assert await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) == cooled_down + + @pytest.mark.asyncio @pytest.mark.parametrize("failure_already_logged", [False, True]) async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None: diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 6a63a502dd3..49d51e99790 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -4531,6 +4531,103 @@ def test_batch_cost_calculator_falls_back_to_text_batch_rate_for_modalities(): assert prompt_cost == pytest.approx(100 * 1e-7) +def _cached_audio_batch_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=1000, audio_tokens=1000, cached_tokens_details=CachedTokensDetails(audio_tokens=1000) + ), + ) + + +def _cached_text_batch_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=1000), + ) + + +@pytest.mark.parametrize( + ("audio_batch_cache_price", "usage", "expected_rate"), + [ + pytest.param(None, _cached_audio_batch_usage(), 1e-7, id="cached-audio-bills-the-standard-audio-cache-price"), + pytest.param(None, _cached_text_batch_usage(), 3e-8, id="cached-text-bills-the-standard-cache-price"), + pytest.param(8e-8, _cached_audio_batch_usage(), 8e-8, id="an-explicit-audio-batch-cache-price-still-wins"), + ], +) +def test_batch_cost_calculator_bills_cached_tokens_at_the_standard_cache_price_when_the_row_has_no_batch_cache_price( + audio_batch_cache_price: float | None, usage: Usage, expected_rate: float +) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 3e-7, + "input_cost_per_audio_token": 1e-6, + "cache_read_input_token_cost": 3e-8, + "cache_read_input_audio_token_cost": 1e-7, + "input_cost_per_token_batches": 1.5e-7, + "input_cost_per_audio_token_batches": 5e-7, + "cache_read_input_audio_token_cost_batches": audio_batch_cache_price, + } + + prompt_cost, _ = batch_cost_calculator( + usage=usage, model="gemini-2.5-flash", custom_llm_provider="vertex_ai", model_info=model_info + ) + + assert prompt_cost == pytest.approx(1000 * expected_rate) + + +def test_batch_cost_calculator_bills_long_context_cached_tokens_at_the_standard_long_context_cache_price() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 1.25e-6, + "input_cost_per_token_above_200k_tokens": 2.5e-6, + "cache_read_input_token_cost": 1.25e-7, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-7, + "input_cost_per_token_batches": 6.25e-7, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-6, + } + usage: Final = Usage( + prompt_tokens=300_000, + completion_tokens=0, + total_tokens=300_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300_000), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, model="gemini-2.5-pro", custom_llm_provider="vertex_ai", model_info=model_info + ) + + assert prompt_cost == pytest.approx(300_000 * 2.5e-7) + + +def _batch_and_standard_cached_audio_cost(model: str, provider: str) -> tuple[str, float, float]: + batch_prompt_cost, _ = batch_cost_calculator(_cached_audio_batch_usage(), model, provider) + standard_prompt_cost, _ = litellm.cost_per_token( + model=model, custom_llm_provider=provider, usage_object=_cached_audio_batch_usage() + ) + return model, batch_prompt_cost, standard_prompt_cost + + +def test_batch_cached_audio_never_costs_more_than_standard_on_any_row_with_a_batch_input_price( + _local_model_cost_map: None, +) -> None: + rows: Final = tuple( + (name, entry["litellm_provider"]) + for name, entry in litellm.model_cost.items() + if isinstance(entry, dict) + and entry.get("input_cost_per_token_batches") is not None + and entry.get("cache_read_input_audio_token_cost") is not None + ) + assert rows, "no cost map row prices both Batch input and Standard cached audio" + + costs: Final = tuple(_batch_and_standard_cached_audio_cost(name, provider) for name, provider in rows) + overpriced: Final = tuple(cost for cost in costs if cost[1] > cost[2]) + + assert overpriced == () + + @pytest.mark.parametrize( ("image_batch_rate", "image_tokens", "expected_completion_cost"), ( @@ -6254,7 +6351,7 @@ def test_batch_cost_calculator_bills_cached_tokens_at_the_flat_batch_cached_rate assert prompt_cost == pytest.approx(100 * 1e-6 + 900 * 1e-7) -def test_batch_cost_calculator_bills_cached_tokens_at_the_batch_input_rate_without_a_cached_batch_rate( +def test_batch_cost_calculator_bills_cached_tokens_at_the_standard_cache_rate_without_a_cached_batch_rate( _tiered_batch_models: None, ) -> None: from litellm.cost_calculator import batch_cost_calculator @@ -6263,7 +6360,7 @@ def test_batch_cost_calculator_bills_cached_tokens_at_the_batch_input_rate_witho usage=_cached_usage(300_048, 300_045, 11), model=_FLAT_CACHE_BATCH_MODEL, custom_llm_provider="openai" ) - assert prompt_cost == pytest.approx(300_048 * 3e-6) + assert prompt_cost == pytest.approx(3 * 3e-6 + 300_045 * 2e-7) def _cache_write_usage(prompt_tokens: int, cache_write_tokens: int, completion_tokens: int) -> Usage: @@ -11548,3 +11645,68 @@ def test_upstream_cost_header_total_on_an_unmapped_alias_is_split_by_the_served_ input_share: Final = (10 * 0.001) / (10 * 0.001 + 20 * 0.002) assert breakdown["input_cost"] == pytest.approx(0.25 * input_share) assert breakdown["output_cost"] == pytest.approx(0.25 * (1 - input_share)) + + +_PER_CALL_RATES: Final = MappingProxyType({"input_cost_per_token": 2e-06, "output_cost_per_token": 4e-06}) + + +def _per_call_priced_reply(model: str, usage_extras: Mapping[str, str]) -> ModelResponse: + usage: Final = Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000) + for field, value in usage_extras.items(): + setattr(usage, field, value) + return ModelResponse(model=model, choices=[Choices(message=Message(content="hi"))], usage=usage) + + +@pytest.mark.parametrize( + ("row", "call_kwargs", "usage_extras", "multiplier"), + [ + pytest.param( + {"litellm_provider": "openai", "mode": "chat", "regional_processing_uplift_multiplier_eu": 1.25}, + { + "model": "openai/lit-eu-residency-test", + "api_key": "fake-key", + "api_base": "https://eu.api.openai.com/v1", + }, + {}, + 1.25, + id="openai-data-residency", + ), + pytest.param( + {"litellm_provider": "vertex_ai", "mode": "chat", "regional_endpoint_uplift_multiplier": 1.5}, + { + "model": "vertex_ai/lit-regional-endpoint-test", + "vertex_project": "lit-project", + "vertex_location": "us-east5", + "vertex_credentials": "{}", + }, + {}, + 1.5, + id="vertex-regional-endpoint", + ), + pytest.param( + {"litellm_provider": "anthropic", "mode": "chat", "provider_specific_entry": {"fast": 2.0}}, + {"model": "anthropic/lit-fast-mode-test", "api_key": "fake-key"}, + {"speed": "fast"}, + 2.0, + id="anthropic-fast-mode", + ), + ], +) +def test_per_call_custom_prices_keep_the_regional_and_speed_multipliers( + monkeypatch: pytest.MonkeyPatch, + row: Mapping[str, object], + call_kwargs: Mapping[str, str], + usage_extras: Mapping[str, str], + multiplier: float, +) -> None: + monkeypatch.setitem(litellm.model_cost, call_kwargs["model"], {**row, **_PER_CALL_RATES}) + _invalidate_model_cost_lowercase_map() + + response: Final = litellm.completion( + messages=[{"role": "user", "content": "hello"}], + mock_response=_per_call_priced_reply(call_kwargs["model"].split("/", 1)[1], usage_extras), + **call_kwargs, + **_PER_CALL_RATES, + ) + + assert response._hidden_params["response_cost"] == pytest.approx((1000 * 2e-06 + 1000 * 4e-06) * multiplier) diff --git a/tests/unit/test_lint_workflow_diff_gates.py b/tests/unit/test_lint_workflow_diff_gates.py index 816ef6ce707..5e6ae9d285a 100644 --- a/tests/unit/test_lint_workflow_diff_gates.py +++ b/tests/unit/test_lint_workflow_diff_gates.py @@ -383,10 +383,20 @@ def _guard_base_argument_and_checkout(run: StepRun) -> tuple[str, str]: return argv[argv.index("--base") + 1], _sha(run.runner.checkout, "HEAD") +def _schema_base_commit_and_its_type_in_the_checkout(run: StepRun) -> tuple[str, str]: + argv: Final = _argv_recorded(run.runner, "uv") + base_commit: Final = argv[argv.index("--base-commit") + 1] + return base_commit, _git(run.runner.checkout, "cat-file", "-t", base_commit).strip() + + def _merge_base(upstream: StaleBaseUpstream, ref: str) -> str: return upstream.merge_base_by_ref[ref] +def _merge_base_as_a_commit(upstream: StaleBaseUpstream, ref: str) -> tuple[str, str]: + return upstream.merge_base_by_ref[ref], "commit" + + def _ui_files_changed_since_merge_base(upstream: StaleBaseUpstream, ref: str) -> frozenset[str]: merge_base: Final = upstream.merge_base_by_ref[ref] changed: Final = _git(upstream.path, "diff", "--name-only", f"--relative={UI_DASHBOARD}", merge_base, upstream.head) @@ -402,6 +412,7 @@ CONSUMERS: Final = MappingProxyType( "test-linting.yml:python": Consumer(_gate_base_exported, _merge_base), "test-linting.yml:ui": Consumer(_ui_files_queued_for_prettier, _ui_files_changed_since_merge_base), "test-unit.yml:ui-unit": Consumer(_ui_files_passed_to_vitest_related, _ui_files_changed_since_merge_base), + "test-unit.yml:unit": Consumer(_schema_base_commit_and_its_type_in_the_checkout, _merge_base_as_a_commit), "spend-logs-guard.yml:spend-logs-guard": Consumer(_guard_base_argument_and_checkout, _merge_base_and_tip), "cost-map-guard.yml:cost-map-guard": Consumer(_guard_base_argument_and_checkout, _merge_base_and_tip), } diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index f815b107648..f5962cb4933 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -25836,6 +25836,441 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} +def _caused_by_an_entra_rejection(error: litellm.AuthenticationError) -> litellm.AuthenticationError: + from litellm.llms.azure.common_utils import EntraRejectionError + + try: + raise error from EntraRejectionError(status_code=401, message=error.message) + except litellm.AuthenticationError as raised: + return raised + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize("caller_credential", ["client_secret", "azure_password", "azure_ad_token"]) +async def test_router_caller_entra_credential_rejection_leaves_the_shared_deployment_routable(caller_credential): + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **{caller_credential: "not-the-credential"}, + ) + + response: Final = await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert response.choices[0].message.content == "hi" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "caller_field", + [ + pytest.param({"client_secret": "environment-secret"}, id="environment_secret"), + pytest.param({"client_id": "environment-client"}, id="environment_client_id"), + pytest.param({"azure_scope": "https://cognitiveservices.azure.com/.default"}, id="default_scope"), + pytest.param({"client_secret": ""}, id="blank_secret"), + ], +) +async def test_router_caller_entra_field_that_changes_the_deployment_credential_leaves_it_routable( + caller_field, monkeypatch +): + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + monkeypatch.setenv("AZURE_CLIENT_ID", "environment-client") + monkeypatch.delenv("AZURE_SCOPE", raising=False) + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + "azure_scope": "api://configured-app/.default", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **caller_field, + ) + + response: Final = await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert response.choices[0].message.content == "hi" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + ("deployment_secret", "router_default_secret"), + [ + pytest.param({"client_secret": "configured-secret"}, None, id="deployment"), + pytest.param({}, {"client_secret": "configured-secret"}, id="router_default_litellm_params"), + ], +) +async def test_router_configured_entra_credential_rejection_cools_down_the_deployment( + deployment_secret, router_default_secret +): + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + **deployment_secret, + }, + "model_info": {"id": "configured-deployment"}, + } + ], + default_litellm_params=router_default_secret, + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000222: The provided client secret keys are expired.", + llm_provider="azure", + model="gpt-4.1", + ), + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "caller_fields", + [ + pytest.param({"azure_scope": "https://cognitiveservices.azure.com/.default"}, id="default_scope"), + pytest.param({"azure_username": ""}, id="empty_username"), + pytest.param({"client_secret": "environment-secret"}, id="environment_secret"), + ], +) +async def test_router_caller_entra_field_that_changes_nothing_still_cools_down_a_rejected_deployment( + caller_fields, monkeypatch +): + monkeypatch.delenv("AZURE_SCOPE", raising=False) + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **caller_fields, + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_caller_repeating_the_secret_an_empty_router_default_selects_still_cools_down(monkeypatch): + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": "configured-deployment"}, + } + ], + default_litellm_params={"client_secret": ""}, + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + client_secret="environment-secret", + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_caller_entra_field_keeps_the_shared_deployment(): + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "configured-key"}, + "model_info": {"id": "shared-deployment"}, + } + ] + ) + + response: Final = await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + client_id="caller-client", + ) + + assert response._hidden_params["model_id"] == "shared-deployment" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + ("litellm_params", "provider"), + [ + pytest.param({"model": "openai/gpt-4.1-mini", "api_key": "revoked-key"}, "openai", id="openai"), + pytest.param( + { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "revoked-key", + }, + "azure", + id="azure_api_key", + ), + ], +) +@pytest.mark.parametrize("caller_entra_field", ["client_id", "azure_scope"]) +async def test_router_auth_failure_entra_did_not_cause_cools_down_despite_a_caller_entra_field( + litellm_params, provider, caller_entra_field +): + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": litellm_params, + "model_info": {"id": "shared-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError( + message="Incorrect API key provided.", llm_provider=provider, model=litellm_params["model"] + ), + **{caller_entra_field: "caller-value"}, + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +async def test_router_fallback_auth_failure_after_an_entra_rejection_still_cools_down_the_fallback(monkeypatch): + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + }, + "model_info": {"id": "entra-deployment"}, + }, + { + "model_name": "key-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa-fallback.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "revoked-key", + }, + "model_info": {"id": "key-deployment"}, + }, + ], + fallbacks=[{"entra-model": ["key-model"]}], + ) + + with respx.mock: + fallback_route: Final = respx.post(url__startswith="https://litellm-qa-fallback.openai.azure.com").mock( + return_value=httpx.Response( + 401, json={"error": {"code": "401", "message": "Access denied due to invalid subscription key."}} + ) + ) + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + client_secret="not-the-credential", + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="key-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert fallback_route.call_count == 1 + assert excinfo.value.all_deployments_in_cooldown is True + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "stream_request", + [ + pytest.param( + lambda router, **caller_fields: router.acompletion( + model="entra-model", messages=[{"role": "user", "content": "hi"}], stream=True, **caller_fields + ), + id="chat_completions", + ), + pytest.param( + lambda router, **caller_fields: router.aanthropic_messages( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + max_tokens=5, + stream=True, + **caller_fields, + ), + id="messages", + ), + ], +) +async def test_router_caller_entra_rejection_on_a_stream_leaves_the_shared_deployment_routable( + monkeypatch: pytest.MonkeyPatch, + stream_request: Callable[..., Awaitable[object]], +) -> None: + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.") + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = Router( + model_list=[ + { + "model_name": "entra-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "azure_ad_token_provider": rejected_token, + }, + "model_info": {"id": "entra-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await stream_request(router, client_secret="not-the-credential") + + with pytest.raises(litellm.AuthenticationError): + await stream_request(router) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router( @@ -29667,11 +30102,21 @@ def test_reset_custom_routing_strategy(): assert "get_available_deployment" not in router._core.__dict__ assert "_core_async_get_available_deployment" not in router._core.__dict__ - assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment + assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment router._reset_custom_routing_strategy() +@pytest.mark.asyncio +async def test_select_picks_through_a_custom_routing_strategy() -> None: + router: Final = _create_custom_routing_router() + router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router)) + + picked: Final = await router.select("azure-model", {"metadata": {}}) + + assert picked["model_info"]["id"] == "very-special-endpoint" + + @pytest.mark.asyncio async def test_azure_router_acompletion_passes_deployment_credentials_for_plain_and_stream( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5ad6eb819ef..3bbb69413d9 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -844,6 +844,34 @@ async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_b await admit_turn() +@pytest.mark.asyncio +async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy(): + router: Final = Router( + model_list=[ + { + "model_name": "ws-least-busy", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("ws-least-busy-a", "ws-least-busy-b") + ], + routing_strategy="least-busy", + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + for _ in range(3): + await admit_turn() + + picks: Final = [ + (await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ + "model_info" + ]["id"] + for _ in range(2) + ] + + assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"] + + @pytest.mark.asyncio async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): router = Router( diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index c00670473af..70231b620be 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -1019,3 +1019,144 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy assert while_a_is_down._hidden_params["model_id"] == "B" assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 assert after_a_recovers._hidden_params["model_id"] == "A" + + +def _two_deployment_router(model: str) -> Router: + return Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": model, "api_key": "bad", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": model, "api_key": "good", "weight": 0}, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + disable_cooldowns=True, + enable_weighted_failover=True, + ) + + +_ENDPOINTS: Final = [ + ("aembedding", "openai/text-embedding-3-small", {"input": ["hi"]}), + ("aimage_generation", "openai/dall-e-3", {"prompt": "a cat"}), + ("atranscription", "openai/whisper-1", {"file": b"audio"}), + ("aspeech", "openai/tts-1", {"input": "hi", "voice": "alloy"}), + ("atext_completion", "openai/gpt-3.5-turbo-instruct", {"prompt": "hi"}), + ("arerank", "cohere/rerank-english-v3.0", {"query": "q", "documents": ["d"]}), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("endpoint", "model", "arguments"), _ENDPOINTS, ids=[e[0] for e in _ENDPOINTS]) +async def test_weighted_failover_moves_every_endpoint_to_another_deployment( + monkeypatch: pytest.MonkeyPatch, endpoint: str, model: str, arguments: dict[str, object] +) -> None: + keys_called: list[object] = [] + answer: Final = object() + + async def provider(**kwargs: object) -> object: + keys_called.append(kwargs.get("api_key")) + if kwargs.get("api_key") == "bad": + raise litellm.InternalServerError(message="A is down", llm_provider="openai", model=model) + return answer + + monkeypatch.setattr(litellm, endpoint, provider) + router: Final = _two_deployment_router(model) + + response: Final = await getattr(router, endpoint)(model="test-model", **arguments) + + assert response is answer + assert keys_called == ["bad", "good"] + + +@pytest.mark.asyncio +async def test_weighted_failover_skips_the_failed_deployment_for_clientside_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + api_bases_called: list[object] = [] + answer: Final = object() + + async def provider(**kwargs: object) -> object: + api_bases_called.append(kwargs.get("api_base")) + if kwargs.get("api_base") == "http://a": + raise litellm.InternalServerError(message="A is down", llm_provider="openai", model="m") + return answer + + monkeypatch.setattr(litellm, "aembedding", provider) + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "http://a", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "http://b", "weight": 0}, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + disable_cooldowns=True, + enable_weighted_failover=True, + ) + + response: Final = await router.aembedding(model="test-model", input=["hi"], api_key="tenant-key") + + assert response is answer + assert api_bases_called == ["http://a", "http://b"] + + +@pytest.mark.parametrize( + ("endpoint", "attempt", "arguments"), + [ + ("embedding", "_embedding", {"input": ["hi"]}), + ("image_generation", "_image_generation", {"prompt": "a cat"}), + ], +) +def test_a_sync_attempt_stamps_the_failed_deployment( + monkeypatch: pytest.MonkeyPatch, endpoint: str, attempt: str, arguments: dict[str, object] +) -> None: + def provider(**kwargs: object) -> object: + raise litellm.InternalServerError(message="down", llm_provider="openai", model="m") + + monkeypatch.setattr(litellm, endpoint, provider) + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/m", "api_key": "k"}, + "model_info": {"id": "A"}, + } + ], + ) + + with pytest.raises(litellm.InternalServerError) as raised: + getattr(router, attempt)(model="test-model", **arguments) + + assert getattr(raised.value, "failed_deployment_id", None) == "A" + + +@pytest.mark.asyncio +async def test_a_failed_pick_stamps_no_deployment() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "k"}, + "model_info": {"id": "A"}, + } + ], + ) + + with pytest.raises(litellm.BadRequestError, match="no healthy deployments") as raised: + await router._aembedding(input=["hi"], model="no-such-model") + + assert getattr(raised.value, "failed_deployment_id", None) is None diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 3d30964129d..dd428b286bd 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1081,6 +1081,11 @@ "count": 1 } }, + "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { + "max-lines": { + "count": 1 + } + }, "src/app/(dashboard)/users/_components/BulkEditUsers.tsx": { "prefer-const": { "count": 1 @@ -1204,6 +1209,11 @@ "count": 1 } }, + "src/components/ModelInfoEditForm.tsx": { + "max-lines": { + "count": 1 + } + }, "src/components/Navbar/BlogDropdown/BlogDropdown.test.tsx": { "max-nested-callbacks": { "count": 10 @@ -1314,6 +1324,11 @@ "count": 1 } }, + "src/components/add_model/ComplexityRouterConfig.tsx": { + "max-lines": { + "count": 1 + } + }, "src/components/add_model/RouterConfigBuilder.tsx": { "react-hooks/set-state-in-effect": { "count": 1 @@ -1339,6 +1354,11 @@ "count": 1 } }, + "src/components/add_model/build_complexity_router_config.ts": { + "max-lines": { + "count": 1 + } + }, "src/components/add_model/cache_control_settings.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2000,6 +2020,9 @@ "src/components/templates/key_edit_view.tsx": { "local/filename-pascal-case": { "count": 1 + }, + "max-lines": { + "count": 1 } }, "src/components/templates/key_info_view.tsx": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx index 1339300b2ba..7f6f7ea1d17 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx @@ -1,5 +1,6 @@ "use client"; import React from "react"; +import { useTranslation } from "react-i18next"; import CodeBlock from "@/components/CodeBlock"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import DocLink from "./DocLink"; @@ -12,6 +13,7 @@ interface ApiRefProps { } const APIReferenceView: React.FC = ({ proxySettings }) => { + const { t } = useTranslation(); let base_url = ""; const customDocBaseUrl = proxySettings?.LITELLM_UI_API_DOC_BASE_URL; if (customDocBaseUrl && customDocBaseUrl.trim()) { @@ -25,14 +27,11 @@ const APIReferenceView: React.FC = ({ proxySettings }) => {
{/* Header row with Docs link on the right */}
-

OpenAI Compatible Proxy: API Reference

+

{t("apiReference.view.title")}

-

- LiteLLM is OpenAI Compatible. This means your API Key works with the OpenAI SDK. Just replace the base_url to - point to your litellm proxy. Example Below{" "} -

+

{t("apiReference.view.description")}

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/DocLink.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/DocLink.tsx index f5a9d5e3a15..cee174d7fdb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/DocLink.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/DocLink.tsx @@ -1,5 +1,6 @@ import React from "react"; import { ExternalLink } from "lucide-react"; +import { useTranslation } from "react-i18next"; import { cn } from "@/lib/cva.config"; @@ -9,21 +10,22 @@ export type DocLinkProps = { }; const DocLink = ({ href, className }: DocLinkProps) => { + const { t } = useTranslation(); return ( - API Reference Docs + {t("apiReference.docLink.label")} - (opens in a new tab) + {t("apiReference.docLink.newTab")} ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx index d7c977b0870..4d8498a1be4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx @@ -1,17 +1,19 @@ "use client"; +import { useTranslation } from "react-i18next"; import APIReferenceView from "./_components/APIReferenceView"; import { DeprecationBanner } from "@/components/DeprecationBanner"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; const APIReferencePage = () => { + const { t } = useTranslation(); const { accessToken } = useAuthorized(); const proxySettings = useProxySettings(accessToken); return ( <> - + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx index 9ffd23884b4..56e2c6dbecc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx @@ -1,4 +1,5 @@ import React, { useId } from "react"; +import { useTranslation } from "react-i18next"; import { Controller } from "react-hook-form"; import { z } from "zod"; import { FormField } from "@/components/shared/form/FormField"; @@ -16,6 +17,7 @@ import { import { Input } from "@/components/ui/input"; import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import i18n from "@/i18n/i18n"; import { useZodForm } from "@/lib/forms/useZodForm"; const budgetShape = { @@ -54,19 +56,23 @@ const numberProps = ) => onChange(event.target.value === "" ? null : event.target.valueAsNumber), }); -const OptionalBadge = () => Optional; +const OptionalBadge = () => { + const { t } = useTranslation(); + return {t("budgets.form.optional")}; +}; function SpendLimitField({ form }: { form: BudgetForm }) { + const { t } = useTranslation(); const amountId = useId(); return ( - Spend limit + {t("budgets.form.spendLimit")} @@ -91,7 +97,7 @@ function SpendLimitField({ form }: { form: BudgetForm }) { value={field.value ?? NO_RESET} onValueChange={(next) => field.onChange(next === NO_RESET ? null : next)} > - + @@ -124,13 +130,14 @@ function RateLimitField({ form, name, label }: { form: BudgetForm; name: NumberF } export function BudgetFormFields({ form, isEditing }: { form: BudgetForm; isEditing: boolean }) { + const { t } = useTranslation(); return ( {({ ref, ...field }) => } @@ -139,15 +146,13 @@ export function BudgetFormFields({ form, isEditing }: { form: BudgetForm; isEdit
- Rate limits + {t("budgets.form.rateLimits")} - - - - - Batches count against the daily limit instead of the per-minute ones. Provider rate limits still apply. - + + + + {t("budgets.form.batchNote")}
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx index bd10837640d..c6bf7e8121e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx @@ -1,4 +1,5 @@ import React, { useId } from "react"; +import { useTranslation } from "react-i18next"; import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { applyBudgetPrecision } from "./budgetPrecision"; import { BudgetFormFields, type BudgetFormValues, useBudgetForm } from "./BudgetFormFields"; @@ -11,6 +12,7 @@ interface BudgetModalProps { setIsModalVisible: React.Dispatch>; } const BudgetModal: React.FC = ({ isModalVisible, setIsModalVisible }) => { + const { t } = useTranslation(); const formId = useId(); const form = useBudgetForm({ budget_id: "" }); const createBudget = useCreateBudget(); @@ -22,14 +24,14 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis const handleCreate = async (formValues: BudgetFormValues) => { try { - toast.info("Making API Call"); + toast.info(t("budgets.createModal.toast.pending")); await createBudget.mutateAsync(applyBudgetPrecision(formValues)); - toast.success("Budget Created"); + toast.success(t("budgets.createModal.toast.created")); form.reset(); setIsModalVisible(false); } catch (error) { console.error("Error creating the budget:", error); - toast.fromError(`Error creating the budget: ${error}`); + toast.fromError(t("budgets.createModal.toast.error", { error: String(error) })); } }; @@ -37,15 +39,17 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis !open && handleCancel()}> - Create budget + + {t("budgets.createModal.title")} +
- }>Cancel + }>{t("budgets.createModal.cancel")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx index 785b09252af..bbc4fed94f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx @@ -6,6 +6,7 @@ import { Page, PageTabs, PageTabsList, PageTabsTrigger } from "@/components/shared/Page"; import { Plus, Wallet } from "lucide-react"; import React, { useCallback, useState } from "react"; +import { useTranslation } from "react-i18next"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { prism } from "react-syntax-highlighter/dist/esm/styles/prism"; @@ -29,6 +30,7 @@ interface BudgetSettingsPageProps { } const BudgetPanel: React.FC = ({ accessToken }) => { + const { t } = useTranslation(); const syntaxTheme = useSyntaxTheme(prism); const [isCreateModelVisible, setIsCreateModelVisible] = useState(false); const [isEditModalVisible, setIsEditModalVisible] = useState(false); @@ -65,10 +67,10 @@ const BudgetPanel: React.FC = ({ accessToken }) => { } try { await deleteBudget.mutateAsync(selectedBudget.budget_id); - toast.success("Budget deleted."); + toast.success(t("budgets.panel.toast.deleted")); } catch (error) { console.error("Error deleting budget:", error); - toast.fromError("Failed to delete budget"); + toast.fromError(t("budgets.panel.toast.deleteFailed")); } finally { setIsDeleteModalVisible(false); setSelectedBudget(null); @@ -85,22 +87,22 @@ const BudgetPanel: React.FC = ({ accessToken }) => { - Budgets + {t("budgets.panel.title")} - Spend, TPM and RPM limits you can assign to customers. + {t("budgets.panel.description")} {canModify && ( <> )} - Budgets - Examples + {t("budgets.panel.tabs.budgets")} + {t("budgets.panel.tabs.examples")} @@ -122,15 +124,15 @@ const BudgetPanel: React.FC = ({ accessToken }) => { /> = ({ accessToken }) => {
-

How to use budget id

+

{t("budgets.panel.examples.intro")}

- Assign Budget to Customer + {t("budgets.panel.examples.tabs.assignBudget")} - Test it (Curl) + {t("budgets.panel.examples.tabs.curl")} - Test it (OpenAI SDK) + {t("budgets.panel.examples.tabs.openaiSdk")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx index 80872c58d28..9d9e32ed86e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.tsx @@ -2,6 +2,7 @@ import { Inbox, ShieldAlert } from "lucide-react"; import React, { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { BUDGET_DURATION_FILTER_OPTIONS, @@ -21,6 +22,7 @@ import { import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import i18n from "@/i18n/i18n"; import { ApiError } from "@/lib/http/client"; import { BUDGET_TABLE_HIDDEN_COLUMNS, getBudgetTableColumns } from "./BudgetTableColumns"; @@ -49,11 +51,19 @@ const formatFilterValue = (columnId: string, value: unknown): string => { } if (columnId === "max_budget") { const { min, max, unlimitedOnly } = (value ?? {}) as MaxBudgetFilterValue; - return unlimitedOnly === true ? "Unlimited only" : `${min ? `$${min}` : "any"} to ${max ? `$${max}` : "any"}`; + return unlimitedOnly === true + ? i18n.t("budgets.table.filterChips.unlimitedOnly") + : i18n.t("budgets.table.filterChips.maxBudgetRange", { + min: min ? `$${min}` : i18n.t("budgets.table.filterChips.any"), + max: max ? `$${max}` : i18n.t("budgets.table.filterChips.any"), + }); } if (columnId === "created_at") { const { from, to } = (value ?? {}) as CreatedAtFilterValue; - return `${from || "any"} to ${to || "any"}`; + return i18n.t("budgets.table.filterChips.createdRange", { + from: from || i18n.t("budgets.table.filterChips.any"), + to: to || i18n.t("budgets.table.filterChips.any"), + }); } return String(value); }; @@ -81,22 +91,24 @@ const normalizeCreatedAt = (draft: CreatedAtFilterValue): CreatedAtFilterValue | }; function EmptyState({ hasQuery }: { hasQuery: boolean }) { + const { t } = useTranslation(); return (
-
{hasQuery ? "No matching budgets" : "No budgets yet"}
+
+ {hasQuery ? t("budgets.table.emptyState.noMatchTitle") : t("budgets.table.emptyState.title")} +
- {hasQuery - ? "No budget matches your search or filters." - : "Create a budget to set spend, TPM and RPM limits for customers."} + {hasQuery ? t("budgets.table.emptyState.noMatchDescription") : t("budgets.table.emptyState.description")}
); } function ErrorState({ error }: { error: Error }) { + const { t } = useTranslation(); const forbidden = error instanceof ApiError && error.status === 403; return (
@@ -104,10 +116,10 @@ function ErrorState({ error }: { error: Error }) {
- {forbidden ? "You do not have access to budgets" : "Could not load budgets"} + {forbidden ? t("budgets.table.errorState.forbiddenTitle") : t("budgets.table.errorState.title")}
- {forbidden ? "Ask a proxy admin to grant you the admin viewer role." : error.message} + {forbidden ? t("budgets.table.errorState.forbiddenDescription") : error.message}
); @@ -141,19 +153,20 @@ function DurationFilter({ selected, onChange }: { selected: string[]; onChange: } function BudgetFilterFields({ get, set }: FilterDraft) { + const { t } = useTranslation(); const maxBudget = (get("max_budget") as MaxBudgetFilterValue | undefined) ?? {}; const created = (get("created_at") as CreatedAtFilterValue | undefined) ?? {}; const unlimitedOnly = maxBudget.unlimitedOnly === true; return ( <> - + set("budget_duration", selected)} /> - +
set("max_budget", normalizeMaxBudget({ ...maxBudget, min: event.target.value }))} - placeholder="Min" - aria-label="Minimum max budget" + placeholder={t("budgets.table.filters.min")} + aria-label={t("budgets.table.filters.minAriaLabel")} data-testid="budget-filter-max-budget-min" /> set("max_budget", normalizeMaxBudget({ ...maxBudget, max: event.target.value }))} - placeholder="Max" - aria-label="Maximum max budget" + placeholder={t("budgets.table.filters.max")} + aria-label={t("budgets.table.filters.maxAriaLabel")} data-testid="budget-filter-max-budget-max" />
@@ -184,23 +197,23 @@ function BudgetFilterFields({ get, set }: FilterDraft) { onCheckedChange={(checked) => set("max_budget", normalizeMaxBudget({ unlimitedOnly: checked === true }))} data-testid="budget-filter-max-budget-unlimited" /> - Unlimited only + {t("budgets.table.filters.unlimitedOnly")} - +
set("created_at", normalizeCreatedAt({ ...created, from: event.target.value }))} - aria-label="Created from" + aria-label={t("budgets.table.filters.createdFrom")} data-testid="budget-filter-created-from" /> set("created_at", normalizeCreatedAt({ ...created, to: event.target.value }))} - aria-label="Created to" + aria-label={t("budgets.table.filters.createdTo")} data-testid="budget-filter-created-to" />
@@ -210,11 +223,12 @@ function BudgetFilterFields({ get, set }: FilterDraft) { } const BudgetTable: React.FC = ({ list, canModify, onEditClick, onDeleteClick }) => { + const { t } = useTranslation(); const [filtersOpen, setFiltersOpen] = useState(false); const columns = useMemo( - () => getBudgetTableColumns({ canModify, onEditClick, onDeleteClick }), - [canModify, onEditClick, onDeleteClick], + () => getBudgetTableColumns({ canModify, onEditClick, onDeleteClick }, t), + [canModify, onEditClick, onDeleteClick, t], ); const hasQuery = list.searchValue.trim() !== "" || list.columnFilters.length > 0; @@ -239,7 +253,7 @@ const BudgetTable: React.FC = ({ list, canModify, onEditClick, columnFilters={list.columnFilters} onColumnFiltersChange={list.onColumnFiltersChange} isLoading={list.isLoading} - loadingMessage="Loading budgets…" + loadingMessage={t("budgets.table.loading")} noDataMessage={emptyMessage} size="compact" toolbar={(table) => ( @@ -248,7 +262,7 @@ const BudgetTable: React.FC = ({ list, canModify, onEditClick, table={table} searchValue={list.searchValue} onSearchChange={list.onSearchChange} - searchPlaceholder="Search by budget ID…" + searchPlaceholder={t("budgets.table.searchPlaceholder")} onOpenFilters={() => setFiltersOpen(true)} onRefresh={list.refetch} isRefreshing={list.isFetching} @@ -259,8 +273,8 @@ const BudgetTable: React.FC = ({ list, canModify, onEditClick, table={table} open={filtersOpen} onOpenChange={setFiltersOpen} - title="Filters" - description="Narrow down your budgets" + title={t("budgets.table.filters.title")} + description={t("budgets.table.filters.description")} > {(draft) => } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index 3ec291c1158..d0538e7ed1c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -1,7 +1,9 @@ "use client"; import { ColumnDef, FilterFn } from "@tanstack/react-table"; +import type { TFunction } from "i18next"; import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; +import { useTranslation } from "react-i18next"; import { DataTableSortHeader } from "@/components/shared/DataTable"; import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; @@ -27,15 +29,17 @@ const serverFilter: FilterFn = () => true; serverFilter.autoRemove = () => false; function RateLimitCell({ value }: { value: number | null | undefined }) { + const { t } = useTranslation(); if (value == null) { - return n/a; + return {t("budgets.table.cells.notApplicable")}; } return {value}; } function BudgetDurationCell({ value }: { value: string | null | undefined }) { + const { t } = useTranslation(); if (!value) { - return Not set; + return {t("budgets.table.cells.notSet")}; } return {getBudgetDurationLabel(value)}; } @@ -47,10 +51,11 @@ interface BudgetRowActionsProps { } function BudgetRowActions({ budget, onEditClick, onDeleteClick }: BudgetRowActionsProps) { + const { t } = useTranslation(); return ( @@ -59,7 +64,7 @@ function BudgetRowActions({ budget, onEditClick, onDeleteClick }: BudgetRowActio onEditClick(budget)}> - Edit budget + {t("budgets.table.rowActions.edit")} onDeleteClick(budget)} > - Delete budget + {t("budgets.table.rowActions.delete")} @@ -87,16 +92,15 @@ interface BudgetTableColumnsDeps { onDeleteClick: (budget: budgetItem) => void; } -export const getBudgetTableColumns = ({ - canModify, - onEditClick, - onDeleteClick, -}: BudgetTableColumnsDeps): ColumnDef[] => [ +export const getBudgetTableColumns = ( + { canModify, onEditClick, onDeleteClick }: BudgetTableColumnsDeps, + t: TFunction, +): ColumnDef[] => [ { id: "budget_id", accessorKey: "budget_id", - meta: { title: "Budget ID" }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.budgetId") }, + header: ({ column }) => , cell: ({ row }) => ( ), @@ -105,32 +109,34 @@ export const getBudgetTableColumns = ({ id: "max_budget", accessorKey: "max_budget", filterFn: serverFilter, - meta: { title: "Max Budget", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.maxBudget"), numeric: true }, + header: ({ column }) => , size: 120, - cell: ({ row }) => , + cell: ({ row }) => ( + + ), }, { id: "tpm_limit", accessorKey: "tpm_limit", - meta: { title: "TPM", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.tpm"), numeric: true }, + header: ({ column }) => , size: 100, cell: ({ row }) => , }, { id: "rpm_limit", accessorKey: "rpm_limit", - meta: { title: "RPM", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.rpm"), numeric: true }, + header: ({ column }) => , size: 100, cell: ({ row }) => , }, { id: "tpd_limit", accessorKey: "tpd_limit", - meta: { title: "TPD (batch)", numeric: true }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.tpd"), numeric: true }, + header: ({ column }) => , size: 110, cell: ({ row }) => , }, @@ -138,9 +144,9 @@ export const getBudgetTableColumns = ({ id: "budget_duration", accessorKey: "budget_duration", filterFn: serverFilter, - meta: { title: "Reset" }, + meta: { title: t("budgets.table.columns.reset") }, sortDescFirst: false, - header: ({ column }) => , + header: ({ column }) => , size: 110, cell: ({ row }) => , }, @@ -148,8 +154,8 @@ export const getBudgetTableColumns = ({ id: "created_at", accessorKey: "created_at", filterFn: serverFilter, - meta: { title: "Created" }, - header: ({ column }) => , + meta: { title: t("budgets.table.columns.created") }, + header: ({ column }) => , size: 160, cell: ({ row }) => , }, @@ -158,7 +164,7 @@ export const getBudgetTableColumns = ({ { id: "actions", meta: { className: "text-right", headerClassName: "text-right" }, - header: () => Actions, + header: () => {t("budgets.table.columns.actions")}, size: 64, enableSorting: false, enableHiding: false, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx index 4aa69c5c15e..81c1f652eab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx @@ -1,4 +1,5 @@ import React, { useEffect, useId } from "react"; +import { useTranslation } from "react-i18next"; import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { applyBudgetPrecision } from "./budgetPrecision"; @@ -23,6 +24,7 @@ interface EditBudgetModalProps { existingBudget: budgetItem; } const EditBudgetModal: React.FC = ({ isModalVisible, setIsModalVisible, existingBudget }) => { + const { t } = useTranslation(); const formId = useId(); const form = useBudgetForm(toFormValues(existingBudget)); const { dirtyFields } = form.formState; @@ -39,15 +41,15 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs const handleUpdate = async (formValues: BudgetFormValues) => { try { - toast.info("Making API Call"); + toast.info(t("budgets.editModal.toast.pending")); const changedValues = pickDirty(formValues, dirtyFields); await updateBudget.mutateAsync(applyBudgetPrecision({ budget_id: formValues.budget_id, ...changedValues })); - toast.success("Budget Updated"); + toast.success(t("budgets.editModal.toast.updated")); form.reset(); setIsModalVisible(false); } catch (error) { console.error("Error updating the budget:", error); - toast.fromError(`Error updating the budget: ${error}`); + toast.fromError(t("budgets.editModal.toast.error", { error: String(error) })); } }; @@ -55,15 +57,15 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs !open && handleCancel()}> - Edit budget + {t("budgets.editModal.title")}
- }>Cancel + }>{t("budgets.editModal.cancel")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx index 9e9c2ad7fc8..6aee2bf7dfd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useState } from "react"; +import { useTranslation } from "react-i18next"; import { CircleAlert } from "lucide-react"; import { z } from "zod"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -17,21 +18,23 @@ import { useZodForm } from "@/lib/forms/useZodForm"; import { toast } from "@/lib/toast"; import { revokeSessionAndClearClientState } from "@/app/(dashboard)/hooks/useLogout"; import { getLoginUrl } from "@/utils/returnUrlUtils"; +import i18n from "@/i18n/i18n"; const changePasswordSchema = z .object({ - currentPassword: z.string().min(1, "Current password is required"), - newPassword: z.string().min(1, "New password is required"), - confirmNewPassword: z.string().min(1, "Confirm your new password"), + currentPassword: z.string().min(1, { error: () => i18n.t("changePassword.schema.currentRequired") }), + newPassword: z.string().min(1, { error: () => i18n.t("changePassword.schema.newRequired") }), + confirmNewPassword: z.string().min(1, { error: () => i18n.t("changePassword.schema.confirmRequired") }), }) .refine((values) => values.newPassword === values.confirmNewPassword, { - message: "New passwords do not match", + error: () => i18n.t("changePassword.schema.mismatch"), path: ["confirmNewPassword"], }); type ChangePasswordValues = z.infer; export function ChangePasswordForm() { + const { t } = useTranslation(); const { accessToken, passwordResetRequired } = useAuthorized(); const form = useZodForm(changePasswordSchema, { defaultValues: { currentPassword: "", newPassword: "", confirmNewPassword: "" }, @@ -49,12 +52,12 @@ export function ChangePasswordForm() { // The session key was minted restricted; only a fresh login lifts it. // Revoke it server-side too (best-effort) so it doesn't sit valid // until the expiry reaper gets to it. - toast.success("Password updated. Please log in with your new password."); + toast.success(t("changePassword.form.toasts.updatedSignIn")); await revokeSessionAndClearClientState(accessToken); window.location.replace(getUiPageUrl(getLoginUrl(getProxyBaseUrl()))); return; } - toast.success("Password updated"); + toast.success(t("changePassword.form.toasts.updated")); form.reset(); } catch (error) { setSubmitError(extractProxyErrorMessage(error)); @@ -67,34 +70,31 @@ export function ChangePasswordForm() {
-

Change Password

-

- Enter your current password and choose a new one. The new password must meet this proxy's password - policy. -

+

{t("changePassword.form.title")}

+

{t("changePassword.form.description")}

{passwordResetRequired && ( - - Your password must be changed before you can use the dashboard: it was either found in a known data - breach or set by an administrator as a temporary password. After updating it, you will be signed out to - log in again. - + {t("changePassword.form.resetRequired")} )}
- + {({ ref, ...field }) => } - + {({ ref, ...field }) => } - + {({ ref, ...field }) => } @@ -108,8 +108,10 @@ export function ChangePasswordForm() {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index b3acb8f9168..a6823d50b21 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useState } from "react"; +import { Trans, useTranslation } from "react-i18next"; import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; @@ -76,20 +77,21 @@ const SpendRow: React.FC<{ label: string; value: string; hint?: string; subdued? ); const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { + const { t } = useTranslation(); const stats = view.stats; const cheaper = stats.saved_pct != null && stats.saved_pct >= 0; const classifierCost = stats.baseline_spend == null ? null : stats.savings_estimated_classifier_cost ?? null; const comparedAll = stats.savings_estimated_turns === stats.turns; return ( - +

- Total estimated savings + {t("costOptimization.autoRouter.hero.totalSavings")}

- {stats.saved_spend == null ? "Unavailable" : usd(stats.saved_spend)} + {stats.saved_spend == null ? t("costOptimization.autoRouter.unavailable") : usd(stats.saved_spend)}

{stats.saved_pct != null && ( = ({ view }) => {
{stats.baseline_spend != null && !comparedAll && (

- Compared on {stats.savings_estimated_turns.toLocaleString()} of {stats.turns.toLocaleString()} requests; - adaptive and quality routers are excluded + {t("costOptimization.autoRouter.hero.comparedOn", { + compared: stats.savings_estimated_turns.toLocaleString(), + total: stats.turns.toLocaleString(), + })}

)} {stats.unattributed_saved_spend != null && (

- Per-router records differ from recorded savings by {usd(Math.abs(stats.unattributed_saved_spend))}, for - example history from before per-router tracking, so the baseline comparison is unavailable + {t("costOptimization.autoRouter.hero.unattributed", { + amount: usd(Math.abs(stats.unattributed_saved_spend)), + })}

)}
{stats.baseline_spend != null && classifierCost == null && (

- Breakdown unavailable because some usage predates classification-cost tracking. + {t("costOptimization.autoRouter.hero.breakdownUnavailable")}

)}
@@ -157,20 +170,24 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { }; const StackedTurnBar: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => { + const { t } = useTranslation(); const segments = buckets.filter((b) => b.turns > 0); return (
{segments.map((b) => (
))}
@@ -185,47 +202,57 @@ const StackedTurnBar: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => { ); }; -const BucketTable: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => ( - - - - Bucket - Turns - - Hit rate - - - - {buckets.map((b) => ( - - - - - - {b.label} - {b.sublabel} - - - - - {b.turns.toLocaleString()} - - -
-
-
- - - {pctLabel(b.hitRatePct)} - +const BucketTable: React.FC<{ buckets: BucketRow[] }> = ({ buckets }) => { + const { t } = useTranslation(); + return ( +
+ + + + {t("costOptimization.autoRouter.caching.columns.bucket")} + + + {t("costOptimization.autoRouter.caching.columns.turns")} + + + + {t("costOptimization.autoRouter.caching.columns.hitRate")} + - ))} - -
-); + + + {buckets.map((b) => ( + + + + + + {b.label} + {b.sublabel} + + + + + {b.turns.toLocaleString()} + + +
+
+
+ + + {pctLabel(b.hitRatePct)} + + + ))} + + + ); +}; const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => { - const buckets = bucketRows(cache); + const { t } = useTranslation(); + const buckets = bucketRows(cache, t); const total = bucketTurnsTotal(cache); const expiredMissPct = expiredMissShare(cache); return ( @@ -233,7 +260,7 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => {
-

Cache hit rate

+

{t("costOptimization.autoRouter.caching.hitRate")}

{pctLabel(cache.hit_rate_pct)}

{expiredMissPct === null ? null : ( @@ -248,13 +275,12 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => { } > - Expired-miss + {t("costOptimization.autoRouter.caching.expiredMiss.label")} {pctLabel(expiredMissPct)} - share of all measured turns that missed cache because a return to an earlier tier came after its TTL - lapsed + {t("costOptimization.autoRouter.caching.expiredMiss.tooltip")} @@ -263,17 +289,27 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => {
-

Share of turns

+

+ {t("costOptimization.autoRouter.caching.shareOfTurns")} +

- {total.toLocaleString()} turns - measured + {total.toLocaleString()} + ), + }} + />

{cache.unordered_turns > 0 && (

- {cache.unordered_turns.toLocaleString()} turns arrived out of order across pods and are not bucketed + {t("costOptimization.autoRouter.caching.unorderedTurns", { + turns: cache.unordered_turns.toLocaleString(), + })}

)}
@@ -291,13 +327,14 @@ interface BenchmarksBodyProps { } const BenchmarksBody: React.FC = ({ isPending, error, data, selectedKey, autoRouters }) => { - if (isPending) return Loading auto-router usage...; + const { t } = useTranslation(); + if (isPending) return {t("costOptimization.autoRouter.body.loading")}; if (error instanceof ApiError && error.status === 403) { - return Auto-router usage is visible to proxy admin roles only; + return {t("costOptimization.autoRouter.body.forbidden")}; } - if (error || !data) return Auto-router usage is unavailable right now; + if (error || !data) return {t("costOptimization.autoRouter.body.unavailable")}; - const view = viewFor(data, selectedKey); + const view = viewFor(data, selectedKey, t); const stats = view.stats; return ( <> @@ -305,30 +342,32 @@ const BenchmarksBody: React.FC = ({ isPending, error, data, -

- Savings and spend count requests on the selected UTC days. Actual spend covers every request on complexity - routers, including LLM classification cost. Baseline is actual spend plus recorded savings, so savings can be - zero or negative. -

+

{t("costOptimization.autoRouter.body.savingsNote")}

-

- Session metrics cover every session that overlaps the range, including its turns outside the range. -

+

{t("costOptimization.autoRouter.body.sessionNote")}

@@ -338,10 +377,8 @@ const BenchmarksBody: React.FC = ({ isPending, error, data,
-

Auto-router prompt caching

-

- every turn falls in exactly one bucket, by what the router did -

+

{t("costOptimization.autoRouter.caching.title")}

+

{t("costOptimization.autoRouter.caching.subtitle")}

@@ -362,21 +399,26 @@ export const AutoRouterUsageView: React.FC = ({ apiKey, userId, }) => { + const { t } = useTranslation(); const { dateValue, onDateChange } = activity; const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey, userId); const [selectedKey, setSelectedKey] = useState(ALL_ROUTERS); const { data: autoRouters } = useAutoRouters(); const groups = data?.groups ?? []; - const selectedLabel = data ? viewFor(data, selectedKey).label : "All auto-routers"; + const selectedLabel = data ? viewFor(data, selectedKey, t).label : t("costOptimization.autoRouter.allRouters"); const rangeLabel = formatRangeLabel(dateValue.from, dateValue.to); return (
-

Auto-router usage

- {rangeLabel &&

{rangeLabel} (UTC)

} +

{t("costOptimization.autoRouter.usageTitle")}

+ {rangeLabel && ( +

+ {t("costOptimization.autoRouter.rangeUtc", { range: rangeLabel })} +

+ )}
@@ -386,7 +428,7 @@ export const AutoRouterUsageView: React.FC = ({ {selectedLabel} - All auto-routers + {t("costOptimization.autoRouter.allRouters")} {groups.map((g) => ( {groupLabel(g, groups)} @@ -398,12 +440,7 @@ export const AutoRouterUsageView: React.FC = ({
- {userId && ( -

- Usage for this user across API keys and JWT-authenticated requests. Older sessions recorded without a user ID - are not included. -

- )} + {userId &&

{t("costOptimization.autoRouter.userScopeNote")}

} = ({ }; const AutoRouterBenchmarksTab: React.FC = ({ accessToken, activity }) => { + const { t } = useTranslation(); const [visitedTabs, setVisitedTabs] = useState(["usage"]); const handleTabChange = (value: unknown) => { @@ -430,10 +468,10 @@ const AutoRouterBenchmarksTab: React.FC = ({ acces - Usage + {t("costOptimization.autoRouter.tabs.usage")} - Shadow Evals + {t("costOptimization.autoRouter.tabs.shadowEvals")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx index 0273fae9ec1..2a5cf572008 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterSummaryTable.tsx @@ -1,3 +1,5 @@ +import { useTranslation } from "react-i18next"; + import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; @@ -10,33 +12,40 @@ interface AutoRouterSummaryTableProps { } export default function AutoRouterSummaryTable({ groups, selectedGroup }: AutoRouterSummaryTableProps) { + const { t } = useTranslation(); const visibleGroups = selectedGroup ? [selectedGroup] : groups; return ( - Router usage and savings - - Selected UTC days. Cost includes LLM calls and classification. Tokens count routed LLM input and output. - + {t("costOptimization.routerSummary.title")} + {t("costOptimization.routerSummary.description")} - +
- Auto-router - Tokens through router - Cost via router - Cost per 1M tokens - Saved vs. premium model - % saved + {t("costOptimization.routerSummary.columns.router")} + + {t("costOptimization.routerSummary.columns.tokens")} + + + {t("costOptimization.routerSummary.columns.cost")} + + + {t("costOptimization.routerSummary.columns.costPerMillion")} + + + {t("costOptimization.routerSummary.columns.saved")} + + {t("costOptimization.routerSummary.columns.savedPct")} {visibleGroups.length === 0 ? ( - No auto-routers in this range + {t("costOptimization.routerSummary.empty")} ) : ( @@ -44,29 +53,32 @@ export default function AutoRouterSummaryTable({ groups, selectedGroup }: AutoRo {groupLabel(group, groups)} - {group.total_tokens == null ? "Unavailable" : group.total_tokens.toLocaleString()} + {group.total_tokens == null + ? t("costOptimization.routerSummary.unavailable") + : group.total_tokens.toLocaleString()} {usd(group.spend)} {group.total_tokens != null && group.total_tokens > 0 ? usd((group.spend * 1_000_000) / group.total_tokens) - : "Unavailable"} + : t("costOptimization.routerSummary.unavailable")} - {group.saved_spend == null ? "Unavailable" : usd(group.saved_spend)} + {group.saved_spend == null + ? t("costOptimization.routerSummary.unavailable") + : usd(group.saved_spend)} - {group.saved_pct == null ? "Unavailable" : pctLabel(group.saved_pct)} + {group.saved_pct == null + ? t("costOptimization.routerSummary.unavailable") + : pctLabel(group.saved_pct)} )) )}
-

- Savings are estimated against each router's premium baseline. Token totals and unit costs are unavailable - for usage recorded before token tracking; unit cost also requires nonzero tokens. -

+

{t("costOptimization.routerSummary.footnote")}

 ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index 7a63ccd1ee4..ab34f09f0ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { ArrowDown, ArrowUp, ArrowUpDown, Info } from "lucide-react"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; @@ -67,6 +68,7 @@ const SortableHead = ({ sort: SortState; onSort: (column: SortColumn) => void; }) => { + const { t } = useTranslation(); const active = sort.column === column; const ActiveArrow = sort.dir === "asc" ? ArrowUp : ArrowDown; const Arrow = active ? ActiveArrow : ArrowUpDown; @@ -76,7 +78,7 @@ const SortableHead = ({
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx index 1b7369c9520..0e286418355 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx @@ -1,6 +1,8 @@ "use client"; import { ColumnDef } from "@tanstack/react-table"; +import type { TFunction } from "i18next"; +import { useTranslation } from "react-i18next"; import { CredentialIdentity } from "@/components/shared/CredentialIdentity"; import { Copy, Info, Loader2, Pencil, RefreshCw, Trash2 } from "lucide-react"; @@ -14,6 +16,7 @@ import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/h import { Switch } from "@/components/ui/switch"; import { getDisplayModelName } from "@/components/view_model/model_name_display"; import { copyToClipboard, formatPerSecondCost } from "@/utils/dataUtils"; +import i18n from "@/i18n/i18n"; export const MODEL_ID_COLUMN_ID = "model_info_id"; export const MODEL_NAME_COLUMN_ID = "model_name"; @@ -56,6 +59,7 @@ const formatShortDate = (value: string | null | undefined): string | null => { }; function ModelInformationCell({ model, displayName }: { model: ModelData; displayName: string }) { + const { t } = useTranslation(); const litellmModelName = model.litellm_model_name || "-"; return ( @@ -85,26 +89,37 @@ function ModelInformationCell({ model, displayName }: { model: ModelData; displa
{model.provider ? : null} - {model.provider || "Unknown provider"} + + {model.provider || t("modelsAndEndpoints.modelsTableColumns.unknownProvider")} +
- Public Model Name + + {t("modelsAndEndpoints.modelsTableColumns.information.publicModelName")} + {displayName}
- LiteLLM Model Name + + {t("modelsAndEndpoints.modelsTableColumns.information.litellmModelName")} + {litellmModelName} @@ -117,15 +132,16 @@ function ModelInformationCell({ model, displayName }: { model: ModelData; displa } function CredentialsHeader() { + const { t } = useTranslation(); return ( - Credentials + {t("modelsAndEndpoints.modelsTableColumns.columns.credentials")} @@ -135,23 +151,25 @@ function CredentialsHeader() {
- Credential types + + {t("modelsAndEndpoints.modelsTableColumns.credentialTypes.title")} +
- Reusable + {t("modelsAndEndpoints.modelsTableColumns.credentialTypes.reusable")} - Credentials saved in LiteLLM that can be added to models repeatedly. + {t("modelsAndEndpoints.modelsTableColumns.credentialTypes.reusableDescription")}
- Manual + {t("modelsAndEndpoints.modelsTableColumns.credentialTypes.manual")} - Credentials added directly during model creation or defined in the config file. + {t("modelsAndEndpoints.modelsTableColumns.credentialTypes.manualDescription")}
@@ -162,11 +180,12 @@ function CredentialsHeader() { } function CredentialsCell({ credentialName, label }: { credentialName: string | undefined; label: string | undefined }) { + const { t } = useTranslation(); if (!credentialName) { return ( - Manual + {t("modelsAndEndpoints.modelsTableColumns.credentialTypes.manual")} ); } @@ -182,10 +201,13 @@ function CredentialsCell({ credentialName, label }: { credentialName: string | u } function CreatedByCell({ model }: { model: ModelData }) { + const { t } = useTranslation(); const isConfigModel = !model.model_info?.db_model; const createdAt = formatShortDate(model.model_info.created_at); - const primary = isConfigModel ? "Defined in config" : model.model_info.created_by || "Unknown"; - const secondaryForDbModel = createdAt ?? "Unknown date"; + const primary = isConfigModel + ? t("modelsAndEndpoints.modelsTableColumns.createdBy.definedInConfig") + : model.model_info.created_by || t("modelsAndEndpoints.modelsTableColumns.createdBy.unknown"); + const secondaryForDbModel = createdAt ?? t("modelsAndEndpoints.modelsTableColumns.createdBy.unknownDate"); return (
@@ -207,6 +229,7 @@ function CostRow({ label, value }: { label: string; value: string }) { } function CostsCell({ model }: { model: ModelData }) { + const { t } = useTranslation(); const { input_cost: inputCost, output_cost: outputCost, output_cost_per_second: perSecond } = model; const hasPerSecond = perSecond != null; const showInput = inputCost != null && (!hasPerSecond || Number(inputCost) > 0); @@ -218,12 +241,25 @@ function CostsCell({ model }: { model: ModelData }) { return ( - {showInput && } - {showOutput && } - {hasPerSecond && } + {showInput && ( + + )} + {showOutput && ( + + )} + {hasPerSecond && ( + + )}
} /> @@ -231,6 +267,7 @@ function CostsCell({ model }: { model: ModelData }) { } function AccessGroupsCell({ accessGroups }: { accessGroups: string[] | null }) { + const { t } = useTranslation(); if (!accessGroups || accessGroups.length === 0) { return -; } @@ -253,7 +290,7 @@ function AccessGroupsCell({ accessGroups }: { accessGroups: string[] | null }) { } trigger={ - +{overflow.length} more + {t("modelsAndEndpoints.modelsTableColumns.accessGroups.more", { hiddenCount: overflow.length })} } /> @@ -281,6 +318,7 @@ function ModelRowActions({ onDeleteClick, onTogglePauseClick, }: ModelRowActionsProps) { + const { t } = useTranslation(); const modelId = model.model_info?.id; const isConfigModel = !model.model_info?.db_model; const isAdmin = userRole === "Admin" && !isViewOnly; @@ -290,17 +328,19 @@ function ModelRowActions({ const resolvePauseTooltip = (): string => { if (isConfigModel) { - return "Config models cannot be paused from the dashboard. Pause is DB-backed."; + return t("modelsAndEndpoints.modelsTableColumns.actions.configPauseTooltip"); } if (!isAdmin) { - return "Only proxy admins can pause or resume a model."; + return t("modelsAndEndpoints.modelsTableColumns.actions.adminOnlyPauseTooltip"); } - return isBlocked ? "Resume model — restore normal routing." : "Pause model — stop routing requests until resumed."; + return isBlocked + ? t("modelsAndEndpoints.modelsTableColumns.actions.resumeTooltip") + : t("modelsAndEndpoints.modelsTableColumns.actions.pauseTooltip"); }; const deleteTooltip = isConfigModel - ? "Config model cannot be deleted on the dashboard. Please delete it from the config file." - : "Delete model"; + ? t("modelsAndEndpoints.modelsTableColumns.actions.configDeleteTooltip") + : t("modelsAndEndpoints.modelsTableColumns.actions.delete"); return (
@@ -319,7 +359,11 @@ function ModelRowActions({ size="sm" checked={!isBlocked} disabled={!isPauseToggleable} - aria-label={isBlocked ? "Resume model" : "Pause model"} + aria-label={ + isBlocked + ? t("modelsAndEndpoints.modelsTableColumns.actions.resume") + : t("modelsAndEndpoints.modelsTableColumns.actions.pause") + } data-testid={`model-pause-toggle-${modelId}`} onCheckedChange={(nextChecked) => { if (isPauseToggleable && onTogglePauseClick && modelId) { @@ -339,7 +383,7 @@ function ModelRowActions({ , +})); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-test", userRole: "Admin" }), +})); + +const renderTabWithBanner = (queryClient: QueryClient) => + render( + + <> + + + + , + ); + +describe("PriceDataManagementTab", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("refetches deprecations after reloading price data", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "old-model", + litellm_provider: null, + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: null, + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [ + { + model_name: "current-model", + litellm_provider: null, + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: null, + }, + ], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(Response.json(deprecatedResponse)) + .mockResolvedValueOnce(Response.json(currentResponse)); + vi.stubGlobal("fetch", fetchMock); + renderTabWithBanner(queryClient); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "reload" })); + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect( + fetchMock.mock.calls.map( + ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, + ), + ).toEqual(["/model/deprecations", "/model/deprecations"]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx index b2beb5beeed..445dc505ce6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx @@ -2,13 +2,16 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; +import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; import { modelCostMapKeys } from "../../hooks/models/useModelCostMap"; import PriceDataManagementTab from "./PriceDataManagementTab"; vi.mock("@/components/price_data_reload", () => ({ default: ({ onReloadSuccess }: { onReloadSuccess: () => void }) => , })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-test", userRole: "Admin" }), +})); const renderTab = (queryClient: QueryClient) => render( @@ -27,13 +30,16 @@ describe("PriceDataManagementTab", () => { const queryClient = new QueryClient(); const liveKey = modelCostMapKeys.list({}); const catalogKey = modelCostMapKeys.list({ filters: { catalog_only: "true" } }); + const deprecationKey = modelDeprecationKeys.all; queryClient.setQueryData(liveKey, {}); queryClient.setQueryData(catalogKey, {}); + queryClient.setQueryData(deprecationKey, {}); renderTab(queryClient); fireEvent.click(screen.getByRole("button", { name: "reload" })); expect(queryClient.getQueryState(liveKey)?.isInvalidated).toBe(true); expect(queryClient.getQueryState(catalogKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(deprecationKey)?.isInvalidated).toBe(true); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx index ce6f8e77d38..4dddc5333ea 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx @@ -1,10 +1,13 @@ import PriceDataReload from "@/components/price_data_reload"; import React from "react"; +import { useTranslation } from "react-i18next"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useQueryClient } from "@tanstack/react-query"; +import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; import { modelCostMapKeys } from "../../hooks/models/useModelCostMap"; const PriceDataManagementTab = () => { + const { t } = useTranslation(); const { accessToken } = useAuthorized(); const queryClient = useQueryClient(); @@ -12,17 +15,18 @@ const PriceDataManagementTab = () => {
-

Price Data Management

-

- Manage model pricing data and configure automatic reload schedules -

+

{t("modelsAndEndpoints.priceDataManagementTab.title")}

+

{t("modelsAndEndpoints.priceDataManagementTab.description")}

{ - queryClient.invalidateQueries({ queryKey: modelCostMapKeys.all }); + void Promise.all([ + queryClient.invalidateQueries({ queryKey: modelCostMapKeys.all }), + queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }), + ]); }} - buttonText="Reload Price Data" + buttonText={t("modelsAndEndpoints.priceDataManagementTab.reloadButton")} size="middle" type="primary" className="w-full" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts new file mode 100644 index 00000000000..66d065e6caa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it } from "vitest"; +import type { ModelDeprecationInfo } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; +import { + bucketByLocalDate, + countDistinctModels, + deprecationBannerTitle, + deprecationEntryText, + localDaysUntil, + parseLocalDate, +} from "./deprecationBannerText"; + +const makeInfo = (overrides: Partial = {}): ModelDeprecationInfo => ({ + model_name: "gpt-4o", + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + ...overrides, +}); + +describe("local deprecation dates", () => { + it("parses date parts as a local calendar date and counts from local midnight", () => { + const parsed = parseLocalDate("2026-11-06"); + + expect([parsed.getFullYear(), parsed.getMonth(), parsed.getDate(), parsed.getHours()]).toEqual([2026, 10, 6, 0]); + expect(localDaysUntil("2026-11-06", new Date(2026, 10, 4, 18))).toBe(2); + }); + + it("buckets entries by their local date and warning window", () => { + const buckets = bucketByLocalDate( + [ + makeInfo({ model_name: "yesterday", deprecation_date: "2026-11-03", status: "imminent" }), + makeInfo({ model_name: "today", deprecation_date: "2026-11-04", status: "deprecated" }), + makeInfo({ model_name: "outside-window", deprecation_date: "2026-12-05", status: "imminent" }), + ], + new Date(2026, 10, 4), + 30, + ); + + expect( + buckets.deprecated.map(({ model_name, days_until_deprecation, status }) => [ + model_name, + days_until_deprecation, + status, + ]), + ).toEqual([["yesterday", -1, "deprecated"]]); + expect( + buckets.imminent.map(({ model_name, days_until_deprecation, status }) => [ + model_name, + days_until_deprecation, + status, + ]), + ).toEqual([["today", 0, "imminent"]]); + }); +}); + +describe("countDistinctModels", () => { + it("counts visible names and providers, not deployment aliases", () => { + expect( + countDistinctModels([ + makeInfo({ model_name: "team-a", team_public_model_name: "shared", litellm_provider: "openai" }), + makeInfo({ model_name: "team-b", team_public_model_name: "shared", litellm_provider: "openai" }), + makeInfo({ model_name: "team-c", team_public_model_name: "shared", litellm_provider: "anthropic" }), + ]), + ).toBe(2); + }); +}); + +describe("deprecationBannerTitle", () => { + it("formats both counts with singular and plural wording", () => { + expect(deprecationBannerTitle(2, 1, 30)).toBe("2 models are deprecated, 1 deprecates within 30 days"); + expect(deprecationBannerTitle(1, 2, 30)).toBe("1 model is deprecated, 2 deprecate within 30 days"); + }); + + it("omits empty counts and uses singular model wording", () => { + expect(deprecationBannerTitle(1, 0, 30)).toBe("1 model is deprecated"); + expect(deprecationBannerTitle(0, 1, 30)).toBe("1 model deprecates within 30 days"); + expect(deprecationBannerTitle(0, 3, 30)).toBe("3 models deprecate within 30 days"); + expect(deprecationBannerTitle(0, 0, 30)).toBe(""); + }); +}); + +describe("deprecationEntryText", () => { + it("formats deprecated entries with provider and successor", () => { + const deprecatedInfo: ModelDeprecationInfo = { + model_name: "gpt-4", + litellm_provider: "openai", + deprecation_date: "2026-06-01", + days_until_deprecation: -152, + status: "deprecated", + successor_model: "gpt-5.6-sol", + }; + + expect(deprecationEntryText(deprecatedInfo)).toBe( + "gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5.6-sol", + ); + }); + + it("omits a null provider and successor", () => { + expect( + deprecationEntryText( + makeInfo({ + litellm_provider: null, + successor_model: null, + }), + ), + ).toBe("gpt-4o deprecates on 2026-11-30, in 26 days"); + }); + + it("uses singular wording for one day", () => { + expect( + deprecationEntryText( + makeInfo({ + days_until_deprecation: 1, + }), + ), + ).toBe("gpt-4o deprecates on 2026-11-30, in 1 day"); + }); + + it("says an imminent model deprecates today", () => { + expect( + deprecationEntryText( + makeInfo({ + days_until_deprecation: 0, + deprecation_date: "2026-11-04", + }), + ), + ).toBe("gpt-4o deprecates today (2026-11-04)"); + }); + + it("uses the team-visible model name when present", () => { + expect( + deprecationEntryText( + makeInfo({ + model_name: "internal-name", + team_public_model_name: "team-name", + }), + ), + ).toBe("team-name deprecates on 2026-11-30, in 26 days"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts new file mode 100644 index 00000000000..0c035c5bdd6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts @@ -0,0 +1,96 @@ +import type { ModelDeprecationInfo } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; + +const shownModelName = (info: ModelDeprecationInfo): string => info.team_public_model_name ?? info.model_name; + +export const countDistinctModels = (infos: readonly ModelDeprecationInfo[]): number => + new Set(infos.map((info) => JSON.stringify([shownModelName(info), info.litellm_provider ?? null]))).size; + +export const parseLocalDate = (value: string): Date => { + const [year, month, day] = value.split("-"); + return new Date(Number(year), Number(month) - 1, Number(day)); +}; + +const startOfLocalDay = (today: Date): Date => new Date(today.getFullYear(), today.getMonth(), today.getDate()); + +export const localDaysUntil = (deprecationDate: string, today: Date): number => + Math.round((parseLocalDate(deprecationDate).getTime() - startOfLocalDay(today).getTime()) / 86_400_000); + +export const bucketByLocalDate = ( + infos: readonly ModelDeprecationInfo[], + today: Date, + warnWithinDays: number, +): { deprecated: ModelDeprecationInfo[]; imminent: ModelDeprecationInfo[] } => { + const inWindow = infos + .map((info): ModelDeprecationInfo => { + const daysUntilDeprecation = localDaysUntil(info.deprecation_date, today); + return { + ...info, + days_until_deprecation: daysUntilDeprecation, + status: daysUntilDeprecation < 0 ? "deprecated" : "imminent", + }; + }) + .filter((info) => info.days_until_deprecation <= warnWithinDays); + + return { + deprecated: inWindow.filter((info) => info.status === "deprecated"), + imminent: inWindow.filter((info) => info.status === "imminent"), + }; +}; + +const hasSameDedupeKey = (candidate: ModelDeprecationInfo, info: ModelDeprecationInfo): boolean => { + const sameDisplayedModel = shownModelName(candidate) === shownModelName(info); + const sameProvider = (candidate.litellm_provider ?? null) === (info.litellm_provider ?? null); + const sameDeprecationDate = candidate.deprecation_date === info.deprecation_date; + const sameSuccessor = (candidate.successor_model ?? null) === (info.successor_model ?? null); + const sameVisibleIdentity = sameDisplayedModel && sameProvider; + const sameDeprecationIdentity = sameDeprecationDate && sameSuccessor; + return sameVisibleIdentity && sameDeprecationIdentity; +}; + +export const dedupeDeprecationEntries = (infos: readonly ModelDeprecationInfo[]): ModelDeprecationInfo[] => + infos.filter((info, index) => infos.findIndex((candidate) => hasSameDedupeKey(candidate, info)) === index); + +export const deprecationBannerFingerprint = (infos: readonly ModelDeprecationInfo[]): string => + JSON.stringify( + infos + .map((info) => + JSON.stringify([ + shownModelName(info), + info.litellm_provider ?? null, + info.deprecation_date, + info.status, + info.successor_model ?? null, + ]), + ) + .sort(), + ); + +export const deprecationBannerTitle = ( + deprecatedCount: number, + imminentCount: number, + warnWithinDays: number, +): string => { + const deprecatedText = + deprecatedCount === 0 ? "" : `${deprecatedCount} ${deprecatedCount === 1 ? "model is" : "models are"} deprecated`; + const imminentModelLabel = imminentCount === 1 ? "model" : "models"; + const imminentSubject = deprecatedCount > 0 ? `${imminentCount}` : `${imminentCount} ${imminentModelLabel}`; + const imminentVerb = imminentCount === 1 ? "deprecates" : "deprecate"; + const imminentText = imminentCount === 0 ? "" : `${imminentSubject} ${imminentVerb} within ${warnWithinDays} days`; + return [deprecatedText, imminentText].filter((text) => text.length > 0).join(", "); +}; + +export const deprecationEntryText = (info: ModelDeprecationInfo): string => { + const model = `${shownModelName(info)}${info.litellm_provider ? ` (${info.litellm_provider})` : ""}`; + const successor = info.successor_model ? `, migrate to ${info.successor_model}` : ""; + + if (info.status === "deprecated") { + return `${model} was deprecated on ${info.deprecation_date}${successor}`; + } + + if (info.days_until_deprecation === 0) { + return `${model} deprecates today (${info.deprecation_date})${successor}`; + } + + const dayLabel = info.days_until_deprecation === 1 ? "day" : "days"; + return `${model} deprecates on ${info.deprecation_date}, in ${info.days_until_deprecation} ${dayLabel}${successor}`; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx index 4d18966625e..7aeff6f44f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx @@ -1,9 +1,9 @@ /* @vitest-environment jsdom */ -import { screen } from "@testing-library/react"; +import { screen, waitForElementToBeRemoved } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "@/components/networking"; -import { renderWithProviders } from "../../../../tests/test-utils"; +import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; import ModelsAndEndpointsPage from "./page"; vi.mock("./panels/AllModelsPanel", () => ({ default: () =>
})); @@ -201,6 +201,7 @@ const createMockTeamData = (overrides = {}) => ({ describe("ModelsAndEndpointsPage ?team drill-in", () => { beforeEach(() => { + testQueryClient.clear(); authState.userRole = "Admin"; authState.isViewOnly = false; can.mockReturnValue(true); @@ -221,6 +222,7 @@ describe("ModelsAndEndpointsPage ?team drill-in", () => { }); afterEach(() => { + vi.unstubAllGlobals(); vi.clearAllMocks(); }); @@ -244,4 +246,53 @@ describe("ModelsAndEndpointsPage ?team drill-in", () => { expect(screen.queryByRole("button", { name: /edit settings/i })).not.toBeInTheDocument(); expect(screen.queryByText("Disable all global guardrails")).not.toBeInTheDocument(); }); + + it("refreshes deprecation banner data when the models list is refreshed", async () => { + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "gpt-4", + litellm_provider: "openai", + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: "gpt-5", + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockImplementationOnce(async (input) => { + const requestUrl = input instanceof Request ? input.url : String(input); + expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); + return Response.json(deprecatedResponse); + }) + .mockImplementationOnce(async (input) => { + const requestUrl = input instanceof Request ? input.url : String(input); + expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); + return Response.json(currentResponse); + }); + vi.stubGlobal("fetch", fetchMock); + + const user = userEvent.setup({ delay: null }); + renderWithProviders(); + + expect( + await screen.findByText("gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5"), + ).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Refresh models" })); + await waitForElementToBeRemoved(() => screen.queryByRole("alert")); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 3cfed7297c8..be8ee1bb8c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -29,6 +29,9 @@ vi.mock("./detailNavigation", () => ({ vi.mock("@/components/molecules/DecisionModelsBanner", () => ({ default: () =>
, })); +vi.mock("./components/ModelDeprecationBanner", () => ({ + default: () =>
, +})); vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, })); @@ -109,6 +112,12 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByRole("region", { name: "Decision models" })).not.toBeInTheDocument(); }); + it("renders the model deprecation banner", () => { + renderPage(); + + expect(screen.getByRole("region", { name: "Model deprecations" })).toBeInTheDocument(); + }); + it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; renderPage(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 9f1ee04c5a2..b4082a14ffe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -4,7 +4,10 @@ import { useMemo, useState } from "react"; import { RefreshCw } from "lucide-react"; import { parseAsString, useQueryStates } from "nuqs"; import { useQueryClient } from "@tanstack/react-query"; +import type { ParseKeys } from "i18next"; +import { useTranslation } from "react-i18next"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { useUrlTab } from "@/hooks/useUrlTab"; @@ -28,6 +31,7 @@ import AccessGroupBudgetsPanel from "@/app/(dashboard)/models-and-endpoints/pane import PriceDataPanel from "@/app/(dashboard)/models-and-endpoints/panels/PriceDataPanel"; import { Button } from "@/components/ui/button"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import ModelDeprecationBanner from "./components/ModelDeprecationBanner"; const MODEL_TAB_SLUGS = [ "add", @@ -49,16 +53,16 @@ const TAB_KEYS = [BASE_TAB_KEY, ...MODEL_TAB_SLUGS] as const; const PREFILL_URL_PARAMS = { provider: parseAsString, litellm_model: parseAsString }; -const TAB_LABELS: Record = { - add: "Add Model", - "auto-routers": "Auto-Routers", - "llm-credentials": "LLM Credentials", - "pass-through": "Pass-Through Endpoints", - health: "Health Status", - "retry-settings": "Model Retry Settings", - "model-group-alias": "Model Group Alias", - "access-group-budgets": "Model Access Group Budgets", - "price-data": "Price Data Reload", +const TAB_LABELS: Record = { + add: "modelsAndEndpoints.page.tabs.add", + "auto-routers": "modelsAndEndpoints.page.tabs.autoRouters", + "llm-credentials": "modelsAndEndpoints.page.tabs.llmCredentials", + "pass-through": "modelsAndEndpoints.page.tabs.passThrough", + health: "modelsAndEndpoints.page.tabs.health", + "retry-settings": "modelsAndEndpoints.page.tabs.retrySettings", + "model-group-alias": "modelsAndEndpoints.page.tabs.modelGroupAlias", + "access-group-budgets": "modelsAndEndpoints.page.tabs.accessGroupBudgets", + "price-data": "modelsAndEndpoints.page.tabs.priceData", }; const renderPanel = (key: string) => { @@ -89,6 +93,7 @@ const renderPanel = (key: string) => { }; export default function ModelsAndEndpointsPage() { + const { t } = useTranslation(); const { accessToken, userRole, userId: userID, premiumUser, isViewOnly } = useAuthorized(); const { data: teams } = useTeams(); const { data: uiSettings } = useUISettings(); @@ -146,25 +151,27 @@ export default function ModelsAndEndpointsPage() { void setPrefillParams({ provider: null, litellm_model: null }); }; - const allModelsLabel = isAdmin ? "Deployed Models" : "Your Models"; + const allModelsLabel = isAdmin + ? t("modelsAndEndpoints.page.tabs.deployedModels") + : t("modelsAndEndpoints.page.tabs.yourModels"); const tabLabel = (slug: "" | ModelTabSlug): React.ReactNode => { if (!slug) return allModelsLabel; if (slug === "auto-routers" || slug === "access-group-budgets") { return ( - {TAB_LABELS[slug]} + {t(TAB_LABELS[slug])} ); } - return TAB_LABELS[slug]; + return t(TAB_LABELS[slug]); }; const handleRefreshClick = () => { setLastRefreshed(new Date().toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })); - queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + void invalidateModelQueries(queryClient); }; - const invalidateModels = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + const invalidateModels = () => invalidateModelQueries(queryClient); if (teamId) { return ( @@ -189,18 +196,17 @@ export default function ModelsAndEndpointsPage() {
-

Model Management

+

{t("modelsAndEndpoints.page.title")}

{isAdmin ? ( -

Add and manage models for the proxy

+

{t("modelsAndEndpoints.page.adminDescription")}

) : ( -

- View your models and manage routers for teams that allow it. -

+

{t("modelsAndEndpoints.page.userDescription")}

)}
{!isViewOnly && } + {modelId ? (
{lastRefreshed && ( - Last Refreshed: {lastRefreshed} + + {t("modelsAndEndpoints.page.lastRefreshed", { time: lastRefreshed })} + )} -
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.tsx index 20eb7b2e64c..d632408c85a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.tsx @@ -3,6 +3,7 @@ import { SortingState } from "@tanstack/react-table"; import { Inbox } from "lucide-react"; import React, { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import { DataTable } from "@/components/shared/DataTable"; @@ -21,20 +22,24 @@ import { getAccessGroupBudgetColumns } from "@/app/(dashboard)/models-and-endpoi const DEFAULT_SORTING: SortingState = [{ id: "access_group", desc: false }]; function EmptyState() { + const { t } = useTranslation(); return (
-
No model access groups yet
+
+ {t("modelsAndEndpoints.accessGroupBudgetsPanel.empty.title")} +
- Put a deployment in an access group from its model settings, then give the group a shared budget here. + {t("modelsAndEndpoints.accessGroupBudgetsPanel.empty.description")}
); } export default function AccessGroupBudgetsPanel() { + const { t } = useTranslation(); const { userRole } = useAuthorized(); const { data: accessGroups, isLoading } = useModelAccessGroups(); const setBudget = useSetModelAccessGroupBudget(); @@ -57,7 +62,7 @@ export default function AccessGroupBudgetsPanel() { { accessGroup, params }, { onSuccess: () => { - toast.success(`Budget saved for "${accessGroup}"`); + toast.success(t("modelsAndEndpoints.accessGroupBudgetsPanel.toasts.saved", { accessGroup })); setEditing(null); }, }, @@ -69,7 +74,7 @@ export default function AccessGroupBudgetsPanel() { const accessGroup = clearing.access_group; clearBudget.mutate(accessGroup, { onSuccess: () => { - toast.success(`Budget cleared for "${accessGroup}"`); + toast.success(t("modelsAndEndpoints.accessGroupBudgetsPanel.toasts.cleared", { accessGroup })); setClearing(null); }, }); @@ -77,10 +82,7 @@ export default function AccessGroupBudgetsPanel() { return (
-

- A model access group can carry one budget that every key granted the group by name draws from together. Keys - that reach the group's models through a wildcard or all-proxy-models are not charged against it. -

+

{t("modelsAndEndpoints.accessGroupBudgetsPanel.description")}

} size="compact" /> @@ -105,12 +107,19 @@ export default function AccessGroupBudgetsPanel() { setClearing(null)} onOk={handleConfirmClear} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx index c9b3defcd28..3d299aa0971 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx @@ -2,8 +2,10 @@ import { useSyncExternalStore } from "react"; import { act, fireEvent, renderWithProviders, screen, waitFor, within } from "../../../../../tests/test-utils"; import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { testQueryClient } from "../../../../../tests/test-utils"; import AddModelPanel from "./AddModelPanel"; +import ModelDeprecationBanner from "../components/ModelDeprecationBanner"; const modelCreateCall = vi.fn(); const mockPtuEnabled = vi.fn(); @@ -147,11 +149,23 @@ const advancedOpenExtras = { const baseModelInfo = { access_groups: undefined, mode: undefined }; +const TEST_TODAY = new Date(2026, 10, 4); + const { api_base: _omitted, ...ALWAYS_MOUNTED_WITHOUT_API_BASE } = alwaysMounted; -const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction) => { +const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction, showBanner = false) => { const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); - renderWithProviders(, { searchParams, onUrlUpdate }); + renderWithProviders( + showBanner ? ( + <> + + + + ) : ( + + ), + { searchParams, onUrlUpdate }, + ); await screen.findByText("Provider"); const openAdvanced = async () => { @@ -186,11 +200,69 @@ const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction) = describe("AddModelPanel submit payload contract", () => { beforeEach(() => { vi.clearAllMocks(); + testQueryClient.clear(); mockUseModelCostMap.mockReturnValue({ data: {} }); mockPtuEnabled.mockReturnValue(false); mockAuthorized.mockReturnValue(PROXY_ADMIN); }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("refreshes deprecation banner data after adding a model", async () => { + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "old-model", + litellm_provider: null, + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: null, + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [ + { + model_name: "current-model", + litellm_provider: null, + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: null, + }, + ], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(Response.json(deprecatedResponse)) + .mockResolvedValueOnce(Response.json(currentResponse)); + vi.stubGlobal("fetch", fetchMock); + + const { fillRequired, submit } = await setup(undefined, undefined, true); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + await fillRequired(); + await submit(); + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect( + fetchMock.mock.calls.map( + ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, + ), + ).toEqual(["/model/deprecations", "/model/deprecations"]); + }); + it("lists catalog models in the model picker, not entries registered at runtime for deployments", async () => { mockUseModelCostMap.mockImplementation((_enabled: boolean, catalogOnly: boolean) => ({ data: catalogOnly diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx index 392a3c3940a..13bec9506e4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx @@ -15,6 +15,7 @@ import { Providers, getPlaceholder, getProviderModels } from "@/components/provi import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; +import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { buildDecisionCatalog } from "@/lib/decisionModels"; @@ -65,7 +66,7 @@ export default function AddModelPanel() { const decisionCatalog = useMemo(() => buildDecisionCatalog(modelCostMapData), [modelCostMapData]); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); - const refresh = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + const refresh = () => invalidateModelQueries(queryClient); const mountedValues = () => projectMountedValues(registry, form.getValues); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelRetrySettingsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelRetrySettingsPanel.tsx index a5c10ca8a8b..50dc4908865 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelRetrySettingsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelRetrySettingsPanel.tsx @@ -1,6 +1,7 @@ "use client"; import { useCallback, useEffect, useState } from "react"; +import { useTranslation } from "react-i18next"; import ModelRetrySettingsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab"; import { getCallbacksCall } from "@/components/networking"; import { useUpdateRetryPolicy } from "@/app/(dashboard)/hooks/routerSettings/useUpdateRetryPolicy"; @@ -23,6 +24,7 @@ interface RouterSettings { } export default function ModelRetrySettingsPanel() { + const { t } = useTranslation(); const { accessToken, userId: userID, userRole } = useAuthorized(); const { availableModelGroups } = useModelDashboardData(); const updateRetryPolicy = useUpdateRetryPolicy(accessToken); @@ -69,7 +71,7 @@ export default function ModelRetrySettingsPanel() { { retry_policy: globalRetryPolicy, model_group_retry_policy: modelGroupRetryPolicy }, { onSuccess: () => { - toast.success("Retry settings saved successfully"); + toast.success(t("modelsAndEndpoints.modelRetrySettingsPanel.toasts.saved")); void fetchRetrySettings().then((routerSettings) => { if (routerSettings) { applyRetrySettings(routerSettings); @@ -77,7 +79,7 @@ export default function ModelRetrySettingsPanel() { }); }, onError: () => { - toast.fromError("Failed to save retry settings"); + toast.fromError(t("modelsAndEndpoints.modelRetrySettingsPanel.toasts.saveFailed")); }, }, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx index f94c43b720a..4e337a88add 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx @@ -2,6 +2,8 @@ import { useId, useState } from "react"; import { Check, ExternalLink, KeyRound } from "lucide-react"; +import type { TFunction } from "i18next"; +import { useTranslation } from "react-i18next"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Button } from "@/components/ui/button"; import { InputGroupButton } from "@/components/ui/input-group"; @@ -13,13 +15,14 @@ import type { MCPServer } from "@/components/mcp_tools/types"; import { fetchClient } from "@/lib/http/api"; import { ApiError } from "@/lib/http/client"; import { toast } from "@/lib/toast"; +import i18n from "@/i18n/i18n"; -const saveErrorMessage = (error: unknown): string => { +const saveErrorMessage = (error: unknown, t: TFunction = i18n.t.bind(i18n)): string => { if (error instanceof ApiError) { const detail = (error.body as { detail?: { error?: string } } | null)?.detail?.error; if (detail) return detail; } - return error instanceof Error && error.message ? error.message : "Failed to connect"; + return error instanceof Error && error.message ? error.message : t("playground.chat.byok.failed"); }; export function ByokConnect({ @@ -33,6 +36,7 @@ export function ByokConnect({ onConnected: () => void; onCancel?: () => void; }) { + const { t } = useTranslation(); const [credential, setCredential] = useState(""); const [save, setSave] = useState(true); const [isSaving, setIsSaving] = useState(false); @@ -41,7 +45,7 @@ export function ByokConnect({ const connect = async () => { if (!credential.trim()) { - toast.error("Enter your API key"); + toast.error(t("playground.chat.byok.enterKey")); return; } setIsSaving(true); @@ -50,11 +54,11 @@ export function ByokConnect({ params: { path: { server_id: server.server_id } }, body: { credential: credential.trim(), save }, }); - toast.success(`Connected to ${name}`); + toast.success(t("playground.chat.byok.connected", { name })); setCredential(""); onConnected(); } catch (error) { - toast.error(saveErrorMessage(error)); + toast.error(saveErrorMessage(error, t)); } finally { setIsSaving(false); } @@ -65,12 +69,14 @@ export function ByokConnect({

- {server.has_user_credential ? `Replace your ${name} API key` : `${name} needs your API key`} + {server.has_user_credential + ? t("playground.chat.byok.replaceTitle", { name }) + : t("playground.chat.byok.connectTitle", { name })}

{server.has_user_credential - ? `The new key replaces the one the gateway uses for ${name}.` - : `The gateway calls ${name} with your own key. Connect it to list and use its tools.`} + ? t("playground.chat.byok.replaceDescription", { name }) + : t("playground.chat.byok.connectDescription", { name })}

{server.byok_description && server.byok_description.length > 0 && ( @@ -84,17 +90,17 @@ export function ByokConnect({ )} - {name} API key + {t("playground.chat.byok.keyLabel", { name })} setCredential(event.target.value)} - placeholder="Paste your API key" + placeholder={t("playground.chat.byok.keyPlaceholder")} /> {server.byok_api_key_help_url && ( - Where to find this key + {t("playground.chat.byok.helpLink")} @@ -102,18 +108,18 @@ export function ByokConnect({ - Save the key for next time + {t("playground.chat.byok.save")}
{onCancel && ( )}
@@ -130,12 +136,13 @@ export function ByokKeyButton({ name: string; onConnected: () => void; }) { + const { t } = useTranslation(); const [open, setOpen] = useState(false); return ( }> - {server.has_user_credential ? "Replace key" : "Connect key"} + {server.has_user_credential ? t("playground.chat.byok.replaceKey") : t("playground.chat.byok.connectKey")} void; }) { + const { t } = useTranslation(); const input = useRef(null); if (!capabilities.attach) return null; - const label = ATTACH_LABEL[capabilities.attach]; + const label = t(ATTACH_LABEL[capabilities.attach]); return ( <> @@ -69,9 +71,15 @@ function AttachButton({ } function SendButton({ draft, disabled }: { draft: DraftState; disabled: boolean }) { + const { t } = useTranslation(); if (draft.isLoading) return ( - + ); @@ -79,7 +87,7 @@ function SendButton({ draft, disabled }: { draft: DraftState; disabled: boolean void draft.send()} > @@ -115,6 +123,7 @@ export function Composer({ settingsControl: ReactNode; sendDisabled: boolean; }) { + const { t } = useTranslation(); const isDirectCall = endpoint === EndpointType.MCP; const isRealtime = endpoint === EndpointType.REALTIME; const message = useRef(null); @@ -138,7 +147,7 @@ export function Composer({ {directTool ? ( ) : ( -

Pick a server and a tool above to fill in its arguments

+

{t("playground.chat.composer.pickTool")}

)}
); @@ -148,9 +157,9 @@ export function Composer({ value={draft.text} onChange={(event) => draft.setText(event.target.value)} onKeyDown={handleKeyDown} - placeholder={capabilities.placeholder} + placeholder={t(capabilities.placeholder)} rows={2} - aria-label="Message" + aria-label={t("playground.chat.composer.message")} className="max-h-60" /> ); @@ -168,21 +177,31 @@ export function Composer({ {capabilities.mcp && mcpControl} {capabilities.vectorStores && ( - } emptyText="No vector stores yet" {...vectorStores} /> + } + emptyText={t("playground.chat.composer.noVectorStores")} + {...vectorStores} + /> )} {capabilities.guardrails && ( - } emptyText="No guardrails yet" {...guardrails} /> + } + emptyText={t("playground.chat.composer.noGuardrails")} + {...guardrails} + /> )} {capabilities.codeInterpreter && ( - Code interpreter + {t("playground.chat.composer.codeInterpreter")} )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx index e3e52be3b7b..42ef5c92f09 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx @@ -1,6 +1,7 @@ "use client"; import { ChevronDown, FileText, Music, X } from "lucide-react"; +import { useTranslation } from "react-i18next"; import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; import { Attachment, @@ -33,16 +34,19 @@ export function EndpointMenu({ value: EndpointType; onValueChange: (value: EndpointType) => void; }) { + const { t } = useTranslation(); const current = ENDPOINT_OPTIONS.find((option) => option.value === value); return ( - }> + } + > {shortEndpoint(current?.label ?? value)} onValueChange(next as EndpointType)}> - Endpoint + {t("playground.chat.controls.endpoint")} {ENDPOINT_OPTIONS.map((option) => ( {option.label} @@ -61,19 +65,24 @@ export function VoiceMenu({ value: OpenAIVoice; onValueChange: (value: OpenAIVoice) => void; }) { + const { t } = useTranslation(); const current = OPEN_AI_VOICE_SELECT_OPTIONS.find((option) => option.value === value); return ( - }> - Voice: {current?.label.split(" - ")[0] ?? value} + } + > + + {t("playground.chat.controls.voiceValue", { voice: current ? t(current.label).split(" - ")[0] : value })} + onValueChange(next as OpenAIVoice)}> - Voice + {t("playground.chat.controls.voice")} {OPEN_AI_VOICE_SELECT_OPTIONS.map((option) => ( - {option.label} + {t(option.label)} ))} @@ -99,6 +108,7 @@ export function DraftAttachments({ files: readonly DraftFile[]; onRemove: (index: number) => void; }) { + const { t } = useTranslation(); return ( {files.map((draft, index) => ( @@ -109,7 +119,10 @@ export function DraftAttachments({ {formatBytes(draft.file.size)} - onRemove(index)}> + onRemove(index)} + > diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx index 31eb395408a..c6b214eb081 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx @@ -16,6 +16,9 @@ import { Timer, Wrench, } from "lucide-react"; +import type { TFunction } from "i18next"; +import { Trans, useTranslation } from "react-i18next"; +import i18n from "@/i18n/i18n"; import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; import ReasoningContent from "@/components/chat_ui/ReasoningContent"; @@ -68,11 +71,20 @@ const fellBack = (message: MessageType): boolean => (message.trace?.attemptedFal export const servedModel = (message: MessageType): string | undefined => message.trace?.modelGroup ?? (message.isError ? undefined : message.model ?? message.trace?.requestedModel); -const replyModel = (message: MessageType): Readonly<{ value: string; detail: string }> | null => { +const replyModel = ( + message: MessageType, + t: TFunction = i18n.t.bind(i18n), +): Readonly<{ value: string; detail: string }> | null => { const served = servedModel(message); - if (served) return { value: served, detail: fellBack(message) ? "Served by, after a fallback" : "Served by" }; + if (served) + return { + value: served, + detail: fellBack(message) + ? t("playground.chat.conversation.servedByAfterFallback") + : t("playground.chat.conversation.servedBy"), + }; const requested = message.trace?.requestedModel ?? message.model; - return requested ? { value: requested, detail: "Requested model" } : null; + return requested ? { value: requested, detail: t("playground.chat.conversation.requestedModel") } : null; }; function Metric({ icon, value, label }: { icon: ReactNode; value: string; label?: string }) { @@ -86,35 +98,60 @@ function Metric({ icon, value, label }: { icon: ReactNode; value: string; label? } function ReplyMetrics({ message }: { message: MessageType }) { + const { t } = useTranslation(); const usage = message.usage; const cost = usage?.cost ?? message.trace?.cost; const cached = usage?.servedFromResponseCache || message.trace?.cacheHit; return (
{message.timeToFirstToken !== undefined && ( - } value={formatSeconds(message.timeToFirstToken)} label="to first token" /> + } + value={formatSeconds(message.timeToFirstToken)} + label={t("playground.chat.conversation.metrics.firstToken")} + /> )} {message.totalLatency !== undefined && ( - } value={formatSeconds(message.totalLatency)} label="total" /> + } + value={formatSeconds(message.totalLatency)} + label={t("playground.chat.conversation.metrics.total")} + /> )} {usage?.promptTokens !== undefined && ( - } value={String(usage.promptTokens)} label="in" /> + } + value={String(usage.promptTokens)} + label={t("playground.chat.conversation.metrics.in")} + /> )} {usage?.completionTokens !== undefined && ( - } value={String(usage.completionTokens)} label="out" /> + } + value={String(usage.completionTokens)} + label={t("playground.chat.conversation.metrics.out")} + /> )} {usage?.reasoningTokens ? ( - } value={String(usage.reasoningTokens)} label="reasoning" /> + } + value={String(usage.reasoningTokens)} + label={t("playground.chat.conversation.metrics.reasoning")} + /> ) : null} {usage?.cacheReadTokens ? ( - } value={String(usage.cacheReadTokens)} label="cached" /> + } + value={String(usage.cacheReadTokens)} + label={t("playground.chat.conversation.metrics.cached")} + /> ) : null} {cost !== undefined && } value={formatCost(cost)} />} {message.toolName && } value={message.toolName} />} - {cached && Cached response} + {cached && {t("playground.chat.conversation.metrics.cachedResponse")}}
); } @@ -152,7 +189,10 @@ function UserTurn({ message }: { message: MessageType }) { - Called {message.mcpToolCall.tool} + {message.mcpToolCall.tool} }} + /> 
                 {JSON.stringify(message.mcpToolCall.arguments, null, 2)}
@@ -190,7 +230,8 @@ function AssistantReply({
   inspected: boolean;
   onInspect: () => void;
 }) {
-  const model = replyModel(message);
+  const { t } = useTranslation();
+  const model = replyModel(message, t);
   const showMcpEvents = isLast && extras.mcpEvents.length > 0 && MCP_EVENT_ENDPOINTS.includes(extras.endpoint);
   const showCode = isLast && extras.codeInterpreter && extras.endpoint === EndpointType.RESPONSES;
   return (
@@ -201,7 +242,7 @@ function AssistantReply({
             
             {fellBack(message) && }
@@ -220,7 +261,7 @@ function AssistantReply({
         )}
         {message.isError ? (
           
-            The request failed
+            {t("playground.chat.conversation.requestFailed")}
             {textOf(message)}
           
         ) : (
@@ -239,17 +280,19 @@ function AssistantReply({
             )}
- {!message.isImage && !message.isAudio && } + {!message.isImage && !message.isAudio && ( + + )} {message.trace && ( - Inspect + {t("playground.chat.conversation.inspect")} )} @@ -260,12 +303,17 @@ function AssistantReply({ } function WaitingTurn({ target }: { target: string | undefined }) { + const { t } = useTranslation(); return (
- {target ? `Waiting for ${target}` : "Waiting for a reply"} + + {target + ? t("playground.chat.conversation.waitingFor", { target }) + : t("playground.chat.conversation.waitingForReply")} +
@@ -273,12 +321,13 @@ function WaitingTurn({ target }: { target: string | undefined }) { } function SwitchMarker({ model }: { model: string }) { + const { t } = useTranslation(); return ( - Switched to {model} + {t("playground.chat.conversation.switchedTo", { model })} ); } @@ -290,16 +339,15 @@ function EmptyConversation({ suggestions: readonly string[]; onPick: (text: string) => void; }) { + const { t } = useTranslation(); return ( - Send a test request - - Pick a model and anything the request should go through, then inspect what the gateway did with it - + {t("playground.chat.conversation.empty.title")} + {t("playground.chat.conversation.empty.description")} {suggestions.length > 0 && ( @@ -340,13 +388,14 @@ export function Conversation({ suggestions: readonly string[]; onPickSuggestion: (text: string) => void; }) { + const { t } = useTranslation(); if (messages.length === 0) return ; const lastIndex = messages.length - 1; const waiting = isLoading && messages[lastIndex]?.role === "user"; return ( - + {messages.flatMap((message, index) => { const switched = switchedModel(messages, index); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx index 3a9919ce38f..d294c859a18 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx @@ -14,6 +14,9 @@ import { Shuffle, Timer, } from "lucide-react"; +import type { TFunction } from "i18next"; +import { useTranslation } from "react-i18next"; +import i18n from "@/i18n/i18n"; import type { MessageType } from "@/components/chat_ui/types"; import { Accordion, AccordionContent, AccordionItem, AccordionTrigger } from "@/components/ui/accordion"; import { Badge } from "@/components/ui/badge"; @@ -80,130 +83,175 @@ const millis = (value: number | undefined): string | undefined => value === undefined ? undefined : `${Math.round(value)}ms`; function Routing({ message }: { message: MessageType }) { + const { t } = useTranslation(); const trace = message.trace; const served = servedModel(message); const fellBack = (trace?.attemptedFallbacks ?? 0) > 0; const reported = Boolean(trace?.callId || trace?.modelGroup || trace?.providerModel); return ( - + {trace?.requestedModel && ( - - + + )} {served && ( - + {fellBack && } - + )} {trace?.providerModel && ( - + )} {trace?.apiBase && ( - + )} {trace?.callId && ( - + {trace.callId} - + )} {trace?.responseId && ( - + {trace.responseId} - + )} , value: count(trace?.attemptedFallbacks) }, - { label: "Retries", icon: , value: count(trace?.attemptedRetries) }, - { label: "Overhead", icon: , value: millis(trace?.overheadMs) }, + { + label: t("playground.chat.inspector.routing.fallbacks"), + icon: , + value: count(trace?.attemptedFallbacks), + }, + { + label: t("playground.chat.inspector.routing.retries"), + icon: , + value: count(trace?.attemptedRetries), + }, + { label: t("playground.chat.inspector.routing.overhead"), icon: , value: millis(trace?.overheadMs) }, ]} /> - {!reported && ( - - The gateway did not share routing details with this page. Open the dashboard from the proxy to see them - - )} + {!reported && {t("playground.chat.inspector.routing.notReported")}} ); } -const noneRan = (policies: readonly string[] | undefined): string | null => { - if (policies === undefined) return "No guardrails ran on this request"; - return policies.length === 0 ? "No guardrails or policies ran on this request" : null; +const noneRan = (policies: readonly string[] | undefined, t: TFunction = i18n.t.bind(i18n)): string | null => { + if (policies === undefined) return t("playground.chat.inspector.guardrails.noGuardrailsRan"); + return policies.length === 0 ? t("playground.chat.inspector.guardrails.noneRan") : null; }; function Guardrails({ message }: { message: MessageType }) { + const { t } = useTranslation(); const guardrails = message.trace?.appliedGuardrails; const policies = message.trace?.appliedPolicies; if (guardrails === undefined) { return ( - - This endpoint does not report which guardrails or policies ran + + {t("playground.chat.inspector.guardrails.unsupported")} ); } - const emptyNote = guardrails.length === 0 ? noneRan(policies) : null; + const emptyNote = guardrails.length === 0 ? noneRan(policies, t) : null; return ( - + {guardrails.map((name) => ( - - Ran + + {t("playground.chat.inspector.guardrails.ran")} ))} {policies?.map((name) => ( - - Ran + + {t("playground.chat.inspector.guardrails.ran")} ))} {emptyNote && {emptyNote}} - {policies === undefined && Policies are only reported when the dashboard is opened from the proxy} + {policies === undefined && {t("playground.chat.inspector.guardrails.policiesUnreported")}} ); } function SpeedAndCost({ message }: { message: MessageType }) { + const { t } = useTranslation(); const usage = message.usage; const cost = usage?.cost ?? message.trace?.cost; return ( - + , value: message.timeToFirstToken === undefined ? undefined : formatSeconds(message.timeToFirstToken), }, { - label: "Total", + label: t("playground.chat.inspector.performance.total"), icon: , value: message.totalLatency === undefined ? undefined : formatSeconds(message.totalLatency), }, - { label: "Cost", icon: , value: cost === undefined ? undefined : formatCost(cost) }, - { label: "Tokens in", icon: , value: count(usage?.promptTokens) }, - { label: "Tokens out", icon: , value: count(usage?.completionTokens) }, - { label: "Cached", icon: , value: count(usage?.cacheReadTokens) }, - { label: "Reasoning", icon: , value: count(usage?.reasoningTokens) }, - { label: "Gateway time", icon: , value: millis(message.trace?.responseMs) }, + { + label: t("playground.chat.inspector.performance.cost"), + icon: , + value: cost === undefined ? undefined : formatCost(cost), + }, + { + label: t("playground.chat.inspector.performance.tokensIn"), + icon: , + value: count(usage?.promptTokens), + }, + { + label: t("playground.chat.inspector.performance.tokensOut"), + icon: , + value: count(usage?.completionTokens), + }, + { + label: t("playground.chat.inspector.performance.cached"), + icon: , + value: count(usage?.cacheReadTokens), + }, + { + label: t("playground.chat.inspector.performance.reasoning"), + icon: , + value: count(usage?.reasoningTokens), + }, + { + label: t("playground.chat.inspector.performance.gatewayTime"), + icon: , + value: millis(message.trace?.responseMs), + }, ]} /> @@ -211,13 +259,12 @@ function SpeedAndCost({ message }: { message: MessageType }) { } function NothingYet() { + const { t } = useTranslation(); return ( - Nothing to inspect yet - - Each reply shows here with the deployment that served it, fallbacks, guardrails, timing and cost - + {t("playground.chat.inspector.empty.title")} + {t("playground.chat.inspector.empty.description")} ); @@ -235,6 +282,7 @@ function Trace({ message }: { message: MessageType | null }) { } function CodeTab({ code }: { code: (sdk: CodeSdk) => string }) { + const { t } = useTranslation(); const [sdk, setSdk] = useState("openai"); const snippet = code(sdk); return ( @@ -246,21 +294,21 @@ function CodeTab({ code }: { code: (sdk: CodeSdk) => string }) { variant="outline" size="sm" spacing={0} - aria-label="SDK" + aria-label={t("playground.chat.inspector.code.sdk")} > - OpenAI SDK + {t("playground.chat.inspector.code.openai")} - Azure SDK + {t("playground.chat.inspector.code.azure")} - +
- Send this request from your own app. Replace YOUR_LITELLM_API_KEY with a virtual key + {t("playground.chat.inspector.code.description")} 
{snippet}
@@ -271,12 +319,13 @@ function CodeTab({ code }: { code: (sdk: CodeSdk) => string }) { } export function InspectorPanel({ message, code }: { message: MessageType | null; code: (sdk: CodeSdk) => string }) { + const { t } = useTranslation(); return (
- Trace - Code + {t("playground.chat.inspector.tabs.trace")} + {t("playground.chat.inspector.tabs.code")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx index e5e0828825d..2a72a5facc9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx @@ -2,6 +2,7 @@ import { useId, useState, type ReactNode } from "react"; import { Blocks, KeyRound, Search } from "lucide-react"; +import { Trans, useTranslation } from "react-i18next"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; @@ -84,9 +85,14 @@ function MasterRow({ onSelect: () => void; summary?: string; }) { + const { t } = useTranslation(); return (
- + ) : undefined; if (!selection.includes(focused)) return ( - Check this server to choose which of its tools the model can call + {t("playground.chat.mcp.checkServer")} ); const state = toolsByServer[focused]; @@ -283,11 +299,11 @@ export function McpDialog({ if (state.failed) return ( - Couldn't load this server's tools + {t("playground.chat.mcp.toolsFailed")} ); const tools = toolNames(focused); - if (tools.length === 0) return This server has no tools; + if (tools.length === 0) return {t("playground.chat.mcp.noTools")}; const allowed = allowedTools(toolLimits, focused, tools); return (
@@ -295,7 +311,7 @@ export function McpDialog({ {replaceKey}
@@ -333,9 +349,9 @@ export function McpDialog({ return ( !open && setReplacingKey(null)}> - }> + }> - MCP + {t("playground.chat.mcp.trigger")} {selection.length > 0 && ( {selection.length} @@ -344,23 +360,28 @@ export function McpDialog({ - MCP servers - - Pick the servers the model can use, then limit their tools. A server allows every tool until you uncheck - some. - + {t("playground.chat.mcp.title")} + {t("playground.chat.mcp.description")} {nothingAvailable ? (
- - Add one on the MCP servers page, or send as a key that can reach one + + }} + />
) : (
- +
{isLoading ? ( @@ -370,11 +391,18 @@ export function McpDialog({ ) : (
+ -
)} @@ -383,7 +411,7 @@ export function McpDialog({
)} - }>Done + }>{t("playground.chat.mcp.done")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx index bbea909711a..3daef04b206 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx @@ -1,6 +1,9 @@ "use client"; import { Copy, Split } from "lucide-react"; +import type { TFunction } from "i18next"; +import { useTranslation } from "react-i18next"; +import i18n from "@/i18n/i18n"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; @@ -8,24 +11,30 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip import { cn } from "@/lib/cva.config"; import { toast } from "@/lib/toast"; -export function copyText(value: string, label: string): void { +export function copyText(value: string, label: string, t: TFunction = i18n.t.bind(i18n)): void { void navigator.clipboard.writeText(value).then( - () => toast.success(`${label} copied`), - () => toast.error("Unable to copy to clipboard"), + () => toast.success(t("playground.chat.modelName.copied", { label })), + () => toast.error(t("playground.chat.modelName.copyFailed")), ); } export function CopyIconButton({ value, label }: { value: string; label: string }) { + const { t } = useTranslation(); return ( copyText(value, label)} /> + )} } + render={ +
{isWide ? ( -