diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index b3a07a6e0ff..4eb6b272c43 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -19,9 +19,6 @@ jobs: build-ui: runs-on: ubuntu-latest timeout-minutes: 10 - defaults: - run: - working-directory: ui/litellm-dashboard steps: - name: Checkout repository @@ -35,18 +32,11 @@ jobs: with: category: ui - - name: Setup Node.js + # Built through the image stage rather than the checkout, because the + # stage copies ui/litellm-dashboard/ alone: an import reaching above the + # dashboard root resolves in a checkout and fails in every image we ship. + # Dockerfile, docker/Dockerfile.non_root and ui/Dockerfile share this + # stage verbatim, so building one covers all three. + - name: Build the dashboard as the shipped images build it if: steps.changes.outputs.decision != 'skip' - uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0 - with: - node-version-file: ui/litellm-dashboard/.nvmrc - cache: "npm" - cache-dependency-path: ui/litellm-dashboard/package-lock.json - - - name: Install dependencies - if: steps.changes.outputs.decision != 'skip' - run: npm ci - - - name: Build - if: steps.changes.outputs.decision != 'skip' - run: npm run build + run: docker build --target ui-builder -f Dockerfile . diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 1b71232bc2e..9b8b132df62 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -74,12 +74,19 @@ jobs: - name: Run Clippy with Bedrock auth run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings + - name: Run Clippy with all gateway features + run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings + - name: Run Rust tests run: cargo test --workspace --locked - name: Run core tests with Bedrock auth run: cargo test -p litellm-core --features bedrock-auth --locked + # Not --all-features: python-config links libpython, which this job does not install. + - name: Run gateway tests with the server feature + run: cargo test -p litellm-ai-gateway --features server --locked + release-wheel: name: release wheel runs-on: ubuntu-latest diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 6da5fc07e80..6bc44995804 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -151,6 +151,7 @@ jobs: tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/batches_endpoints + tests/test_litellm/proxy/container_endpoints tests/test_litellm/proxy/fine_tuning_endpoints tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a37cb194757..9a9b1138a7d 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14074 }, "reportArgumentType": { - "limit": 2215 + "limit": 2214 }, "reportAssignmentType": { "limit": 319 @@ -42,7 +42,7 @@ "limit": 12 }, "reportIndexIssue": { - "limit": 25 + "limit": 24 }, "reportInvalidTypeForm": { "limit": 34 diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 72f7f74bcf6..be6b9093f53 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -428,3 +428,16 @@ envFrom: {{- end }} {{- end }} {{- end -}} + +{{/* +ingress-nginx's admission webhook rejects a dot in an Exact or Prefix path +(strict-validate-path-type) and serves ImplementationSpecific as a plain +prefix location, so a dotted path takes that type there. +*/}} +{{- define "litellm.ingress.pathType" -}} +{{- if and (eq .controller "nginx") (contains "." .path) -}} +ImplementationSpecific +{{- else -}} +{{- .pathType -}} +{{- end -}} +{{- end -}} diff --git a/helm/litellm/templates/ingress.yaml b/helm/litellm/templates/ingress.yaml index f77ef537b02..732564b280f 100644 --- a/helm/litellm/templates/ingress.yaml +++ b/helm/litellm/templates/ingress.yaml @@ -5,6 +5,10 @@ {{- $gatewayPort := .Values.gateway.service.port -}} {{- $backendPort := .Values.backend.service.port -}} {{- $uiPort := .Values.ui.service.port -}} +{{- $controller := .Values.ingress.controller | default "alb" -}} +{{- if not (has $controller (list "alb" "nginx")) }} +{{- fail (printf "ingress.controller: unknown controller %q, expected one of alb, nginx" $controller) }} +{{- end }} {{/* Backends addressable from ingress.extraPaths, keyed by the `service` field. */}} @@ -27,10 +31,11 @@ /litellm-asset-prefix, so without /*.txt they fall to the backend catch-all → 404 → client-side navigation never settles and the login flow spins in an infinite redirect loop (/ ⇄ /ui/login). ui/nginx.conf already serves *.txt - from the export; the rule only routes the request to it. Needs an ingress - controller whose ImplementationSpecific path is a wildcard pattern - (AWS ALB: `*` = 0+ chars); this chart targets the AWS Load Balancer - Controller. + from the export; the rule only routes the request to it. It needs an + ingress controller whose ImplementationSpecific path is a wildcard pattern + (AWS ALB: `*` = 0+ chars), so it is rendered for ingress.controller=alb + only: ingress-nginx serves ImplementationSpecific as a literal prefix + location, where /*.txt can never match. */}} {{- $uiPaths := list (dict "path" "/" "pathType" "Exact") @@ -38,8 +43,10 @@ (dict "path" "/litellm-asset-prefix" "pathType" "Prefix") (dict "path" "/_next" "pathType" "Prefix") (dict "path" "/ui" "pathType" "Prefix") - (dict "path" "/*.txt" "pathType" "ImplementationSpecific") -}} +{{- if eq $controller "alb" }} +{{- $uiPaths = append $uiPaths (dict "path" "/*.txt" "pathType" "ImplementationSpecific") }} +{{- end }} {{/* Gateway data-plane prefixes — must mirror gateway/routes/allowlist.py. Versioned paths are listed explicitly to avoid routing management routes @@ -83,12 +90,6 @@ adding to it. */}} {{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}} -{{- range $uiPaths }} -{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" .path .pathType) }} -{{- end }} -{{- range $gatewayPrefixes }} -{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|Prefix" .) }} -{{- end }} apiVersion: networking.k8s.io/v1 kind: Ingress metadata: @@ -115,8 +116,10 @@ spec: paths: # --- UI (Next.js static export) --- {{- range $uiPaths }} + {{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" .path "pathType" .pathType) }} + {{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" .path $pathType) }} - path: {{ .path }} - pathType: {{ .pathType }} + pathType: {{ $pathType }} backend: service: name: {{ $uiName }} @@ -134,8 +137,10 @@ spec: port: number: {{ $gatewayPort }} {{- range $gatewayPrefixes }} + {{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" . "pathType" "Prefix") }} + {{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" . $pathType) }} - path: {{ . }} - pathType: Prefix + pathType: {{ $pathType }} backend: service: name: {{ $gatewayName }} @@ -147,10 +152,11 @@ spec: Rendered after every built-in path so an entry can never take precedence over a default, and before the backend catch-all. Position only decides the match on controllers that honour manifest - order: the AWS Load Balancer Controller this chart targets sorts - Exact paths first and Prefix paths longest-first, but keeps + order: the AWS Load Balancer Controller (ingress.controller=alb) + sorts Exact paths first and Prefix paths longest-first, but keeps ImplementationSpecific paths in manifest order, which is what the - /*.txt rule above already depends on. + /*.txt rule above already depends on. ingress-nginx ignores order + and serves the longest matching location. */}} {{- range $idx, $extra := .Values.ingress.extraPaths }} {{- if not (kindIs "map" $extra) }} @@ -164,10 +170,11 @@ spec: {{- if not $target }} {{- fail (printf "ingress.extraPaths[%d] (path %s): unknown service %q, expected one of backend, gateway, ui" $idx $extra.path $service) }} {{- end }} - {{- $pathType := $extra.pathType | default "Prefix" }} - {{- if not (has $pathType (list "Prefix" "Exact" "ImplementationSpecific")) }} - {{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $pathType) }} + {{- $requestedPathType := $extra.pathType | default "Prefix" }} + {{- if not (has $requestedPathType (list "Prefix" "Exact" "ImplementationSpecific")) }} + {{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $requestedPathType) }} {{- end }} + {{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" $extra.path "pathType" $requestedPathType) }} {{- if eq $extra.path "/" }} {{- fail (printf "ingress.extraPaths[%d]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture" $idx) }} {{- end }} diff --git a/helm/litellm/tests/ingress_controller_tests.yaml b/helm/litellm/tests/ingress_controller_tests.yaml new file mode 100644 index 00000000000..40790ba674a --- /dev/null +++ b/helm/litellm/tests/ingress_controller_tests.yaml @@ -0,0 +1,205 @@ +suite: test ingress.controller +templates: + - ingress.yaml +values: + - ./values/required.yaml +tests: + - it: keeps the AWS Load Balancer Controller path types by default + set: + ingress.enabled: true + asserts: + - contains: + path: spec.rules[0].http.paths + content: + path: /favicon.ico + pathType: Exact + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 + - contains: + path: spec.rules[0].http.paths + content: + path: /eu.assemblyai + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + - contains: + path: spec.rules[0].http.paths + content: + path: /*.txt + pathType: ImplementationSpecific + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 + + - it: renders no dotted Exact or Prefix path for ingress-nginx, whose admission webhook rejects them + set: + ingress.enabled: true + ingress.controller: nginx + asserts: + - notMatchRegexRaw: + pattern: 'path: /\S*\.\S*\n\s+pathType: (Exact|Prefix)\n' + - contains: + path: spec.rules[0].http.paths + content: + path: /favicon.ico + pathType: ImplementationSpecific + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 + - contains: + path: spec.rules[0].http.paths + content: + path: /eu.assemblyai + pathType: ImplementationSpecific + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + + - it: drops the /*.txt wildcard for ingress-nginx and keeps every other route as is + set: + ingress.enabled: true + ingress.controller: nginx + asserts: + - notContains: + path: spec.rules[0].http.paths + content: + path: /*.txt + any: true + - contains: + path: spec.rules[0].http.paths + content: + path: /ui + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 + - contains: + path: spec.rules[0].http.paths + content: + path: /test + pathType: Exact + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + - equal: + path: spec.rules[0].http.paths[-1] + value: + path: / + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-backend + port: + number: 4001 + + - it: rejects an extraPaths entry that repeats a built-in path at the pathType ingress-nginx renders it with + set: + ingress.enabled: true + ingress.controller: nginx + ingress.extraPaths: + - path: /favicon.ico + service: ui + pathType: ImplementationSpecific + asserts: + - failedTemplate: + errorMessage: "ingress.extraPaths[0]: path /favicon.ico with pathType ImplementationSpecific is already routed by this chart, and a duplicate would take it over rather than add to it" + + - it: rejects a controller it has no path types for + set: + ingress.enabled: true + ingress.controller: traefik + asserts: + - failedTemplate: + errorMessage: 'ingress.controller: unknown controller "traefik", expected one of alb, nginx' + + - it: rejects an extraPaths entry that repeats a built-in path once ingress-nginx normalizes its pathType + set: + ingress.enabled: true + ingress.controller: nginx + ingress.extraPaths: + - path: /favicon.ico + service: ui + asserts: + - failedTemplate: + errorMessage: "ingress.extraPaths[0]: path /favicon.ico with pathType ImplementationSpecific is already routed by this chart, and a duplicate would take it over rather than add to it" + + - it: renders a dotted extraPaths entry as ImplementationSpecific for ingress-nginx + set: + ingress.enabled: true + ingress.controller: nginx + ingress.extraPaths: + - path: /eu.assemblyai.custom + service: gateway + - path: /robots.txt + service: ui + pathType: Exact + asserts: + - notMatchRegexRaw: + pattern: 'path: "?/\S*\.\S*"?\n\s+pathType: (Exact|Prefix)\n' + - contains: + path: spec.rules[0].http.paths + content: + path: /eu.assemblyai.custom + pathType: ImplementationSpecific + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + - contains: + path: spec.rules[0].http.paths + content: + path: /robots.txt + pathType: ImplementationSpecific + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 + + - it: keeps the requested pathType of a dotted extraPaths entry for the AWS Load Balancer Controller + set: + ingress.enabled: true + ingress.extraPaths: + - path: /eu.assemblyai.custom + service: gateway + - path: /robots.txt + service: ui + pathType: Exact + asserts: + - contains: + path: spec.rules[0].http.paths + content: + path: /eu.assemblyai.custom + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + - contains: + path: spec.rules[0].http.paths + content: + path: /robots.txt + pathType: Exact + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 378c3b7a618..461330ba491 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -10,6 +10,18 @@ imagePullSecrets: [] ingress: enabled: false className: "" + # Which ingress controller serves this Ingress. Controllers disagree on the + # pathTypes they accept, so this picks the pathType of the dotted paths, the + # built-in ones and any dotted extraPaths entry alike: + # alb AWS Load Balancer Controller (default): Exact and Prefix paths plus + # the /*.txt wildcard that routes the UI's RSC payloads. + # nginx ingress-nginx: its admission webhook rejects a dot in an Exact or + # Prefix path (strict-validate-path-type, on by default from v1.12.0 + # until v1.12.6 / v1.13.2 allowed dots again), so /favicon.ico and + # /eu.assemblyai render as ImplementationSpecific, which nginx serves + # as a plain prefix location. /*.txt is dropped: nginx has no + # wildcard pathType, so that rule could never match there. + controller: alb annotations: {} host: "" # optional; if set, becomes the rule's host tls: [] @@ -26,7 +38,8 @@ ingress: # # path required; the HTTP path to route # service which component serves it: gateway (default), backend, or ui - # pathType Prefix (default), Exact, or ImplementationSpecific + # pathType Prefix (default), Exact, or ImplementationSpecific; a dotted + # path renders as ImplementationSpecific when controller is nginx # # The target component only answers paths its own route allowlist keeps, so # a path here still has to be one that component serves. diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 2283814ab35..b51de9609d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -18,10 +18,15 @@ recoverable one. constant: it grows with the number of pending migrations, so a fresh database that has to replay every migration this package ships overruns a per-command budget sized for the short bookkeeping commands, on a laptop as much as on a -slow CI runner. The Python ``prisma`` wrapper spawns Node and the schema engine -as separate children, so killing the wrapper on timeout leaves them running: -the retry then contends with that orphan for Prisma's advisory lock and cannot -finish any sooner. Migrate deploy therefore runs under its own budget. +slow CI runner. Migrate deploy therefore runs under its own budget. + +The Python ``prisma`` wrapper spawns Node, which spawns the Rust schema +engine, so killing only the wrapper on timeout leaves the engine running with +no parent: it keeps mutating the database after the proxy has given up, holds +Prisma's advisory lock so every retry and every later boot queues behind it, +and dies mid-migration once its pipes close, leaving a half-applied ledger row. +Every Prisma command therefore runs in a process group of its own, and a +timeout kills the whole group. All three budgets are overridable so an operator can widen them without a release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, @@ -35,10 +40,12 @@ the deploy override says otherwise. import math import os import shutil +import signal import subprocess +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import IO, Optional, Union from litellm_proxy_extras._logging import logger @@ -167,6 +174,49 @@ def heal_incomplete_nodeenv_cache() -> bool: return True +def _kill_process_group(process: "subprocess.Popen[str]") -> None: + if os.name == "nt": + process.kill() + return + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + return + + +def run_prisma( + argv: Sequence[str], + *, + timeout: float, + env: Mapping[str, str], + stdout: Union[IO[str], int, None] = subprocess.PIPE, + stderr: Optional[int] = subprocess.PIPE, +) -> "subprocess.CompletedProcess[str]": + """Run one Prisma CLI command in its own process group, bounded by ``timeout``. + + Raises ``subprocess.TimeoutExpired`` once the budget is spent, after killing + the command together with every process it spawned, and + ``subprocess.CalledProcessError`` on a non-zero exit. Output is captured as + text unless ``stdout``/``stderr`` say otherwise. + """ + with subprocess.Popen( + argv, + env=env, + stdout=stdout, + stderr=stderr, + text=True, + start_new_session=True, + ) as process: + try: + out, err = process.communicate(timeout=timeout) + except BaseException: + _kill_process_group(process) + raise + if process.returncode: + raise subprocess.CalledProcessError(process.returncode, process.args, out, err) + return subprocess.CompletedProcess(process.args, process.returncode, out, err) + + def ensure_prisma_toolchain( prisma_command: str, prisma_env: dict[str, str] ) -> ToolchainBootstrap: @@ -179,14 +229,7 @@ def ensure_prisma_toolchain( timeout = prisma_bootstrap_timeout() logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout) try: - subprocess.run( - [prisma_command, BOOTSTRAP_ARG], - timeout=timeout, - check=True, - capture_output=True, - text=True, - env=prisma_env, - ) + run_prisma([prisma_command, BOOTSTRAP_ARG], timeout=timeout, env=prisma_env) except subprocess.TimeoutExpired: logger.warning( "Preparing the Prisma CLI toolchain timed out after %ss. Raise %s " diff --git a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py index 157d595404e..3a5865a54cd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py +++ b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py @@ -16,7 +16,7 @@ import tempfile from pathlib import Path from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout +from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout, run_prisma REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL" @@ -66,7 +66,7 @@ def apply_replica_identity_full( with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir: sql_path = Path(tmp_dir) / "replica_identity_full.sql" sql_path.write_text(REPLICA_IDENTITY_FULL_SQL) - subprocess.run( + run_prisma( [ prisma_command, "db", @@ -77,9 +77,6 @@ def apply_replica_identity_full( schema_path, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=prisma_env, ) except subprocess.CalledProcessError as e: diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index a4ea4789b49..71e7e9c683b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -10,6 +10,7 @@ from dataclasses import dataclass, replace from pathlib import Path from typing import Optional +from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger from litellm_proxy_extras.replica_identity import ( REPLICA_IDENTITY_FULL_ENV_VAR, @@ -231,7 +232,7 @@ class ProxyExtrasDBManager: # 1. Generate migration SQL file by comparing empty state to current db state logger.info("Generating baseline migration...") migration_file = init_dir / "migration.sql" - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -242,14 +243,13 @@ class ProxyExtrasDBManager: "--script", ], stdout=open(migration_file, "w"), - check=True, timeout=prisma_command_timeout(), env=prisma_env, ) # 3. Mark the migration as applied since it represents current state logger.info("Marking baseline migration as applied...") - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -257,7 +257,6 @@ class ProxyExtrasDBManager: "--applied", "0_init", ], - check=True, timeout=prisma_command_timeout(), env=prisma_env, ) @@ -286,7 +285,7 @@ class ProxyExtrasDBManager: """Mark a specific migration as rolled back""" # Set up environment for offline mode if configured prisma_env = _get_prisma_env() - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -295,8 +294,6 @@ class ProxyExtrasDBManager: migration_name, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, env=prisma_env, ) @@ -348,11 +345,9 @@ class ProxyExtrasDBManager: def _resolve_specific_migration(migration_name: str): """Mark a specific migration as applied""" prisma_env = _get_prisma_env() - subprocess.run( + prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], timeout=prisma_command_timeout(), - check=True, - capture_output=True, env=prisma_env, ) @@ -436,7 +431,7 @@ class ProxyExtrasDBManager: try: logger.info("Generating migration diff between DB and schema.prisma...") with open(diff_sql_path, "w") as f: - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -447,7 +442,6 @@ class ProxyExtrasDBManager: schema_path, "--script", ], - check=True, timeout=prisma_command_timeout(), stdout=f, env=_get_prisma_env(), @@ -470,7 +464,7 @@ class ProxyExtrasDBManager: migration_files = sorted(Path(migrations_dir).glob("*/migration.sql")) for mig_file in migration_files: try: - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "db", @@ -481,9 +475,6 @@ class ProxyExtrasDBManager: schema_path, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"Applied migration: {mig_file.parent.name}") @@ -516,7 +507,7 @@ class ProxyExtrasDBManager: applied_ok = False try: logger.info("Running prisma db execute to apply the migration diff...") - result = subprocess.run( + result = prisma_toolchain.run_prisma( [ _get_prisma_command(), "db", @@ -527,9 +518,6 @@ class ProxyExtrasDBManager: schema_path, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"prisma db execute stdout: {result.stdout}") @@ -558,7 +546,7 @@ class ProxyExtrasDBManager: for migration_name in migration_names: try: logger.info(f"Resolving migration: {migration_name}") - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -567,9 +555,6 @@ class ProxyExtrasDBManager: migration_name, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.debug(f"Resolved migration: {migration_name}") @@ -762,11 +747,12 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - subprocess.run( + prisma_toolchain.run_prisma( [_get_prisma_command(), "db", "push", "--accept-data-loss"], timeout=prisma_command_timeout(), - check=True, env=_get_prisma_env(), + stdout=None, + stderr=None, ) return True except ( @@ -789,12 +775,9 @@ class ProxyExtrasDBManager: try: while not budget.exhausted: try: - result = subprocess.run( + result = prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "deploy"], timeout=deploy_timeout, - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -1031,12 +1014,9 @@ class ProxyExtrasDBManager: logger.info("Running prisma migrate deploy") try: # Set migrations directory for Prisma - result = subprocess.run( + result = prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "deploy"], timeout=prisma_migrate_deploy_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -1108,7 +1088,7 @@ class ProxyExtrasDBManager: f"Found failed migration: {failed_migration}, marking as rolled back" ) # Mark the failed migration as rolled back - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -1117,9 +1097,6 @@ class ProxyExtrasDBManager: failed_migration, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info( @@ -1244,10 +1221,12 @@ class ProxyExtrasDBManager: if ProxyExtrasDBManager.spend_logs_is_partitioned(): raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR) # Use prisma db push with increased timeout - subprocess.run( + prisma_toolchain.run_prisma( [_get_prisma_command(), "db", "push", "--accept-data-loss"], timeout=prisma_command_timeout(), - check=True, + stdout=None, + stderr=None, + env=_get_prisma_env(), ) return True except subprocess.TimeoutExpired: diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 406f07eb792..040d67d25e4 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -42,7 +42,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): "Error: P3018\nMigration name: 20250326162113_baseline\n" "Database error code: 42501\npermission denied for schema public" ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="permission"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -60,7 +60,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n" 'Reason: syntax error at or near "BRKN" LINE 42' ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -124,7 +124,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path): def fake_resolve(*args, **kwargs): resolve_called["n"] += 1 - monkeypatch.setattr("subprocess.run", fake_run) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", fake_run) monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set @@ -139,7 +139,7 @@ def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_pat (tmp_path / "schema.prisma").write_text("// stub") stderr = "db push error" - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="prisma db push failed"): ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) @@ -209,7 +209,7 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error( "Error: P3009\nMigration `20260101000000_some_migration` failed\n" "relation already exists" ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises( RuntimeError, match="Failed to mark migration .* as applied" ): @@ -228,7 +228,7 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): stdout = "Applied migration.\n" stderr = "" - monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult()) resolve_called = {"n": 0} monkeypatch.setattr( @@ -296,7 +296,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): "_resolve_specific_migration", lambda name: pytest.fail("a deadlocked migration must never be marked applied"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR)) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -309,7 +309,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None) with patch( - "subprocess.run", + "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR), ): with pytest.raises(RuntimeError, match="after 4 attempts"): @@ -343,7 +343,7 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ "_resolve_specific_migration", lambda name: pytest.fail("a deadlocked migration must never be marked applied"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -372,7 +372,7 @@ def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path "_resolve_specific_migration", lambda name: pytest.fail("a deadlocked migration must never be marked applied"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -395,7 +395,7 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): "_roll_back_migration", lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -417,7 +417,7 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): lambda name: 'ERROR: syntax error at or near "BRKN"', ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -427,7 +427,7 @@ def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): waiter as victim) is retried, not fatal.""" _stub_v2_env(monkeypatch, tmp_path) monkeypatch.setattr( - "subprocess.run", _succeed_after(1, "Database error: deadlock detected") + "litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, "Database error: deadlock detected") ) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -446,7 +446,10 @@ def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path): """v2: the advisory-lock waiter that times out while a peer's retry holds the lock retries instead of dying.""" _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR)) + monkeypatch.setattr( + "litellm_proxy_extras.prisma_toolchain.run_prisma", + _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR), + ) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -456,7 +459,7 @@ def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_pa """v2: a plain P1002 (database unreachable) stays fatal.""" _stub_v2_env(monkeypatch, tmp_path) stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out." - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md index 857a744e014..ae8ae5a6870 100644 --- a/litellm-rust/ADDING_A_PROVIDER.md +++ b/litellm-rust/ADDING_A_PROVIDER.md @@ -26,4 +26,4 @@ variants of it. The test for a good abstraction is that adding the next provider is a few declarative lines, not a new file of duplicated flow. Only diverge from the base when behavior is genuinely different, and say so explicitly in the PR. -**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`. +**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run the commands under "Checks" in [CLAUDE.md](CLAUDE.md). diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index 3dcf1853efc..d9c944529df 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -174,10 +174,14 @@ for changes under `litellm-rust/`. ```bash cd litellm-rust cargo fmt --check +cargo clippy --workspace --all-targets -- -D warnings +cargo clippy -p litellm-core --all-targets --features bedrock-auth -- -D warnings # the ai-gateway binary + server code is behind the `server` feature -cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings -cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings +cargo clippy -p litellm-ai-gateway --all-targets --all-features -- -D warnings cargo test --workspace +cargo test -p litellm-core --features bedrock-auth +# the `auth`, `routes`, `state` and `realtime` tests only exist under `server` +cargo test -p litellm-ai-gateway --features server ``` When a Rust path is exposed through Python, add Python parity tests that compare diff --git a/litellm-rust/README.md b/litellm-rust/README.md index a0d79c6f0a5..e43dc7ea6ad 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -49,11 +49,6 @@ function per top-level route, mirroring the core entrypoints. ## Checks -Run these before pushing Rust changes. GitHub Actions runs the same checks for -changes under `litellm-rust/`. - -```bash -cargo fmt --check -cargo clippy --workspace --all-targets -- -D warnings -cargo test --workspace -``` +Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust +changes. That list is the single source of truth and matches what GitHub Actions +runs for changes under `litellm-rust/`. diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index 4a689cb9579..c0a29ab14bc 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -49,11 +49,5 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` ## Checks before push -25. Run, and keep green: - ```bash - cd litellm-rust - cargo fmt --check - cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings - cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings - cargo test --workspace - ``` +25. Run, and keep green, the commands under "Checks" in `litellm-rust/CLAUDE.md`. + That list is the single source of truth and matches what GitHub Actions runs. diff --git a/litellm-rust/crates/ai-gateway/Dockerfile b/litellm-rust/crates/ai-gateway/Dockerfile index adf6fca0741..2bc3c05ad7e 100644 --- a/litellm-rust/crates/ai-gateway/Dockerfile +++ b/litellm-rust/crates/ai-gateway/Dockerfile @@ -36,12 +36,12 @@ FROM chef AS builder # whenever only gateway source changes. COPY --from=planner /build/litellm-rust/recipe.json recipe.json RUN cargo chef cook --locked --release \ - -p litellm-ai-gateway --features python-config \ + -p litellm-ai-gateway --features server,python-config \ --recipe-path recipe.json # Now copy the real sources and build the gateway binary. Deps are already cooked # above, so this step only recompiles the gateway crate. COPY litellm-rust/ . -RUN cargo build --locked --release -p litellm-ai-gateway --features python-config +RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config # ---- Runtime ---------------------------------------------------------------- # python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3 diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 5cbb47220be..1675e6f1b16 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -100,7 +100,7 @@ Worker tuning, rarely needed: `LITELLM_LOG_CHANNEL_CAPACITY` (4096), ## Build & run with Docker -The image is built `--features python-config` and installs litellm **from this +The image is built `--features server,python-config` and installs litellm **from this repo's source** (the config reader is newer than any PyPI release), so the build **context is the repo root**: @@ -135,10 +135,10 @@ docker run --rm -p 4001:4001 \ ```bash # config.yaml mode — needs litellm importable in the active python env LITELLM_CONFIG_PATH=./crates/ai-gateway/config.yaml \ - cargo run --release -p litellm-ai-gateway --features python-config + cargo run --release -p litellm-ai-gateway --features server,python-config # env stand-in mode — no python, no config -cargo run --release -p litellm-ai-gateway +cargo run --release -p litellm-ai-gateway --features server ``` ## Deploy on Render diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs index c32e727de54..c0d72e90b77 100644 --- a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -106,16 +106,16 @@ impl RealTimeStreaming { /// `litellm_call_id`, replacing the gateway-generated fallback. fn on_session(&mut self, event: &RealtimeEvent) { let session = event.data.get("session").and_then(Value::as_object); - if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) { - if !id.is_empty() { - self.id = id.to_string(); - self.litellm_call_id = id.to_string(); - } + if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) + && !id.is_empty() + { + self.id = id.to_string(); + self.litellm_call_id = id.to_string(); } - if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) { - if !model.is_empty() { - self.model = model.to_string(); - } + if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) + && !model.is_empty() + { + self.model = model.to_string(); } } @@ -323,6 +323,32 @@ mod tests { assert_eq!(streaming.dropped(), 0); } + #[test] + fn blank_session_id_and_model_keep_the_gateway_fallbacks() { + let mut streaming = RealTimeStreaming::new( + Vec::new(), + "call_fallback".to_string(), + "gpt-realtime".to_string(), + RequestMetadata::default(), + ); + + streaming.observe(&event( + r#"{"type":"session.created","session":{"id":"","model":""}}"#, + )); + let payload = streaming.build_payload(); + assert_eq!(payload.id, "call_fallback"); + assert_eq!(payload.litellm_call_id, "call_fallback"); + assert_eq!(payload.model, "gpt-realtime"); + + streaming.observe(&event( + r#"{"type":"session.updated","session":{"id":"sess_002","model":""}}"#, + )); + let payload = streaming.build_payload(); + assert_eq!(payload.id, "sess_002"); + assert_eq!(payload.litellm_call_id, "sess_002"); + assert_eq!(payload.model, "gpt-realtime"); + } + #[test] fn payload_serializes_with_camelcase_times_and_realtime_call_type() { let mut streaming = RealTimeStreaming::new( diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index b8ee77c4269..f7bbb37dff4 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -50,18 +50,17 @@ where provider_model, params.api_key.as_deref(), params.api_base.as_deref(), - ) { - if let Some(handoff) = pool.take(&key) { - return crate::io::realtime::realtime_warm( - provider_model, - handoff, - idle_timeout, - observe, - client_in, - client_out, - ) - .await; - } + ) && let Some(handoff) = pool.take(&key) + { + return crate::io::realtime::realtime_warm( + provider_model, + handoff, + idle_timeout, + observe, + client_in, + client_out, + ) + .await; } // Cold path: fresh dial (the original behavior). diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 42a86763b6d..fc27d3a118a 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Callable, Coroutine from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final @@ -83,6 +84,47 @@ class ServiceLogging(CustomLogger): return open_telemetry_logger return None + @staticmethod + def _sync_dispatch_loop() -> asyncio.AbstractEventLoop | None: + """The event loop a blocking caller can dispatch on, or ``None`` if it has none.""" + try: + loop: Final = asyncio.get_event_loop() + except RuntimeError: + return None + return None if loop.is_closed() else loop + + @staticmethod + async def _emit_guarded(hook: Callable[[], Coroutine[object, object, None]]) -> None: + """Emit one service event, absorbing anything the callbacks raise. + + Monitoring must not break the call it monitors. Sync callers are the ones that + swallow their own service failures (a Redis batch read returns an empty dict), + so an exception from a misconfigured callback would replace a Redis outage with + a callback error and skip the caller's fallback handling. + """ + try: + await hook() + except Exception as e: + verbose_logger.exception("Error emitting service event - %s", e) + + @staticmethod + def _dispatch_from_sync(hook: Callable[[], Coroutine[object, object, None]]) -> None: + """Run an async service hook from a blocking caller, whatever event loop it holds. + + Takes a factory rather than a coroutine so the hook is built on the path that + runs it, and only ever once. + """ + loop: Final = ServiceLogging._sync_dispatch_loop() + try: + if loop is None: + asyncio.run(ServiceLogging._emit_guarded(hook)) + elif loop.is_running(): + loop.create_task(ServiceLogging._emit_guarded(hook)) + else: + loop.run_until_complete(ServiceLogging._emit_guarded(hook)) + except Exception as e: + verbose_logger.exception("Error dispatching service event - %s", e) + def service_success_hook( self, service: ServiceTypes, @@ -99,54 +141,45 @@ class ServiceLogging(CustomLogger): if self.mock_testing: self.mock_testing_sync_success_hook += 1 - try: - # Try to get the current event loop - loop: Final = asyncio.get_event_loop() - # Check if the loop is running - if loop.is_running(): - # If we're in a running loop, create a task - loop.create_task( - self.async_service_success_hook( - service=service, - duration=duration, - call_type=call_type, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - else: - # Loop exists but not running, we can use run_until_complete - loop.run_until_complete( - self.async_service_success_hook( - service=service, - duration=duration, - call_type=call_type, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - except RuntimeError: - # No event loop exists, create a new one and run - asyncio.run( - self.async_service_success_hook( - service=service, - duration=duration, - call_type=call_type, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) + self._dispatch_from_sync( + lambda: self.async_service_success_hook( + service=service, + duration=duration, + call_type=call_type, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, ) + ) - def service_failure_hook(self, service: ServiceTypes, duration: float, error: Exception, call_type: str): + def service_failure_hook( + self, + service: ServiceTypes, + duration: float, + error: Exception, + call_type: str, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: float | datetime | None = None, + ): """ - [TODO] Not implemented for sync calls yet. V0 is focused on async monitoring (used by proxy). + Handles both sync and async monitoring by checking for existing event loop. """ if self.mock_testing: self.mock_testing_sync_failure_hook += 1 + self._dispatch_from_sync( + lambda: self.async_service_failure_hook( + service=service, + duration=duration, + error=error, + call_type=call_type, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + ) + ) + async def async_service_success_hook( self, service: ServiceTypes, diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 598c9e67faf..df67ba08416 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -8,10 +8,9 @@ Has 4 primary methods: - async_get_cache """ -import asyncio import time import traceback -from concurrent.futures import ThreadPoolExecutor +from collections.abc import Sequence from threading import Lock from typing import TYPE_CHECKING, Any, Final @@ -188,31 +187,38 @@ class DualCache(BaseCache): local_only: bool = False, **kwargs, ): - received_args: Final = locals() - received_args.pop("self") - - def run_in_new_loop(): - """Run the coroutine in a new event loop within this thread.""" - new_loop: Final = asyncio.new_event_loop() - try: - asyncio.set_event_loop(new_loop) - return new_loop.run_until_complete(self.async_batch_get_cache(**received_args)) - finally: - new_loop.close() - asyncio.set_event_loop(None) - try: - # First, try to get the current event loop - _ = asyncio.get_running_loop() - # If we're already in an event loop, run in a separate thread - # to avoid nested event loop issues - with ThreadPoolExecutor(max_workers=1) as executor: - future: Final = executor.submit(run_in_new_loop) - return future.result() + in_memory_result: Final = ( + self.in_memory_cache.batch_get_cache(keys, **kwargs) if self.in_memory_cache is not None else None + ) + result: Final = in_memory_result if in_memory_result is not None else tuple(None for _ in keys) - except RuntimeError: - # No running event loop, we can safely run in this thread - return run_in_new_loop() + if None not in result or self.redis_cache is None or local_only: + return result + + sublist_keys, previous_access_times = self._reserve_redis_batch_keys(time.time(), keys, result) + if len(sublist_keys) == 0: + return result + + try: + redis_result: Final = self.redis_cache.batch_get_cache( + key_list=sublist_keys, parent_otel_span=parent_otel_span + ) + except Exception: + # Do not throttle subsequent callers if the Redis read fails. + self._rollback_redis_batch_key_reservations(previous_access_times) + raise + + if self.in_memory_cache is not None: + for key, value in redis_result.items(): + if value is not None: + self.in_memory_cache.set_cache(key, value, **self._backfill_kwargs(kwargs)) + + return list( # mutable-ok: public list contract + redis_result.get(key) if value is None else value for key, value in zip(keys, result) + ) + except Exception: + verbose_logger.error(traceback.format_exc()) async def async_get_cache( self, @@ -251,7 +257,7 @@ class DualCache(BaseCache): self, current_time: float, keys: list[str], - result: list[Any], + result: Sequence[Any], ) -> tuple[list[str], dict[str, float | None]]: """ Atomically choose keys to fetch from Redis and reserve their access time. diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2b04a075114..58733b384b9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -78,10 +78,18 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... +_BREAKER_GUARD_FRAME_NAMES: Final = frozenset( + {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} +) + + def _get_call_stack_info(num_frames: int = 2) -> str: """ Get the function names from the previous 1-2 functions in the call stack. + Frames belonging to this module's circuit-breaker guards are skipped so the + reported callers stay the real ones even on guarded methods. + Args: num_frames: Number of previous frames to include (default: 2) @@ -102,11 +110,11 @@ def _get_call_stack_info(num_frames: int = 2) -> str: return "unknown" function_names: Final = [] - for _ in range(num_frames): - if frame is None: - break - func_name = frame.f_code.co_name - function_names.append(func_name) + while frame is not None and len(function_names) < num_frames: + if frame.f_code.co_name in _BREAKER_GUARD_FRAME_NAMES and frame.f_globals.get("__name__") == __name__: + frame = frame.f_back + continue + function_names.append(frame.f_code.co_name) frame = frame.f_back if not function_names: @@ -241,6 +249,23 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int: + """Reject the call if the breaker is open, else return the swallowed-failure count to compare against.""" + if breaker.is_open(): + raise Exception(f"Redis circuit breaker is open — skipping {name}") + return _swallowed_redis_failures.get() + + +def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None: + """Record success only when nothing failed while the call ran. + + Several Redis methods catch their own connection errors and return a default, so a + method that returned is not on its own proof of a healthy Redis. + """ + if _swallowed_redis_failures.get() == swallowed_before: + breaker.record_success() + + async def _run_under_circuit_breaker( breaker: RedisCircuitBreaker, name: str, @@ -249,20 +274,33 @@ async def _run_under_circuit_breaker( """Run one Redis coroutine under a circuit breaker. Shared by the method decorator and the Lua script executor so both feed the same - health signal. Success is recorded only when nothing failed while ``call`` ran, - because several Redis methods catch their own connection errors and return a default. + health signal. """ - if breaker.is_open(): - raise Exception(f"Redis circuit breaker is open — skipping {name}") - swallowed_before: Final = _swallowed_redis_failures.get() + swallowed_before: Final = _enter_circuit_breaker(breaker, name) try: result: Final = await call() except Exception as e: if _is_redis_health_failure(e): breaker.record_failure() raise - if _swallowed_redis_failures.get() == swallowed_before: - breaker.record_success() + _exit_circuit_breaker(breaker, swallowed_before) + return result + + +def _run_under_circuit_breaker_sync( + breaker: RedisCircuitBreaker, + name: str, + call: Callable[[], _RedisCallResult], +) -> _RedisCallResult: + """Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path.""" + swallowed_before: Final = _enter_circuit_breaker(breaker, name) + try: + result: Final = call() + except Exception as e: + if _is_redis_health_failure(e): + breaker.record_failure() + raise + _exit_circuit_breaker(breaker, swallowed_before) return result @@ -288,6 +326,14 @@ def _redis_circuit_breaker_guard(method): return wrapper +def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) -> Callable[..., _RedisCallResult]: + return functools.wraps(method)( + lambda self, *args, **kwargs: _run_under_circuit_breaker_sync( + self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs) + ) + ) + + class RedisCache(BaseCache): # if users don't provider one, use the default litellm cache @@ -1146,14 +1192,13 @@ class RedisCache(BaseCache): """ key_value_dict = {} _key_list: Final = [key for key in key_list if key is not None] + start_time: Final = time.time() try: - _keys: Final = [] - for cache_key in _key_list: - cache_key = self.check_and_fix_namespace(key=cache_key or "") - _keys.append(cache_key) - start_time: Final = time.time() + swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") + _keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list] results: Final = self._run_redis_mget_operation(keys=_keys) + _exit_circuit_breaker(self._circuit_breaker, swallowed_before) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( @@ -1178,7 +1223,18 @@ class RedisCache(BaseCache): return decoded_results except Exception as e: + failed_at: Final = time.time() + self.service_logger_obj.service_failure_hook( + service=ServiceTypes.REDIS, + duration=failed_at - start_time, + error=e, + call_type=f"batch_get_cache <- {_get_call_stack_info()}", + start_time=start_time, + end_time=failed_at, + parent_otel_span=parent_otel_span, + ) verbose_logger.error("Error occurred in batch get cache - %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @_redis_circuit_breaker_guard diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 7368de1e968..d70f947469a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -932,7 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: - if role == "user" or role == "system" or role == "tool": + if role in ("user", "system", "developer", "tool"): return {"type": "input_text", "text": content} else: return {"type": "output_text", "text": content} diff --git a/litellm/constants.py b/litellm/constants.py index be13d9aac5f..f5acadc32ab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_ran DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) +AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"}) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( @@ -39,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", + "heuristic_v2_router_limit", } ) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c8e5610b4a9..9bb613654e4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -831,10 +831,10 @@ class CustomGuardrail(CustomLogger): # should run guardrail litellm_guardrails: Final = request_data.get("guardrails") if litellm_guardrails is None or not isinstance(litellm_guardrails, list): - return response + return None if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: - return response + return None # CHECK IF GUARDRAIL REJECTS THE REQUEST result: Final = await self.async_post_call_success_hook( @@ -850,7 +850,7 @@ class CustomGuardrail(CustomLogger): ) if not self._is_valid_response_type(result): - return response + return None return result diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index e012d35b8f3..12ff38ce4ba 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -6,7 +6,8 @@ It searches the vector store for relevant context and appends it to the messages """ from collections.abc import Awaitable, Callable -from typing import TYPE_CHECKING, Any, Final, cast +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm import litellm.vector_stores @@ -24,10 +25,35 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.utils import PrismaClient + from litellm.router import Router else: LiteLLMLoggingObj = Any +class ProxyRuntime(Protocol): + def llm_router(self) -> "Router | None": ... + + def prisma_client(self) -> "PrismaClient | None": ... + + +@dataclass(frozen=True, slots=True) +class ProxyServerRuntime: + def llm_router(self) -> "Router | None": + try: + from litellm.proxy.proxy_server import llm_router + except ImportError: + return None + return llm_router + + def prisma_client(self) -> "PrismaClient | None": + try: + from litellm.proxy.proxy_server import prisma_client + except ImportError: + return None + return prisma_client + + class VectorStorePreCallHook(CustomLogger): CONTENT_PREFIX_STRING = "Context:\n\n" """ @@ -39,8 +65,9 @@ class VectorStorePreCallHook(CustomLogger): 3. Appends the search results as context to the messages """ - def __init__(self): + def __init__(self, proxy_runtime: ProxyRuntime | None = None): super().__init__() + self.proxy_runtime: Final[ProxyRuntime] = proxy_runtime or ProxyServerRuntime() async def async_get_chat_completion_prompt( self, @@ -79,21 +106,8 @@ class VectorStorePreCallHook(CustomLogger): if litellm.vector_store_registry is None: return model, messages, non_default_params - # Get prisma_client for database fallback - prisma_client = None - llm_router = None - try: - from litellm.proxy.proxy_server import ( - llm_router as _llm_router, - ) - from litellm.proxy.proxy_server import ( - prisma_client as _prisma_client, - ) - - prisma_client = _prisma_client - llm_router = _llm_router - except ImportError: - pass + prisma_client: Final = self.proxy_runtime.prisma_client() + llm_router: Final = self.proxy_runtime.llm_router() # Use database fallback to ensure synchronization across instances vector_stores_to_run: list[ @@ -136,15 +150,23 @@ class VectorStorePreCallHook(CustomLogger): Callable[..., Awaitable[VectorStoreSearchResponse]], litellm.vector_stores.asearch, ) - search_response = await search_function( - **{ - "vector_store_id": vector_store_id, - "query": query, - "custom_llm_provider": custom_llm_provider, - "metadata": request_metadata, - **litellm_params_for_vector_store, - }, - ) + try: + search_response = await search_function( + **{ + "vector_store_id": vector_store_id, + "query": query, + "custom_llm_provider": custom_llm_provider, + "metadata": request_metadata, + **litellm_params_for_vector_store, + }, + ) + except Exception as search_error: + verbose_logger.warning( + "Vector store search failed for vector_store_id=%s, continuing without its context: %s", + vector_store_id, + search_error, + ) + continue verbose_logger.debug("search_response: %s", search_response) @@ -153,7 +175,7 @@ class VectorStorePreCallHook(CustomLogger): # Process search results and append as context modified_messages = self._append_search_results_to_messages( - messages=messages, search_response=search_response + messages=modified_messages, search_response=search_response ) # Get the number of results for logging diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 4d043701f40..87f007ca1d5 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -18,9 +18,11 @@ caller's identity metadata, minus two things that must never be forwarded as-is: from __future__ import annotations from collections.abc import Mapping +from types import MappingProxyType from typing import Final from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES +from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) @@ -142,6 +144,19 @@ def forwarded_internal_call_metadata( } +def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]: + kwargs: Final = request_kwargs or MappingProxyType({}) + return MappingProxyType( + {k: v for k in ("litellm_session_id", "litellm_trace_id") if isinstance(v := kwargs.get(k), str)} + ) + + +def effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None: + return initialize_standard_callback_dynamic_params(dict(request_kwargs) if request_kwargs else None).get( + "turn_off_message_logging" + ) + + def sanitized_forwardable_call_metadata( parent_metadata: Mapping[str, object], call_origin: InternalCallOrigin, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f54eeca5178..51d6858b7dc 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -414,6 +414,11 @@ def _resolve_vertex_location_for_cost( return VertexBase.get_vertex_region(configured_location, model) +def _provider_response_id(source: object) -> str | None: + candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) + return candidate if isinstance(candidate, str) and candidate else None + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -429,6 +434,7 @@ class Logging(LiteLLMLoggingBaseClass): custom_pricing: bool = False stream_options = None litellm_request_debug: bool = False + streamed_anthropic_message_id: str | None = None def __init__( self, @@ -2136,7 +2142,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["cache_hit"] = cache_hit if self.call_type == CallTypes.anthropic_messages.value: - result = self._handle_anthropic_messages_response_logging(result=result) + result = self._anthropic_messages_logged_response(result=result) elif ( self.call_type == CallTypes.generate_content.value or self.call_type == CallTypes.agenerate_content.value @@ -3806,6 +3812,23 @@ class Logging(LiteLLMLoggingBaseClass): ) return None + def record_streamed_anthropic_message_id(self, message_id: str) -> None: + self.streamed_anthropic_message_id = message_id + + def _anthropic_messages_logged_response(self, result: Any) -> ModelResponse: + """ + The ModelResponse a /v1/messages spend_logs row is built from. + + A streaming call bridged onto the Responses API is the one case where the `msg_` id the + caller was served is minted locally rather than issued upstream, so it is absent from the + response the row would otherwise be keyed on and has to be carried over here. + """ + logged: Final = self._handle_anthropic_messages_response_logging(result=result) + streamed_message_id: Final = self.streamed_anthropic_message_id + if streamed_message_id is None: + return logged + return logged.model_copy(update={"id": streamed_message_id}) + def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse: """ Handles logging for Anthropic messages responses. @@ -3832,11 +3855,12 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(result, ResponsesAPIResponse): return self._translate_responses_api_response_to_model_response(result) + provider_response_id: Final = _provider_response_id(result) httpx_response: Final = self.model_call_details.get("httpx_response", None) if httpx_response and isinstance(httpx_response, httpx.Response): result = litellm.AnthropicConfig().transform_response( raw_response=httpx_response, - model_response=litellm.ModelResponse(), + model_response=litellm.ModelResponse(id=provider_response_id), model=self.model, messages=[], logging_obj=self, @@ -3859,7 +3883,7 @@ class Logging(LiteLLMLoggingBaseClass): status_code=200, headers={}, ), - model_response=litellm.ModelResponse(), + model_response=litellm.ModelResponse(id=provider_response_id), json_mode=None, speed=self.optional_params.get("speed") if self.optional_params else None, ) @@ -3882,7 +3906,7 @@ class Logging(LiteLLMLoggingBaseClass): return LiteLLMResponsesTransformationHandler().transform_response( model=self.model, raw_response=result, - model_response=litellm.ModelResponse(), + model_response=litellm.ModelResponse(id=_provider_response_id(result)), logging_obj=self, request_data={}, messages=[], @@ -3897,7 +3921,7 @@ class Logging(LiteLLMLoggingBaseClass): "usage-only ModelResponse to keep the spend_logs row.", str(e), ) - model_response: Final = litellm.ModelResponse() + model_response: Final = litellm.ModelResponse(id=_provider_response_id(result)) model_response.model = self.model usage: Final = getattr(result, "usage", None) if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage): 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 5504756ceb8..bf99035a6b1 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 @@ -5,6 +5,8 @@ Helper utilities for tracking the cost of built-in tools. from collections.abc import Mapping from typing import Final, Literal +from pydantic import ValidationError + import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -13,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, + ResponsesToolUsage, WebSearchOptions, ) from litellm.types.utils import ( @@ -32,6 +35,17 @@ def _output_item_type(output_item: object) -> str | None: return item_type if isinstance(item_type, str) else None +def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: + tool_usage: Final = getattr(response_object, "tool_usage", None) + if tool_usage is None: + return None + try: + web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search + except ValidationError: + return None + return None if web_search is None else web_search.num_requests + + def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: details: Final = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): @@ -182,15 +196,19 @@ class StandardBuiltInToolCostTracking: Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by get_cost_for_web_search_request and never reach here. This path prices per call, so it must count - the web_search_call items. Chat-completions responses only expose url_citation annotations with no - count, so they floor to a single billable search. + the web_search_call items, unless the response reports the billable count itself + (Bedrock's tool_usage.web_search.num_requests, which excludes open_page fetches). Chat-completions + responses only expose url_citation annotations with no count, so they floor to a single billable search. """ - if isinstance(response_object, ResponsesAPIResponse): - count = sum( - 1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call" - ) - return max(count, 1) - return 1 + if not isinstance(response_object, ResponsesAPIResponse): + return 1 + reported: Final = _reported_web_search_requests(response_object) + if reported is not None: + return reported + count: Final = sum( + 1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call" + ) + return max(count, 1) @staticmethod def _handle_file_search_cost( diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index b62226a6a19..390abf41955 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -11,7 +11,7 @@ from typing import Final from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH -_REDACTED: Final = "REDACTED" +REDACTED: Final = "REDACTED" def _build_secret_patterns() -> "re.Pattern[str]": @@ -89,7 +89,7 @@ _SECRET_RE: Final = _build_secret_patterns() def redact_string(value: str) -> str: """Scrub known secret/credential patterns from *value* and return the result.""" - return _SECRET_RE.sub(_REDACTED, value) + return _SECRET_RE.sub(REDACTED, value) _UNIX_SYSTEM_PATH: Final = r"/(?:etc|var|opt|usr|home|root|private|Users|tmp|mnt|srv)/[^\s'\"\)\]}>,]+" @@ -110,7 +110,7 @@ def redact_internal_details(value: str) -> str: on top of redact_string(). For client-facing messages only: server logs keep this detail.""" marker_index: Final = value.find(_TRACEBACK_MARKER) without_traceback: Final = value[:marker_index].rstrip() if marker_index != -1 else value - return _INTERNAL_DETAIL_RE.sub(_REDACTED, redact_string(without_traceback)) + return _INTERNAL_DETAIL_RE.sub(REDACTED, redact_string(without_traceback)) def redact_structured_value(key: str | None, value: str) -> str: @@ -126,4 +126,4 @@ def redact_structured_value(key: str | None, value: str) -> str: if scrubbed != value or key is None: return scrubbed rendered: Final = f"'{key}': '{value}'" - return _REDACTED if redact_string(rendered) != rendered else value + return REDACTED if redact_string(rendered) != rendered else value diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 3d60c1bda12..15b2c879224 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -1,9 +1,10 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any, Final from pydantic import BaseModel -from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER +from litellm.litellm_core_utils.secret_redaction import REDACTED class SensitiveDataMasker: @@ -214,6 +215,46 @@ def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dic return masked +def redact_credentials_in_payload(data: Mapping[str, object]) -> Mapping[str, object]: + """Return a copy of ``data`` where every value under a credential-named key is + replaced by the shared ``REDACTED`` marker, nested mappings are recursed into, + and every other value is preserved by identity. + + Sensitive-key detection is delegated to the shared :class:`SensitiveDataMasker`, + so the credential names stay in one place. Unlike + :func:`mask_credentials_in_payload`, no prefix or suffix of the secret survives + and non-string secrets are covered too, which is what a payload rendered + straight to stdout needs. ``None`` is preserved so an unset credential still + reads as unset, and lists and tuples are rebuilt element by element so a + credential nested inside one is caught as well. The walk is bounded only to stop + runaway recursion, and a container sitting at that bound is replaced wholesale + rather than passed through, so burying a credential deeper than the walk goes + hides it instead of exposing it. + """ + return _redact_mapping(data, 0) + + +def _redact_mapping(data: Mapping[str, object], depth: int) -> Mapping[str, object]: + return {key: _redact_entry(key, value, depth) for key, value in data.items()} + + +def _redact_entry(key: str, value: object, depth: int) -> object: + if value is not None and _default_masker.is_sensitive_key(key): + return REDACTED + if not isinstance(value, (Mapping, list, tuple)): + return value + if depth >= DEFAULT_MAX_RECURSE_DEPTH: + return REDACTED + if isinstance(value, Mapping): + return _redact_mapping(value, depth + 1) + return _redact_sequence(value, depth + 1) + + +def _redact_sequence(values: Sequence[object], depth: int) -> Sequence[object]: + redacted: Final = tuple(_redact_entry("", item, depth) for item in values) + return redacted if isinstance(values, tuple) else list(redacted) + + # Usage example: """ masker = SensitiveDataMasker() diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 355faa41e68..6ae17bac6ff 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2337,6 +2337,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + if self.stream_options is None: + usage: Final = calculate_total_usage(chunks=self.chunks) + processed_chunk._hidden_params["usage"] = usage # pyright: ignore[reportPrivateUsage] # sync parity # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 9d61701d26d..87a29ca50ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -21,6 +21,7 @@ from litellm.llms.anthropic.experimental_pass_through.context_management import ) from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, + litellm_logging_obj_from_kwargs, local_model_name, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -621,6 +622,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tool_name_mapping=tool_name_mapping, polyfill_result=polyfill_result, is_async=True, + litellm_logging_obj=litellm_logging_obj_from_kwargs(kwargs), ) if transformed_stream is not None: return transformed_stream @@ -755,6 +757,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tool_name_mapping=tool_name_mapping, polyfill_result=polyfill_result, is_async=False, + litellm_logging_obj=litellm_logging_obj_from_kwargs(kwargs), ) if transformed_stream is not None: return transformed_stream diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index cc5879df56d..78ff83cafbf 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.types.llms.anthropic import ( from litellm.types.utils import AdapterCompletionStreamWrapper, Delta if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.types.utils import ModelResponseStream @@ -287,12 +288,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): applied_edits: list[AppliedEdit] | None = None, compaction_block: CompactionBlock | None = None, iterations_usage: list[UsageIteration] | None = None, + litellm_logging_obj: "LiteLLMLoggingObject | None" = None, ): # Wrap the upstream stream so chunks that carry both content and a # finish_reason (fake-streamed providers) are split into two — see # _CombinedChunkSplitter. super().__init__(_CombinedChunkSplitter(completion_stream)) self.model = model + self._message_id: str = f"msg_{uuid.uuid4()}" + if litellm_logging_obj is not None: + litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) # Mapping of truncated tool names to original names (for OpenAI's 64-char limit) self.tool_name_mapping = tool_name_mapping or {} # Polyfill applied_edits on final message_delta. @@ -507,7 +512,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): { "type": "message_start", "message": { - "id": f"msg_{uuid.uuid4()}", + "id": self._message_id, "type": "message", "role": "assistant", "content": [], @@ -741,7 +746,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): { "type": "message_start", "message": { - "id": f"msg_{uuid.uuid4()}", + "id": self._message_id, "type": "message", "role": "assistant", "content": [], diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 199a8ab77e7..573a461e89e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -174,6 +174,7 @@ from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage from .streaming_iterator import AnthropicStreamWrapper if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.types.llms.anthropic import ContentBlockContentBlockDict ToolResultContent: TypeAlias = str | list[ToolMessageContentPart] @@ -264,6 +265,7 @@ class AnthropicAdapter: tool_name_mapping: dict[str, str] | None = None, polyfill_result: PolyfillResult | None = None, is_async: bool = True, + litellm_logging_obj: "LiteLLMLoggingObject | None" = None, ) -> AsyncIterator[bytes] | Iterator[bytes] | None: """ Translate OpenAI streaming response to Anthropic format. @@ -290,6 +292,7 @@ class AnthropicAdapter: applied_edits=applied_edits, compaction_block=compaction_block, iterations_usage=iterations_usage, + litellm_logging_obj=litellm_logging_obj, ) # Return the SSE-wrapped version for proper event formatting. if is_async: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index ec0560016da..0445c23ed8c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -20,7 +20,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( ) from litellm.types.llms.openai import ResponsesAPIResponse -from ..utils import local_model_name +from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper from .transformation import LiteLLMAnthropicToResponsesAPIAdapter @@ -186,7 +186,9 @@ class LiteLLMMessagesToResponsesAPIHandler: if stream: wrapper: Final = AnthropicResponsesStreamWrapper( - responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")) + responses_stream=result, + model=local_model_name(model, kwargs.get("custom_llm_provider")), + litellm_logging_obj=litellm_logging_obj_from_kwargs(responses_kwargs), ) return wrapper.async_anthropic_sse_wrapper() @@ -266,7 +268,9 @@ class LiteLLMMessagesToResponsesAPIHandler: if stream: wrapper: Final = AnthropicResponsesStreamWrapper( - responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")) + responses_stream=result, + model=local_model_name(model, kwargs.get("custom_llm_provider")), + litellm_logging_obj=litellm_logging_obj_from_kwargs(responses_kwargs), ) return wrapper.async_anthropic_sse_wrapper() diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 292d2622c7f..a97ce18d179 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -4,7 +4,7 @@ import json import traceback from collections import deque from collections.abc import AsyncIterator, Mapping -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from litellm import verbose_logger from litellm._uuid import uuid @@ -12,6 +12,9 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUs from .transformation import LiteLLMAnthropicToResponsesAPIAdapter +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject + class AnthropicResponsesStreamWrapper: """ @@ -31,10 +34,13 @@ class AnthropicResponsesStreamWrapper: self, responses_stream: Any, model: str, + litellm_logging_obj: "LiteLLMLoggingObject | None" = None, ) -> None: self.responses_stream = responses_stream self.model = model self._message_id: str = f"msg_{uuid.uuid4()}" + if litellm_logging_obj is not None: + litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) self._current_block_index: int = -1 # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: dict[str, int] = {} diff --git a/litellm/llms/anthropic/experimental_pass_through/utils.py b/litellm/llms/anthropic/experimental_pass_through/utils.py index 716a4f54778..55fe9c47faf 100644 --- a/litellm/llms/anthropic/experimental_pass_through/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/utils.py @@ -1,11 +1,14 @@ import os from collections.abc import Mapping from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final import litellm from litellm.types.utils import ModelInfo +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject + OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH: Final = 64 _EFFORT_DEGRADATION_CHAIN: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( @@ -24,6 +27,14 @@ def prompt_cache_key_from_user_id(user_id: object) -> str | None: return str(user_id)[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None +def litellm_logging_obj_from_kwargs(kwargs: Mapping[str, object]) -> "LiteLLMLoggingObject | None": + """The logging object the bridged call logs through, when the caller supplied one.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + candidate: Final = kwargs.get("litellm_logging_obj") + return candidate if isinstance(candidate, Logging) else None + + def local_model_name(model: str, custom_llm_provider: object) -> str: """The id the provider itself knows, for reporting back to the caller in ``message_start``.""" return model.removeprefix(f"{custom_llm_provider}/") if isinstance(custom_llm_provider, str) else model diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 6fdd277a04f..3189f5b57ac 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -7,6 +7,7 @@ from litellm.exceptions import UnsupportedParamsError from litellm.llms.openai.chat.gpt_5_transformation import ( OpenAIGPT5Config, _get_effort_level, + is_gpt_reasoning_series_name, ) from litellm.types.llms.openai import AllMessageValues @@ -35,26 +36,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - """Check if the Azure model string refers to a gpt-5 variant. - - Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix - used for manual routing. - """ - # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07, - # …) are regular chat models: they support temperature and tool_choice but NOT - # reasoning_effort. They must NOT be routed through the GPT-5 reasoning path. - # - # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning - # models and must stay on the GPT-5 path. The distinguishing feature is that - # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5" - # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version - # number (i.e. "gpt-5.-chat"). - # - # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather - # than a substring check) makes this boundary explicit and avoids any ambiguity - # if future model names coincidentally contain "gpt-5-chat" as an interior run. - _normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "azure/" - return ("gpt-5" in model and not _normalized.startswith("gpt-5-chat")) or "gpt5_series" in model + return is_gpt_reasoning_series_name(model) or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> list[str]: """Get supported parameters for Azure OpenAI GPT-5 models. diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0ac0662205a..880a51eb584 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, ) from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.chat.gpt_5_transformation import GPT_REASONING_SERIES_MARKERS from litellm.types.llms.azure import ( API_VERSION_MONTH_SUPPORTED_RESPONSE_FORMAT, API_VERSION_YEAR_SUPPORTED_RESPONSE_FORMAT, @@ -139,7 +140,7 @@ class AzureOpenAIConfig(BaseConfig): name family needs the rename, including the ``gpt-5-chat*`` models that are excluded from the reasoning path by https://github.com/BerriAI/litellm/issues/13781. """ - return "gpt-5" in model or "gpt5_series" in model + return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) or "gpt5_series" in model def _is_response_format_supported_model(self, model: str) -> bool: """ diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 7fe9d3dec52..f2d405e9a17 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig @@ -207,20 +208,18 @@ class AzureAIStudioConfig(OpenAIConfig): message["content"] = texts return stripped_messages - def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: - try: - if "/" in model: - model = model.split("/", 1)[1] - if ( - model in litellm.open_ai_chat_completion_models - or model in litellm.open_ai_text_completion_models - or model in litellm.open_ai_embedding_models - ): - return True + def _is_foundry_model_inference_base(self, api_base: str) -> bool: + return is_foundry_model_inference_base(api_base) - except Exception: + def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: + if api_base is None or self._is_foundry_model_inference_base(api_base): return False - return False + stripped_model: Final = model.split("/", 1)[1] if "/" in model else model + return ( + stripped_model in litellm.open_ai_chat_completion_models + or stripped_model in litellm.open_ai_text_completion_models + or stripped_model in litellm.open_ai_embedding_models + ) def _get_openai_compatible_provider_info( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 26a90157455..aa34bab5b2e 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from typing import Final, Literal +from urllib.parse import urlparse import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter @@ -10,6 +11,14 @@ from litellm.types.router import GenericLiteLLMParams AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"] +def is_foundry_model_inference_base(api_base: str) -> bool: + parsed: Final = urlparse(api_base) + host: Final = parsed.hostname + if host is None or not host.endswith(".services.ai.azure.com"): + return False + return "/openai/deployments" not in parsed.path + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 65c3997c099..c65edbf56e6 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -1,8 +1,10 @@ from typing import Final +from urllib.parse import urlsplit, urlunsplit from openai import OpenAI import litellm +from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -16,6 +18,16 @@ from litellm.utils import convert_to_model_response_object from .cohere_transformation import AzureAICohereConfig +def _foundry_models_route_base(api_base: str | None) -> str | None: + if api_base is None or not is_foundry_model_inference_base(api_base): + return api_base + parts: Final = urlsplit(api_base) + path: Final = parts.path.rstrip("/") + if path.endswith("/models"): + return api_base + return urlunsplit((parts.scheme, parts.netloc, f"{path}/models", parts.query, parts.fragment)) + + class AzureAIEmbedding(OpenAIChatCompletion): def _process_response( self, @@ -214,6 +226,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): assemble result in-order, and return """ + resolved_api_base: Final = _foundry_models_route_base(api_base) if aembedding is True: return self.async_embedding( model, @@ -223,7 +236,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response, optional_params, api_key, - api_base, + resolved_api_base, client, ) @@ -245,7 +258,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response=model_response, optional_params=optional_params, api_key=api_key, - api_base=api_base, + api_base=resolved_api_base, client=client, ) @@ -262,7 +275,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response, optional_params, api_key, - api_base, + resolved_api_base, client=(client if client is not None and isinstance(client, OpenAI) else None), aembedding=aembedding, shared_session=shared_session, diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1e634ced29b..c3da992a904 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -7,7 +7,7 @@ import urllib.parse from collections.abc import Callable from datetime import datetime from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload import httpx from pydantic import BaseModel, ValidationError @@ -48,12 +48,24 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile( SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"}) -class Boto3CredentialsInfo(BaseModel): - credentials: Credentials +class BedrockRequestTarget(BaseModel): aws_region_name: str aws_bedrock_runtime_endpoint: str | None +class Boto3CredentialsInfo(BedrockRequestTarget): + credentials: Credentials + + +class BearerRequestTarget(BedrockRequestTarget): + credentials: None = None + + +def bedrock_bearer_token(api_key: str | None) -> str | None: + token: Final = api_key if api_key is not None else get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + return token or None + + class _WebIdentityTokenClaims(BaseModel): aud: str | list[str] | None = None iss: str | None = None @@ -1387,9 +1399,26 @@ class BaseAWSLLM: else: return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" + @overload def _get_boto_credentials_from_optional_params( - self, optional_params: dict, model: str | None = None - ) -> Boto3CredentialsInfo: + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + model: str | None = None, + bearer_token: None = None, + ) -> Boto3CredentialsInfo: ... + + @overload + def _get_boto_credentials_from_optional_params( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + model: str | None = None, + *, + bearer_token: str, + ) -> BearerRequestTarget: ... + + def _get_boto_credentials_from_optional_params( + self, optional_params: dict, model: str | None = None, bearer_token: str | None = None + ) -> Boto3CredentialsInfo | BearerRequestTarget: """ Get boto3 credentials from optional params @@ -1420,6 +1449,12 @@ class BaseAWSLLM: ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) + if bearer_token is not None: + return BearerRequestTarget( + aws_region_name=aws_region_name, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + ) + credentials: Final[Credentials] = self.get_credentials( aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, @@ -1432,7 +1467,6 @@ class BaseAWSLLM: aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, ) - return Boto3CredentialsInfo( credentials=credentials, aws_region_name=aws_region_name, @@ -1451,14 +1485,9 @@ class BaseAWSLLM: api_key: str | None = None, supports_bearer_token: bool = True, ) -> AWSPreparedRequest: - if not supports_bearer_token: - aws_bearer_token: str | None = None - elif api_key is not None: - aws_bearer_token = api_key - else: - aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + aws_bearer_token: Final = bedrock_bearer_token(api_key) if supports_bearer_token else None - if aws_bearer_token: + if aws_bearer_token is not None: try: from botocore.awsrequest import AWSRequest except ImportError: @@ -1555,13 +1584,9 @@ class BaseAWSLLM: Returns: Tuple[dict, Optional[str]]: A tuple containing the headers and the json str body of the request """ - if api_key is not None: - aws_bearer_token: str | None = api_key - else: - aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + aws_bearer_token: Final = bedrock_bearer_token(api_key) - # If aws bearer token is set, use it directly in the header - if aws_bearer_token: + if aws_bearer_token is not None: headers = headers or {} headers["Content-Type"] = "application/json" headers["Authorization"] = f"Bearer {aws_bearer_token}" diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7d5f99ca893..a75124325ae 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -from ..base_aws_llm import BaseAWSLLM, Credentials +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token from ..common_utils import BedrockError, _get_all_bedrock_regions from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -349,17 +349,21 @@ class BedrockConverseLLM(BaseAWSLLM): litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls - credentials: Final[Credentials | None] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bedrock_bearer_token(api_key) is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) ### SET RUNTIME ENDPOINT ### diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 91c3a363c31..04c6ec86a13 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -149,19 +149,15 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): - Temperature and parameter validation """ - # Filter out AWS credentials using the existing method from BaseAWSLLM - self._get_boto_credentials_from_optional_params(optional_params, model) + inference_params: Final = {k: v for k, v in optional_params.items() if k not in self.aws_authentication_params} - # Strip routing prefixes to get the actual model ID clean_model_id: Final = self._get_model_id(model) - # Use Moonshot's transform_request which handles message transformation - # and tool_choice="required" workaround return MoonshotChatConfig.transform_request( self, model=clean_model_id, messages=messages, - optional_params=optional_params, + optional_params=inference_params, litellm_params=litellm_params, headers=headers, ) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index c34ca7750e2..5fb86d476f4 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,7 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, get_args +from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -42,14 +42,25 @@ if TYPE_CHECKING: class BedrockEmbedding(BaseAWSLLM): + @overload + def _load_credentials( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + bearer_token: None = None, + ) -> tuple[Credentials, str]: ... + + @overload + def _load_credentials( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + bearer_token: str, + ) -> tuple[None, str]: ... + def _load_credentials( self, optional_params: dict, - ) -> tuple[Any, str]: - try: - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + bearer_token: str | None = None, + ) -> tuple[Credentials | None, str]: ## CREDENTIALS ## # pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key: Final = optional_params.pop("aws_secret_access_key", None) @@ -78,17 +89,21 @@ class BedrockEmbedding(BaseAWSLLM): if aws_region_name is None: aws_region_name = "us-west-2" - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bearer_token is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) return credentials, aws_region_name @@ -233,7 +248,7 @@ class BedrockEmbedding(BaseAWSLLM): client: HTTPHandler | None, timeout: float | httpx.Timeout | None, batch_data: list[dict], - credentials: Any, + credentials: Credentials | None, extra_headers: dict | None, endpoint_url: str, aws_region_name: str, @@ -301,7 +316,7 @@ class BedrockEmbedding(BaseAWSLLM): client: AsyncHTTPHandler | None, timeout: float | httpx.Timeout | None, batch_data: list[dict], - credentials: Any, + credentials: Credentials | None, extra_headers: dict | None, endpoint_url: str, aws_region_name: str, @@ -383,7 +398,9 @@ class BedrockEmbedding(BaseAWSLLM): litellm_params: dict, api_key: str | None = None, ) -> EmbeddingResponse: - credentials, aws_region_name = self._load_credentials(optional_params) + credentials, aws_region_name = self._load_credentials( + optional_params, bearer_token=bedrock_bearer_token(api_key) + ) ### TRANSFORMATION ### unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 9d8631c7c26..5c517f2049c 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import ImageResponse -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token from ..common_utils import BedrockError if TYPE_CHECKING: @@ -198,7 +198,9 @@ class BedrockImageEdit(BaseAWSLLM): Returns: BedrockImageEditPreparedRequest: The prepared request object """ - boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model) + boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params( + optional_params, model, bearer_token=bedrock_bearer_token(api_key) + ) # Use the existing ARN-aware provider detection method bedrock_provider: Final = self.get_bedrock_invoke_provider(model) diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index 6fac14a0dc3..c78e3c147cb 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import ImageResponse -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token from ..common_utils import BedrockError if TYPE_CHECKING: @@ -220,7 +220,9 @@ class BedrockImageGeneration(BaseAWSLLM): prepped (httpx.Request): The prepared request object body (bytes): The request body """ - boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model) + boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params( + optional_params, model, bearer_token=bedrock_bearer_token(api_key) + ) # Use the existing ARN-aware provider detection method bedrock_provider: Final = self.get_bedrock_invoke_provider(model) diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index dd20a8c2ed4..de72735e4ec 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -139,7 +139,7 @@ def _build_query_params( return {name: value if isinstance(value, str) else str(value) for name, value in supplied if value is not None} -def _error_message_from_response(response: httpx.Response) -> str: +def error_message_from_response(response: httpx.Response) -> str: try: body: Final = response.json() except ValueError: @@ -153,6 +153,16 @@ def _error_message_from_response(response: httpx.Response) -> str: return response.text +def raise_for_error_status(response: httpx.Response, container_provider_config: "BaseContainerConfig") -> None: + if not httpx.codes.is_error(response.status_code): + return + raise container_provider_config.get_error_class( + error_message=error_message_from_response(response), + status_code=response.status_code, + headers=response.headers, + ) + + def _transform_response( response: httpx.Response, returns_binary: bool, @@ -163,7 +173,7 @@ def _transform_response( if httpx.codes.is_error(response.status_code): raise BaseLLMException( status_code=response.status_code, - message=_error_message_from_response(response), + message=error_message_from_response(response), headers=dict(response.headers), ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 71a598a6fe7..f281c249c72 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -99,9 +100,12 @@ from litellm.types.containers.main import ( ) from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopPlan, AgenticLoopRequestPatch, AgenticLoopSafetyError, + converted_stream_requested, + is_interception_internal_key, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -2760,6 +2764,7 @@ class BaseLLMHTTPHandler: ) if self._has_agentic_completion_hook(logging_obj): + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = run_async_function( self._call_agentic_completion_hooks, response=initial_response, @@ -2770,10 +2775,19 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=dict(litellm_params), + kwargs=agentic_kwargs, api_surface="responses", ) - return final_response if final_response is not None else initial_response + result: Final = final_response if final_response is not None else initial_response + if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): + return self._wrap_responses_response_as_fake_stream( + result=result, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + return result return initial_response @@ -2939,6 +2953,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -2948,15 +2963,12 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=dict(litellm_params), + kwargs=agentic_kwargs, api_surface="responses", ) result: Final = final_response if final_response is not None else initial_response - interception_converted_stream: Final = litellm_params.get( - "_code_interpreter_interception_converted_stream" - ) or litellm_params.get("_websearch_interception_converted_stream") - if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"): + if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): return self._wrap_responses_response_as_fake_stream( result=result, model=model, @@ -5420,8 +5432,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup: Final = { k: v for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and not k.startswith("_compression_interception") + if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) and k != "_code_interpreter_interception_converted_stream" and k not in internal_keys and k not in optional_params @@ -8753,17 +8764,19 @@ class BaseLLMHTTPHandler: json=data, timeout=timeout, ) - - return container_provider_config.transform_container_create_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_create_response( + raw_response=response, + logging_obj=logging_obj, + ) async def async_container_create_handler( self, @@ -8829,17 +8842,19 @@ class BaseLLMHTTPHandler: json=data, timeout=timeout, ) - - return container_provider_config.transform_container_create_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_create_response( + raw_response=response, + logging_obj=logging_obj, + ) def container_list_handler( self, @@ -8919,17 +8934,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_list_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_list_response( + raw_response=response, + logging_obj=logging_obj, + ) async def async_container_list_handler( self, @@ -8996,17 +9013,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_list_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_list_response( + raw_response=response, + logging_obj=logging_obj, + ) def container_retrieve_handler( self, @@ -9084,17 +9103,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_retrieve_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_retrieve_response( + raw_response=response, + logging_obj=logging_obj, + ) async def async_container_retrieve_handler( self, @@ -9161,17 +9182,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_retrieve_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_retrieve_response( + raw_response=response, + logging_obj=logging_obj, + ) def container_delete_handler( self, @@ -9249,17 +9272,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_delete_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_delete_response( + raw_response=response, + logging_obj=logging_obj, + ) async def async_container_delete_handler( self, @@ -9326,17 +9351,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_delete_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_delete_response( + raw_response=response, + logging_obj=logging_obj, + ) def container_file_list_handler( self, @@ -9418,17 +9445,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_file_list_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_file_list_response( + raw_response=response, + logging_obj=logging_obj, + ) async def async_container_file_list_handler( self, @@ -9497,17 +9526,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_file_list_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_file_list_response( + raw_response=response, + logging_obj=logging_obj, + ) def container_file_content_handler( self, @@ -9583,17 +9614,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_file_content_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_file_content_response( + raw_response=response, + logging_obj=logging_obj, + ) async def async_container_file_content_handler( self, @@ -9659,17 +9692,19 @@ class BaseLLMHTTPHandler: headers=headers, params=params or None, ) - - return container_provider_config.transform_container_file_content_response( - raw_response=response, - logging_obj=logging_obj, - ) - except Exception as e: raise self._handle_error( e=e, provider_config=container_provider_config, ) + raise_for_error_status( + response=response, + container_provider_config=container_provider_config, + ) + return container_provider_config.transform_container_file_content_response( + raw_response=response, + logging_obj=logging_obj, + ) ###### VECTOR STORE HANDLER ###### @staticmethod diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index dd5bee1fe8b..d8eb1f9f8d7 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,11 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ -from dataclasses import dataclass +from dataclasses import dataclass, replace +from datetime import datetime from typing import Final from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.litellm_core_utils.llm_cost_calc.utils import ( + apply_off_peak_pricing, parse_completion_tokens_details, parse_prompt_tokens_details, ) @@ -32,6 +34,19 @@ class TokenBreakdown: return self.text_tokens + self.cached_tokens + self.cache_creation_tokens +@dataclass(frozen=True, slots=True) +class TokenRates: + input_rate: float + cache_read_rate: float + cache_creation_rate: float + output_rate: float + reasoning_rate: float | None + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: prompt_details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = prompt_details["cache_hit_tokens"] @@ -57,69 +72,75 @@ def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) -> return float(value) -def _calculate_prompt_cost( - breakdown: TokenBreakdown, - model_info: ModelInfo, - tier: dict | None, -) -> float: - if tier is not None: - return ( - (breakdown.text_tokens * tier_rate(tier, "input_cost_per_token")) - + (breakdown.cached_tokens * tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token")) - + ( - breakdown.cache_creation_tokens - * tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") - ) - ) - - input_cost: Final = float(model_info.get("input_cost_per_token") or 0.0) - cache_read_cost: Final = _flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token") - cache_creation_cost: Final = _flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token") - - return ( - (breakdown.text_tokens * input_cost) - + (breakdown.cached_tokens * cache_read_cost) - + (breakdown.cache_creation_tokens * cache_creation_cost) +def _flat_rates(model_info: ModelInfo) -> TokenRates: + reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token") + return TokenRates( + input_rate=float(model_info.get("input_cost_per_token") or 0.0), + cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"), + output_rate=float(model_info.get("output_cost_per_token") or 0.0), + reasoning_rate=None if reasoning_rate is None else float(reasoning_rate), ) -def _calculate_completion_cost( - breakdown: TokenBreakdown, - model_info: ModelInfo, - tier: dict | None, -) -> float: +def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: # A tier that declares output rates keeps the request on them, all-or-nothing. A tier table # spelling out only input rates would serve every completion for free, so there the model's # own output rates stand in - tier_declares_output: Final = tier is not None and "output_cost_per_token" in tier - output_cost: Final = ( - tier_rate(tier, "output_cost_per_token") - if tier_declares_output - else float(model_info.get("output_cost_per_token") or 0.0) - ) - tier_declares_reasoning: Final = tier is not None and "output_cost_per_reasoning_token" in tier - model_reasoning_rate: Final = None if tier_declares_output else model_info.get("output_cost_per_reasoning_token") - reasoning_cost: Final = ( - tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") - if tier_declares_reasoning - else float(model_reasoning_rate) - if model_reasoning_rate is not None - else output_cost + flat_rates: Final = _flat_rates(model_info) + tier_declares_output: Final = "output_cost_per_token" in tier + tier_declares_reasoning: Final = "output_cost_per_reasoning_token" in tier + return TokenRates( + input_rate=tier_rate(tier, "input_cost_per_token"), + cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token"), + output_rate=tier_rate(tier, "output_cost_per_token") if tier_declares_output else flat_rates.output_rate, + reasoning_rate=( + tier_rate(tier, "output_cost_per_reasoning_token") + if tier_declares_reasoning + else None + if tier_declares_output + else flat_rates.reasoning_rate + ), ) - return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) + +def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: + input_rate, output_rate, cache_read_rate = apply_off_peak_pricing( + model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate + ) + return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate) -def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]: +def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]: + prompt_cost: Final = ( + (breakdown.text_tokens * rates.input_rate) + + (breakdown.cached_tokens * rates.cache_read_rate) + + (breakdown.cache_creation_tokens * rates.cache_creation_rate) + ) + completion_cost: Final = (breakdown.completion_tokens * rates.output_rate) + ( + breakdown.reasoning_tokens * rates.billed_reasoning_rate + ) + return prompt_cost, completion_cost + + +def cost_per_token( + model: str, + usage: Usage, + custom_llm_provider: str = "dashscope", + current_time: datetime | None = None, +) -> tuple[float, float]: """ Calculate cost per token for Dashscope models. - Supports both tiered and flat pricing with cached and reasoning tokens. + Supports both tiered and flat pricing with cached and reasoning tokens, and swaps in the + model's off_peak_pricing rates while one of its windows is open. Args: model: Model name without provider prefix usage: LiteLLM Usage block custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases + current_time: The moment the request is billed at; defaults to now, UTC Returns: Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd) @@ -133,8 +154,7 @@ def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashsco if tiered_pricing else None ) + standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier) + rates: Final = _off_peak_rates(model_info, current_time, standard_rates) - prompt_cost: Final = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier) - completion_cost: Final = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier) - - return prompt_cost, completion_cost + return _bill(breakdown, rates) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 0223be300b0..b02f953425d 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -61,6 +61,14 @@ def _get_effort_level(value: str | dict | None) -> str | None: return None +GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6") + + +def is_gpt_reasoning_series_name(model: str) -> bool: + normalized: Final = model.split("/")[-1] + return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat") + + class OpenAIGPT5Config(OpenAIGPTConfig): """Configuration for gpt-5 models including GPT-5-Codex variants. @@ -73,21 +81,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07, - # …) are regular chat models: they support temperature and tool_choice but NOT - # reasoning_effort. They must NOT be routed through the GPT-5 reasoning path. - # - # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning - # models and must stay on the GPT-5 path. The distinguishing feature is that - # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5" - # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version - # number (i.e. "gpt-5.-chat"). - # - # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather - # than a substring check) makes this boundary explicit and avoids any ambiguity - # if future model names coincidentally contain "gpt-5-chat" as an interior run. - _normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "openai/" - return "gpt-5" in model and not _normalized.startswith("gpt-5-chat") + return is_gpt_reasoning_series_name(model) @classmethod def is_model_gpt_5_search_model(cls, model: str) -> bool: @@ -122,6 +116,8 @@ class OpenAIGPT5Config(OpenAIGPTConfig): def is_model_gpt_5_4_plus_model(cls, model: str) -> bool: """Check if the model is gpt-5.4 or newer (5.4, 5.5, 5.6, etc., including pro).""" model_name: Final = model.split("/")[-1] + if model_name.startswith("gpt-6"): + return True if not model_name.startswith("gpt-5."): return False try: diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index bcd4ea43243..2db6d78a218 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -11,6 +11,7 @@ import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional +from urllib.parse import urlsplit import httpx import openai @@ -43,6 +44,14 @@ _OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(OpenAI) _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(AzureOpenAI) +_OPENAI_API_HOST: Final[str] = "api.openai.com" + + +def is_openai_backed_api_base(api_base: str) -> bool: + hostname: Final = urlsplit(api_base).hostname + return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}")) + + class OpenAIError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 1cfc6e06ee9..edc8d64d9c2 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -2,7 +2,6 @@ import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast -from urllib.parse import urlparse import httpx @@ -55,6 +54,7 @@ from .common_utils import ( OpenAIError, build_output_token_limit_response, drop_params_from_unprocessable_entity_error, + is_openai_backed_api_base, is_output_token_limit_error, ) from .workload_identity import resolve_openai_workload_identity_config @@ -1190,10 +1190,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ if stream_options is not None: return {"stream_options": stream_options} - else: - # by default litellm will include usage for openai endpoints - if api_base is None or urlparse(api_base).hostname == "api.openai.com": - return {"stream_options": {"include_usage": True}} + if api_base is None or is_openai_backed_api_base(api_base): + return {"stream_options": {"include_usage": True}} return {} # Embedding diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a0db7aadb9e..ecc89c5f135 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -33,8 +33,9 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import accumulate from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from pydantic import BaseModel, TypeAdapter @@ -42,6 +43,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.base_llm.guardrail_translation.base_translation import ( @@ -74,6 +76,7 @@ from litellm.types.llms.openai import ( OutputTextDoneEvent, ResponseAPIUsage, ResponseCompletedEvent, + ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -115,6 +118,199 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call_output": "output", "message": "content"} +) + +_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {} + + +def _item_rewrite_field(item: Mapping[str, object]) -> str | None: + item_type: Final = item.get("type") + if item_type is None: + return "content" if "content" in item else None + if not isinstance(item_type, str): + return None + return _PATCHABLE_ITEM_FIELDS.get(item_type) + + +def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + field: Final = _item_rewrite_field(item) + if field is None or not isinstance(rewritten, Mapping): + return None + rewritten_content: Final = rewritten.get("content") + if isinstance(item.get(field), str) and isinstance(rewritten_content, str): + return {**item, field: rewritten_content} # mutable-ok: request input items must stay JSON-plain dicts + rewritten_row: Final = cast("AllMessageValues", rewritten) # cast-ok: guardrails hand back chat-shaped rows + converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [rewritten_row] # mutable-ok: converter signature takes a list + ) + if len(converted_items) != 1 or not isinstance(converted_items[0], Mapping): + return None + first_converted: Final = cast("Mapping[str, object]", converted_items[0]) # cast-ok: isinstance-checked above + converted_value: Final = first_converted.get(field) + if converted_value is None: + return None + return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts + + +def _is_function_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") + + +def _last_message_role(messages: Sequence[object]) -> str | None: + if not messages: + return None + last: Final = messages[-1] + role: Final = last.get("role") if isinstance(last, Mapping) else getattr(last, "role", None) + return role if isinstance(role, str) else None + + +def _provenance_unit_bounds( + raw_input: Sequence[object], + solo_conversions: Sequence[Sequence[object]], +) -> tuple[tuple[int, int], ...]: + trailing_roles: Final = tuple( + accumulate( + (_last_message_role(messages) for messages in solo_conversions), + lambda previous, current: current if current is not None else previous, + ) + ) + start_indexes: Final = tuple( + index + for index in range(len(raw_input)) + if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + ) + return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) + + +def _input_item_provenance( + raw_input: Sequence[object], + expected_messages: Sequence[object], +) -> tuple[Mapping[int, int], frozenset[int]] | None: + if not all(isinstance(item, Mapping) for item in raw_input): + return None + solo_conversions: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", [item]), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + for item in raw_input + ) + full_conversion: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input)), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + if full_conversion != tuple(expected_messages): + return None + units: Final = _provenance_unit_bounds(raw_input, solo_conversions) + unit_messages: Final = tuple( + tuple(solo_conversions[start]) + if end - start == 1 + else tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input[start:end])), # cast-ok: checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + for start, end in units + ) + if tuple(message for messages in unit_messages for message in messages) != full_conversion: + return None + boundaries: Final = tuple(accumulate((len(messages) for messages in unit_messages), initial=0)) + item_for_message: Final = MappingProxyType( + { + message_index: start + for unit_index, (start, end) in enumerate(units) + if end - start == 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) + } + ) + tainted: Final = frozenset( + message_index + for unit_index, (start, end) in enumerate(units) + if end - start > 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) + ) + return item_for_message, tainted + + +class _RequestFields(NamedTuple): + input: tuple[object, ...] + instructions: str | None + + +class _ExtractedInputs(NamedTuple): + inputs: GenericGuardrailAPIInputs + task_mappings: tuple[tuple[int, int | None], ...] + + +def _patched_request_fields( + raw_input: object, + instructions: object, + original_messages: Sequence[object], + structured_messages: Sequence[object], +) -> _RequestFields | None: + if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): + return None + offset: Final = 1 if instructions else 0 + provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:]) + if provenance is None: + return None + item_for_message, tainted = provenance + changed: Final = tuple( + (index, rewritten) + for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages)) + if original != rewritten + ) + instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset) + rewritten_instructions: Final = ( + instruction_rewrites[0].get("content") + if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping) + else instructions + ) + instructions_value: Final = rewritten_instructions if isinstance(rewritten_instructions, str) else None + if rewritten_instructions is not None and instructions_value is None: + return None + body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset) + if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes): + return None + replacements: Final = MappingProxyType( + { + item_for_message[message_index]: _rewritten_input_item( + cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings + rewritten, + ) + for message_index, rewritten in body_changes + } + ) + if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()): + return None + return _RequestFields( + input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)), + instructions=instructions_value, + ) + + +def _patch_or_convert_request_fields( + raw_input: object, + instructions: object, + original_messages: Sequence[object], + structured_messages: Sequence[AllMessageValues], +) -> _RequestFields | None: + if not isinstance(structured_messages, list): + return None + patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages) + if patched is not None: + return patched + input_items, converted_instructions = ( + LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages) + ) + return _RequestFields(input=tuple(input_items), instructions=converted_instructions) + + def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) @@ -162,9 +358,8 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Final[str | ResponseInputParam | None] = data.get("input") - if input_data is None: + if not isinstance(input_data, (str, list)): return data - structured_messages: Final = self.get_structured_messages(data) raw_tools: Final = data.get("tools") original_tools: Final[tuple[Mapping[str, object], ...]] = ( @@ -173,94 +368,93 @@ class OpenAIResponsesHandler(BaseTranslation): flattened_tool_groups: Final = tuple( form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) - flattened_tools: Final = tuple( - cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list - for group in flattened_tool_groups - for tool in group - ) - tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list - copy.deepcopy(flattened_tools) - ) - - # Handle simple string input - if isinstance(input_data, str): - inputs = GenericGuardrailAPIInputs(texts=[input_data]) - if tools_to_check: - inputs["tools"] = tools_to_check - if structured_messages: - inputs["structured_messages"] = structured_messages - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=data, - input_type="request", - logging_obj=litellm_logging_obj, - ) - guardrailed_texts = guardrailed_inputs.get("texts", []) - data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data - self._apply_guardrailed_tools_to_data( - data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") - ) - verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") - return data - - # Handle list input (ResponseInputParam) - if not isinstance(input_data, list): + extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) + if not extracted.inputs.get("texts"): return data + if structured_messages: + extracted.inputs["structured_messages"] = structured_messages + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=extracted.inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + self._apply_guardrailed_tools_to_data( + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") + ) + written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) + if written_back is not None: + data["input"] = list(written_back.input) # mutable-ok: JSON body + if written_back.instructions is None: + data.pop("instructions", None) + else: + data["instructions"] = written_back.instructions # rebind-ok: data is an out-param + elif isinstance(input_data, str): + guardrailed_texts: Final = guardrailed_inputs.get("texts") or () + data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param + else: + await self._apply_guardrail_responses_to_input( + messages=input_data, + responses=guardrailed_inputs.get("texts") or (), + task_mappings=extracted.task_mappings, + ) + verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) + return data + def _extract_guardrail_inputs( + self, + data: Mapping[str, object], + input_data: "str | ResponseInputParam", + flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]], + ) -> _ExtractedInputs: texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] task_mappings: Final[list[tuple[int, int | None]]] = [] - - # Step 1: Extract all text content, images, and tools - for msg_idx, message in enumerate(input_data): - self._extract_input_text_and_images( - message=message, - msg_idx=msg_idx, - texts_to_check=texts_to_check, - images_to_check=images_to_check, - task_mappings=task_mappings, + tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list + copy.deepcopy( + tuple( + cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list + for group in flattened_tool_groups + for tool in group + ) ) + ) + if isinstance(input_data, str): + texts_to_check.append(input_data) + else: + for msg_idx, message in enumerate(input_data): + self._extract_input_text_and_images( + message=message, + msg_idx=msg_idx, + texts_to_check=texts_to_check, + images_to_check=images_to_check, + task_mappings=task_mappings, + ) + inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tools_to_check: + inputs["tools"] = tools_to_check + model: Final = data.get("model") + if isinstance(model, str): + inputs["model"] = model + return _ExtractedInputs(inputs=inputs, task_mappings=tuple(task_mappings)) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: - inputs = GenericGuardrailAPIInputs(texts=texts_to_check) - if images_to_check: - inputs["images"] = images_to_check - if tools_to_check: - inputs["tools"] = tools_to_check - if structured_messages: - inputs["structured_messages"] = structured_messages - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=data, - input_type="request", - logging_obj=litellm_logging_obj, - ) - - guardrailed_texts = guardrailed_inputs.get("texts", []) - self._apply_guardrailed_tools_to_data( - data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") - ) - - # Step 3: Map guardrail responses back to original input structure - await self._apply_guardrail_responses_to_input( - messages=input_data, - responses=guardrailed_texts, - task_mappings=task_mappings, - ) - - verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", input_data) - - return data + @staticmethod + def _written_back_request_fields( + data: Mapping[str, object], + structured_messages: Sequence[AllMessageValues] | None, + guardrailed_inputs: GenericGuardrailAPIInputs, + ) -> _RequestFields | None: + guardrailed: Final = guardrailed_inputs.get("structured_messages") + if guardrailed is None or guardrailed is structured_messages: + return None + return _patch_or_convert_request_fields( + data.get("input"), + data.get("instructions"), + structured_messages or (), + guardrailed, + ) def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from Responses API request (tools[].name for function @@ -331,8 +525,8 @@ class OpenAIResponsesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam - responses: list[str], - task_mappings: list[tuple[int, int | None]], + responses: Sequence[str], + task_mappings: Sequence[tuple[int, int | None]], ) -> None: """ Apply guardrail responses back to input messages. diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 01313e95878..b97521b90c2 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import * @@ -88,7 +89,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): parts: Final = model.split("/") if len(parts) > 1 and parts[0] not in ("openai",): return False - return "gpt-5" in model and "gpt-5-chat" not in model + return is_gpt_reasoning_series_name(model) @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index a36c920dda0..970759479fe 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -998,6 +998,16 @@ def replace_project_and_location_in_route(requested_route: str, vertex_project: return modified_route +def _api_version_for_route(requested_route: str) -> Literal["v1", "v1beta1"]: + return "v1beta1" if "cachedContent" in requested_route else "v1" + + +def _with_api_version(requested_route: str) -> str: + if not requested_route.startswith("/projects/"): + return requested_route + return f"/{_api_version_for_route(requested_route)}{requested_route}" + + def construct_target_url( base_url: str, requested_route: str, @@ -1017,18 +1027,19 @@ def construct_target_url( new_base_url: Final = httpx.URL(base_url) if "locations" in requested_route: # contains the target project id + location - if vertex_project and vertex_location: - requested_route = replace_project_and_location_in_route(requested_route, vertex_project, vertex_location) - return new_base_url.copy_with(path=requested_route) + targeted_route: Final = ( + replace_project_and_location_in_route(requested_route, vertex_project, vertex_location) + if vertex_project and vertex_location + else requested_route + ) + return new_base_url.copy_with(path=_with_api_version(targeted_route)) """ - Add endpoint version (e.g. v1beta for cachedContent, v1 for rest) - Add default project id - Add default location """ - vertex_version: Literal["v1", "v1beta1"] = "v1" - if "cachedContent" in requested_route: - vertex_version = "v1beta1" + vertex_version: Literal["v1", "v1beta1"] = _api_version_for_route(requested_route) # Check if the requested route starts with a version # e.g. /v1beta1/publishers/google/models/gemini-3-pro-preview:streamGenerateContent diff --git a/litellm/main.py b/litellm/main.py index 0bca4a7350e..2929790f2bd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -26,6 +26,7 @@ from copy import deepcopy from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from urllib.parse import urlsplit from litellm._logging import _redact_string from litellm._uuid import uuid @@ -60,6 +61,7 @@ if TYPE_CHECKING: from litellm.types.utils import TokenCountResponse from litellm.constants import ( + AZURE_OPENAI_AUDIO_PROVIDERS, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, ) @@ -984,6 +986,12 @@ def mock_completion( _OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1" +_OPENAI_API_HOST: Final = "api.openai.com" + + +def _is_openai_backed_api_base(api_base: str) -> bool: + hostname: Final = urlsplit(api_base).hostname + return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}")) def _resolve_openai_api_base(api_base: str | None) -> str: @@ -1053,7 +1061,7 @@ def responses_api_bridge_check( # natively by Chat Completions with reasoning on, so custom-only requests stay on # chat and keep their native custom tool_call response shape. # - The UNSET-effort arm only fires against endpoints known to enforce that - # constraint (the default OpenAI endpoint, or Azure OpenAI where api_base is + # constraint (any api.openai.com host, or Azure OpenAI where api_base is # always set): chat-only OpenAI-compatible backends registered under the openai # provider with a custom api_base and gpt-5.4+ model names serve tools without # reasoning fine and have no /responses route, so they keep pre-existing @@ -1068,14 +1076,15 @@ def responses_api_bridge_check( reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None else: reasoning_active = reasoning_effort != "none" - # The reasoning+tools constraint is enforced only by the real OpenAI endpoint (and Azure OpenAI). - # Resolve the effective base arg>global>env>default exactly as the chat handler does, so a custom - # base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread as the default and - # bridged to a /responses route it lacks. A whitespace-only base collapses to the default too. - resolved_api_base: Final = _resolve_openai_api_base(api_base) - on_constraint_enforcing_endpoint: Final = custom_llm_provider == "azure" or resolved_api_base.strip() in ( - "", - _OPENAI_DEFAULT_API_BASE, + # The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com + # host (the default URL or a PrivateLink hostname such as .privatelink.api.openai.com) and + # by Azure OpenAI. Resolve the effective base arg>global>env>default exactly as the chat handler + # does, so a custom base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread + # as the default and bridged to a /responses route it lacks. A whitespace-only base collapses to + # the default too. + resolved_api_base: Final = _resolve_openai_api_base(api_base).strip() + on_constraint_enforcing_endpoint: Final = ( + custom_llm_provider == "azure" or resolved_api_base == "" or _is_openai_backed_api_base(resolved_api_base) ) if ( custom_llm_provider in ("openai", "azure") @@ -7769,7 +7778,7 @@ def transcription( provider=LlmProviders(custom_llm_provider), ) - if custom_llm_provider == "azure" and provider_config is None: + if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None: # azure configs api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -8056,7 +8065,10 @@ def speech( custom_llm_provider=custom_llm_provider, ) response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None - if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers: + if custom_llm_provider == "openai" or ( + custom_llm_provider in litellm.openai_compatible_providers + and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS + ): if voice is None or not (isinstance(voice, str)): raise litellm.BadRequestError( message="'voice' is required to be passed as a string for OpenAI TTS", @@ -8110,7 +8122,7 @@ def speech( aspeech=aspeech, shared_session=shared_session, ) - elif custom_llm_provider == "azure": + elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: # Check if this is Azure Speech Service (Cognitive Services TTS) if model.startswith("speech/"): from litellm.llms.azure.text_to_speech.transformation import ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9b358a1cedd..44c4f10ec38 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10305,6 +10305,24 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/grok-4-6-comes-to-microsoft-foundry-models-built-for-long-horizon-reasoning-and-/4547578", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, @@ -29277,6 +29295,75 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + "cache_creation_input_token_cost_flex": 6.25e-06, + "cache_creation_input_token_cost_priority": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_read_input_token_cost_flex": 5e-07, + "cache_read_input_token_cost_priority": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "input_cost_per_token_batches": 5e-06, + "input_cost_per_token_flex": 5e-06, + "input_cost_per_token_priority": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "output_cost_per_token_batches": 2.5e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 0.0001, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -52911,6 +52998,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -52945,6 +53037,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53007,6 +53104,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53195,6 +53297,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53226,6 +53333,11 @@ "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "output_cost_per_token": 1.65e-05, "output_cost_per_token_above_272k_tokens": 2.475e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 94bca9460dd..81dd9057c9b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -448,6 +448,17 @@ def _append_query_params(url: str, params: dict[str, str]) -> str: return urlunparse(parsed._replace(query=urlencode(query_params))) +def _resolve_mcp_server_by_name_or_id(lookup: str, client_ip: str | None) -> MCPServer | None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + by_name: Final = global_mcp_server_manager.get_mcp_server_by_name(lookup, client_ip=client_ip) + if by_name is not None: + return by_name + return global_mcp_server_manager.get_mcp_server_by_id(lookup, client_ip=client_ip) + + def _resolve_oauth2_server_for_root_endpoints( client_ip: str | None = None, ) -> MCPServer | None: @@ -1766,10 +1777,6 @@ async def authorize( resource: str | None = None, ): # Redirect to real OAuth provider with PKCE support - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id): if is_proxy_api_resource(request, resource): return await native_client_authorize( @@ -1797,9 +1804,7 @@ async def authorize( lookup_name: Final[str | None] = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = ( - global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None - ) + mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if lookup_name else None if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: @@ -1855,10 +1860,6 @@ async def token_endpoint( 3. Return the token 4. Return a virtual key in this response """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - if mcp_server_name is None and is_gateway_dcr_client_id(client_id): from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load master_key, @@ -1882,7 +1883,7 @@ async def token_endpoint( lookup_name: Final = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) + mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: @@ -2288,10 +2289,6 @@ async def _build_oauth_protected_resource_response( Returns: OAuth protected resource metadata dict """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - request_base_url: Final = get_request_base_url(request) client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) explicitly_named: Final = mcp_server_name is not None @@ -2304,7 +2301,7 @@ async def _build_oauth_protected_resource_response( mcp_server: MCPServer | None = None if mcp_server_name: - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) # Build resource URL based on the pattern if mcp_server_name: @@ -2562,10 +2559,6 @@ def _build_oauth_authorization_server_response( registry lookups; unlike :func:`_build_oauth_protected_resource_response` it does not need to await any upstream IO. """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - request_base_url: Final = get_request_base_url(request) client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) explicitly_named: Final = mcp_server_name is not None @@ -2583,7 +2576,7 @@ def _build_oauth_authorization_server_response( mcp_server: MCPServer | None = None if mcp_server_name: - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth authorization server") @@ -2709,10 +2702,6 @@ async def oauth_authorization_server_legacy(request: Request, mcp_server_name: s @router.post("/{mcp_server_name}/register") @router.post("/register") async def register_client(request: Request, mcp_server_name: str | None = None): - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - # Get the correct base URL considering X-Forwarded-* headers request_base_url: Final = get_request_base_url(request) @@ -2748,7 +2737,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None): ) return dummy_return - mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server: Final = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) if mcp_server is None: return dummy_return return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a772c569bfa..36e1078dc76 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6147,13 +6147,13 @@ class MCPServerManager: internal_networks = IPAddressUtils.parse_internal_networks(general_settings.get("mcp_internal_ip_ranges")) return IPAddressUtils.is_internal_ip(client_ip, internal_networks) - def get_mcp_server_by_id(self, server_id: str) -> MCPServer | None: - """ - Get the MCP Server from the server id - """ + def get_mcp_server_by_id(self, server_id: str, client_ip: str | None = None) -> MCPServer | None: + """Get the MCP Server from the server id.""" registry: Final = self.get_registry() for server in registry.values(): if server.server_id == server_id: + if not self._is_server_accessible_from_ip(server, client_ip): + return None return server return None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 98464d3a127..5d5a25e7cd6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4239,6 +4239,8 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): access_group_id: str access_group_name: str models: tuple[str, ...] + mcp_server_ids: tuple[str, ...] = () + agent_ids: tuple[str, ...] = () class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index d0ac94d3710..e4dd77e2f82 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -5,10 +5,13 @@ Handles agent permission checking for keys and teams using object_permission_id. Follows the same pattern as MCP permission handling. """ +import asyncio +from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from typing import Final, TypeAlias from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_ObjectPermissionTable, @@ -443,15 +446,47 @@ class AgentRequestHandler: return [] -async def accessible_agents(user_api_key_auth: UserAPIKeyAuth) -> tuple[AgentResponse, ...]: - """Every registry agent for proxy admins, else the agents the key's and team's grants reach.""" +def _granted_ids(access: AgentAccess) -> frozenset[str]: + match access: + case UnrestrictedAgentAccess(): + return frozenset() + case RestrictedAgentAccess(agent_ids): + return agent_ids + + +ResolveAgentAccess: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[AgentAccess]] +EffectiveAuthContexts: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[Sequence[UserAPIKeyAuth]]] + + +async def _granted_agent_ids( + user_api_key_auth: UserAPIKeyAuth, + resolve_access: ResolveAgentAccess, + effective_contexts: EffectiveAuthContexts, +) -> frozenset[str]: + """Union of the explicit grants reachable from the key, its team, or (for a dashboard session) + the user's real teams and user row. No grant anywhere yields the empty set, unlike the + open-by-default ``resolve_agent_access`` that guards direct access.""" + accesses: Final = await asyncio.gather( + *(resolve_access(auth_context) for auth_context in await effective_contexts(user_api_key_auth)) + ) + return frozenset().union(*(_granted_ids(access) for access in accesses)) + + +async def accessible_agents( + user_api_key_auth: UserAPIKeyAuth, + all_agents: tuple[AgentResponse, ...] | None = None, + resolve_access: ResolveAgentAccess | None = None, + effective_contexts: EffectiveAuthContexts = build_effective_auth_contexts, +) -> tuple[AgentResponse, ...]: + """Every registry agent for proxy admins, else only the agents the caller was granted.""" from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - all_agents: Final = global_agent_registry.get_agent_list() + agents: Final = global_agent_registry.get_agent_list() if all_agents is None else all_agents if user_api_key_auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value): - return all_agents - match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_auth): - case UnrestrictedAgentAccess(): - return all_agents - case RestrictedAgentAccess(allowed_agent_ids): - return tuple(agent for agent in all_agents if agent.agent_id in allowed_agent_ids) + return agents + allowed_agent_ids: Final = await _granted_agent_ids( + user_api_key_auth, + AgentRequestHandler.resolve_agent_access if resolve_access is None else resolve_access, + effective_contexts, + ) + return tuple(agent for agent in agents if agent.agent_id in allowed_agent_ids) diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 677f1a0fdda..55bb1e3925a 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -16,6 +16,10 @@ if TYPE_CHECKING: from litellm.proxy._types import EnterpriseLicenseData +AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" +HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." + + class LicenseCheck: """ - Check if license in env @@ -149,6 +153,19 @@ class LicenseCheck: return False return team_count > _max_teams_in_license + def heuristic_v2_router_limit(self) -> int | None: + """ + How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the + signed license lists the auto_router feature, otherwise one. A license verified through + the API carries no feature list, so it does not lift the limit either. + """ + if self.airgapped_license_data is None: + return 1 + allowed_features: Final = self.airgapped_license_data.get("allowed_features") + if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features: + return None + return 1 + def verify_license_without_api_request(self, public_key, license_key): try: from cryptography.hazmat.primitives import hashes @@ -179,19 +196,21 @@ class LicenseCheck: # Decode and parse the data license_data: Final = json.loads(message.decode()) - self.airgapped_license_data = EnterpriseLicenseData(**license_data) - # debug information provided in license data verbose_proxy_logger.debug("License data: %s", license_data) # Check expiration date expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") if expiration_date < datetime.now(): + self.airgapped_license_data = None return False, "License has expired" + self.airgapped_license_data = EnterpriseLicenseData(**license_data) + return True except Exception as e: + self.airgapped_license_data = None verbose_proxy_logger.debug( "litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s", e, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fc83c1ddeed..6542842f5e4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -467,6 +467,42 @@ def _getattr_object(value: object, name: str, default: object = None) -> object: return getattr(value, name, default) +_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( + { + status.HTTP_401_UNAUTHORIZED: "authentication_error", + status.HTTP_403_FORBIDDEN: "permission_error", + status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error", + } +) + + +def _error_status_code(exc: object, default: int) -> int: + """The HTTP status an exception carries, or ``default`` when it carries none.""" + carried: Final = _getattr_object(exc, "status_code") + return carried if isinstance(carried, int) and not isinstance(carried, bool) else default + + +def _openai_error_type(exc: object, status_code: int) -> str: + """OpenAI types ``error.type`` as a required string, so an exception carrying none + falls back to the type its status code stands for.""" + carried: Final = _getattr_object(exc, "type") + if isinstance(carried, str): + return carried + mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) + if mapped is not None: + return mapped + if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR: + return "invalid_request_error" + return "internal_server_error" + + +def _openai_error_param(exc: object) -> str | None: + """OpenAI types ``error.param`` as nullable, so an exception carrying none + serializes as JSON ``null``.""" + carried: Final = _getattr_object(exc, "param") + return carried if isinstance(carried, str) else None + + class _UpstreamHttpResponse(Protocol): @property def status_code(self) -> int: ... @@ -540,11 +576,12 @@ def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, s message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) + error_status: Final = _error_status_code(exc, status.HTTP_400_BAD_REQUEST) return ProxyException( message=message, - type=getattr(exc, "type", "None"), - param=getattr(exc, "param", "None"), - code=getattr(exc, "status_code", status.HTTP_400_BAD_REQUEST), + type=_openai_error_type(exc, error_status), + param=_openai_error_param(exc), + code=error_status, provider_specific_fields=merged_fields, headers=headers, ) @@ -827,25 +864,22 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: are byte-identical. """ # Preserve status code from HTTPException (e.g. guardrail blocks) - error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) + error_status: Final = _error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) - # Built in one statement then given its one optional key, rather than spread - # conditionally: the spread form costs two extra dict constructions, which - # type-discipline-budget.json's LIT002 ceiling has no room for. error_obj: Final = { "message": message, - "type": getattr(exc, "type", "None"), - "param": getattr(exc, "param", "None"), + "type": _openai_error_type(exc, error_status), + "param": _openai_error_param(exc), "code": str(error_status), } - if merged_fields: - error_obj["provider_specific_fields"] = merged_fields - return error_status, error_obj + if not merged_fields: + return error_status, error_obj + return error_status, {**error_obj, "provider_specific_fields": merged_fields} def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: @@ -922,7 +956,7 @@ async def create_response( "error": { "message": _CLIENT_DISCONNECT_DETAIL, "type": "client_disconnect", - "param": "None", + "param": None, "code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED), } }, @@ -3417,8 +3451,8 @@ class ProxyBaseLLMRequestProcessing: _code = status.HTTP_500_INTERNAL_SERVER_ERROR raise ProxyException( message=redact_internal_details_from_client_message(getattr(e, "message", error_msg)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), + type=_openai_error_type(e, _code), + param=_openai_error_param(e), openai_code=getattr(e, "code", None), code=_code, provider_specific_fields=getattr(e, "provider_specific_fields", None), @@ -3628,11 +3662,12 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e + stream_error_status: Final = _error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) proxy_exception: Final = ProxyException( message=redact_internal_details_from_client_message(getattr(e, "message", str(e))), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=_openai_error_type(e, stream_error_status), + param=_openai_error_param(e), + code=stream_error_status, ) stream_completed = True yield serialize_error(proxy_exception) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 96621b08ba1..552d1ea434f 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -1,10 +1,12 @@ import json import re -from collections.abc import Collection -from typing import Any, Final +from collections.abc import Collection, Mapping +from types import MappingProxyType, UnionType +from typing import Any, Final, Union, get_args, get_origin import orjson from fastapi import Request, UploadFile, status +from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB @@ -40,6 +42,65 @@ def _is_json_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) == "application/json" +def _numeric_form_type(annotation: object) -> type[int] | type[float] | None: + """The scalar to parse an ``int``/``float``-typed field as, else ``None``.""" + unwrapped: Final = get_args(annotation)[0] if get_origin(annotation) is ReadOnly else annotation + candidates: Final = ( + tuple(arg for arg in get_args(unwrapped) if arg is not type(None)) + if get_origin(unwrapped) in (Union, UnionType) + else (unwrapped,) + ) + if len(candidates) != 1: + return None + if candidates[0] is int: + return int + if candidates[0] is float: + return float + return None + + +def numeric_form_fields(annotations: Mapping[str, object]) -> Mapping[str, type[int] | type[float]]: + """ + The numeric fields of a request schema, mapped to the scalar to parse them as. + + Only a bare ``int``/``float`` or an optional one qualifies, so container and + literal fields are left alone and ``bool`` is excluded on purpose. + """ + return MappingProxyType( + { + name: scalar + for name, annotation in annotations.items() + if (scalar := _numeric_form_type(annotation)) is not None + } + ) + + +def _numeric_form_value(value: object, scalar: type[int] | type[float]) -> object: + if not isinstance(value, str): + return value + try: + return scalar(value) + except ValueError: + return value + + +def coerce_numeric_form_fields( + parsed_body: Mapping[str, object], + numeric_fields: Mapping[str, type[int] | type[float]], +) -> Mapping[str, object]: + """ + Parse the numeric fields of a form-encoded body back into numbers. + + ``request.form()`` yields every field as a string, so a provider that puts the + value in a JSON body would send a string where its API requires a number. A + value that will not parse is left as-is for the provider to reject as before. + """ + return { + name: _numeric_form_value(value, numeric_fields[name]) if name in numeric_fields else value + for name, value in parsed_body.items() + } + + async def _read_request_body(request: Request | None) -> dict: """ Safely read the request body and parse it as JSON. diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index eaa3db336a9..f3a2abc1225 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -15,10 +15,11 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, - filter_container_list_response, get_container_forwarding_params, + list_owned_containers, record_container_owner, ) @@ -173,6 +174,9 @@ async def list_containers( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + after: str | None = None, + limit: int | None = None, + order: str | None = None, ): """ Container list endpoint for retrieving a list of containers. @@ -206,55 +210,54 @@ async def list_containers( version, ) - # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = {"query_params": query_params, "model": query_params.get("model")} - - # Extract custom_llm_provider using priority chain custom_llm_provider: Final = ( get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) or "openai" ) + data: Final[dict[str, Any]] = { + "query_params": query_params, + "model": query_params.get("model"), + "order": order, + "custom_llm_provider": custom_llm_provider, + } - # Add custom_llm_provider to data - data["custom_llm_provider"] = custom_llm_provider + async def fetch_page(page_after: str | None, page_limit: int | None) -> object: + processor: Final = ProxyBaseLLMRequestProcessing(data={**data, "after": page_after, "limit": page_limit}) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="alist_containers", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) - # Process request using ProxyBaseLLMRequestProcessing - processor: Final = ProxyBaseLLMRequestProcessing(data=data) - try: - response: Final = await processor.base_process_llm_request( - request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="alist_containers", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - ) - except Exception as e: - raise await processor._handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - version=version, - ) - - # Ownership filtering runs OUTSIDE the LLM-exception scope: a DB error - # in the ownership lookup is not an LLM-API error and shouldn't be - # translated to a provider-shaped failure (which would also fire the - # post_call_failure_hook for what is in fact a successful upstream call). - return await filter_container_list_response( - response=response, + if is_proxy_admin(user_api_key_dict): + return await fetch_page(after, limit) + return await list_owned_containers( + fetch_page=fetch_page, + after=after, + limit=limit, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index aaee1d3e264..cc742f0520b 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -6,7 +6,9 @@ FastAPI route handlers for ALL container file endpoints. """ import json +from collections.abc import Mapping, Sequence from pathlib import Path +from types import MappingProxyType from typing import Any, Final from fastapi import APIRouter, Depends, Request, Response @@ -56,6 +58,7 @@ def _create_handler_for_path_params( route_type: str, returns_binary: bool = False, is_multipart: bool = False, + query_param_names: Sequence[str] = (), ): """ Dynamically create a handler with the correct path parameter signature. @@ -114,6 +117,7 @@ def _create_handler_for_path_params( user_api_key_dict=user_api_key_dict, route_type=route_type, path_params={"container_id": container_id}, + query_param_names=query_param_names, ) return handler_container_id @@ -133,6 +137,7 @@ def _create_handler_for_path_params( user_api_key_dict=user_api_key_dict, route_type=route_type, path_params={"container_id": container_id, "file_id": file_id}, + query_param_names=query_param_names, ) return handler_container_file @@ -150,6 +155,7 @@ def _create_handler_for_path_params( user_api_key_dict=user_api_key_dict, route_type=route_type, path_params={}, + query_param_names=query_param_names, ) return handler_no_params @@ -351,12 +357,17 @@ async def _process_multipart_upload_request( ) +def _declared_query_params(query_params: Mapping[str, str], query_param_names: Sequence[str]) -> Mapping[str, str]: + return MappingProxyType({name: query_params[name] for name in query_param_names if name in query_params}) + + async def _process_request( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, route_type: str, path_params: dict[str, str], + query_param_names: Sequence[str] = (), ): """Common request processing logic.""" from litellm.proxy.proxy_server import ( @@ -376,6 +387,7 @@ async def _process_request( query_params: Final = dict(request.query_params) data: Final[dict[str, Any]] = { "query_params": query_params, + **_declared_query_params(query_params, query_param_names), **path_params, } @@ -452,7 +464,13 @@ def register_container_file_endpoints(router: APIRouter) -> None: is_multipart = endpoint_config.get("is_multipart", False) # Create handler with correct signature for path params - handler = _create_handler_for_path_params(path_params, route_type, returns_binary, is_multipart) + handler = _create_handler_for_path_params( + path_params, + route_type, + returns_binary, + is_multipart, + query_param_names=endpoint_config.get("query_params", ()), + ) # Register routes route_method = getattr(router, method) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e3088771c82..14480232a4a 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,9 +1,10 @@ import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias from fastapi import HTTPException +from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -46,6 +47,12 @@ _CONTAINER_STORED_ID_CACHE: Final = InMemoryCache(max_size_in_memory=10000, defa # different users with different scopes get disjoint cache entries. _ALLOWED_CONTAINER_IDS_CACHE: Final = InMemoryCache(max_size_in_memory=2048, default_ttl=60) +DEFAULT_CONTAINER_LIST_LIMIT: Final = 20 +OWNED_CONTAINER_LIST_PAGE_SIZE: Final = 100 +OWNED_CONTAINER_LIST_MAX_PAGES: Final = 5 + +FetchContainerListPage: TypeAlias = Callable[[str | None, int | None], Awaitable[object]] + def _allowed_container_ids_cache_key(owner_scopes: Sequence[str]) -> str: """JSON-encode the sorted scope list — using a separator like ``|`` @@ -337,27 +344,23 @@ def _get_container_list_data(response: object) -> Sequence[object] | None: return data if isinstance(data, list) else None -def _set_container_list_data(response: Any, data: list[object], removed_filtered_items: bool = False) -> object: +def _get_has_more(response: object) -> bool: if isinstance(response, dict): - response["data"] = data - if data: - response["first_id"] = _get_response_id(data[0]) - response["last_id"] = _get_response_id(data[-1]) - else: - response["first_id"] = None - response["last_id"] = None - response["has_more"] = False - if removed_filtered_items: - response["has_more"] = False - return response + return response.get("has_more") is True + return getattr(response, "has_more", None) is True - response.data = data - response.first_id = _get_response_id(data[0]) if data else None - response.last_id = _get_response_id(data[-1]) if data else None - if not data and hasattr(response, "has_more"): - response.has_more = False - if removed_filtered_items and hasattr(response, "has_more"): - response.has_more = False + +def _with_container_list_page(response: object, data: Sequence[object], has_more: bool) -> object: + page: Final = { + "data": list(data), + "first_id": _get_response_id(data[0]) if data else None, + "last_id": _get_response_id(data[-1]) if data else None, + "has_more": has_more, + } + if isinstance(response, dict): + return {**response, **page} + if isinstance(response, BaseModel): + return response.model_copy(update=page) return response @@ -366,16 +369,16 @@ async def _get_allowed_container_ids( ) -> AbstractSet[str]: owner_scopes: Final = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: - return set() + return frozenset() cache_key: Final = _allowed_container_ids_cache_key(owner_scopes) cached: Final = _ALLOWED_CONTAINER_IDS_CACHE.get_cache(cache_key) if cached is not None: - return set(cached) + return frozenset(cached) prisma_client: Final = await _get_prisma_client() if prisma_client is None: - return set() + return frozenset() table: Final = ManagedObjectRepository(prisma_client).table rows: Final[Sequence[prisma_models.LiteLLM_ManagedObjectTable]] = await table.find_many( @@ -384,34 +387,69 @@ async def _get_allowed_container_ids( "created_by": {"in": owner_scopes}, } ) - allowed_ids: Final = {row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None} - # ``InMemoryCache.get_cache`` attempts ``json.loads`` on the stored - # value; passing a set would round-trip through that path - # unnecessarily. Store as a list and rehydrate above. - _ALLOWED_CONTAINER_IDS_CACHE.set_cache(cache_key, list(allowed_ids)) + allowed_ids: Final = frozenset( + row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None + ) + _ALLOWED_CONTAINER_IDS_CACHE.set_cache(cache_key, tuple(allowed_ids)) return allowed_ids -async def filter_container_list_response( - response: object, +def _is_owned_container(item: object, allowed_container_ids: AbstractSet[str], custom_llm_provider: str) -> bool: + container_id: Final = _get_response_id(item) + if container_id is None: + return False + original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) + return _container_model_object_id(original_container_id, resolved_provider) in allowed_container_ids + + +async def _collect_owned_containers( + fetch_page: FetchContainerListPage, + after: str | None, + needed: int, + allowed_container_ids: AbstractSet[str], + custom_llm_provider: str, + pages_left: int, + collected: tuple[object, ...], +) -> tuple[object, tuple[object, ...]]: + page: Final = await fetch_page(after, OWNED_CONTAINER_LIST_PAGE_SIZE) + page_data: Final = _get_container_list_data(page) or () + owned: Final = collected + tuple( + item for item in page_data if _is_owned_container(item, allowed_container_ids, custom_llm_provider) + ) + upstream_last_id: Final = _get_response_id(page_data[-1]) if page_data else None + if len(owned) >= needed or upstream_last_id is None or pages_left <= 1 or not _get_has_more(page): + return page, owned + return await _collect_owned_containers( + fetch_page=fetch_page, + after=upstream_last_id, + needed=needed, + allowed_container_ids=allowed_container_ids, + custom_llm_provider=custom_llm_provider, + pages_left=pages_left - 1, + collected=owned, + ) + + +async def list_owned_containers( + fetch_page: FetchContainerListPage, + after: str | None, + limit: int | None, user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, ) -> object: - if is_proxy_admin(user_api_key_dict): - return response - - data: Final = _get_container_list_data(response) - if data is None: - return response - allowed_container_ids: Final = await _get_allowed_container_ids(user_api_key_dict) - filtered: Final[list[object]] = [] - for item in data: - container_id = _get_response_id(item) - if container_id is None: - continue - original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) - if _container_model_object_id(original_container_id, resolved_provider) in allowed_container_ids: - filtered.append(item) - - return _set_container_list_data(response, filtered, removed_filtered_items=len(filtered) != len(data)) + page_limit: Final = limit if limit is not None else DEFAULT_CONTAINER_LIST_LIMIT + last_page, owned = await _collect_owned_containers( + fetch_page=fetch_page, + after=after, + needed=page_limit + 1, + allowed_container_ids=allowed_container_ids, + custom_llm_provider=custom_llm_provider, + pages_left=OWNED_CONTAINER_LIST_MAX_PAGES, + collected=(), + ) + return _with_container_list_page( + last_page, + owned[:page_limit], + has_more=len(owned) > page_limit or _get_has_more(last_page), + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 30526d30dc5..7f2616c2fb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -41,7 +41,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -56,7 +56,6 @@ from litellm.proxy.guardrails.anthropic_sse import ( is_raw_sse_stream, model_response_text, ) -from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import ( BedrockChecksConfigModel, BedrockGuardrailStreamingParams, @@ -713,9 +712,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # logic becomes shared across providers. #### CALL HOOKS - proxy only #### - def _load_credentials( - self, - ): + def _load_credentials(self, bearer_token: str | None = None): try: from botocore.credentials import Credentials except ImportError: @@ -737,17 +734,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): aws_region_name=aws_region_name, ) - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_session_name=aws_session_name, - aws_profile_name=aws_profile_name, - aws_role_name=aws_role_name, - aws_web_identity_token=aws_web_identity_token, - aws_sts_endpoint=aws_sts_endpoint, - aws_external_id=aws_external_id, + credentials: Final[Credentials | None] = ( + None + if bearer_token is not None + else self.get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) ) return credentials, aws_region_name @@ -779,13 +780,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): proxy_endpoint_url = f"{proxy_endpoint_url}{request_path}" encoded_data: Final = json.dumps(data).encode("utf-8") - # first check api-key, if none, fall back to sigV4 - if api_key is not None: - aws_bearer_token: str | None = api_key - else: - aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + aws_bearer_token: Final = bedrock_bearer_token(api_key) - if aws_bearer_token: + if aws_bearer_token is not None: try: from botocore.awsrequest import AWSRequest except ImportError: @@ -916,7 +913,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials() + credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) allow_chunking: Final = not self._content_uses_contextual_grounding(content) completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator @@ -958,7 +955,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, content: Sequence[BedrockContentItem], base_request_data: Mapping[str, object], - credentials: "Credentials", + credentials: "Credentials | None", aws_region_name: str, api_key: str | None, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper @@ -1096,7 +1093,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, content: Sequence[BedrockContentItem], base_request_data: Mapping[str, object], - credentials: "Credentials", + credentials: "Credentials | None", aws_region_name: str, api_key: str | None, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper @@ -1146,7 +1143,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, content: Sequence[BedrockContentItem], base_request_data: Mapping[str, object], - credentials: "Credentials", + credentials: "Credentials | None", aws_region_name: str, api_key: str | None, request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper @@ -1873,9 +1870,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does. return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials() - body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} api_key: Final[str | None] = request_data.get("api_key") if request_data else None + credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} prepared_request: Final = self._prepare_request( credentials=credentials, diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 1fc3c06e6bf..93d859066b0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -916,9 +916,10 @@ class CompresrGuardrail(CustomGuardrail): def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None: """Compressed content mirrored into the Responses `texts` channel. - The chat/Anthropic handlers round-trip ``structured_messages``; the - Responses translation cannot rebuild its input from chat messages and - instead writes back through ``texts``. This matches by value, so a + The chat/Anthropic/Responses handlers round-trip + ``structured_messages``; translations without that round-trip write + back through ``texts``, so the compressed content is mirrored there + too. This matches by value, so a replacement is applied only when it is unambiguous: one compression per text, and every occurrence in ``texts`` accounted for by a compressed target. Anything else is left uncompressed rather than risk a wrong or diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index fc881a60f43..9d993384461 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -50,6 +50,9 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel BYPASS_HEADER: Final = "x-headroom-bypass" +_STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( + (CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses) +) HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 @@ -725,6 +728,10 @@ class HeadroomGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs + if request_data.get("background"): + verbose_proxy_logger.debug("Headroom: background request; skipping compression") + return inputs + structured_messages: Final = inputs.get("structured_messages") if not _is_object_list(structured_messages) or not structured_messages: return inputs @@ -826,9 +833,9 @@ class HeadroomGuardrail(CustomGuardrail): ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) effective: Final = base_result if base_result is not None else kwargs - if call_type not in (CallTypes.completion, CallTypes.acompletion): + if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: return base_result - if not effective.get("stream"): + if not effective.get("stream") or effective.get("background"): return base_result if not has_headroom_retrieve_tool(effective.get("tools")): return base_result diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c6b8df1b493..9029e926b35 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -121,9 +121,6 @@ def _a2a_jsonrpc_error_chunk(exc: HTTPException, request_id: str | None) -> Mapp } -endpoint_guardrail_translation_mappings = None - - def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: """Populate data['litellm_metadata'] from user_api_key_dict if absent.""" if "litellm_metadata" not in data: @@ -164,7 +161,6 @@ class UnifiedLLMGuardrails(CustomLogger): Use this if you want to MODIFY the input """ - global endpoint_guardrail_translation_mappings from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) @@ -186,18 +182,15 @@ class UnifiedLLMGuardrails(CustomLogger): ) return data - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + mappings: Final = load_guardrail_translation_mappings() try: - if CallTypes(call_type) not in endpoint_guardrail_translation_mappings: + if CallTypes(call_type) not in mappings: return data except ValueError: return data # handle unmapped call types - endpoint_translation: Final = _as_endpoint_translation( - endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - ) + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) _ensure_litellm_metadata(data, user_api_key_dict) @@ -222,8 +215,6 @@ class UnifiedLLMGuardrails(CustomLogger): This can NOT modify the input, only used to reject or accept a call before going to LLM API """ - global endpoint_guardrail_translation_mappings - verbose_proxy_logger.debug("Running UnifiedLLMGuardrails moderation hook") guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None) @@ -241,14 +232,11 @@ class UnifiedLLMGuardrails(CustomLogger): ) return data - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() - if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: + mappings: Final = load_guardrail_translation_mappings() + if call_type is not None and CallTypes(call_type) not in mappings: return data - endpoint_translation: Final = _as_endpoint_translation( - endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - ) + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) _ensure_litellm_metadata(data, user_api_key_dict) @@ -271,7 +259,6 @@ class UnifiedLLMGuardrails(CustomLogger): Uses Enkrypt AI guardrails to check the response for policy violations, PII, and injection attacks """ - global endpoint_guardrail_translation_mappings # Local import avoids a module-level cyclic import with # litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException @@ -319,10 +306,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + mappings: Final = load_guardrail_translation_mappings() - if CallTypes(call_type) not in endpoint_guardrail_translation_mappings: + if CallTypes(call_type) not in mappings: verbose_proxy_logger.warning( "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; " "skipping post-call scanning.", @@ -332,9 +318,7 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = _as_endpoint_translation( - endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - ) + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) try: response = await endpoint_translation.process_output_response( @@ -906,8 +890,6 @@ class UnifiedLLMGuardrails(CustomLogger): sampling_rate=1 means every chunk, sampling_rate=5 means every 5th chunk, etc. """ - global endpoint_guardrail_translation_mappings - # Local import avoids a module-level cyclic import with # litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException @@ -978,9 +960,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield item return - # Initialize translation mappings if needed - if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + mappings: Final = load_guardrail_translation_mappings() # Streaming text transformation (incremental_diff) diverges enough from the # block_only path that it runs as its own iterator. It requires a route we @@ -989,7 +969,7 @@ class UnifiedLLMGuardrails(CustomLogger): if streaming_transform_mode == "incremental_diff": transform_call_type: Final = self._resolve_transform_call_type( user_api_key_dict=user_api_key_dict, - mappings=endpoint_guardrail_translation_mappings, + mappings=mappings, ) if transform_call_type is not None: async for transformed_item in self._run_incremental_transform_stream( @@ -1000,7 +980,7 @@ class UnifiedLLMGuardrails(CustomLogger): call_type=transform_call_type, sampling_rate=sampling_rate, end_of_stream_only=end_of_stream_only, - mappings=endpoint_guardrail_translation_mappings, + mappings=mappings, ): yield transformed_item return @@ -1037,7 +1017,7 @@ class UnifiedLLMGuardrails(CustomLogger): call_type = _infer_call_type(call_type=None, completion_response=item) # If call type not supported, just pass through all chunks - if call_type is None or CallTypes(call_type) not in endpoint_guardrail_translation_mappings: + if call_type is None or CallTypes(call_type) not in mappings: yield item async for remaining_item in response: yield remaining_item @@ -1049,7 +1029,7 @@ class UnifiedLLMGuardrails(CustomLogger): # moderation runs below. if end_of_stream_only: if not buffer_until_moderated: - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation = mappings[CallTypes(call_type)]() stream_has_ended = hasattr( endpoint_translation, "_check_streaming_has_ended" ) and endpoint_translation._check_streaming_has_ended(responses_so_far) @@ -1063,7 +1043,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation = mappings[CallTypes(call_type)]() scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) if _is_redundant_scan(scan_key, last_scan_key): verbose_proxy_logger.debug( @@ -1143,14 +1123,14 @@ class UnifiedLLMGuardrails(CustomLogger): yield item # Stream has ended - do final processing with all collected chunks - if call_type is not None and CallTypes(call_type) in endpoint_guardrail_translation_mappings: + if call_type is not None and CallTypes(call_type) in mappings: verbose_proxy_logger.debug( "Processing final streaming response with all %s chunks for guardrail %s", len(responses_so_far), guardrail_to_apply.guardrail_name, ) - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation = mappings[CallTypes(call_type)]() # When buffering, snapshot the original chunks before moderation. # A shallow copy suffices: end-of-stream diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 003576ba555..0dc50cd6196 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -654,23 +654,16 @@ class InMemoryGuardrailHandler: source: Literal["db", "config"] = "db", ) -> None: """ - Update a guardrail in memory - - - updates the guardrail in memory - - updates the guardrail params in litellm.callback_manager + Update a guardrail in memory: a changed name or litellm_params rebuilds the + live callback from the new row (fail-closed: an invalid row keeps the + previous instance and raises), anything else only refreshes the stored row """ - self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail - self._sources[guardrail_id] = source - - tracked_callbacks: Final = self._tracked_callbacks(guardrail_id) - if not tracked_callbacks: + updated_guardrail: Final = cast(Guardrail, {**guardrail, "guardrail_id": guardrail_id}) + if self._has_guardrail_params_changed(guardrail_id, updated_guardrail): + self.reinitialize_guardrail(guardrail=updated_guardrail, source=source) return - updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) - tracked_callbacks[0].update_in_memory_litellm_params(litellm_params=updated_litellm_params) - for sibling_callback in tracked_callbacks[1:]: - sibling_stage = sibling_callback.event_hook - sibling_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) - sibling_callback.event_hook = sibling_stage + self.IN_MEMORY_GUARDRAILS[guardrail_id] = updated_guardrail + self._sources[guardrail_id] = source def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 21d12c8f720..c4c15c40d1e 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -5,10 +5,11 @@ This hook uses the DBSpendUpdateWriter to batch-write response IDs to the databa instead of writing immediately on each request. """ -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Callable, Mapping from typing import TYPE_CHECKING, Any, Final, cast from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger @@ -32,6 +33,44 @@ _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" _RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) +_RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _response_payload(response_obj: object) -> Mapping[str, object] | None: + try: + return _RESPONSE_PAYLOAD_ADAPTER.validate_python(response_obj) + except ValidationError: + return None + + +def _rewrite_advertised_id( + event: BaseLiteLLMOpenAIResponseObject, + rewrite: Callable[[str], str], +) -> BaseLiteLLMOpenAIResponseObject: + event_id: Final = getattr(event, "id", None) + if isinstance(event_id, str) and event_id.startswith("resp_"): + setattr(event, "id", rewrite(event_id)) + return event + + nested: Final = getattr(event, "response", None) + if isinstance(nested, ResponsesAPIResponse): + setattr(nested, "id", rewrite(nested.id)) + setattr(event, "response", nested) + return event + + payload: Final = _response_payload(nested) + if payload is None: + return event + + payload_id: Final = payload.get("id") + if not isinstance(payload_id, str): + return event + + rewritten: Final = {**payload, "id": rewrite(payload_id)} # mutable-ok: pydantic cannot serialize a frozen map + setattr(event, "response", rewritten) + return event + + def _is_responses_api_create_route(request_route: str | None) -> bool: if request_route is None: return False @@ -196,10 +235,6 @@ class ResponsesIDSecurity(CustomLogger): user_api_key_dict: "UserAPIKeyAuth", request_cache: dict[str, str] | None = None, ) -> BaseLiteLLMOpenAIResponseObject: - # encrypt the response id using the symmetric key - # encrypt the response id, and encode the user id and response id in base64 - - # Check if signing key is available signing_key: Final = self._get_signing_key() if signing_key is None: verbose_proxy_logger.debug( @@ -210,43 +245,22 @@ class ResponsesIDSecurity(CustomLogger): ) return response - response_id: Final = getattr(response, "id", None) - response_obj: Final = getattr(response, "response", None) + def encrypt(original_id: str) -> str: + cached: Final = request_cache.get(original_id) if request_cache is not None else None + if cached is not None: + return cached - if response_id and isinstance(response_id, str) and response_id.startswith("resp_"): - # Check request-scoped cache first (for streaming consistency) - if request_cache is not None and response_id in request_cache: - setattr(response, "id", request_cache[response_id]) - else: - encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( - response_id, - user_api_key_dict.user_id or "", - user_api_key_dict.team_id or "", - ) + managed_id: Final = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + original_id, + user_api_key_dict.user_id or "", + user_api_key_dict.team_id or "", + ) + encrypted_id: Final = f"resp_{encrypt_value_helper(value=managed_id)}" + if request_cache is not None: + request_cache[original_id] = encrypted_id + return encrypted_id - encoded_user_id_and_response_id = encrypt_value_helper(value=encrypted_response_id) - encrypted_id = f"resp_{encoded_user_id_and_response_id}" - if request_cache is not None: - request_cache[response_id] = encrypted_id - setattr(response, "id", encrypted_id) - - elif response_obj and isinstance(response_obj, ResponsesAPIResponse): - # Check request-scoped cache first (for streaming consistency) - if request_cache is not None and response_obj.id in request_cache: - setattr(response_obj, "id", request_cache[response_obj.id]) - else: - encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( - response_obj.id, - user_api_key_dict.user_id or "", - user_api_key_dict.team_id or "", - ) - encoded_user_id_and_response_id = encrypt_value_helper(value=encrypted_response_id) - encrypted_id = f"resp_{encoded_user_id_and_response_id}" - if request_cache is not None: - request_cache[response_obj.id] = encrypted_id - setattr(response_obj, "id", encrypted_id) - setattr(response, "response", response_obj) - return response + return _rewrite_advertised_id(response, encrypt) async def async_post_call_success_hook( self, diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 83caa92ede5..06f99e4ae9c 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -2,7 +2,7 @@ import asyncio import io import traceback from collections.abc import Sequence -from typing import Final +from typing import Final, get_type_hints import orjson from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status @@ -16,11 +16,18 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.http_parsing_utils import ( + coerce_numeric_form_fields, + numeric_form_fields, +) from litellm.proxy.route_llm_request import route_request +from litellm.types.images.main import ImageEditRequestParams from litellm.types.llms.openai import ChatCompletionUserMessage router: Final = APIRouter() +IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams)) + async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO: """ @@ -279,7 +286,12 @@ async def image_edit_api( ######################################################### # Read request body and convert UploadFiles to BytesIO ######################################################### - data: Final = await _read_request_body(request=request) + data: Final = dict( + coerce_numeric_form_fields( + parsed_body=await _read_request_body(request=request), + numeric_fields=IMAGE_EDIT_NUMERIC_FORM_FIELDS, + ) + ) image_files: Final = await batch_to_bytesio(image) mask_files: Final = await batch_to_bytesio(mask) if image_files: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 93423ca5a1a..d066f1e9138 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,6 +17,7 @@ import json import traceback from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Any, Final, Literal, Protocol, cast, overload import fastapi @@ -77,6 +78,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, UserListResponse, + UserSearchWhere, UserUpdateResult, ) from litellm.types.proxy.management_endpoints.scim_v2 import ( @@ -2080,6 +2082,22 @@ async def _authorize_user_list_request( return ",".join(allowed_org_ids) +_NO_SEARCH_WHERE: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _user_search_where(search: str | None) -> Mapping[str, object]: + """Prisma predicate for `/user/list?search=`: user_id or user_email contains it, case-insensitive.""" + if not search: + return _NO_SEARCH_WHERE + search_where: Final[UserSearchWhere] = { + "OR": ( + {"user_id": {"contains": search, "mode": "insensitive"}}, + {"user_email": {"contains": search, "mode": "insensitive"}}, + ) + } + return search_where + + @router.get( "/user/list", tags=["Internal User management"], @@ -2091,6 +2109,10 @@ async def get_users( user_ids: str | None = fastapi.Query(default=None, description="Get list of users by user_ids"), sso_user_ids: str | None = fastapi.Query(default=None, description="Get list of users by sso_user_id"), user_email: str | None = fastapi.Query(default=None, description="Filter users by partial email match"), + search: str | None = fastapi.Query( + default=None, + description="Combined search: matches users whose 'user_id' or 'user_email' contains the value (case-insensitive).", + ), team: str | None = fastapi.Query(default=None, description="Filter users by team id"), page: int = fastapi.Query(default=1, ge=1, description="Page number"), page_size: int = fastapi.Query(default=25, ge=1, le=100, description="Number of items per page"), @@ -2121,6 +2143,8 @@ async def get_users( Get list of users by sso_ids. Comma separated list of sso_ids. user_email: Optional[str] Filter users by partial email match + search: Optional[str] + Combined search: matches users whose user_id or user_email contains the value (case-insensitive) team: Optional[str] Filter users by team id. Will match if user has this team in their teams array. page: int @@ -2197,7 +2221,11 @@ async def get_users( where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_id_list}}} ## Filter any none fastapi.Query params - e.g. where_conditions: {'user_email': {'contains': Query(None), 'mode': 'insensitive'}, 'teams': {'has': Query(None)}} - where_conditions = {k: v for k, v in where_conditions.items() if v is not None} + where: Final[Mapping[str, object]] = { + key: value + for key, value in (*where_conditions.items(), *_user_search_where(search).items()) + if value is not None + } # Build order_by conditions @@ -2206,14 +2234,14 @@ async def get_users( ) users: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await UserRepository(prisma_client).table.find_many( - where=where_conditions, + where=where, skip=skip, take=page_size, order=(order_by if order_by else {"created_at": "desc"}), # Default to created_at desc if no sort specified ) # Get total count of user rows - total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where) # Get key count for each user user_key_counts: Final = await get_user_key_counts(prisma_client, [user.user_id for user in users]) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 613e726f89d..82ee33cbc39 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,10 +13,11 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager from json import JSONDecodeError from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator @@ -49,6 +50,7 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UserAPIKeyAuth, ) +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -96,6 +98,9 @@ from litellm.router_strategy.complexity_router import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -153,6 +158,8 @@ class _ProxyModelTable(Protocol): def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + def create(self, *, data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + def update( self, *, where: Mapping[str, object], data: Mapping[str, object] ) -> Awaitable[_ProxyModelRow | None]: ... @@ -166,6 +173,9 @@ class _TxModelTables(Protocol): litellm_proxymodeltable: _ProxyModelTable +_RowT = TypeVar("_RowT") + + class _ExistingModelRow(Protocol): @property def litellm_params(self) -> Mapping[str, object]: ... @@ -269,6 +279,66 @@ def _raise_on_strategy_router_write_violation( ) +HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301 +_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_HEURISTIC_V2_DB_ROWS_SQL: Final = """ +SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable" +WHERE model_id <> $1 + AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END) + -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2' +""" + + +def _effective_complexity_router_config( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> object: + """The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one.""" + incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config + if incoming is not None or existing_params is None: + return incoming + return existing_params.complexity_router_config + + +@asynccontextmanager +async def _heuristic_v2_slot( + prisma_client: PrismaClient, *, effective_config: object, model_id: str | None +) -> AsyncGenerator[_ProxyModelTable, None]: + """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled. + + A write that leaves the row on classifier_type heuristic_v2 under a limited license runs + inside one transaction that takes an advisory lock in its own statement before counting + (a statement's snapshot predates anything it locks), so pods cannot both pass the count: + the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged + against the license limit and the write is refused with a 403 before it happens. The row + being edited keeps its own slot through ``model_id``. Every other write, and every write on + an unlimited license, goes through the repository table with no lock. Only the row write + itself may run inside: anything that needs a second connection (the team model bookkeeping) + must wait until the transaction has committed and the lock is released. The transaction + writes bypass the repository's publish-on-write, so the config change is published once + after commit, the way delete_team_models does. + """ + from litellm.proxy.proxy_server import _license_check, llm_router + + limit: Final = _license_check.heuristic_v2_router_limit() + if limit is None or not uses_heuristic_v2_classifier(effective_config): + yield _proxy_model_table(prisma_client) + return + async with prisma_client.db.tx() as tx_ctx: + tables: Final[_TxModelTables] = tx_ctx + await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY) + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "") + db_held: Final = rows[0].get("held") if rows else 0 + config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows) + violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit) + if violation is not None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}" + ) + yield tables.litellm_proxymodeltable + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + + ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" _REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm") @@ -720,22 +790,29 @@ async def patch_model( ) requested_model_name: Final = patch_data.model_name + stored_model_name: str | None = None + + async def write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRow | None: + nonlocal stored_model_name + stored_model_name = update_data.get("model_name") + update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + update_data["updated_at"] = cast(str, get_utc_datetime()) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + patch_data.litellm_params, db_model.litellm_params + ), + model_id=model_id, + ) as table: + return await table.update(where={"model_id": model_id}, data=update_data) + # Handle team model updates with proper alias management - update_data: Final = await _update_team_model_in_db( + updated_model: Final = await _update_team_model_in_db( db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, - ) - - # Add metadata about update - update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name - update_data["updated_at"] = cast(str, get_utc_datetime()) - - # Perform partial update - updated_model: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": model_id}, - data=update_data, + write_row=write_row, ) if updated_model is None: @@ -746,7 +823,6 @@ async def patch_model( param=None, ) - stored_model_name: Final = update_data.get("model_name") if ( stored_model_name is not None and stored_model_name == requested_model_name @@ -980,7 +1056,8 @@ async def _add_model_to_db( prisma_client: PrismaClient, new_encryption_key: str | None = None, should_create_model_in_db: bool = True, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": # encrypt litellm params # _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model @@ -998,18 +1075,20 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id _create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above - if should_create_model_in_db: - model_response = await ModelRepository(prisma_client).table.create(data=_create_data) - else: - model_response = LiteLLM_ProxyModelTable(**_data) - return model_response + if not should_create_model_in_db: + return LiteLLM_ProxyModelTable(**_data) + if slot is None: + return await _proxy_model_table(prisma_client).create(data=_create_data) + async with slot as table: + return await table.create(data=_create_data) async def _add_team_model_to_db( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": """ If 'team_id' is provided, @@ -1040,6 +1119,7 @@ async def _add_team_model_to_db( model_params=model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, + slot=slot, ) if original_model_name: @@ -1060,7 +1140,8 @@ async def _update_team_model_in_db( patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> PrismaCompatibleUpdateDBModel: + write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], +) -> _RowT: """ Handle team model updates with proper alias management. @@ -1068,6 +1149,9 @@ async def _update_team_model_in_db( - Creates unique internal model_name and team alias - Adds model to team object - Preserves team_public_model_name for external reference + + The row is written through ``write_row`` before the team's model list is touched, so a + refused or failed write leaves the team as it was (the create path orders itself the same way). """ # Validate team_id if present in patch_data from litellm.proxy.proxy_server import premium_user @@ -1079,9 +1163,7 @@ async def _update_team_model_in_db( premium_user=premium_user, ) - # Validated before any write, beside the premium check the create path already runs - # here. The team ACL is updated below and autocommits, so a validator that raises - # further down would leave the team mutated and the deployment row never written. + # Validated before the row write, beside the premium check the create path already runs here. # # The merged view is what gets stored, so that is what has to satisfy the invariants. # Validating the patch alone rejected a partial edit of an already valid deployment: @@ -1101,7 +1183,7 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return update_db_model(db_model=db_model, updated_patch=patch_data) + return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -1120,11 +1202,14 @@ async def _update_team_model_in_db( db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None is_new_team_assignment: Final = db_team_id != patch_team_id + # Team rows keep their internal UUID-based model_name; the public name lives in model_info + patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None + row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + if is_new_team_assignment: await _setup_new_team_model_assignment( team_id=patch_team_id, public_model_name=public_model_name, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, ) else: @@ -1132,12 +1217,11 @@ async def _update_team_model_in_db( team_id=patch_team_id, public_model_name=public_model_name, db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) - return update_db_model(db_model=db_model, updated_patch=patch_data) + return row def _get_public_model_name( @@ -1189,13 +1273,9 @@ def _get_public_model_name( async def _setup_new_team_model_assignment( team_id: str, public_model_name: str, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Set up a new team model with unique name and team membership.""" - unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}" - patch_data.model_name = unique_model_name - + """Register a newly team-assigned model's public name on the team.""" await team_model_add( data=TeamModelAddRequest( team_id=team_id, @@ -1385,7 +1465,6 @@ async def _update_existing_team_model_assignment( team_id: str, public_model_name: str, db_model: Deployment, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient | None, ) -> None: @@ -1409,9 +1488,6 @@ async def _update_existing_team_model_assignment( old_public_name: Final = db_model.model_info.team_public_model_name if db_model.model_info else None if old_public_name and public_model_name != old_public_name: - # Clear user-supplied public name from patch before any early return so the - # caller does not overwrite the internal UUID-based model_name in the DB. - patch_data.model_name = None if prisma_client is None: verbose_proxy_logger.warning( "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries" @@ -1459,10 +1535,6 @@ async def _update_existing_team_model_assignment( # else: old_public_name == public_model_name (no rename needed) # No team_model_add/delete calls required; public name is already registered - # Always clear patch_data.model_name to prevent caller from overwriting - # the internal UUID-based model_name in the DB with the user-supplied public name - patch_data.model_name = None - class ModelManagementAuthChecks: """ @@ -1878,18 +1950,19 @@ async def add_new_model( reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None) try: _original_litellm_model_name: Final = model_params.model_name - if model_params.model_info.team_id is None: - model_response = await _add_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - else: - model_response = await _add_team_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) + add_model: Final = ( + _add_model_to_db if model_params.model_info.team_id is None else _add_team_model_to_db + ) + model_response = await add_model( + model_params=priced_model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + slot=_heuristic_v2_slot( + prisma_client, + effective_config=priced_model_params.litellm_params.complexity_router_config, + model_id=priced_model_params.model_info.id, + ), + ) reload_outcome = await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) @@ -1903,6 +1976,8 @@ async def add_new_model( passed_model_info=priced_model_params.model_info, ) except Exception as e: + if isinstance(e, HTTPException): + raise verbose_proxy_logger.exception("Exception in add_new_model: %s", e) else: @@ -2070,10 +2145,17 @@ async def update_model( "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, **({} if renamed_to is None else {"model_name": renamed_to}), } - model_response: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": _model_id}, - data=_data, - ) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + model_params.litellm_params, deployment.litellm_params + ), + model_id=_model_id, + ) as table: + model_response: Final = await table.update( + where={"model_id": _model_id}, + data=_data, + ) if renamed_to is not None: await sync_access_groups_for_renamed_model( prisma_client=prisma_client, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 714cf252e69..90d7539b38d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -10,6 +10,7 @@ All /team management endpoints """ import asyncio +import copy import json import math import traceback @@ -2189,6 +2190,26 @@ async def update_team( if field in updated_kv } + _writes_metadata_backed_field: Final = any( + field in updated_kv + for field in ( + *LiteLLM_ManagementEndpoint_MetadataFields, + *LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) + ) + if isinstance(existing_team_row.metadata, dict): + if "metadata" not in updated_kv and (_team_member_fields_in_request or _writes_metadata_backed_field): + updated_kv["metadata"] = copy.deepcopy(existing_team_row.metadata) + elif isinstance(updated_kv.get("metadata"), dict): + updated_kv["metadata"] = { + **updated_kv["metadata"], + **{ + key: existing_team_row.metadata[key] + for key in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS + if key in existing_team_row.metadata + }, + } + if _team_member_fields_in_request and TeamMemberBudgetHandler.should_create_budget( team_member_budget=data.team_member_budget, team_member_rpm_limit=data.team_member_rpm_limit, @@ -4318,6 +4339,8 @@ async def _resolve_team_access_group_resources( access_group_id=group.access_group_id, access_group_name=group.access_group_name, models=tuple(group.access_model_names or ()), + mcp_server_ids=tuple(group.access_mcp_server_ids or ()), + agent_ids=tuple(group.access_agent_ids or ()), ) for group in resolved_groups ), diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b48b8d81494..29f216fd450 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1730,6 +1730,16 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: return headers +def _is_vertex_anthropic_count_tokens_route(endpoint: str) -> bool: + return endpoint.rsplit("/", 1)[-1].split(":", 1)[0] == "count-tokens" + + +def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]: + if not _is_vertex_anthropic_count_tokens_route(endpoint): + return headers + return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"}) + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: @@ -2128,7 +2138,7 @@ async def _base_vertex_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=target, - custom_headers=headers, + custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a36a365f39a..0acc7b1b584 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -107,6 +107,7 @@ class AnthropicPassthroughLoggingHandler: start_time=start_time, end_time=end_time, logging_obj=logging_obj, + response_id=optional_str(response_body.get("id")), ) return { @@ -148,8 +149,9 @@ class AnthropicPassthroughLoggingHandler: return model @staticmethod - def _extract_model_from_anthropic_chunks( + def _extract_message_start_field( all_chunks: Sequence[str | bytes], + field: str, ) -> str | None: for raw in all_chunks: text = raw.decode("utf-8") if isinstance(raw, bytes) else raw @@ -163,11 +165,23 @@ class AnthropicPassthroughLoggingHandler: if not isinstance(data, dict): continue if data.get("type") == "message_start": - model = (data.get("message") or {}).get("model") - if model: - return model + value = (data.get("message") or {}).get(field) + if isinstance(value, str) and value: + return value return None + @staticmethod + def _extract_model_from_anthropic_chunks( + all_chunks: Sequence[str | bytes], + ) -> str | None: + return AnthropicPassthroughLoggingHandler._extract_message_start_field(all_chunks, "model") + + @staticmethod + def _extract_response_id_from_anthropic_chunks( + all_chunks: Sequence[str | bytes], + ) -> str | None: + return AnthropicPassthroughLoggingHandler._extract_message_start_field(all_chunks, "id") + @staticmethod def _stream_was_interrupted( all_chunks: Sequence[str | bytes], @@ -251,6 +265,7 @@ class AnthropicPassthroughLoggingHandler: start_time: datetime, end_time: datetime, logging_obj: LiteLLMLoggingObj, + response_id: str | None = None, ): """ Create the standard logging object for Anthropic passthrough @@ -312,8 +327,7 @@ class AnthropicPassthroughLoggingHandler: json.dumps(kwargs, indent=4, default=str), ) - # set litellm_call_id to logging response object - litellm_model_response.id = logging_obj.litellm_call_id + litellm_model_response.id = response_id or logging_obj.litellm_call_id litellm_model_response.model = model logging_obj.model_call_details["model"] = model if not logging_obj.model_call_details.get("custom_llm_provider"): @@ -413,6 +427,7 @@ class AnthropicPassthroughLoggingHandler: start_time=start_time, end_time=end_time, logging_obj=litellm_logging_obj, + response_id=AnthropicPassthroughLoggingHandler._extract_response_id_from_anthropic_chunks(all_chunks), ) return { diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index b6cbd2d7889..06957fd1c0f 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -29,6 +29,12 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.path_utils import safe_filename +from litellm.proxy.prompts.prompt_registry import ( + DEFAULT_PROMPT_ENVIRONMENT, + get_base_prompt_id, + get_version_number, + prompt_environment_or_default, +) from litellm.repositories.table_repositories import PromptRepository from litellm.types.prompts.init_prompts import ( ListPromptsResponse, @@ -102,165 +108,20 @@ def _prompt_table(prisma_client: "PrismaClient") -> _PromptTableActions: return PromptRepository(prisma_client).table -def get_base_prompt_id(prompt_id: str) -> str: - """ - Extract the base prompt ID by stripping the version suffix if present. - - Args: - prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v1" or "jack_success_v1") - - Returns: - Base prompt ID without version suffix (e.g., "jack_success") - - Examples: - >>> get_base_prompt_id("jack_success.v1") - "jack_success" - >>> get_base_prompt_id("jack_success_v1") - "jack_success" - >>> get_base_prompt_id("jack_success") - "jack_success" - """ - # Try dot separator first (.v) - if ".v" in prompt_id: - return prompt_id.split(".v")[0] - # Try underscore separator (_v) - if "_v" in prompt_id: - return prompt_id.split("_v")[0] - return prompt_id - - -def get_version_number(prompt_id: str) -> int: - """ - Extract the version number from a versioned prompt ID. - - Args: - prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v2" or "jack_success_v2") - - Returns: - Version number (defaults to 1 if no version suffix or invalid format) - - Examples: - >>> get_version_number("jack_success.v2") - 2 - >>> get_version_number("jack_success_v2") - 2 - >>> get_version_number("jack_success") - 1 - """ - # Try dot separator first (.v) - if ".v" in prompt_id: - version_str = prompt_id.split(".v")[1] - try: - return int(version_str) - except ValueError: - pass - - # Try underscore separator (_v) - if "_v" in prompt_id: - version_str = prompt_id.split("_v")[1] - try: - return int(version_str) - except ValueError: - pass - - return 1 - - -def construct_versioned_prompt_id(prompt_id: str, version: int | None = None) -> str: - """ - Construct a versioned prompt ID from a base prompt_id and version number. - - Args: - prompt_id: Base prompt ID (e.g., "jack_success") - version: Version number (if None, returns the base prompt_id unchanged) - - Returns: - Versioned prompt ID (e.g., "jack_success.v4") - - Examples: - >>> construct_versioned_prompt_id("jack_success", 4) - "jack_success.v4" - >>> construct_versioned_prompt_id("jack_success", None) - "jack_success" - >>> construct_versioned_prompt_id("jack_success.v2", 4) - "jack_success.v4" - """ - if version is None: - return prompt_id - - # Strip any existing version suffix first - base_id: Final = get_base_prompt_id(prompt_id) - return f"{base_id}.v{version}" - - -def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Mapping[str, object]) -> str: - """ - Find the latest version of a prompt from available prompt IDs. - - Args: - prompt_id: Base prompt ID or versioned prompt ID (e.g., "jack_success" or "jack_success.v2") - all_prompt_ids: Dictionary of all available prompt IDs (keys are prompt IDs) - - Returns: - The prompt ID with the highest version number, or the original prompt_id if no versions exist - - Examples: - >>> all_ids = {"jack.v1": {}, "jack.v2": {}, "jack.v3": {}} - >>> get_latest_version_prompt_id("jack", all_ids) - "jack.v3" - >>> get_latest_version_prompt_id("jack.v1", all_ids) - "jack.v3" - >>> all_ids = {"simple": {}} - >>> get_latest_version_prompt_id("simple", all_ids) - "simple" - """ - base_id: Final = get_base_prompt_id(prompt_id=prompt_id) - - # Find all versions of this prompt - matching_versions: Final = [] - for stored_prompt_id in all_prompt_ids: - if get_base_prompt_id(prompt_id=stored_prompt_id) == base_id: - version_num = get_version_number(prompt_id=stored_prompt_id) - matching_versions.append((version_num, stored_prompt_id)) - - # Use the highest version number - if matching_versions: - matching_versions.sort(reverse=True) - return matching_versions[0][1] - else: - # No versioned prompts found, use the base ID as-is - return prompt_id - - def get_latest_prompt_versions(prompts: list[PromptSpec]) -> list[PromptSpec]: """ - Filter a list of prompts to return only the latest version of each unique prompt. - - Args: - prompts: List of PromptSpec objects - - Returns: - List of PromptSpec objects with only the latest version of each prompt + Filter prompts down to the latest version per (base prompt id, environment). """ - latest_prompts: Final[dict[str, PromptSpec]] = {} - - for prompt in prompts: - base_id = get_base_prompt_id(prompt_id=prompt.prompt_id) - version = get_version_number(prompt_id=prompt.prompt_id) - - # Keep the prompt with the highest version number - if base_id not in latest_prompts: - latest_prompts[base_id] = prompt - else: - existing_version = get_version_number(prompt_id=latest_prompts[base_id].prompt_id) - if version > existing_version: - latest_prompts[base_id] = prompt - + sorted_prompts: Final = sorted(prompts, key=lambda prompt: get_version_number(prompt_id=prompt.prompt_id)) + latest_prompts: Final = { + (get_base_prompt_id(prompt_id=prompt.prompt_id), prompt_environment_or_default(prompt.environment)): prompt + for prompt in sorted_prompts + } return list(latest_prompts.values()) async def get_next_version_for_prompt( - prisma_client: "PrismaClient", prompt_id: str, environment: str = "development" + prisma_client: "PrismaClient", prompt_id: str, environment: str = DEFAULT_PROMPT_ENVIRONMENT ) -> int: """ Get the next version number for a prompt in a specific environment. @@ -403,11 +264,14 @@ async def list_prompts( if key_metadata is not None: prompts: Final = cast(list[str] | None, key_metadata.get("prompts", None)) if prompts is not None: - all_prompts = [ - IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[prompt_id] - for prompt_id in prompts - if prompt_id in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS + allowed_prompt_ids: Final = frozenset(prompts) + allowed_prompts: Final = [ + spec + for spec in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.values() + if spec.prompt_id in allowed_prompt_ids + or get_base_prompt_id(prompt_id=spec.prompt_id) in allowed_prompt_ids ] + all_prompts = get_latest_prompt_versions(prompts=allowed_prompts) if environment: all_prompts = [p for p in all_prompts if p.environment == environment] prompt_list: Final = [] @@ -576,7 +440,7 @@ def _get_prompt_template(prompt_spec: PromptSpec, base_prompt_id: str) -> Prompt metadata=parsed.get("metadata"), ) else: - prompt_callback: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(prompt_spec.prompt_id) + prompt_callback: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=prompt_spec) if prompt_callback is not None: integration_name: Final = prompt_callback.integration_name if integration_name == "dotprompt": @@ -690,15 +554,10 @@ async def get_prompt_info( if env_prompts: prompt_spec = create_versioned_prompt_spec(db_prompt=env_prompts[0]) - # Fallback: use in-memory registry (no environment filter) - if prompt_spec is None and environment is None: - prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if prompt_spec is None: - latest_prompt_id: Final = get_latest_version_prompt_id( - prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, - ) - prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) + if prompt_spec is None: + prompt_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec( + prompt_id, version=requested_version, environment=environment + ) if prompt_spec is None: raise HTTPException( @@ -785,7 +644,7 @@ async def create_prompt( environment: Final = ( request.prompt_info.environment if request.prompt_info and request.prompt_info.environment - else "development" + else DEFAULT_PROMPT_ENVIRONMENT ) # Get next version number @@ -885,7 +744,7 @@ async def update_prompt( environment: Final = ( request.prompt_info.environment if request.prompt_info and request.prompt_info.environment - else "development" + else DEFAULT_PROMPT_ENVIRONMENT ) # Check if any version of this prompt exists (in any environment) @@ -897,9 +756,7 @@ async def update_prompt( detail=f"Prompt with ID {base_prompt_id} not found", ) - # Check if it's a config prompt - existing_in_memory: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if existing_in_memory and existing_in_memory.prompt_info.prompt_type == "config": + if IN_MEMORY_PROMPT_REGISTRY.has_config_prompt(base_prompt_id=base_prompt_id): raise HTTPException( status_code=400, detail="Cannot update config prompts.", @@ -988,40 +845,26 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - # Try to get prompt directly first - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - - # If not found, try to find the latest version - if existing_prompt is None: - latest_prompt_id: Final = get_latest_version_prompt_id( - prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, - ) - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) - # Use the resolved prompt_id for deletion - prompt_id = latest_prompt_id + base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) + existing_prompt: Final = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec(prompt_id, environment=environment) if existing_prompt is None: raise HTTPException(status_code=404, detail=f"Prompt with ID {prompt_id} not found") - if existing_prompt.prompt_info.prompt_type == "config": + if IN_MEMORY_PROMPT_REGISTRY.has_config_prompt(base_prompt_id=base_prompt_id): raise HTTPException( status_code=400, detail="Cannot delete config prompts.", ) - # Get the base prompt ID (without version suffix) for database deletion - base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) - - # Build delete filter; scope to environment if provided - delete_where: Final[dict[str, str]] = {"prompt_id": base_prompt_id} - if environment: - delete_where["environment"] = environment - - # Delete versions from the database (scoped to environment if provided) + delete_where: Final[dict[str, str]] = { + "prompt_id": base_prompt_id, + **({"environment": environment} if environment else {}), + } await _prompt_table(prisma_client).delete_many(where=delete_where) - - IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id, environment=environment or None) + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id( + base_prompt_id=base_prompt_id, environment=environment or None + ) env_msg: Final = f" from {environment}" if environment else "" return {"message": f"Prompt {base_prompt_id} deleted successfully{env_msg}"} @@ -1093,7 +936,7 @@ async def patch_prompt( try: # Resolve the target row: find the latest version in the given environment base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) - env: Final = environment or "development" + env: Final = prompt_environment_or_default(environment) requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None # Build query to find the exact row by composite unique key @@ -1117,11 +960,7 @@ async def patch_prompt( target_row: Final = db_rows[0] - # Check if prompt exists in memory - versioned_id: Final = f"{base_prompt_id}.v{target_row.version}" - existing_prompt: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(versioned_id) - - if existing_prompt and existing_prompt.prompt_info.prompt_type == "config": + if IN_MEMORY_PROMPT_REGISTRY.has_config_prompt(base_prompt_id=base_prompt_id): raise HTTPException( status_code=400, detail="Cannot update config prompts.", diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index addfb3f80d5..7803352e9e7 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -1,6 +1,6 @@ import importlib import os -from collections.abc import Callable +from collections.abc import Callable, Sequence from pathlib import Path from typing import Final @@ -14,6 +14,87 @@ from litellm.types.prompts.init_prompts import ( prompt_initializer_registry = {} +DEFAULT_PROMPT_ENVIRONMENT: Final = "development" +PROMPT_ENVIRONMENT_SERVE_PRECEDENCE: Final = ("production", "staging", "development") + + +def get_base_prompt_id(prompt_id: str) -> str: + """ + Extract the base prompt ID by stripping the version suffix if present. + + Examples: + >>> get_base_prompt_id("jack_success.v1") + "jack_success" + >>> get_base_prompt_id("jack_success_v1") + "jack_success" + >>> get_base_prompt_id("jack_success") + "jack_success" + """ + if ".v" in prompt_id: + return prompt_id.split(".v")[0] + if "_v" in prompt_id: + return prompt_id.split("_v")[0] + return prompt_id + + +def get_version_number(prompt_id: str) -> int: + """ + Extract the version number from a versioned prompt ID (defaults to 1). + + Examples: + >>> get_version_number("jack_success.v2") + 2 + >>> get_version_number("jack_success_v2") + 2 + >>> get_version_number("jack_success") + 1 + """ + if ".v" in prompt_id: + version_str = prompt_id.split(".v")[1] + try: + return int(version_str) + except ValueError: + pass + + if "_v" in prompt_id: + version_str = prompt_id.split("_v")[1] + try: + return int(version_str) + except ValueError: + pass + + return 1 + + +def prompt_environment_or_default(environment: str | None) -> str: + return environment or DEFAULT_PROMPT_ENVIRONMENT + + +def registry_key_for_prompt(prompt: PromptSpec) -> str: + return f"{prompt.prompt_id}::{prompt_environment_or_default(prompt.environment)}" + + +def parse_prompt_version(raw_version: object) -> int | None: + if isinstance(raw_version, bool): + return None + if isinstance(raw_version, int): + return raw_version + if isinstance(raw_version, str) and raw_version.isdigit(): + return int(raw_version) + return None + + +def _spec_version(prompt: PromptSpec) -> int: + return prompt.version if prompt.version is not None else get_version_number(prompt_id=prompt.prompt_id) + + +def _default_serve_environment(prompts: Sequence[PromptSpec]) -> str: + present: Final = frozenset(prompt_environment_or_default(prompt.environment) for prompt in prompts) + ladder_pick: Final = next((env for env in PROMPT_ENVIRONMENT_SERVE_PRECEDENCE if env in present), None) + if ladder_pick is not None: + return ladder_pick + return min(present) if present else DEFAULT_PROMPT_ENVIRONMENT + def get_prompt_initializer_from_integrations(): """ @@ -113,17 +194,16 @@ class InMemoryPromptRegistry: """ import litellm - prompt_id: Final = prompt.prompt_id - if prompt_id in self.IN_MEMORY_PROMPTS: - verbose_proxy_logger.debug("prompt_id already exists in IN_MEMORY_PROMPTS") - return self.IN_MEMORY_PROMPTS[prompt_id] + registry_key: Final = registry_key_for_prompt(prompt) + if registry_key in self.IN_MEMORY_PROMPTS: + verbose_proxy_logger.debug("prompt already exists in IN_MEMORY_PROMPTS") + return self.IN_MEMORY_PROMPTS[registry_key] parsed_prompt, custom_prompt_callback = self._build_prompt_callback(prompt=prompt) litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback) - # store references to the prompt in memory - self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt - self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback + self.IN_MEMORY_PROMPTS[registry_key] = parsed_prompt + self.prompt_id_to_custom_prompt[registry_key] = custom_prompt_callback return parsed_prompt @@ -166,68 +246,93 @@ class InMemoryPromptRegistry: import litellm parsed_prompt, new_callback = self._build_prompt_callback(prompt=prompt) - stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None) - self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None) + registry_key: Final = registry_key_for_prompt(parsed_prompt) + stale_callback: Final = self.prompt_id_to_custom_prompt.pop(registry_key, None) + self.IN_MEMORY_PROMPTS.pop(registry_key, None) if stale_callback is not None: litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback) litellm.logging_callback_manager.add_litellm_callback(new_callback) - self.IN_MEMORY_PROMPTS[prompt.prompt_id] = parsed_prompt - self.prompt_id_to_custom_prompt[prompt.prompt_id] = new_callback + self.IN_MEMORY_PROMPTS[registry_key] = parsed_prompt + self.prompt_id_to_custom_prompt[registry_key] = new_callback return parsed_prompt def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None: - existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id) + existing: Final = self.IN_MEMORY_PROMPTS.get(registry_key_for_prompt(prompt)) if existing is None: return self.initialize_prompt(prompt=prompt) if existing.litellm_params == prompt.litellm_params and existing.prompt_info == prompt.prompt_info: return existing return self.reload_prompt(prompt=prompt) - def get_prompt_by_id(self, prompt_id: str) -> PromptSpec | None: + def resolve_prompt_spec( + self, + prompt_id: str, + version: int | None = None, + environment: str | None = None, + ) -> PromptSpec | None: """ - Get a prompt by its ID from memory - """ - return self.IN_MEMORY_PROMPTS.get(prompt_id) + Resolve a prompt spec by base prompt id, optional version, and optional environment. - def get_prompt_callback_by_id(self, prompt_id: str) -> CustomPromptManagement | None: + With no environment, resolves within the default serve environment + (production > staging > development > alphabetical first present). + With no version, resolves to the highest version in the chosen environment. """ - Get a prompt callback by its ID from memory - """ - return self.prompt_id_to_custom_prompt.get(prompt_id) + base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) + base_matches: Final = tuple( + spec + for spec in self.IN_MEMORY_PROMPTS.values() + if get_base_prompt_id(prompt_id=spec.prompt_id) == base_prompt_id + ) + if not base_matches: + return None + resolved_environment: Final = ( + environment if environment is not None else _default_serve_environment(base_matches) + ) + env_matches: Final = tuple( + spec for spec in base_matches if prompt_environment_or_default(spec.environment) == resolved_environment + ) + if not env_matches: + return None + if version is not None: + return next((spec for spec in env_matches if _spec_version(spec) == version), None) + return max(env_matches, key=_spec_version) - def remove_prompt(self, prompt_id: str) -> None: + def get_prompt_callback_for_prompt(self, prompt: PromptSpec) -> CustomPromptManagement | None: + return self.prompt_id_to_custom_prompt.get(registry_key_for_prompt(prompt)) + + def has_config_prompt(self, base_prompt_id: str) -> bool: + return any( + spec.prompt_info.prompt_type == "config" + for spec in self.IN_MEMORY_PROMPTS.values() + if get_base_prompt_id(prompt_id=spec.prompt_id) == base_prompt_id + ) + + def remove_prompt(self, registry_key: str) -> None: import litellm - self.IN_MEMORY_PROMPTS.pop(prompt_id, None) - stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt_id, None) + self.IN_MEMORY_PROMPTS.pop(registry_key, None) + stale_callback: Final = self.prompt_id_to_custom_prompt.pop(registry_key, None) if stale_callback is not None: litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback) def delete_prompts_by_base_id(self, base_prompt_id: str, environment: str | None = None) -> list[str]: """ - Delete all prompts matching the given base prompt ID from memory, along with their - registered callbacks; scoped to one environment when given. + Delete matching prompts from memory, along with their registered callbacks, + scoped to one environment when given. - Args: - base_prompt_id: The base prompt ID (without version suffix) - environment: When set, only delete prompts deployed to this environment - - Returns: - List of prompt IDs that were deleted + Returns the registry keys that were deleted. """ - from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id - - prompts_to_delete: Final = [ - pid - for pid, prompt in self.IN_MEMORY_PROMPTS.items() - if get_base_prompt_id(prompt_id=pid) == base_prompt_id - and (environment is None or prompt.environment == environment) + keys_to_delete: Final = [ + key + for key, spec in self.IN_MEMORY_PROMPTS.items() + if get_base_prompt_id(prompt_id=spec.prompt_id) == base_prompt_id + and (environment is None or prompt_environment_or_default(spec.environment) == environment) ] - for pid in prompts_to_delete: - self.remove_prompt(prompt_id=pid) + for key in keys_to_delete: + self.remove_prompt(registry_key=key) - return prompts_to_delete + return keys_to_delete IN_MEMORY_PROMPT_REGISTRY: Final = InMemoryPromptRegistry() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..c43cc510990 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -120,6 +120,8 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, validate_complexity_router_config_placement, ) from litellm.types.utils import ( @@ -301,7 +303,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -4316,6 +4318,19 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object]) raise ValueError(f"model {model.get('model_name', '')!r}: {violation}") +def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: + """ + Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows. + + Checked here rather than left to router registration for the same reason as the two + validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so + the router's own refusal would turn the extra router into a silently missing model. + """ + violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit) + if violation is not None: + raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}") + + def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place """ Stamps `model_info.id` from the raw litellm_params before plugin resolution swaps @@ -5721,6 +5736,7 @@ class ProxyConfig: model_list: Final = config.get("model_list", None) if model_list: router_params["model_list"] = model_list + validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit()) print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" ) @@ -5810,6 +5826,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6270,6 +6287,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: @@ -7571,7 +7589,7 @@ class ProxyConfig: return create_versioned_prompt_spec(db_prompt=db_prompt) async def _init_prompts_in_db(self, prisma_client: PrismaClient): - from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY, registry_key_for_prompt from litellm.types.prompts.init_prompts import PromptSpec def parse_row(db_prompt: object) -> PromptSpec | None: @@ -7586,21 +7604,12 @@ class ProxyConfig: return None try: - prompt_ids_loaded_before_db_read: Final = frozenset(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS) + registry_keys_loaded_before_db_read: Final = frozenset(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS) prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many() parsed_specs: Final[tuple[PromptSpec, ...]] = tuple( spec for row in prompts_in_db if (spec := parse_row(row)) is not None ) - newest_spec_per_id: Final[Mapping[str, PromptSpec]] = MappingProxyType( - { - spec.prompt_id: spec - for spec in sorted( - parsed_specs, - key=lambda s: s.updated_at.timestamp() if s.updated_at else float("-inf"), - ) - } - ) - for prompt_spec in newest_spec_per_id.values(): + for prompt_spec in parsed_specs: try: IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts @@ -7612,15 +7621,16 @@ class ProxyConfig: # An unparsable row still exists in the DB, so skip the sweep rather than unload its in-memory copy every_row_parsed: Final = len(parsed_specs) == len(prompts_in_db) if every_row_parsed: - deleted_db_prompt_ids: Final = tuple( - prompt_id - for prompt_id in prompt_ids_loaded_before_db_read - if (loaded_spec := IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.get(prompt_id)) is not None + db_registry_keys: Final = frozenset(registry_key_for_prompt(spec) for spec in parsed_specs) + deleted_db_registry_keys: Final = tuple( + registry_key + for registry_key in registry_keys_loaded_before_db_read + if (loaded_spec := IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.get(registry_key)) is not None and loaded_spec.prompt_info.prompt_type == "db" - and prompt_id not in newest_spec_per_id + and registry_key not in db_registry_keys ) - for deleted_prompt_id in deleted_db_prompt_ids: - IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id=deleted_prompt_id) + for deleted_registry_key in deleted_db_registry_keys: + IN_MEMORY_PROMPT_REGISTRY.remove_prompt(registry_key=deleted_registry_key) except Exception as e: verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index e144ff965ae..c8c6c505375 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -761,6 +761,7 @@ async def rag_query( model=model, messages=messages, retrieval_config=merged_retrieval_config, + vector_store_params=store_data, rerank=rerank, stream=stream, router=llm_router, diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 295faaa980a..2a50d5170f0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -173,6 +173,9 @@ class _SessionSpendRow(TypedDict): session_cache_hit_count: ReadOnly[int] session_llm_count: ReadOnly[int] session_agent_count: ReadOnly[int] + session_total_prompt_tokens: ReadOnly[int] + session_total_completion_tokens: ReadOnly[int] + session_total_tokens: ReadOnly[int] session_models: ReadOnly[Sequence[str]] @@ -188,6 +191,9 @@ class _SessionSpendStats(NamedTuple): session_cache_hit_count: int session_llm_count: int session_agent_count: int + session_total_prompt_tokens: int + session_total_completion_tokens: int + session_total_tokens: int session_models: Sequence[str] session_models_truncated: bool @@ -4287,8 +4293,8 @@ async def _build_ui_spend_logs_response( Build the paginated response for the UI spend-logs endpoint. When ``enrich_session_counts`` is ``True`` (the default for the v1/UI - endpoint), each row is enriched with ``session_total_count`` plus spend - and call-type aggregates so the frontend knows which sessions are + endpoint), each row is enriched with ``session_total_count`` plus spend, + token and call-type aggregates so the frontend knows which sessions are expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)`` query serves every referenced session, keyed per api key so two callers reusing a session id never see each other's totals. Rows without a @@ -4356,7 +4362,10 @@ async def _build_ui_spend_logs_response( COUNT(*) FILTER ( WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} )::int AS session_llm_count, - COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count + COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count, + COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens, + COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens, + COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens FROM "LiteLLM_SpendLogs" WHERE session_id = ANY($1::text[]) AND api_key = ANY($2::text[]) @@ -4389,6 +4398,9 @@ async def _build_ui_spend_logs_response( session_cache_hit_count=int(row.get("session_cache_hit_count") or 0), session_llm_count=int(row.get("session_llm_count") or 0), session_agent_count=int(row.get("session_agent_count") or 0), + session_total_prompt_tokens=int(row.get("session_total_prompt_tokens") or 0), + session_total_completion_tokens=int(row.get("session_total_completion_tokens") or 0), + session_total_tokens=int(row.get("session_total_tokens") or 0), session_models=models[:_SESSION_MODELS_LIMIT], session_models_truncated=len(models) > _SESSION_MODELS_LIMIT, ) @@ -4418,6 +4430,9 @@ async def _build_ui_spend_logs_response( row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count row_dict["session_llm_count"] = session_stats.session_llm_count row_dict["session_agent_count"] = session_stats.session_agent_count + row_dict["session_total_prompt_tokens"] = session_stats.session_total_prompt_tokens + row_dict["session_total_completion_tokens"] = session_stats.session_total_completion_tokens + row_dict["session_total_tokens"] = session_stats.session_total_tokens row_dict["session_models"] = session_stats.session_models row_dict["session_models_truncated"] = session_stats.session_models_truncated enriched.append(row_dict) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cab2bd6d9db..1f453d3b1ba 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1478,28 +1478,27 @@ class ProxyLogging: ) -> None: """Process prompt template if applicable.""" - from litellm.proxy.prompts.prompt_endpoints import ( - construct_versioned_prompt_id, - get_latest_version_prompt_id, - ) from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.utils import get_non_default_completion_params - if prompt_version is None: - lookup_prompt_id = get_latest_version_prompt_id( - prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, - ) - else: - lookup_prompt_id = construct_versioned_prompt_id(prompt_id=prompt_id, version=prompt_version) - - custom_logger: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(lookup_prompt_id) - prompt_spec: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(lookup_prompt_id) + raw_prompt_environment: Final = data.get("prompt_environment", None) + prompt_environment: Final = raw_prompt_environment if isinstance(raw_prompt_environment, str) else None + prompt_spec: Final = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec( + prompt_id, + version=prompt_version, + environment=prompt_environment, + ) + custom_logger: Final = ( + IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=prompt_spec) + if prompt_spec is not None + else None + ) litellm_prompt_id: str | None = None if prompt_spec is not None: litellm_prompt_id = prompt_spec.litellm_params.prompt_id data.pop("prompt_id", None) + data.pop("prompt_environment", None) if custom_logger and prompt_spec is not None: is_responses_call: Final = call_type == "aresponses" @@ -1542,6 +1541,7 @@ class ProxyLogging: data.pop("prompt_variables", None) data.pop("prompt_label", None) data.pop("prompt_version", None) + data.pop("prompt_environment", None) def _process_guardrail_metadata(self, data: dict) -> None: """Process guardrails from metadata and add to applied_guardrails.""" @@ -1750,7 +1750,6 @@ class ProxyLogging: litellm_logging_obj: Final = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)) prompt_id: Final[str | None] = data.get("prompt_id", None) - prompt_version: Final[int | None] = data.get("prompt_version", None) ## PROMPT TEMPLATE CHECK ## @@ -1760,11 +1759,13 @@ class ProxyLogging: and prompt_id is not None and (call_type == "completion" or call_type == "acompletion" or call_type == "aresponses") ): + from litellm.proxy.prompts.prompt_registry import parse_prompt_version + await self._process_prompt_template( data=data, litellm_logging_obj=litellm_logging_obj, prompt_id=prompt_id, - prompt_version=prompt_version, + prompt_version=parse_prompt_version(data.get("prompt_version", None)), call_type=call_type, ) @@ -7531,6 +7532,9 @@ def create_model_info_response( max_input_tokens = configured_input if configured_output is not None: max_output_tokens = configured_output + configured_mode: Final = llm_router.get_configured_mode(model_id) + if isinstance(configured_mode, str): + base["mode"] = configured_mode if max_input_tokens is not None: base["max_input_tokens"] = max_input_tokens diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index c928398a87f..fe4732c6492 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -30,7 +30,10 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user -from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.proxy.vector_store_endpoints.utils import ( + can_user_access_vector_store, + filter_listable_vector_stores, +) from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.types.vector_stores import ( @@ -390,11 +393,10 @@ async def list_vector_stores( # Filter vector stores based on access control accessible_vector_stores: Final = [] - for vs in vector_store_map.values(): - if await _check_vector_store_access(vs, user_api_key_dict): - redacted = LiteLLM_ManagedVectorStore(**vs) - redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) - accessible_vector_stores.append(redacted) + for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict): + redacted = LiteLLM_ManagedVectorStore(**vs) + redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) + accessible_vector_stores.append(redacted) total_count: Final = len(accessible_vector_stores) total_pages: Final = (total_count + page_size - 1) // page_size diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 93f1510bf22..6e94a5a88ac 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,11 +1,17 @@ import json import re +from collections.abc import Iterable +from types import MappingProxyType from typing import Any, Final, Literal from fastapi import HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + is_ui_session_credential, + resolve_ui_session_team_ids, +) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LitellmUserRoles, @@ -160,10 +166,16 @@ async def can_user_access_vector_store( if _is_proxy_admin(user_api_key_dict): return True - vector_store_team_id: Final = vector_store.get("team_id") - if vector_store_team_id is None: + if vector_store.get("team_id") is None: return True + return await _is_vector_store_granted(vector_store, user_api_key_dict) + + +async def _is_vector_store_granted( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: vector_store_id: Final = vector_store.get("vector_store_id") or "" key_object_permission = user_api_key_dict.object_permission @@ -178,12 +190,70 @@ async def can_user_access_vector_store( if _object_permission_allows_vector_store(team_object_permission, vector_store_id): return True - if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id: - return True + return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id") + +async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + team: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return user_api_key_dict.model_copy( + update=MappingProxyType( + { + "team_id": team_id, + "team_object_permission": team.object_permission, + "team_object_permission_id": team.object_permission_id, + } + ) + ) + + +async def _vector_store_listing_auth_contexts( + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[UserAPIKeyAuth, ...]: + if not is_ui_session_credential(user_api_key_dict): + return (user_api_key_dict,) + session_key_context: Final = user_api_key_dict.model_copy( + update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None}) + ) + team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict) + team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids]) + return (session_key_context, *team_contexts) + + +async def _is_vector_store_granted_to_any( + vector_store: LiteLLM_ManagedVectorStore, + auth_contexts: tuple[UserAPIKeyAuth, ...], +) -> bool: + for auth_context in auth_contexts: + if await _is_vector_store_granted(vector_store, auth_context): + return True return False +async def filter_listable_vector_stores( + vector_stores: Iterable[LiteLLM_ManagedVectorStore], + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[LiteLLM_ManagedVectorStore, ...]: + """Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants.""" + if _is_proxy_admin(user_api_key_dict): + return tuple(vector_stores) + + auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict) + return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)]) + + async def get_litellm_managed_vector_store( vector_store_id: str, ) -> LiteLLM_ManagedVectorStore | None: diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 94bfc305a6a..1f63152632e 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -11,7 +11,7 @@ __all__ = ["aingest", "aquery", "ingest", "query"] import asyncio import contextvars -from collections.abc import Coroutine, Iterator +from collections.abc import Coroutine, Iterator, Mapping from contextlib import contextmanager from functools import partial from types import MappingProxyType @@ -66,6 +66,10 @@ _FORWARDABLE_RETRIEVAL_CONFIG_KEYS: Final = frozenset( } ) +_SEARCH_ARGS_SET_BY_PIPELINE: Final = frozenset( + {"vector_store_id", "query", "max_num_results", "custom_llm_provider", "router"} +) + def get_ingestion_class(provider: str) -> type[BaseRAGIngestion]: """ @@ -225,6 +229,7 @@ async def _execute_query_pipeline( retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, + vector_store_params: Mapping[str, object] | None = None, **kwargs, ) -> ModelResponse: """ @@ -241,11 +246,19 @@ async def _execute_query_pipeline( # 2. Search vector store # Forward allowlisted provider retrieval_config extras (region, embedding - # model, bucket, credential refs) to the search call; kwargs win on conflict. + # model, bucket, credential refs) to the search call; the managed store's + # params win on conflict. provider_search_params: Final = MappingProxyType( {k: v for k, v in retrieval_config.items() if k in _FORWARDABLE_RETRIEVAL_CONFIG_KEYS} ) - forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs}) + store_search_params: Final = MappingProxyType( + { + k: v + for k, v in (vector_store_params.items() if vector_store_params else ()) + if k not in _SEARCH_ARGS_SET_BY_PIPELINE + } + ) + forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs, **store_search_params}) with _suppressed_sub_call_billing(): search_response: Final = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], @@ -339,6 +352,7 @@ async def aquery( retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, + vector_store_params: Mapping[str, object] | None = None, **kwargs, ) -> ModelResponse: """ @@ -356,6 +370,7 @@ async def aquery( retrieval_config=retrieval_config, rerank=rerank, stream=stream, + vector_store_params=vector_store_params, **kwargs, ) @@ -386,6 +401,7 @@ def query( retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, + vector_store_params: Mapping[str, object] | None = None, **kwargs, ) -> ModelResponse | Coroutine[None, None, ModelResponse]: """ @@ -402,6 +418,7 @@ def query( retrieval_config=retrieval_config, rerank=rerank, stream=stream, + vector_store_params=vector_store_params, **kwargs, ) else: @@ -412,6 +429,7 @@ def query( retrieval_config=retrieval_config, rerank=rerank, stream=stream, + vector_store_params=vector_store_params, **kwargs, ) ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d4b9f4e8cce..3862aec445f 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -8,6 +8,7 @@ from typing import Any, Final, Literal, cast import litellm from litellm.constants import ( + AZURE_OPENAI_AUDIO_PROVIDERS, REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, request_timeout, @@ -400,7 +401,7 @@ async def _arealtime( litellm_metadata=_build_litellm_metadata(kwargs), query_params=query_params, ) - elif _custom_llm_provider == "azure": + elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # set API KEY api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY") diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index bc25f4fffb1..d7f8cd8f8bd 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -562,6 +562,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): hidden_params: Final = getattr(chunk, "_hidden_params", None) if hidden_params is not None: chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params + if ( + chunk_dict.get("usage") is None + and isinstance(hidden_params, dict) + and hidden_params.get("usage") is not None + ): + chunk_dict["usage"] = hidden_params["usage"] return chunk_dict def create_reasoning_summary_text_done_event( diff --git a/litellm/router.py b/litellm/router.py index 2b8b342d253..f33dfbba7bf 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,7 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -117,6 +117,9 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -211,6 +214,7 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, + HeuristicV2RouterLimit, LiteLLM_Params, MockRouterTestingParams, ModelGroupInfo, @@ -590,16 +594,20 @@ set_live_deployment_replay(_replay_live_router_model_cost) # Kwargs that carry no signal about the failed attempt, so log_retry drops them from a -# breadcrumb entirely: the request payload and the router-internal walk state. Credentials are -# handled separately by mask_credentials_in_payload, which scrubs credential-named values from -# whatever kwargs remain rather than trying to enumerate every credential-bearing key here. +# breadcrumb entirely: the request payload, the proxy's snapshot of the inbound request (its body +# aliases the live request metadata, earlier breadcrumbs included, so copying it would nest every +# breadcrumb inside the next one), and the router-internal walk state. Credentials are handled +# separately by mask_credentials_in_payload, which scrubs credential-named values from whatever +# kwargs remain rather than trying to enumerate every credential-bearing key here. RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( ( "messages", "original_function", "attempted_targets", + "proxy_server_request", ) ) +RETRY_BREADCRUMB_LIMIT: Final = 4 class Router: @@ -683,6 +691,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, + heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -759,6 +768,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments + self.heuristic_v2_router_limit = heuristic_v2_router_limit self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks @@ -958,7 +968,6 @@ class Router: self.total_calls: defaultdict = defaultdict(int) # dict to store total calls made to each model self.fail_calls: defaultdict = defaultdict(int) # dict to store fail_calls made to each model self.success_calls: defaultdict = defaultdict(int) # dict to store success_calls made to each model - self.previous_models: list = [] # list to store failed calls (passed in as metadata to next call) # make Router.chat.completions.create compatible for openai.chat.completions.create default_litellm_params = default_litellm_params or {} @@ -8137,35 +8146,31 @@ class Router: """ When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing """ - try: - _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - # Log failed model as the previous model - previous_model: Final = { + _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" + request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var] + attempt_kwargs: Final = MappingProxyType( + {k: v for k, v in kwargs.items() if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS} + ) + attempt_metadata: Final = MappingProxyType( + {k: v for k, v in request_metadata.items() if k != "previous_models"} + ) + previous_model: Final = MappingProxyType( + { "exception_type": type(e).__name__, "exception_string": str(e), + **attempt_kwargs, + _metadata_var: attempt_metadata, } - for ( - k, - v, - ) in kwargs.items(): # log everything in kwargs except the old previous_models value - prevent nesting - if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS: - previous_model[k] = v - elif k == _metadata_var and isinstance(v, dict): - previous_model[_metadata_var] = {} - for metadata_k, metadata_v in kwargs[_metadata_var].items(): - if metadata_k != "previous_models": - previous_model[k][metadata_k] = metadata_v - - # check current size of self.previous_models, if it's larger than 3, remove the first element - if len(self.previous_models) > 3: - self.previous_models.pop(0) - - scrubbed_previous_model: Final = mask_credentials_in_payload(previous_model) - self.previous_models.append(scrubbed_previous_model) - kwargs[_metadata_var]["previous_models"] = self.previous_models - return kwargs - except Exception as e: - raise e + ) + earlier_breadcrumbs: Final = request_metadata.get("previous_models") + kept_breadcrumbs: Final[tuple[object, ...]] = ( + tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] + if isinstance(earlier_breadcrumbs, (list, tuple)) + else () + ) + breadcrumbs: Final = (*kept_breadcrumbs, mask_credentials_in_payload(previous_model)) + kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict + return kwargs def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int: """ @@ -8796,6 +8801,30 @@ class Router: """ return classify_strategy_router_model(litellm_params.model) == "complexity" + def config_deployments(self) -> Iterator[Mapping[str, object]]: + """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset).""" + for deployment in self.model_list: + if not isinstance(deployment, Mapping): + continue + model_info = deployment.get("model_info") + if not (isinstance(model_info, Mapping) and model_info.get("db_model")): + yield deployment + + def heuristic_v2_router_limit_violation(self) -> str | None: + """ + Why one more heuristic_v2 router cannot join this router, or None when it can. + + Judged against every deployment currently on the model_list; an upsert pops the row being + edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is + resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which + is the SDK default, and the proxy injects a resolver backed by its license. + """ + limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None + others: Final = count_heuristic_v2_routers( + deployment for deployment in self.model_list if isinstance(deployment, Mapping) + ) + return heuristic_v2_limit_violation(held=others + 1, limit=limit) + def init_complexity_router_deployment(self, deployment: Deployment): """ Initialize the complexity-router deployment. @@ -8813,6 +8842,10 @@ class Router: ) complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config + if uses_heuristic_v2_classifier(complexity_router_config): + limit_violation: Final = self.heuristic_v2_router_limit_violation() + if limit_violation is not None: + raise ValueError(limit_violation) default_model: str | None = deployment.litellm_params.complexity_router_default_model @@ -9636,8 +9669,16 @@ class Router: raise e def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: + """Put a deployment back the way it was before a failed upsert popped it. + + A rollback re-admits state that was already serving, so it does not go through the + heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first + registered, judging the rollback would drop a serving router over an unrelated failed edit. + """ if previous_deployment is None or self.has_model_id(model_id): return + limit_resolver: Final = self.heuristic_v2_router_limit + self.heuristic_v2_router_limit = None try: self.add_deployment(deployment=previous_deployment) verbose_router_logger.info( @@ -9652,6 +9693,8 @@ class Router: model_id, restore_error, ) + finally: + self.heuristic_v2_router_limit = limit_resolver @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: @@ -9981,6 +10024,17 @@ class Router: coerce_token_limit(model_info.get("max_output_tokens")), ) + def get_configured_mode(self, model_name: str) -> "str | None": + """Return the mode explicitly configured for a concrete deployment.""" + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + + mode: Final = deployment.model_info.get("mode") + if isinstance(mode, str) and mode.strip(): + return mode + return None + def get_configured_display_name(self, model_name: str) -> "str | None": """ Return the display_name explicitly configured in a concrete deployment's diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index c77745a498d..6b443026f61 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -2,23 +2,41 @@ Auto-Routing Strategy that works with a Semantic Router Config """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional +from pydantic import BaseModel, ConfigDict + from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import ( + effective_turn_off_message_logging, + forwarded_internal_call_metadata, + parent_session_kwargs, +) +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from semantic_router.routers.base import Route from litellm.router import Router + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder from litellm.types.router import PreRoutingHookResponse else: Router = Any PreRoutingHookResponse = Any Route = Any SemanticRouter = Any + LiteLLMRouterEncoder = Any + + +class _CallerMetadata(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + metadata: Mapping[str, object] | None = None + litellm_metadata: Mapping[str, object] | None = None class AutoRouter(CustomLogger): @@ -50,6 +68,8 @@ class AutoRouter(CustomLogger): """ from semantic_router.routers import SemanticRouter + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + self.auto_router_config_path: str | None = auto_router_config_path self.auto_router_config: str | None = auto_router_config self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE @@ -59,6 +79,11 @@ class AutoRouter(CustomLogger): self.embedding_model: str = embedding_model self.max_input_chars: int = max_input_chars self.litellm_router_instance: Router = litellm_router_instance + self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder( + litellm_router_instance=litellm_router_instance, + model_name=embedding_model, + max_input_chars=max_input_chars, + ) def _load_semantic_routing_routes(self) -> list[Route]: from semantic_router.routers import SemanticRouter @@ -129,9 +154,6 @@ class AutoRouter(CustomLogger): from semantic_router.routers import SemanticRouter from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) from litellm.types.router import PreRoutingHookResponse resolved_messages: Final = ( @@ -149,34 +171,47 @@ class AutoRouter(CustomLogger): ####################### routelayer = SemanticRouter( routes=self.loaded_routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.litellm_router_instance, - model_name=self.embedding_model, - max_input_chars=self.max_input_chars, - ), + encoder=self.encoder, auto_sync=self.auto_sync_value, ) self.routelayer = routelayer message_content: Final = self._extract_text_from_messages(resolved_messages) - route_name: Final = self._matched_route_name(routelayer, message_content) + route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs) return PreRoutingHookResponse( model=route_name or self.default_model, messages=messages, ) - def _matched_route_name(self, routelayer: "SemanticRouter", text: str) -> str | None: + async def _matched_route_name( + self, routelayer: "SemanticRouter", text: str, request_kwargs: Mapping[str, object] + ) -> str | None: """Name of the route `text` matches, or None when nothing matched or the match failed. - The route layer embeds `text` to compare it against the routes, and that embedding call can + `text` is embedded here rather than by `routelayer(text=...)` so the caller's metadata reaches + `aembedding()` and the embedding's spend lands on the key/team that sent the request; + SemanticRouter has no way to pass kwargs through to its encoder. That embedding call can fail (context limit, timeout, provider error). Choosing a model is a routing decision, so a failure here falls back to the default model rather than failing the user's request. """ from semantic_router.schema import RouteChoice try: - route_choice: Final = routelayer(text=text) + caller: Final = _CallerMetadata.model_validate(request_kwargs) + query_vector: Final = ( + await self.encoder.aencode_queries( + [text], + metadata=forwarded_internal_call_metadata(caller.metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN), + litellm_metadata=forwarded_internal_call_metadata( + caller.litellm_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN + ), + proxy_server_request={"body": {"model": self.embedding_model, "input": [text]}}, + turn_off_message_logging=effective_turn_off_message_logging(request_kwargs), + **parent_session_kwargs(request_kwargs), + ) + )[0] + route_choice: Final = await routelayer.acall(vector=query_vector) except Exception as e: # noqa: BLE001 -- the embedding call behind the route layer can fail many ways (context limit, timeout, provider/network error); none of them may fail the request verbose_router_logger.warning( "AutoRouter: semantic routing failed (%s), falling back to default model %s", e, self.default_model diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 6deba5aa1cf..665ff69ab47 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -1,6 +1,7 @@ #### What this does #### # identifies lowest tpm deployment import random +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -350,9 +351,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): model_group: str, healthy_deployments: list, tpm_keys: list, - tpm_values: list | None, + tpm_values: Sequence | None, rpm_keys: list, - rpm_values: list | None, + rpm_values: Sequence | None, messages: list[dict[str, str]] | None = None, input: str | list | None = None, ) -> dict | None: diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index a8aa543d735..2efbfb5782e 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under ``ignore_invalid_deployments``. """ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -163,6 +163,38 @@ def strategy_router_dependencies( ) +def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool: + """Whether this complexity config classifies with the bundled heuristic_v2 model.""" + return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2" + + +def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool: + """Whether this deployment is a complexity router that classifies with heuristic_v2.""" + return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and ( + uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config")) + ) + + +def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int: + """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers.""" + return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params")))) + + +def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: + """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits. + + ``limit`` None means unlimited. The message is shared by every enforcement point (config + load, model writes, router registration) and stays SDK-neutral: it names the cap and what + the caller can change; the proxy appends how its license lifts the cap. + """ + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make " + f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router." + ) + + def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None: """Reject a complexity config the router would refuse to build a deployment from. diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 39d3e25aacb..6f3ea8eb78a 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -82,6 +82,7 @@ class DeploymentAffinityCheck(CustomLogger): """ CACHE_KEY_PREFIX = "deployment_affinity:v1" + USER_ID_AFFINITY_PREFIX: Final = "user_id:" def __init__( self, @@ -253,15 +254,6 @@ class DeploymentAffinityCheck(CustomLogger): hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" - @staticmethod - def _get_user_key_from_metadata_dict(metadata: dict) -> str | None: - # NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the - # OpenAI `user` parameter, which is an end-user identifier). - user_key: Final = metadata.get("user_api_key_hash") - if user_key is None: - return None - return str(user_key) - @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") @@ -285,22 +277,30 @@ class DeploymentAffinityCheck(CustomLogger): return metadata_dicts @staticmethod - def _get_user_key_from_request_kwargs(request_kwargs: dict) -> str | None: + def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None: + value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None) + return None if value is None else str(value) + + @classmethod + def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None: """ Extract a stable affinity key from request kwargs. - Source (proxy): `metadata.user_api_key_hash` + Source (proxy): `metadata.user_api_key_hash` for virtual-key callers. JWT-authenticated + callers carry no key hash, so their `metadata.user_api_key_user_id` stands in for it, + namespaced under `USER_ID_AFFINITY_PREFIX` so a user id can never alias a key hash. Note: the OpenAI `user` parameter is an end-user identifier and is intentionally not used for deployment affinity. """ - # Check metadata dicts (Proxy usage) - for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): - user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata) - if user_key is not None: - return user_key - - return None + metadata_dicts: Final = cls._iter_metadata_dicts(request_kwargs) + user_api_key_hash: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_hash") + if user_api_key_hash is not None: + return user_api_key_hash + user_id: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_user_id") + if user_id is None: + return None + return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}" @staticmethod def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: @@ -533,9 +533,9 @@ class DeploymentAffinityCheck(CustomLogger): return typed_healthy_deployments verbose_router_logger.debug( - "DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s", + "DeploymentAffinityCheck: caller affinity hit -> deployment=%s user_key=%s", model_id, - self._shorten_for_logs(user_key), + self._shorten_for_logs(self._hash_user_key(user_key)), ) return [deployment] @@ -626,7 +626,7 @@ class DeploymentAffinityCheck(CustomLogger): deployment_model_name, model_id, self.ttl_seconds, - self._shorten_for_logs(user_key), + self._shorten_for_logs(self._hash_user_key(user_key)), ) else: verbose_router_logger.debug( diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 9a714e1724e..5de58a20242 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from pydantic import BaseModel, Field @@ -29,6 +30,13 @@ def is_interception_internal_key( return any(key.startswith(prefix) for prefix in prefixes) +CONVERTED_STREAM_KEYS: Final = frozenset(f"{prefix}_converted_stream" for prefix in INTERCEPTION_INTERNAL_PREFIXES) + + +def converted_stream_requested(params: Mapping[str, object]) -> bool: + return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS) + + class AgenticLoopSafetyError(ValueError): """ Raised when an agentic-loop safety rail refuses a rerun. diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 32d88da0085..b33ff954c35 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -66,6 +66,7 @@ from pydantic import ( ConfigDict, Discriminator, Field, + NonNegativeInt, PrivateAttr, SerializerFunctionWrapHandler, field_serializer, @@ -1321,6 +1322,18 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class WebSearchToolUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + num_requests: NonNegativeInt + + +class ResponsesToolUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + web_search: WebSearchToolUsage | None = None + + ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] """ The status of the response generation. diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index df0a090cdb0..6973f1d1f12 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,6 +1,8 @@ -from typing import Any, Final +from collections.abc import Mapping +from typing import Any, Final, Literal from pydantic import BaseModel, field_validator +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import ( LiteLLM_UserTableWithKeyCount, @@ -9,6 +11,17 @@ from litellm.proxy._types import ( ) +class InsensitiveContains(TypedDict): + contains: ReadOnly[str] + mode: ReadOnly[Literal["insensitive"]] + + +class UserSearchWhere(TypedDict): + """Prisma filter behind `/user/list?search=`: user_id or user_email contains the term, case-insensitive.""" + + OR: ReadOnly[tuple[Mapping[Literal["user_id", "user_email"], InsensitiveContains], ...]] + + class UserListResponse(BaseModel): """ Response model for the user list endpoint diff --git a/litellm/types/proxy/model_listing.py b/litellm/types/proxy/model_listing.py index b59c0f2cf19..24cfa85eee4 100644 --- a/litellm/types/proxy/model_listing.py +++ b/litellm/types/proxy/model_listing.py @@ -11,8 +11,8 @@ class ModelInfoMetadata(TypedDict): class ModelInfoResponse(TypedDict): """OpenAI-compatible model object. `mode`, `max_input_tokens`, and - `max_output_tokens` are attached when the cost map knows them; `metadata` - is present only when the endpoint is called with include_metadata=true. + `max_output_tokens` are attached when the cost map or deployment config + knows them; `metadata` is present only with include_metadata=true. """ id: str diff --git a/litellm/types/router.py b/litellm/types/router.py index 4f4df1a8d2e..7ebd50f1328 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -885,6 +885,18 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... +class HeuristicV2RouterLimit(Protocol): + """ + Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited. + + The Router calls it on every registration and limit query instead of caching the answer, so the + proxy can keep the limit on its license object (re-verified on config load) rather than hand + over a snapshot. + """ + + def __call__(self) -> int | None: ... + + class LiteLLM_RouterFileObject(TypedDict, total=False): """ Tracking the litellm params hash, used for mapping the file id to the right model diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0bdbb83fbb4..34c47d3f201 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3612,6 +3612,7 @@ all_litellm_params = ( "litellm_system_prompt", "provider_specific_header", "prompt_version", + "prompt_environment", "api_base", "force_timeout", "logger_fn", diff --git a/litellm/utils.py b/litellm/utils.py index 239673a20d9..585b5dbe1a8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,6 +83,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, ) +from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload _CachingHandlerResponse = None _LLMCachingHandler = None @@ -543,6 +544,14 @@ def print_verbose( pass +def _print_verbose_is_active() -> bool: + """Whether print_verbose would reach either of its two consumers, so a call site can skip + building a payload nothing would read. _is_debugging_on() is not the same predicate: it reads + litellm._logging.set_verbose, while print_verbose's print reads litellm.set_verbose, and + assigning the documented litellm.set_verbose = True rebinds only the latter.""" + return litellm.set_verbose is True or verbose_logger.isEnabledFor(logging.DEBUG) + + ####### CLIENT ################### # make it easy to log if completion/embedding runs succeeded or failed + see what happened | Non-Blocking def custom_llm_setup(): @@ -1284,16 +1293,18 @@ async def async_post_call_success_deployment_hook( except ValueError: typed_call_type = None # unknown call type + modified_response = response + CustomLogger: Final = _get_cached_custom_logger() for callback in litellm.callbacks: if isinstance(callback, CustomLogger): result = await callback.async_post_call_success_deployment_hook( - request_data, cast(LLMResponseTypes, response), typed_call_type + request_data, cast(LLMResponseTypes, modified_response), typed_call_type ) if result is not None: - return result + modified_response = result - return response + return modified_response async def async_post_call_failure_deployment_hook( @@ -4707,7 +4718,8 @@ def get_optional_params( openai_params=list(DEFAULT_CHAT_COMPLETION_PARAM_VALUES.keys()), additional_drop_params=additional_drop_params, ) - print_verbose(f"Final returned optional params: {optional_params}") + if _print_verbose_is_active(): + print_verbose(f"Final returned optional params: {redact_credentials_in_payload(optional_params)}") optional_params = _apply_openai_param_overrides( optional_params=optional_params, non_default_params=non_default_params, @@ -7462,7 +7474,8 @@ def print_args_passed_to_litellm(original_function, args, kwargs): return args_str: Final = ", ".join(map(repr, args)) - kwargs_str: Final = ", ".join(f"{key}={value!r}" for key, value in kwargs.items()) + redacted_kwargs: Final = redact_credentials_in_payload(kwargs) + kwargs_str: Final = ", ".join(f"{key}={value!r}" for key, value in redacted_kwargs.items()) print_verbose( "\n", ) # new line before diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9b358a1cedd..44c4f10ec38 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10305,6 +10305,24 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/grok-4-6-comes-to-microsoft-foundry-models-built-for-long-horizon-reasoning-and-/4547578", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, @@ -29277,6 +29295,75 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + "cache_creation_input_token_cost_flex": 6.25e-06, + "cache_creation_input_token_cost_priority": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_read_input_token_cost_flex": 5e-07, + "cache_read_input_token_cost_priority": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "input_cost_per_token_batches": 5e-06, + "input_cost_per_token_flex": 5e-06, + "input_cost_per_token_priority": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "output_cost_per_token_batches": 2.5e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 0.0001, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -52911,6 +52998,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -52945,6 +53037,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53007,6 +53104,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53195,6 +53297,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53226,6 +53333,11 @@ "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "output_cost_per_token": 1.65e-05, "output_cost_per_token_above_272k_tokens": 2.475e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index be2b30fc189..24c0ff6b181 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -42,7 +42,7 @@ "limit": 52 }, "B010": { - "limit": 190 + "limit": 187 }, "B018": { "limit": 2 @@ -144,7 +144,7 @@ "limit": 1 }, "PLR1704": { - "limit": 3 + "limit": 1 }, "PLR1714": { "limit": 253 @@ -240,13 +240,13 @@ "limit": 96 }, "TRY201": { - "limit": 403 + "limit": 401 }, "TRY203": { - "limit": 111 + "limit": 109 }, "TRY300": { - "limit": 854 + "limit": 852 }, "UP028": { "limit": 2 diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index 20a8592db20..a56d3dc077e 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -168,6 +168,35 @@ def completed_responses_object(result: StreamingResponse) -> ResponsesObject | N return completed[-1] if completed else None +class AnthropicMessageObject(BaseModel): + id: str + + +class AnthropicStreamEvent(BaseModel): + """One SSE frame of a native Anthropic stream. Only `message_start` carries the + message, so it stays optional and the deltas validate as themselves.""" + + type: str + message: AnthropicMessageObject | None = None + + +def anthropic_message_id(result: StreamingResponse) -> str | None: + """The `msg_...` id the caller was served, which is what the spend row is keyed by + on this route: off the `message_start` frame when streaming, off the body when not.""" + if not result.is_streaming: + return AnthropicMessageObject.model_validate_json(result.body).id + events = ( + AnthropicStreamEvent.model_validate_json(payload) + for payload in result.stream_events + ) + started = tuple( + event.message + for event in events + if event.type == "message_start" and event.message is not None + ) + return started[0].id if started else None + + class OpenAIResponsesBody(BaseModel): model: str input: str diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index 7e6a8b25155..50ea8f4b4df 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -2,7 +2,8 @@ Each test sends a NATIVE provider request through the proxy's passthrough route and verifies the proxy still logged a costed SpendLogs row -(call_type="pass_through_endpoint"), correlated by the x-litellm-call-id header. +(call_type="pass_through_endpoint"), correlated by the id the caller was served: +the x-litellm-call-id header on gemini, the `msg_...` message id on anthropic. Covered: gemini ("gemini-2.5-flash") + anthropic ("claude-haiku-4-5"), streaming + non-streaming, plus native tool calls. See LLM_TRANSLATION_COVERAGE_MATRIX.md. @@ -14,7 +15,7 @@ A passthrough call returning non-2xx fails hard (never a skip); once it returns import pytest from e2e_config import CHEAP_OPENAI_MODEL, unique_marker -from e2e_http import StreamingResponse, require_successful_call, unwrap +from e2e_http import require_successful_call, unwrap from lifecycle import ResourceManager from models import KeyGenerateBody, SpendLogRow from passthrough_client import ( @@ -24,6 +25,7 @@ from passthrough_client import ( JsonSchema, JsonSchemaProperty, PassthroughClient, + anthropic_message_id, completed_responses_object, ) @@ -33,18 +35,18 @@ REALTIME_MODEL = "gpt-realtime-2" pytestmark = pytest.mark.e2e -def _fetch_cost_breakdown(client: PassthroughClient, result: StreamingResponse) -> SpendLogRow: +def _fetch_cost_breakdown(client: PassthroughClient, request_id: str | None) -> SpendLogRow: """The passthrough call's logged row, polled until it carries a cost. Asserts (not skips) that a 2xx passthrough call produced a costed row - the whole point of passthrough spend tracking. """ - assert result.call_id, "passthrough response had no x-litellm-call-id header" + assert request_id, "passthrough response carried no id to correlate its spend row by" rows = client.proxy.poll_logs_for_request_id( - result.call_id, + request_id, predicate=lambda rs: (rs[0].spend or 0) > 0, ) - assert rows, f"no SpendLogs row for passthrough call_id {result.call_id}" + assert rows, f"no SpendLogs row for passthrough request_id {request_id}" row = rows[0] assert row.call_type == "pass_through_endpoint" assert (row.spend or 0) > 0, f"passthrough call was not costed: {row}" @@ -64,7 +66,7 @@ def test_gemini_passthrough_nonstreaming_logs_cost( ) require_successful_call(result) - row = _fetch_cost_breakdown(client, result) + row = _fetch_cost_breakdown(client, result.call_id) assert row.custom_llm_provider == "gemini" assert "gemini" in (row.model or "") assert tag in (row.request_tags or []), f"tags not logged: {row.request_tags}" @@ -107,7 +109,7 @@ def test_gemini_passthrough_streaming_logs_cost( require_successful_call(result) assert result.chunks > 0, "streaming passthrough produced no events" - row = _fetch_cost_breakdown(client, result) + row = _fetch_cost_breakdown(client, result.call_id) assert row.custom_llm_provider == "gemini" @@ -137,7 +139,7 @@ def test_gemini_passthrough_tool_call_logs_cost( require_successful_call(result) assert "functionCall" in result.body, "gemini did not emit a tool call" - row = _fetch_cost_breakdown(client, result) + row = _fetch_cost_breakdown(client, result.call_id) assert row.custom_llm_provider == "gemini" @@ -150,7 +152,7 @@ def test_anthropic_passthrough_nonstreaming_logs_cost( result = client.anthropic_message(scoped_key, "claude-haiku-4-5", "Say hello") require_successful_call(result) - row = _fetch_cost_breakdown(client, result) + row = _fetch_cost_breakdown(client, anthropic_message_id(result)) assert row.custom_llm_provider == "anthropic" assert "claude" in (row.model or "") @@ -164,7 +166,7 @@ def test_anthropic_passthrough_streaming_logs_cost( require_successful_call(result) assert result.chunks > 0, "streaming passthrough produced no events" - row = _fetch_cost_breakdown(client, result) + row = _fetch_cost_breakdown(client, anthropic_message_id(result)) assert row.custom_llm_provider == "anthropic" @@ -190,7 +192,7 @@ def test_anthropic_passthrough_tool_call_logs_cost( require_successful_call(result) assert "tool_use" in result.body, "anthropic did not emit a tool call" - row = _fetch_cost_breakdown(client, result) + row = _fetch_cost_breakdown(client, anthropic_message_id(result)) assert row.custom_llm_provider == "anthropic" diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 22558faedad..9ffb57924b6 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -567,7 +567,7 @@ class TestResolveAllMigrationsLedger: return _FakeCompleted() return _FakeCompleted() - monkeypatch.setattr(utils_module.subprocess, "run", fake_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run) ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma") return calls @@ -604,9 +604,9 @@ class TestPartitionedSpendLogsPushGuard: import litellm_proxy_extras.utils as utils_module def fail_run(cmd, **kwargs): - raise AssertionError(f"subprocess.run should not be called, got: {cmd}") + raise AssertionError(f"run_prisma should not be called, got: {cmd}") - monkeypatch.setattr(utils_module.subprocess, "run", fail_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fail_run) def test_v1_db_push_fails_fast_with_guidance(self, monkeypatch): monkeypatch.setattr( diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index ef8d6c55148..f8f23ea015a 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -2879,7 +2879,6 @@ def response_format_tests(response: litellm.ModelResponse): "model", [ "bedrock/mistral.mistral-large-2407-v1:0", - "bedrock/cohere.command-r-plus-v1:0", "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "mistral.mistral-7b-instruct-v0:2", "meta.llama3-8b-instruct-v1:0", diff --git a/tests/local_testing/test_dual_cache.py b/tests/local_testing/test_dual_cache.py index e60fa5f3746..43b10a9557a 100644 --- a/tests/local_testing/test_dual_cache.py +++ b/tests/local_testing/test_dual_cache.py @@ -240,3 +240,35 @@ async def test_dual_cache_delete(is_async): result = dual_cache.get_cache(test_key) assert result is None + + +@pytest.mark.asyncio +async def test_dual_cache_concurrent_sync_and_async_redis_reads(): + """Sync and async batch reads share one Redis backend in one process, and sync reads never open an async connection""" + redis_cache = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + dual_cache = DualCache(redis_cache=redis_cache) + + run_id = str(uuid.uuid4()) + sync_keys = [f"sync_{run_id}_{index}" for index in range(5)] + async_keys = [f"async_{run_id}_{index}" for index in range(5)] + in_loop_keys = [f"in_loop_{run_id}_{index}" for index in range(3)] + survivor_key = f"survivor_{run_id}" + expected = {key: {"key": key} for key in [*sync_keys, *async_keys, *in_loop_keys, survivor_key]} + for key, value in expected.items(): + await redis_cache.async_set_cache(key, value, ttl=60) + + concurrent_results = await asyncio.gather( + *(asyncio.to_thread(dual_cache.batch_get_cache, keys=[key]) for key in sync_keys), + *(dual_cache.async_batch_get_cache(keys=[key]) for key in async_keys), + ) + assert list(concurrent_results) == [[expected[key]] for key in [*sync_keys, *async_keys]] + + with patch.object( + redis_cache, + "async_batch_get_cache", + side_effect=AssertionError("sync batch reads must not call async Redis"), + ): + in_loop_results = [dual_cache.batch_get_cache(keys=[key]) for key in in_loop_keys] + + assert in_loop_results == [[expected[key]] for key in in_loop_keys] + assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]] diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 07d693af447..bf39d3155b7 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1168,7 +1168,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode): "model, region", [ # ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"], - # ["bedrock/cohere.command-r-plus-v1:0", None], ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None], # ["mistral.mistral-7b-instruct-v0:2", None], # ["meta.llama3-8b-instruct-v1:0", None], @@ -1271,7 +1270,7 @@ def test_bedrock_claude_3_streaming(): "model", [ "claude-haiku-4-5-20251001", - "cohere.command-r-plus-v1:0", # bedrock + "bedrock/mistral.mistral-7b-instruct-v0:2", "gpt-3.5-turbo", ], ) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index abae26e02cd..7e338dafb86 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,6 +1,13 @@ +import time +from collections.abc import Iterator +from typing import Final + import httpx -from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError import pytest +from openai import APIStatusError, BadRequestError, NotFoundError, OpenAI, Stream +from openai.types.responses import ResponseStreamEvent + +BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS: Final = 90 def generate_key(): @@ -153,43 +160,48 @@ def test_cancel_response(): raise e +def admitted_response_id(chunk: ResponseStreamEvent) -> str | None: + response: Final = getattr(chunk, "response", None) + return None if response is None else response.id + + +def events_until_admission(stream: Stream[ResponseStreamEvent], started: float) -> Iterator[ResponseStreamEvent]: + for chunk in stream: + print("stream chunk=", chunk) + yield chunk + if admitted_response_id(chunk) is not None: + return + if time.monotonic() - started > BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS: + return + + def test_cancel_streaming_response(): - try: - client = get_test_client() - from litellm.types.llms.openai import ResponsesAPIResponse + client: Final = get_test_client() + started: Final = time.monotonic() + stream: Final = client.responses.create( + model="gpt-5.5", + input="count from 1 to 500, one number per line", + stream=True, + background=True, + timeout=BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS, + ) - stream = client.responses.create( - model="gpt-5.5", - input="just respond with the word 'ping'", - stream=True, - background=True, + with stream: + events: Final = tuple(events_until_admission(stream, started)) + + elapsed: Final = time.monotonic() - started + keepalive_events: Final = sum(1 for chunk in events if chunk.type == "keepalive") + response_id: Final = next((rid for rid in map(admitted_response_id, events) if rid is not None), None) + if response_id is None and keepalive_events: + pytest.skip( + f"OpenAI held the background stream in keepalive for {elapsed:.0f}s " + f"({keepalive_events} keepalive events) without creating the response" ) + assert response_id is not None, f"no response event within {elapsed:.0f}s of streaming a background response" - collected_chunks = [] - response_id = None - for chunk in stream: - print("stream chunk=", chunk) - collected_chunks.append(chunk) - # Extract response ID from the first chunk that has it - if ( - response_id is None - and hasattr(chunk, "response") - and hasattr(chunk.response, "id") - ): - response_id = chunk.response.id - - assert len(collected_chunks) > 0 - - # cancel the response if we got a response ID - if response_id: - cancel_response = client.responses.cancel(response_id) - print("CANCEL streaming response=", cancel_response) - assert hasattr(cancel_response, "id") - except Exception as e: - if "Cannot cancel a completed response" in str(e): - pass - else: - raise e + cancel_response: Final = client.responses.cancel(response_id) + print("CANCEL streaming response=", cancel_response) + assert cancel_response.status == "cancelled" def test_cancel_invalid_response_id(): diff --git a/tests/pass_through_tests/test_anthropic_passthrough.py b/tests/pass_through_tests/test_anthropic_passthrough.py index d42e06937dc..0452b171f9e 100644 --- a/tests/pass_through_tests/test_anthropic_passthrough.py +++ b/tests/pass_through_tests/test_anthropic_passthrough.py @@ -50,9 +50,9 @@ async def test_anthropic_basic_completion_with_headers(): anthropic_api_output_tokens = ( reported_usage.get("output_tokens", None) if reported_usage else None ) - litellm_call_id = response_headers.get("x-litellm-call-id") + anthropic_message_id = response_json.get("id") - print(f"LiteLLM Call ID: {litellm_call_id}") + print(f"Anthropic message ID: {anthropic_message_id}") # Wait for spend to be logged await asyncio.sleep(15) @@ -64,7 +64,7 @@ async def test_anthropic_basic_completion_with_headers(): print(f"Attempt {attempt + 1}/{max_retries} to check spend logs") async with session.get( - f"http://0.0.0.0:4000/spend/logs?request_id={litellm_call_id}", + f"http://0.0.0.0:4000/spend/logs?request_id={anthropic_message_id}", headers={"Authorization": "Bearer sk-1234"}, ) as spend_response: print("text spend response") @@ -84,25 +84,25 @@ async def test_anthropic_basic_completion_with_headers(): print("Waiting 10 seconds before retry...") await asyncio.sleep(10) - # Spend data might be unavailable (auth error, slow DB write, etc.) - if ( - spend_data is None - or not isinstance(spend_data, list) - or len(spend_data) == 0 - or not isinstance(spend_data[0], dict) - or "request_id" not in spend_data[0] - ): - print(f"Spend data not available or is error response: {spend_data}") - print("Skipping spend assertions (DB write may be slow in CI)") + if not isinstance(spend_data, list): + print(f"Spend endpoint answered with an error response: {spend_data}") + print("Skipping spend assertions (spend logs unreachable in CI)") return + assert spend_data, ( + f"GET /spend/logs?request_id={anthropic_message_id} found no row for the id " + "the caller received" + ) + log_entry = spend_data[0] # Basic existence checks assert isinstance(log_entry, dict), "Log entry should be a dictionary" # Request metadata assertions - assert log_entry["request_id"] == litellm_call_id, "Request ID should match" + assert ( + log_entry["request_id"] == anthropic_message_id + ), "Request ID should be the message id the caller received" assert ( log_entry["call_type"] == "pass_through_endpoint" ), "Call type should be pass_through_endpoint" @@ -182,8 +182,6 @@ async def test_anthropic_streaming_with_headers(): assert response.status == 200, "Response should be successful" response_headers = response.headers print(f"Response headers: {response_headers}") - litellm_call_id = response_headers.get("x-litellm-call-id") - print(f"LiteLLM Call ID: {litellm_call_id}") collected_output = [] async for line in response.content: @@ -194,13 +192,18 @@ async def test_anthropic_streaming_with_headers(): print("Collected output:", "".join(collected_output)) anthropic_api_usage_chunks = [] + anthropic_message_id = None for chunk in collected_output: chunk_json = json.loads(chunk) + if chunk_json.get("type") == "message_start": + anthropic_message_id = chunk_json.get("message", {}).get("id") if "usage" in chunk_json: anthropic_api_usage_chunks.append(chunk_json["usage"]) elif "message" in chunk_json and "usage" in chunk_json["message"]: anthropic_api_usage_chunks.append(chunk_json["message"]["usage"]) + print(f"Anthropic message ID: {anthropic_message_id}") + print( "anthropic_api_usage_chunks", json.dumps(anthropic_api_usage_chunks, indent=4, default=str), @@ -232,7 +235,7 @@ async def test_anthropic_streaming_with_headers(): print(f"Attempt {attempt + 1}/{max_retries} to check spend logs") async with session.get( - f"http://0.0.0.0:4000/spend/logs?request_id={litellm_call_id}", + f"http://0.0.0.0:4000/spend/logs?request_id={anthropic_message_id}", headers={"Authorization": "Bearer sk-1234"}, ) as spend_response: spend_data = await spend_response.json() @@ -250,25 +253,25 @@ async def test_anthropic_streaming_with_headers(): print("Waiting 10 seconds before retry...") await asyncio.sleep(10) - # Spend data might be unavailable (auth error, slow DB write, etc.) - if ( - spend_data is None - or not isinstance(spend_data, list) - or len(spend_data) == 0 - or not isinstance(spend_data[0], dict) - or "request_id" not in spend_data[0] - ): - print(f"Spend data not available or is error response: {spend_data}") - print("Skipping spend assertions (DB write may be slow in CI)") + if not isinstance(spend_data, list): + print(f"Spend endpoint answered with an error response: {spend_data}") + print("Skipping spend assertions (spend logs unreachable in CI)") return + assert spend_data, ( + f"GET /spend/logs?request_id={anthropic_message_id} found no row for the id " + "the caller received" + ) + log_entry = spend_data[0] # Basic existence checks assert isinstance(log_entry, dict), "Log entry should be a dictionary" # Request metadata assertions - assert log_entry["request_id"] == litellm_call_id, "Request ID should match" + assert ( + log_entry["request_id"] == anthropic_message_id + ), "Request ID should be the message id the caller received" assert ( log_entry["call_type"] == "pass_through_endpoint" ), "Call type should be pass_through_endpoint" diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py index 7577570be48..d3ffcf2445e 100644 --- a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py +++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py @@ -9,14 +9,16 @@ be forced by a sequential script: it needs one request to be genuinely mid-fligh other commits. A mocked prisma cannot arbitrate that either, since the property under test is whether Postgres's own advisory lock actually serializes the two requests. -These tests pin the interleaving the same way test_access_group_team_sync.py does: a second -real connection holds the team's advisory lock in its own transaction, so the function under -test is provably blocked on it rather than hoping a sleep lands in the right gap. +These tests pin the interleaving without a timing assumption: a second real connection holds +the team's advisory lock in its own transaction, and the test then waits for Postgres itself +to report the endpoint queued behind that exact lock. A sleep can only guess whether the +endpoint has reached the lock yet; pg_locks answers it. """ import asyncio import json import os +import time import uuid from contextlib import asynccontextmanager from datetime import timedelta @@ -39,6 +41,47 @@ _DELETE_SEEDED = 'DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = $1' _DELETE_USER = 'DELETE FROM "LiteLLM_UserTable" WHERE user_id = $1' _DELETE_TEAM = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = $1' _LOCK_SQL = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" +_HELD_LOCK_KEY_SQL = ( + "SELECT classid::bigint AS classid, objid::bigint AS objid FROM pg_locks " + "WHERE locktype = 'advisory' AND granted AND pid = pg_backend_pid()" +) +_LOCK_WAITER_SQL = ( + "SELECT count(*)::int AS waiters FROM pg_locks " + "WHERE locktype = 'advisory' AND NOT granted " + "AND classid::bigint = $1 AND objid::bigint = $2" +) +_LOCK_WAIT_TIMEOUT_SECONDS = 20.0 +_LOCK_POLL_SECONDS = 0.01 + + +async def _hold_team_lock(held, team_id: str) -> tuple[int, int]: + """Take the team's advisory lock and return its pg_locks key. + + Reading the key back off our own backend avoids re-deriving hashtext()'s signed + 32-bit split here, and pins the watcher to this lock rather than to any advisory + lock another xdist worker happens to hold on the same database.""" + await held.query_raw(_LOCK_SQL, team_id) + rows = await held.query_raw(_HELD_LOCK_KEY_SQL) + assert len(rows) == 1, f"expected exactly one advisory lock on the blocking connection, got {rows}" + return rows[0]["classid"], rows[0]["objid"] + + +async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what: str) -> None: + """Block until Postgres reports `task` queued behind the held lock. + + This is the assertion that the endpoint serializes on the team's advisory lock, and it + is what a fixed sleep was standing in for: the endpoint is only provably waiting once a + non-granted advisory lock on the same key exists. `watcher` must be a connection that is + not itself blocked, so it can observe the queue.""" + classid, objid = lock_key + deadline = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if task.done(): + raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() + if (await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid))[0]["waiters"]: + return + await asyncio.sleep(_LOCK_POLL_SECONDS) + raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s") def _race_ids() -> tuple[str, str]: @@ -110,10 +153,8 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def add_member(): - lock_acquired.set() await _add_team_members_to_team( data=TeamMemberAddRequest( team_id=team_id, @@ -128,11 +169,9 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(add_member()) - await lock_acquired.wait() - await asyncio.sleep(0.2) - assert not task.done(), "member_add did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "member_add") # the delete wins the race: strip the team row while the lock is held await held.execute_raw(_DELETE_TEAM, team_id) @@ -185,10 +224,8 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def run_delete(): - lock_acquired.set() return await team_member_delete( data=TeamMemberDeleteRequest(team_id=team_id, user_id=user_id), user_api_key_dict=_admin_auth(), @@ -196,11 +233,9 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await lock_acquired.wait() - await asyncio.sleep(0.2) - assert not task.done(), "member_delete did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "member_delete") # member_add wins the race: it adds `other_user` while holding the lock await held.litellm_teamtable.update( @@ -265,10 +300,8 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def run_delete(): - lock_acquired.set() return await delete_team( data=DeleteTeamRequest(team_ids=[team_id]), http_request=MagicMock(), @@ -278,11 +311,9 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await lock_acquired.wait() - await asyncio.sleep(0.3) - assert not task.done(), "delete_team did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "delete_team") # member_add wins the race: write the reference while holding the lock await held.litellm_usertable.upsert( diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 733870f3239..0ed33193a9b 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -18,6 +18,7 @@ import ast import json import logging import os +import signal import sys import time from collections.abc import Callable @@ -47,6 +48,7 @@ FAKE_PRISMA = """#!{python} import json import os import pathlib +import subprocess import sys import time @@ -66,6 +68,9 @@ with log_path.open("a") as log: time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: if earlier_same_command == 0: + if os.environ.get("FAKE_PRISMA_GRANDCHILD_PIDFILE"): + grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid)) time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) @@ -272,6 +277,43 @@ def test_migrate_deploy_stops_at_its_own_timeout( assert elapsed < 30 +def _process_is_gone(pid: int, within_seconds: float) -> bool: + deadline = time.monotonic() + within_seconds + while time.monotonic() < deadline: + try: + os.kill(pid, 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The real CLI forks Node and a schema engine; a timeout must not leave them running.""" + _, log_path = toolchain_env + pidfile = tmp_path / "grandchild.pid" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") + monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(pidfile)) + + with pytest.raises(RuntimeError, match="insufficient permissions"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + grandchild_pid = int(pidfile.read_text()) + try: + assert len(_deploy_calls(log_path)) == 2 + assert _process_is_gone(grandchild_pid, within_seconds=5) + finally: + try: + os.kill(grandchild_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + def test_db_push_timeout_hint_names_the_per_command_budget( toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 47be139eb5e..ded3be26630 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -61,6 +61,137 @@ async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_ assert "shared_b" not in dual_cache.last_redis_batch_access_time +def _redis_mock_for_sync_batch(redis_result: dict) -> MagicMock: + mock_redis = MagicMock(spec=RedisCache) + mock_redis.batch_get_cache.return_value = redis_result + return mock_redis + + +def _assert_sync_batch_used_blocking_client(dual_cache: DualCache, mock_redis: MagicMock) -> None: + with patch("asyncio.new_event_loop", side_effect=AssertionError("sync path must not create an event loop")): + result = dual_cache.batch_get_cache(keys=["lit6729_key"]) + + assert result == ["redis_value"] + mock_redis.batch_get_cache.assert_called_once_with(key_list=["lit6729_key"], parent_otel_span=None) + mock_redis.async_batch_get_cache.assert_not_called() + mock_redis.init_async_client.assert_not_called() + assert dual_cache.in_memory_cache.get_cache("lit6729_key") == "redis_value" + + +@pytest.mark.asyncio +async def test_dual_cache_batch_get_cache_uses_sync_redis_client_inside_running_loop(): + """ + Regression test for LIT-6729: sync batch_get_cache ran async_batch_get_cache on a + throwaway event loop, reusing an async Redis client created on another loop and + corrupting its connection pool. The sync path must use the blocking client, never + the async one, and never create an event loop, even when called from a coroutine + (e.g. async_raise_no_deployment_exception -> get_min_cooldown). + """ + mock_redis = _redis_mock_for_sync_batch({"lit6729_key": "redis_value"}) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis) + + _assert_sync_batch_used_blocking_client(dual_cache, mock_redis) + + +def test_dual_cache_batch_get_cache_uses_sync_redis_client_without_running_loop(): + mock_redis = _redis_mock_for_sync_batch({"lit6729_key": "redis_value"}) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis) + + _assert_sync_batch_used_blocking_client(dual_cache, mock_redis) + + +def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): + mock_redis = _redis_mock_for_sync_batch({"miss_key": "from_redis"}) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis) + dual_cache.in_memory_cache.set_cache("hit_key", "from_memory") + + result = dual_cache.batch_get_cache(keys=["hit_key", "miss_key"]) + + assert result == ["from_memory", "from_redis"] + mock_redis.batch_get_cache.assert_called_once_with(key_list=["miss_key"], parent_otel_span=None) + + +def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): + mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) + + first = dual_cache.batch_get_cache(keys=["absent_key"]) + second = dual_cache.batch_get_cache(keys=["absent_key"]) + + assert first == [None] + assert second == [None] + mock_redis.batch_get_cache.assert_called_once() + + +def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): + mock_redis = MagicMock(spec=RedisCache) + mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) + + first_result = dual_cache.batch_get_cache(keys=["shared_a"]) + second_result = dual_cache.batch_get_cache(keys=["shared_a"]) + + assert first_result is None + assert second_result is None + assert mock_redis.batch_get_cache.call_count == 2 + assert "shared_a" not in dual_cache.last_redis_batch_access_time + + +def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): + mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) + dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() + + result = dual_cache.batch_get_cache(keys=["throttled_key"]) + + assert result == [None] + mock_redis.batch_get_cache.assert_not_called() + + +def test_dual_cache_sync_batch_redis_backfill_injects_default_in_memory_ttl(): + """Sync batch_get_cache's Redis-to-memory backfill must honor + default_in_memory_ttl, same as the async path.""" + in_memory_cache = InMemoryCache(default_ttl=600) + mock_redis = _redis_mock_for_sync_batch({"batch_backfill_key": "redis_value"}) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + redis_cache=mock_redis, + default_in_memory_ttl=60, + ) + + before = time.time() + result = dual_cache.batch_get_cache(keys=["batch_backfill_key"]) + after = time.time() + + assert result == ["redis_value"] + expiry = in_memory_cache.ttl_dict["batch_backfill_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +def test_dual_cache_batch_get_cache_forwards_explicit_ttl_to_backfill(): + """An explicit ttl kwarg must reach the in-memory backfill flat, not nested + under a 'kwargs' key the way the old locals()-forwarding path sent it.""" + in_memory_cache = InMemoryCache(default_ttl=600) + mock_redis = _redis_mock_for_sync_batch({"explicit_ttl_key": "redis_value"}) + dual_cache = DualCache(in_memory_cache=in_memory_cache, redis_cache=mock_redis) + + before = time.time() + result = dual_cache.batch_get_cache(keys=["explicit_ttl_key"], ttl=5) + after = time.time() + + assert result == ["redis_value"] + expiry = in_memory_cache.ttl_dict["explicit_ttl_key"] + assert expiry >= before + 5 + assert expiry <= after + 5 + + @pytest.mark.asyncio async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl(): """ diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 487a64797d1..71be8730df1 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1,10 +1,10 @@ import asyncio -from unittest.mock import MagicMock, patch +from collections.abc import Iterator +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unittest.mock import AsyncMock - +from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import RedisCache @@ -17,6 +17,17 @@ def redis_no_ping(): yield +@pytest.fixture +def sync_batch_redis_cache(redis_no_ping): + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ) as get_client: + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.side_effect = OSError("redis unavailable") + get_client.assert_called_once() + yield cache + + @pytest.mark.parametrize( ("namespace", "key", "expected"), [ @@ -504,6 +515,173 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no await call_method(cache) +def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_redis_cache): + """An open breaker must preserve the sync batch read's dictionary fallback.""" + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} + + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} + + +@pytest.fixture +def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: + service_logger = ServiceLogging(mock_testing=True) + failing_client = MagicMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=failing_client + ): + cache = RedisCache(host="127.0.0.1", port=6379, service_logger_obj=service_logger) + yield cache, service_logger + + +@pytest.mark.asyncio +async def test_sync_batch_get_cache_reports_a_failed_read_from_a_running_loop( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], +): + """A swallowed Redis failure must still be reported as a service failure event. + + The routing strategies call this blocking read from inside the request's event loop, + and the read hides the Redis error by returning an empty dict. Without an emitted + failure event, litellm_redis_failed_requests_total stops moving during a Redis + outage while the success path keeps reporting, so the dashboards read healthy. + """ + cache, service_logger = sync_batch_cache_with_service_logger + + assert cache.batch_get_cache(key_list=["lit6729"]) == {} + await asyncio.sleep(0.05) + + assert service_logger.mock_testing_sync_failure_hook == 1 + assert service_logger.mock_testing_async_failure_hook == 1 + + +def test_sync_batch_get_cache_reports_a_failed_read_from_a_worker_thread( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], +): + """The same report must reach the async hook when the caller has no event loop at all.""" + from concurrent.futures import ThreadPoolExecutor + + cache, service_logger = sync_batch_cache_with_service_logger + + with ThreadPoolExecutor(max_workers=1) as pool: + assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {} + + assert service_logger.mock_testing_async_failure_hook == 1 + + +def test_sync_batch_get_cache_reports_a_failed_read_on_an_idle_event_loop( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], +): + """The report must also go out when the caller holds an open loop that is not running.""" + cache, service_logger = sync_batch_cache_with_service_logger + loop = asyncio.new_event_loop() + try: + asyncio.set_event_loop(loop) + assert cache.batch_get_cache(key_list=["lit6729"]) == {} + finally: + asyncio.set_event_loop(None) + loop.close() + + assert service_logger.mock_testing_async_failure_hook == 1 + + +def test_sync_batch_get_cache_survives_a_service_callback_that_raises( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], + monkeypatch: pytest.MonkeyPatch, +): + """A failing service callback must not replace the swallowed Redis failure. + + A misconfigured callback raises while emitting (a datadog callback with no + DD_API_KEY raises at construction), and the failure event is emitted from inside + the except block that swallows the Redis error. If that exception escapes, a Redis + outage surfaces to routing as a callback error and the circuit breaker never + records the failed read. + """ + from concurrent.futures import ThreadPoolExecutor + + import litellm + + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + + cache, service_logger = sync_batch_cache_with_service_logger + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + monkeypatch.setattr( + service_logger, + "init_prometheus_services_logger_if_none", + AsyncMock(side_effect=Exception("callback is misconfigured")), + ) + + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + with ThreadPoolExecutor(max_workers=1) as pool: + assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {} + + assert cache.batch_get_cache(key_list=["lit6729"]) == {} + + +def test_call_stack_info_skips_breaker_guard_frames(): + """Guarded methods must still report their real callers in service-log call_type. + + The breaker guards put their own frames between a method body and its caller, so + without skipping them every guarded method logged the guard machinery instead of + who actually issued the Redis call. + """ + from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _get_call_stack_info, + _redis_circuit_breaker_guard_sync, + ) + + class Guarded: + _circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + + @_redis_circuit_breaker_guard_sync + def probe(self): + return _get_call_stack_info() + + def caller_one(): + return Guarded().probe() + + def caller_two(): + return caller_one() + + assert caller_two() == "caller_one <- caller_two" + + +def test_call_stack_info_skips_guard_frames_when_deployed_without_sources(monkeypatch): + """Guard-frame skipping must survive a bytecode-only deployment. + + Shipping `.pyc` files without their `.py` sources leaves the module's `__file__` pointing + at the compiled file while every frame still carries the compile-time source path, so a + check comparing those two paths stops skipping and the service log then names the guard + machinery instead of the real caller. + """ + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _get_call_stack_info, + _redis_circuit_breaker_guard_sync, + ) + + monkeypatch.setattr(redis_cache_module, "__file__", redis_cache_module.__file__ + "c") + + class Guarded: + _circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + + @_redis_circuit_breaker_guard_sync + def probe(self): + return _get_call_stack_info() + + def caller_one(): + return Guarded().probe() + + def caller_two(): + return caller_one() + + assert caller_two() == "caller_one <- caller_two" + + @pytest.mark.asyncio async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ping): """A reachable Redis must keep the breaker closed, however many earlier calls failed. @@ -580,7 +758,6 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): async def swallows_a_failure(): await asyncio.sleep(0.02) _record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable")) - return None async def succeeds_while_the_other_fails(): await asyncio.sleep(0.05) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 21b60d7a216..6590718878d 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1295,6 +1295,19 @@ def test_text_plus_tool_calls_sequence(): # ============================================================================= +def test_developer_message_content_uses_input_text(): + handler = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [{"role": "developer", "content": "Always answer in French."}] + ) + + assert instructions is None + assert input_items == [ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "Always answer in French."}]} + ] + + def test_tool_message_output_uses_input_text_not_output_text(): """ Test that tool message content uses input_text type, not output_text. diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 176405bb9ca..38988d65c04 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -1,7 +1,7 @@ import json import sys from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException @@ -242,110 +242,156 @@ async def test_should_not_reassign_existing_container_to_different_owner(monkeyp table.update.assert_not_awaited() -@pytest.mark.asyncio -async def test_should_filter_container_list_to_owned_records(monkeypatch): +def _owned_containers_in_db(monkeypatch, *model_object_ids: str) -> AsyncMock: table = AsyncMock() - table.find_many.return_value = [ - SimpleNamespace(model_object_id="container:openai:cntr_owned"), - ] - prisma_client = SimpleNamespace( - db=SimpleNamespace(litellm_managedobjecttable=table) - ) + table.find_many.return_value = [SimpleNamespace(model_object_id=object_id) for object_id in model_object_ids] monkeypatch.setattr( ownership, "_get_prisma_client", - AsyncMock(return_value=prisma_client), + AsyncMock(return_value=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))), ) - auth = UserAPIKeyAuth(user_id="user-1") - response = ContainerListResponse( + return table + + +def _upstream(pages_by_after): + calls = [] + + async def fetch_page(after, limit): + calls.append((after, limit)) + return pages_by_after[after] + + return fetch_page, calls + + +def _page(*container_ids: str, has_more: bool) -> ContainerListResponse: + return ContainerListResponse( object="list", - data=[_container("cntr_owned"), _container("cntr_other")], - has_more=True, + data=[_container(container_id) for container_id in container_ids], + has_more=has_more, ) - filtered = await ownership.filter_container_list_response( - response=response, - user_api_key_dict=auth, + +async def _list_owned(fetch_page, after=None, limit=None): + return await ownership.list_owned_containers( + fetch_page=fetch_page, + after=after, + limit=limit, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), custom_llm_provider="openai", ) - assert [item.id for item in filtered.data] == ["cntr_owned"] - assert filtered.first_id == "cntr_owned" - assert filtered.last_id == "cntr_owned" - assert filtered.has_more is False + +@pytest.mark.asyncio +async def test_should_page_upstream_until_owned_containers_fill_the_limit(monkeypatch): + table = _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + fetch_page, calls = _upstream( + { + None: _page("cntr_other_1", "cntr_other_2", has_more=True), + "cntr_other_2": _page("cntr_owned", has_more=False), + } + ) + + listed = await _list_owned(fetch_page, limit=1) + + assert [item.id for item in listed.data] == ["cntr_owned"] + assert listed.first_id == "cntr_owned" + assert listed.last_id == "cntr_owned" + assert listed.has_more is False + assert calls == [(None, 100), ("cntr_other_2", 100)] where = table.find_many.await_args.kwargs["where"] assert where["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE assert where["created_by"]["in"] == ["user-1", "user:user-1"] @pytest.mark.asyncio -async def test_should_clear_has_more_when_filtered_container_list_is_empty( - monkeypatch, -): - table = AsyncMock() - table.find_many.return_value = [ - SimpleNamespace(model_object_id="container:openai:cntr_owned"), - ] - prisma_client = SimpleNamespace( - db=SimpleNamespace(litellm_managedobjecttable=table) - ) - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=prisma_client), - ) - auth = UserAPIKeyAuth(user_id="user-1") - response = ContainerListResponse( - object="list", - data=[_container("cntr_other")], - has_more=True, - ) +async def test_should_trim_owned_containers_to_the_limit_without_mutating_the_upstream_page(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned_1", "container:openai:cntr_owned_2") + upstream_page = _page("cntr_owned_1", "cntr_other", "cntr_owned_2", has_more=False) + fetch_page, calls = _upstream({None: upstream_page}) - filtered = await ownership.filter_container_list_response( - response=response, - user_api_key_dict=auth, - custom_llm_provider="openai", - ) + listed = await _list_owned(fetch_page, limit=1) - assert filtered.data == [] - assert filtered.first_id is None - assert filtered.last_id is None - assert filtered.has_more is False + assert [item.id for item in listed.data] == ["cntr_owned_1"] + assert listed.first_id == "cntr_owned_1" + assert listed.last_id == "cntr_owned_1" + assert listed.has_more is True + assert calls == [(None, 100)] + assert [item.id for item in upstream_page.data] == ["cntr_owned_1", "cntr_other", "cntr_owned_2"] + assert upstream_page.has_more is False @pytest.mark.asyncio -async def test_should_clear_dict_has_more_when_filtered_container_list_is_empty( - monkeypatch, -): - table = AsyncMock() - table.find_many.return_value = [ - SimpleNamespace(model_object_id="container:openai:cntr_owned"), - ] - prisma_client = SimpleNamespace( - db=SimpleNamespace(litellm_managedobjecttable=table) +async def test_should_start_paging_from_the_requested_cursor(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned_2") + fetch_page, calls = _upstream({"cntr_owned_1": _page("cntr_other", "cntr_owned_2", has_more=False)}) + + listed = await _list_owned(fetch_page, after="cntr_owned_1", limit=1) + + assert [item.id for item in listed.data] == ["cntr_owned_2"] + assert listed.has_more is False + assert calls == [("cntr_owned_1", 100)] + + +@pytest.mark.asyncio +async def test_should_default_to_twenty_owned_containers_per_page(monkeypatch): + owned_ids = tuple(f"cntr_owned_{index}" for index in range(21)) + _owned_containers_in_db(monkeypatch, *(f"container:openai:{container_id}" for container_id in owned_ids)) + fetch_page, _ = _upstream({None: _page(*owned_ids, has_more=False)}) + + listed = await _list_owned(fetch_page) + + assert [item.id for item in listed.data] == list(owned_ids[:20]) + assert listed.last_id == "cntr_owned_19" + assert listed.has_more is True + + +@pytest.mark.asyncio +async def test_should_stop_after_five_upstream_pages_and_keep_has_more(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + fetch_page, calls = _upstream( + { + None: _page("cntr_other_0", has_more=True), + **{f"cntr_other_{index}": _page(f"cntr_other_{index + 1}", has_more=True) for index in range(6)}, + } ) - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=prisma_client), - ) - auth = UserAPIKeyAuth(user_id="user-1") - response = { + + listed = await _list_owned(fetch_page, limit=1) + + assert listed.data == [] + assert listed.first_id is None + assert listed.last_id is None + assert listed.has_more is True + assert len(calls) == 5 + + +@pytest.mark.asyncio +async def test_should_stop_when_upstream_has_no_more_pages(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + fetch_page, calls = _upstream({None: _page("cntr_other", has_more=False)}) + + listed = await _list_owned(fetch_page, limit=1) + + assert listed.data == [] + assert listed.has_more is False + assert calls == [(None, 100)] + + +@pytest.mark.asyncio +async def test_should_build_dict_pages_without_mutating_the_upstream_page(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + upstream_page = {"object": "list", "data": [{"id": "cntr_other"}, {"id": "cntr_owned"}], "has_more": False} + fetch_page, _ = _upstream({None: upstream_page}) + + listed = await _list_owned(fetch_page, limit=1) + + assert listed == { "object": "list", - "data": [{"id": "cntr_other"}], - "has_more": True, + "data": [{"id": "cntr_owned"}], + "first_id": "cntr_owned", + "last_id": "cntr_owned", + "has_more": False, } - - filtered = await ownership.filter_container_list_response( - response=response, - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - - assert filtered["data"] == [] - assert filtered["first_id"] is None - assert filtered["last_id"] is None - assert filtered["has_more"] is False + assert [item["id"] for item in upstream_page["data"]] == ["cntr_other", "cntr_owned"] @pytest.mark.asyncio @@ -647,7 +693,7 @@ async def test_should_return_response_when_owner_recording_raises_unexpected( @pytest.mark.asyncio -async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): +async def test_should_list_owned_containers_inside_list_endpoint(monkeypatch): from litellm.proxy.container_endpoints import endpoints proxy_server_stub = SimpleNamespace( @@ -665,42 +711,37 @@ async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): ) monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) - response = ContainerListResponse( - object="list", - data=[_container("cntr_provider")], - has_more=False, - ) - - class FakeProcessor: - def __init__(self, data): - pass - - async def base_process_llm_request(self, **kwargs): - return response - - async def _handle_llm_api_exception(self, **kwargs): - raise kwargs["e"] - - filter_response = AsyncMock(return_value=response) - monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) - monkeypatch.setattr( - endpoints, - "filter_container_list_response", - filter_response, + upstream_page = _page("cntr_provider", has_more=False) + processor_cls = MagicMock( + side_effect=lambda data: SimpleNamespace(base_process_llm_request=AsyncMock(return_value=upstream_page)) ) + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", processor_cls) + list_owned = AsyncMock(return_value=upstream_page) + monkeypatch.setattr(endpoints, "list_owned_containers", list_owned) result = await endpoints.list_containers( request=SimpleNamespace(query_params={}, headers={}), fastapi_response=SimpleNamespace(), user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + after="cntr_prev", + limit=2, + order="desc", ) - assert result == response - filter_response.assert_awaited_once_with( - response=response, - user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), - custom_llm_provider="openai", - ) + assert result == upstream_page + kwargs = list_owned.await_args.kwargs + assert kwargs["after"] == "cntr_prev" + assert kwargs["limit"] == 2 + assert kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_id="user-1") + assert kwargs["custom_llm_provider"] == "openai" + processor_cls.assert_not_called() + + assert await kwargs["fetch_page"]("cntr_page_cursor", 100) == upstream_page + forwarded = processor_cls.call_args.kwargs["data"] + assert forwarded["after"] == "cntr_page_cursor" + assert forwarded["limit"] == 100 + assert forwarded["order"] == "desc" + assert forwarded["custom_llm_provider"] == "openai" @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 7d70b9a8862..6edc2c9bf77 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1091,8 +1091,8 @@ class TestCustomGuardrailPassthroughSupport: call_type=CallTypes.allm_passthrough_route, ) - # When result is None, should return the original response - assert result == mock_response + # None means the guardrail did not modify the response (LIT-5863 contract) + assert result is None @pytest.mark.asyncio async def test_async_post_call_success_deployment_hook_with_none_call_type(self): @@ -1120,8 +1120,8 @@ class TestCustomGuardrailPassthroughSupport: call_type=None, ) - # Should return the original response when result is None - assert result == mock_response + # None means the guardrail did not modify the response (LIT-5863 contract) + assert result is None def test_is_valid_response_type_with_none(self): """ @@ -2436,3 +2436,73 @@ class TestLoggingOnlyApplyGuardrail: assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] entries = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] assert [e["guardrail_status"] for e in entries] == ["success", "success"] + + +class TestCustomGuardrailPostCallSuccessDeploymentHook: + """Regression tests for LIT-5863: this hook answering the unmodified response instead of + None made the utils.py dispatcher treat the guardrail as having modified the response, + which starved every later callback in litellm.callbacks (notably the lazily-appended + VectorStorePreCallHook that attaches provider_specific_fields["search_results"]).""" + + @pytest.mark.asyncio + async def test_returns_none_when_request_has_no_guardrails(self): + from litellm.types.utils import ModelResponse + + guardrail = CustomGuardrail(guardrail_name="test-guardrail") + response = ModelResponse() + + assert ( + await guardrail.async_post_call_success_deployment_hook( + request_data={}, response=response, call_type=CallTypes.acompletion + ) + is None + ) + assert ( + await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": "not-a-list"}, response=response, call_type=CallTypes.acompletion + ) + is None + ) + + @pytest.mark.asyncio + async def test_returns_none_when_guardrail_should_not_run(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + guardrail = CustomGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + response = ModelResponse() + + result = await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": ["test-guardrail"]}, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is None + + @pytest.mark.asyncio + async def test_returns_modified_response_when_guardrail_runs(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + replacement = ModelResponse() + + class ReplacingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return replacement + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + + result = await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": ["test-guardrail"]}, + response=ModelResponse(), + call_type=CallTypes.acompletion, + ) + + assert result is replacement diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py new file mode 100644 index 00000000000..ae5cffd8ab0 --- /dev/null +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -0,0 +1,287 @@ +import logging +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import Protocol + +import pytest + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + ProxyServerRuntime, + VectorStorePreCallHook, +) +from litellm.types.llms.openai import AllMessageValues +from litellm.types.vector_stores import ( + VectorStoreResultContent, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) +from litellm.vector_stores.vector_store_registry import ( + LiteLLM_ManagedVectorStore, + VectorStoreRegistry, +) + + +def _search_response(text: str) -> VectorStoreSearchResponse: + return VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query="what is litellm?", + data=[ + VectorStoreSearchResult( + score=1.0, + content=[VectorStoreResultContent(text=text, type="text")], + ) + ], + ) + + +@dataclass +class RecordingRouter: + failing_vector_store_ids: frozenset[str] = frozenset() + calls: list[dict[str, object]] = field(default_factory=list) + + async def avector_store_search(self, **kwargs: object) -> VectorStoreSearchResponse: + self.calls.append(kwargs) + vector_store_id = str(kwargs["vector_store_id"]) + if vector_store_id in self.failing_vector_store_ids: + raise litellm.BadRequestError( + message=f"no healthy deployments for {vector_store_id}", + model="text-embedding-3-small", + llm_provider="openai", + ) + return _search_response(f"context from {vector_store_id}") + + +@dataclass(frozen=True) +class FakeProxyRuntime: + router: RecordingRouter | None + + def llm_router(self) -> RecordingRouter | None: + return self.router + + def prisma_client(self) -> None: + return None + + +class RecordingHandler(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + +class RegisterStores(Protocol): + def __call__(self, *vector_store_ids: str, custom_llm_provider: str = "bedrock") -> None: ... + + +@pytest.fixture +def registry_with(monkeypatch: pytest.MonkeyPatch) -> RegisterStores: + def _register(*vector_store_ids: str, custom_llm_provider: str = "bedrock") -> None: + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + vector_stores=[ + LiteLLM_ManagedVectorStore(vector_store_id=vector_store_id, custom_llm_provider=custom_llm_provider) + for vector_store_id in vector_store_ids + ], + ), + ) + + return _register + + +@pytest.fixture +def warnings() -> Iterator[list[logging.LogRecord]]: + handler = RecordingHandler() + verbose_logger.addHandler(handler) + yield handler.records + verbose_logger.removeHandler(handler) + + +class FakeLoggingObj: + def __init__(self, metadata: dict[str, str]) -> None: + self.model_call_details: dict[str, object] = {"litellm_params": {"metadata": metadata}} + + +async def _run_hook( + hook: VectorStorePreCallHook, + vector_store_ids: list[str], + logging_obj: FakeLoggingObj, +) -> tuple[str, list[AllMessageValues], dict[str, object]]: + return await hook.async_get_chat_completion_prompt( + model="chat-model", + messages=[{"role": "user", "content": "what is litellm?"}], + non_default_params={"vector_store_ids": vector_store_ids}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + litellm_logging_obj=logging_obj, + ) + + +@pytest.mark.asyncio +async def test_hook_searches_through_the_injected_router_with_the_request_metadata( + registry_with: RegisterStores, +) -> None: + """Regression (LIT-6752): the hook must reach the Router through its injected runtime, not a proxy_server import.""" + registry_with("vs-router") + router = RecordingRouter() + logging_obj = FakeLoggingObj({"user_api_key_team_id": "team-a"}) + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=router)), + ["vs-router"], + logging_obj, + ) + + assert router.calls == [ + { + "vector_store_id": "vs-router", + "query": "what is litellm?", + "custom_llm_provider": "bedrock", + "metadata": {"user_api_key_team_id": "team-a"}, + } + ] + assert messages[0]["content"] == "Context:\n\ncontext from vs-router\n\n" + + +@pytest.mark.asyncio +async def test_hook_falls_back_to_the_sdk_when_the_runtime_has_no_router( + registry_with: RegisterStores, + warnings: list[logging.LogRecord], +) -> None: + registry_with("vs-sdk", custom_llm_provider="lit6752-not-a-provider") + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)), + ["vs-sdk"], + FakeLoggingObj({"user_api_key_team_id": "team-a"}), + ) + + assert messages == [{"role": "user", "content": "what is litellm?"}] + assert len(warnings) == 1 + assert ( + warnings[0] + .getMessage() + .startswith("Vector store search failed for vector_store_id=vs-sdk, continuing without its context: ") + ) + assert "is not a valid LlmProviders" in warnings[0].getMessage() + + +@pytest.mark.asyncio +async def test_every_healthy_vector_store_contributes_its_own_context(registry_with: RegisterStores) -> None: + """Regression (LIT-6752): each store appended its context to the original messages, so only the last one survived.""" + registry_with("vs-one", "vs-two") + router = RecordingRouter() + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=router)), + ["vs-one", "vs-two"], + FakeLoggingObj({}), + ) + + assert [message["content"] for message in messages] == [ + "Context:\n\ncontext from vs-one\n\n", + "Context:\n\ncontext from vs-two\n\n", + "what is litellm?", + ] + + +@pytest.mark.asyncio +async def test_a_failing_vector_store_warns_with_its_id_and_the_other_stores_still_answer( + registry_with: RegisterStores, + warnings: list[logging.LogRecord], +) -> None: + """Regression (LIT-6752): one unreachable store must not silently drop every other store's context.""" + registry_with("vs-broken", "vs-healthy") + router = RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})) + logging_obj = FakeLoggingObj({"user_api_key_team_id": "team-a"}) + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=router)), + ["vs-broken", "vs-healthy"], + logging_obj, + ) + + search_results = logging_obj.model_call_details["search_results"] + + assert [call["vector_store_id"] for call in router.calls] == ["vs-broken", "vs-healthy"] + assert messages[0]["content"] == "Context:\n\ncontext from vs-healthy\n\n" + assert isinstance(search_results, list) + assert len(search_results) == 1 + assert [record.getMessage() for record in warnings] == [ + "Vector store search failed for vector_store_id=vs-broken, continuing without its context: " + "litellm.BadRequestError: no healthy deployments for vs-broken" + ] + + +@pytest.mark.asyncio +async def test_the_only_vector_store_failing_leaves_the_messages_untouched( + registry_with: RegisterStores, + warnings: list[logging.LogRecord], +) -> None: + registry_with("vs-broken") + original_messages = [{"role": "user", "content": "what is litellm?"}] + + _, messages, _ = await _run_hook( + VectorStorePreCallHook( + proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))) + ), + ["vs-broken"], + FakeLoggingObj({}), + ) + + assert messages == original_messages + assert [(record.levelname, record.getMessage()) for record in warnings] == [ + ( + "WARNING", + "Vector store search failed for vector_store_id=vs-broken, continuing without its context: " + "litellm.BadRequestError: no healthy deployments for vs-broken", + ) + ] + + +@pytest.mark.asyncio +async def test_the_default_hook_reaches_the_proxy_router_through_its_runtime( + registry_with: RegisterStores, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-6752): a hook built with no arguments must still search through the proxy's own Router.""" + from litellm.proxy import proxy_server + + registry_with("vs-default") + router = RecordingRouter() + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "prisma_client", None) + + _, messages, _ = await _run_hook( + VectorStorePreCallHook(), + ["vs-default"], + FakeLoggingObj({"user_api_key_team_id": "team-a"}), + ) + + assert [call["vector_store_id"] for call in router.calls] == ["vs-default"] + assert messages[0]["content"] == "Context:\n\ncontext from vs-default\n\n" + + +def test_the_default_runtime_follows_the_proxy_globals(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + runtime = ProxyServerRuntime() + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + + assert runtime.llm_router() is None + assert runtime.prisma_client() is None + + router = RecordingRouter() + prisma = object() + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + assert runtime.llm_router() is router + assert runtime.prisma_client() is prisma diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index b7f0ca1efe1..0b6832d4bef 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1662,6 +1662,54 @@ def test_generic_cost_per_token_gpt56_cyber( assert completion_cost == pytest.approx(completion_tokens * output_rate) +@pytest.mark.parametrize( + "service_tier,tier_multiplier", + [(None, 1.0), ("flex", 0.5), ("priority", 2.0), ("fast", 2.0)], +) +@pytest.mark.parametrize( + "prompt_tokens,input_side_multiplier,output_multiplier", + [(100000, 1.0, 1.0), (300000, 2.0, 1.5)], +) +def test_generic_cost_per_token_gpt_6_astra_price_sheet( + _local_model_cost_map, + service_tier, + tier_multiplier, + prompt_tokens, + input_side_multiplier, + output_multiplier, +): + """gpt-6-astra launch price sheet: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens. + + Above 272K prompt tokens the input-side rates double and the output rate is 1.5x on the whole + request. Flex is half the applicable rate and fast mode, billed as priority, is double it. + """ + cached_tokens = 50000 + cache_write_tokens = 40000 + text_tokens = prompt_tokens - cached_tokens - cache_write_tokens + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="gpt-6-astra", + usage=usage, + custom_llm_provider="openai", + service_tier=service_tier, + ) + + input_side = tier_multiplier * input_side_multiplier + assert prompt_cost == pytest.approx( + input_side * (text_tokens * 1e-5 + cached_tokens * 1e-6 + cache_write_tokens * 1.25e-5) + ) + assert completion_cost == pytest.approx(tier_multiplier * output_multiplier * completion_tokens * 5e-5) + + @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost", [ diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index fd795ffcc96..6cc3dcceebc 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,5 @@ import os +from collections.abc import Mapping, Sequence import pytest @@ -6,7 +7,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) -from litellm.types.llms.openai import FileSearchTool, WebSearchOptions +from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams @@ -928,3 +929,125 @@ def test_web_search_gate_reads_server_side_tool_usage_details_without_citations( standard_built_in_tools_params=None, ) assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL + + +_BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + "bedrock_mantle/openai.gpt-5.5", + "bedrock_mantle/openai.gpt-5.4", +) + +_BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +def _responses_with_web_search( + model: str, actions: Sequence[Mapping[str, str]], tool_usage: Mapping[str, object] | None = None +) -> ResponsesAPIResponse: + payload = { + "id": "resp_1", + "created_at": 1756900000, + "model": model.split("/", 1)[-1], + "object": "response", + "status": "completed", + "output": [ + {"type": "web_search_call", "id": f"ws_{i}", "status": "completed", "action": action} + for i, action in enumerate(actions) + ], + } + return ResponsesAPIResponse.model_validate( + payload if tool_usage is None else {**payload, "tool_usage": tool_usage} + ) + + +def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_provider: str) -> float: + from litellm.types.utils import Usage + + return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=None, + ) + + +@pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS) +def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model): + """Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike.""" + pricing = litellm.get_model_info(model)["search_context_cost_per_query"] + assert pricing == { + "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + "search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + } + + response = _responses_with_web_search( + model, + actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], + tool_usage={"web_search": {"num_requests": 2}}, + ) + for cost_model in (model, model.split("/", 1)[1]): + cost = _web_search_cost(cost_model, response, "bedrock_mantle") + assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}" + ) + + +@pytest.mark.parametrize("num_requests", [1, 0]) +def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests): + """A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items.""" + model = "bedrock_mantle/openai.gpt-5.6-sol" + response = _responses_with_web_search( + model, + actions=[ + {"type": "search", "query": "litellm"}, + {"type": "open_page", "url": "https://docs.litellm.ai/"}, + ], + tool_usage={"web_search": {"num_requests": num_requests}}, + ) + + cost = _web_search_cost(model, response, "bedrock_mantle") + + assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"{num_requests} reported web search requests must bill {num_requests} x " + f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" + ) + + +@pytest.mark.parametrize( + "tool_usage", + [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}], +) +def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage): + """Without a usable reported count the per-call path keeps counting web_search_call items.""" + model = "bedrock_mantle/openai.gpt-5.6-sol" + response = _responses_with_web_search( + model, + actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], + tool_usage=tool_usage, + ) + + cost = _web_search_cost(model, response, "bedrock_mantle") + + assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x " + f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" + ) + + +def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map): + """OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count.""" + response = _responses_with_web_search( + "gpt-5.6", + actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}], + tool_usage={ + "image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + "web_search": {"num_requests": 1}, + }, + ) + + cost = _web_search_cost("gpt-5.6", response, "openai") + + assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}" diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index f6b8a93c472..fadc4ca49e9 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -312,3 +312,121 @@ def test_mask_credentials_in_payload_masks_only_sensitive_string_leaves(): assert masked != plaintext assert masked.startswith(plaintext[:4]) assert masked.endswith(plaintext[-4:]) + + +def test_redact_credentials_in_payload_leaves_no_fragment_of_the_secret(): + """A payload rendered straight to stdout cannot afford the partial reveal + mask_credentials_in_payload leaves, so every credential-named value is replaced + whole, nested header dicts included, while ordinary params survive verbatim.""" + from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload + + fake_key = "sk-fake-lit6823-0000000000000000" + fake_token = "fake-azure-ad-token-0000" + result = redact_credentials_in_payload( + { + "api_key": fake_key, + "azure_ad_token": fake_token, + "aws_secret_access_key": "fake-aws-secret-0000", + "vertex_credentials": {"private_key": "fake-pem"}, + "extra_headers": {"Authorization": "Bearer fake-bearer-0000", "x-request-id": "abc123"}, + "model": "gpt-4o-mini", + "max_tokens": 17, + "temperature": 0.25, + "api_base": None, + } + ) + + assert fake_key not in str(result) + assert fake_token not in str(result) + assert "fake-aws-secret-0000" not in str(result) + assert "fake-pem" not in str(result) + assert "fake-bearer-0000" not in str(result) + assert result["api_key"] == "REDACTED" + assert result["extra_headers"]["Authorization"] == "REDACTED" + assert result["extra_headers"]["x-request-id"] == "abc123" + assert result["model"] == "gpt-4o-mini" + assert result["max_tokens"] == 17 + assert result["temperature"] == 0.25 + assert result["api_base"] is None + + +def test_redact_credentials_in_payload_reaches_credentials_nested_in_sequences(): + """Free-form kwargs like extra_body and metadata routinely carry lists of dicts, so a + credential hiding one level inside a list or tuple must be replaced too, while the + surrounding container keeps its type and every ordinary element stays verbatim.""" + from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload + + result = redact_credentials_in_payload( + { + "extra_body": {"providers": [{"name": "openai", "api_key": "sk-fake-lit6823-in-a-list"}]}, + "metadata": {"upstreams": ({"aws_secret_access_key": "fake-aws-in-a-tuple"},)}, + "messages": [{"role": "user", "content": "hello"}], + } + ) + + assert "sk-fake-lit6823-in-a-list" not in str(result) + assert "fake-aws-in-a-tuple" not in str(result) + assert result["extra_body"]["providers"][0]["api_key"] == "REDACTED" + assert result["extra_body"]["providers"][0]["name"] == "openai" + assert isinstance(result["extra_body"]["providers"], list) + assert result["metadata"]["upstreams"][0]["aws_secret_access_key"] == "REDACTED" + assert isinstance(result["metadata"]["upstreams"], tuple) + assert result["messages"] == [{"role": "user", "content": "hello"}] + + +@pytest.mark.parametrize("wrap", ["mapping", "sequence"]) +def test_redact_credentials_in_payload_hides_containers_at_the_recursion_limit(wrap): + """The recursion limit exists to bound the walk, not to grant an exemption, so a caller who + buries a credential deeper than the limit must get the container hidden rather than handed + back verbatim. Nesting through lists costs depth twice as fast as nesting through mappings, + so both shapes are pushed well past the limit here.""" + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload + + fake_key = "sk-fake-lit6835-past-the-limit" + node = {"api_key": fake_key} + for _ in range(2 * DEFAULT_MAX_RECURSE_DEPTH + 1): + node = {"extra_body": node} if wrap == "mapping" else {"providers": [node]} + + result = redact_credentials_in_payload({**node, "max_tokens": 17}) + + assert fake_key not in str(result) + assert "REDACTED" in str(result) + assert result["max_tokens"] == 17 + + +def test_redact_credentials_in_payload_leaves_a_realistic_tool_schema_intact(): + """The bound must not eat ordinary payloads: a tool whose JSON schema nests an array of + objects inside a nested object is what agent traffic looks like, and the verbose line is + useless if those leaves come back as REDACTED.""" + from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload + + tool = { + "type": "function", + "function": { + "name": "search_orders", + "parameters": { + "type": "object", + "properties": { + "filters": { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": { + "type": "object", + "properties": {"sku": {"type": "string"}, "qty": {"type": "integer"}}, + }, + } + }, + } + }, + }, + }, + } + + result = redact_credentials_in_payload({"model": "gpt-4o-mini", "tools": [tool], "api_key": "sk-fake-lit6835"}) + + assert "REDACTED" not in str(result["tools"]) + assert result["tools"][0] == tool + assert result["api_key"] == "REDACTED" diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 4b2a3f9247c..0aa73833677 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4761,6 +4761,43 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX" +@pytest.mark.asyncio +async def test_async_fake_stream_final_chunk_carries_hidden_usage(logging_obj: Logging): + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.types.utils import ModelResponse + + model_response = ModelResponse( + id="chatcmpl-fake-stream", + model="my-random-model", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "hello world"}, + "finish_reason": "stop", + } + ], + ) + model_response.usage = Usage(prompt_tokens=1234, completion_tokens=7, total_tokens=1241) + + wrapper = CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=model_response), + model="my-random-model", + custom_llm_provider="anthropic", + logging_obj=logging_obj, + ) + + final_chunk = None + async for chunk in wrapper: + final_chunk = chunk + + assert final_chunk is not None + hidden_usage = final_chunk._hidden_params.get("usage") + assert hidden_usage is not None + assert hidden_usage.prompt_tokens == 1234 + assert hidden_usage.completion_tokens == 7 + assert hidden_usage.total_tokens == 1241 + + class TestStableStreamingResponseId: """ All chunks of one streamed response must share the same top-level id diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py new file mode 100644 index 00000000000..7cd789529c8 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py @@ -0,0 +1,111 @@ +""" +Streaming ``/v1/messages`` against a model that is neither Anthropic nor OpenAI is served by +translating the call onto ``/v1/chat/completions``, and the ``msg_`` id the caller is streamed +is minted right here. It is the only request id such a caller ever sees, so the spend row has +to be keyed on that same value rather than on the provider's own completion id. +""" + +import datetime +import json + +import pytest +import respx + +import litellm +from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, +) +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) + +MESSAGES = [{"role": "user", "content": "hello"}] + +GROQ_CHAT_URL = "https://api.groq.com/openai/v1/chat/completions" + +CHAT_SSE_BODY = ( + b'data: {"id":"chatcmpl-lit6825","object":"chat.completion.chunk","created":1,' + b'"model":"kimi-k2","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},' + b'"finish_reason":null}]}\n\n' + b'data: {"id":"chatcmpl-lit6825","object":"chat.completion.chunk","created":1,' + b'"model":"kimi-k2","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],' + b'"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}\n\n' + b"data: [DONE]\n\n" +) + + +def _logging_obj(call_id: str): + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="kimi-k2", + messages=MESSAGES, + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id=call_id, + function_id="1234", + ) + + +def _streamed_message_id(raw_events: list[bytes]) -> str: + events = [json.loads(chunk.decode().split("data: ", 1)[1]) for chunk in raw_events] + message_start = next(e for e in events if e["type"] == "message_start") + return message_start["message"]["id"] + + +@pytest.fixture(autouse=True) +def _intercept_groq(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("GROQ_API_KEY", "gsk-lit6825-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(GROQ_CHAT_URL).respond( + status_code=200, + headers={"Content-Type": "text/event-stream"}, + content=CHAT_SSE_BODY, + ) + + +@pytest.mark.asyncio +async def test_async_streaming_hands_the_logging_object_the_message_id_the_caller_is_streamed(): + logging_obj = _logging_obj("6825beef-0000-4000-8000-000000000010") + + sse = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=1024, + messages=MESSAGES, + model="groq/kimi-k2", + stream=True, + custom_llm_provider="groq", + litellm_logging_obj=logging_obj, + ) + streamed_id = _streamed_message_id([chunk async for chunk in sse]) + + assert streamed_id.startswith("msg_") + assert logging_obj.streamed_anthropic_message_id == streamed_id + + +def test_sync_streaming_hands_the_logging_object_the_message_id_the_caller_is_streamed(): + logging_obj = _logging_obj("6825beef-0000-4000-8000-000000000011") + + sse = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=1024, + messages=MESSAGES, + model="groq/kimi-k2", + stream=True, + custom_llm_provider="groq", + litellm_logging_obj=logging_obj, + ) + streamed_id = _streamed_message_id(list(sse)) + + assert streamed_id.startswith("msg_") + assert logging_obj.streamed_anthropic_message_id == streamed_id + + +def test_concurrent_streams_are_keyed_on_their_own_message_id(): + """Two callers streaming at once must not be handed, or logged under, one another's id.""" + first = AnthropicStreamWrapper(completion_stream=iter([]), model="kimi-k2") + second = AnthropicStreamWrapper(completion_stream=iter([]), model="kimi-k2") + + assert first._message_id != second._message_id + assert _streamed_message_id(list(first.anthropic_sse_wrapper())) == first._message_id + assert _streamed_message_id(list(second.anthropic_sse_wrapper())) == second._message_id diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py index 589dc64f9b9..3383813245a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py @@ -1,9 +1,11 @@ +import datetime import json import os import sys from unittest.mock import AsyncMock, patch import pytest +import respx sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) @@ -15,6 +17,18 @@ from litellm.llms.anthropic.experimental_pass_through.responses_adapters.handler MESSAGES = [{"role": "user", "content": "hello"}] +RESPONSES_SSE_BODY = ( + b"event: response.created\n" + b'data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_lit6825",' + b'"object":"response","created_at":1,"status":"in_progress","model":"gpt-5.6-luna","output":[],' + b'"parallel_tool_calls":true,"tool_choice":"auto","tools":[]}}\n\n' + b"event: response.completed\n" + b'data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_lit6825",' + b'"object":"response","created_at":1,"status":"completed","model":"gpt-5.6-luna","output":[],' + b'"parallel_tool_calls":true,"tool_choice":"auto","tools":[],' + b'"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7}}}\n\n' +) + def test_build_responses_kwargs_derives_prompt_cache_key_from_user_id(): responses_kwargs = _build_responses_kwargs( @@ -82,3 +96,47 @@ async def test_streaming_message_start_reports_the_provider_local_model(requeste message_start = next(e for e in events if e["type"] == "message_start") assert message_start["message"]["model"] == expected_reported_model + + +@pytest.mark.asyncio +async def test_streaming_hands_the_logging_object_the_message_id_the_caller_is_streamed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + """ + The bridge mints the ``msg_`` id itself, and it is the only request id a streaming + /v1/messages caller ever sees, so the spend row has to be keyed on that same value. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setenv("OPENAI_API_KEY", "sk-lit6825-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").respond( + status_code=200, + headers={"Content-Type": "text/event-stream"}, + content=RESPONSES_SSE_BODY, + ) + + logging_obj = Logging( + model="gpt-5.6-luna", + messages=MESSAGES, + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id="6825beef-0000-4000-8000-000000000003", + function_id="1234", + ) + + sse = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + stream=True, + custom_llm_provider="openai", + litellm_logging_obj=logging_obj, + ) + events = [json.loads(chunk.decode().split("data: ", 1)[1]) async for chunk in sse] + + message_start = next(e for e in events if e["type"] == "message_start") + assert message_start["message"]["id"].startswith("msg_") + assert logging_obj.streamed_anthropic_message_id == message_start["message"]["id"] diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index e06cae97283..bd0f16a695b 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -336,3 +336,15 @@ class TestAzureResolvesTheDeclaredDefaultEffort: drop_params=True, ) assert ("temperature" in mapped) is temperature_survives + + +def test_azure_gpt_6_astra_takes_the_reasoning_series_request_shape(): + params = litellm.get_optional_params( + model="gpt-6-astra", + custom_llm_provider="azure", + max_tokens=100, + reasoning_effort="max", + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + assert params["reasoning_effort"] == "max" diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 11a727c9635..33fbb4e8fc7 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -31,6 +31,46 @@ async def test_get_openai_compatible_provider_info(): assert custom_llm_provider == "azure" +@pytest.mark.parametrize( + "model, api_base, expected_provider", + [ + ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com/models", "azure_ai"), + ("azure_ai/gpt-5.4-nano", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/gpt-4o", "https://my-resource.openai.azure.com", "azure"), + ( + "azure_ai/gpt-4o", + "https://my-resource.services.ai.azure.com/openai/deployments/gpt-4o/chat/completions" + "?api-version=2024-08-01-preview", + "azure", + ), + ("azure_ai/mistral-large-latest", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/mistral-large-latest", "https://my-resource.openai.azure.com", "azure_ai"), + ], +) +def test_foundry_base_keeps_azure_ai_provider(model: str, api_base: str, expected_provider: str): + """Regression for #38276: a Foundry .services.ai.azure.com base must not be reclassified as azure.""" + config = AzureAIStudioConfig() + ( + _, + _, + custom_llm_provider, + ) = config._get_openai_compatible_provider_info( + model=model, + api_base=api_base, + api_key="my-key", + custom_llm_provider="azure_ai", + ) + assert custom_llm_provider == expected_provider + + +def test_is_azure_openai_model_without_api_base_keeps_azure_ai(): + """Metadata lookups (get_model_info, supports_* checks) carry no api_base and must not flip the provider.""" + config = AzureAIStudioConfig() + assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base=None) is False + assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base="https://my-res.openai.azure.com") is True + + def test_azure_ai_validate_environment(): config = AzureAIStudioConfig() headers = config.validate_environment( diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py new file mode 100644 index 00000000000..0401629171d --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py @@ -0,0 +1,69 @@ +import httpx +import pytest +import respx + +from litellm import embedding +from litellm.llms.azure_ai.embed.handler import _foundry_models_route_base + +EMBEDDING_PAYLOAD = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, +} + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ( + "https://my-foundry.services.ai.azure.com", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com/", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com?api-version=2024-05-01-preview", + "https://my-foundry.services.ai.azure.com/models?api-version=2024-05-01-preview", + ), + ( + "https://my-foundry.services.ai.azure.com/models", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small", + "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small", + ), + ( + "https://my-resource.openai.azure.com", + "https://my-resource.openai.azure.com", + ), + ( + "https://Mistral-serverless.eastus2.models.ai.azure.com", + "https://Mistral-serverless.eastus2.models.ai.azure.com", + ), + (None, None), + ], +) +def test_foundry_models_route_base(api_base, expected): + assert _foundry_models_route_base(api_base) == expected + + +@respx.mock +def test_azure_ai_embedding_calls_foundry_models_route(): + route = respx.post("https://my-foundry.services.ai.azure.com/models/embeddings").mock( + return_value=httpx.Response(200, json=EMBEDDING_PAYLOAD) + ) + + response = embedding( + model="azure_ai/text-embedding-3-small", + input=["hello world"], + api_base="https://my-foundry.services.ai.azure.com", + api_key="fake-key", + ) + + assert route.called + assert response.data is not None + assert len(response.data) == 1 diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py new file mode 100644 index 00000000000..531f334e460 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py @@ -0,0 +1,66 @@ +import pytest + +from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, +) + +AWS_AUTH_PARAMS = { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + "aws_region_name": "us-west-2", + "aws_session_name": "session", + "aws_role_name": "arn:aws:iam::000000000000:role/example", + "aws_web_identity_token": "web-identity", + "aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com", + "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com", + "aws_external_id": "external", +} + + +def test_transform_request_never_resolves_aws_credentials(): + """A broken credential chain must not stop the request body from being built.""" + config = AmazonMoonshotConfig() + + transformed = config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"aws_profile_name": "litellm-profile-that-does-not-exist", "max_tokens": 16}, + litellm_params={}, + headers={}, + ) + + assert transformed["model"] == "moonshot.kimi-k2-thinking" + assert transformed["max_tokens"] == 16 + assert "aws_profile_name" not in transformed + + +@pytest.mark.parametrize("aws_param", sorted(AWS_AUTH_PARAMS)) +def test_transform_request_keeps_aws_params_out_of_the_body(aws_param: str): + config = AmazonMoonshotConfig() + + transformed = config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params={aws_param: AWS_AUTH_PARAMS[aws_param]}, + litellm_params={}, + headers={}, + ) + + assert aws_param not in transformed + + +def test_transform_request_leaves_the_caller_aws_params_in_place_for_signing(): + """sign_request reads the aws_* keys off optional_params after transform_request runs.""" + config = AmazonMoonshotConfig() + optional_params = dict(AWS_AUTH_PARAMS) + + config.transform_request( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "Hello"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert optional_params == AWS_AUTH_PARAMS diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 21e3239f623..c4d6896b17b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -513,3 +513,27 @@ def test_the_rust_opt_in_needs_no_sigv4_principal(): assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys() assert params["aws_region_name"] == "us-east-1" assert seen["call"][0]["api_key"] == "bedrock-bearer-token" + + +@pytest.mark.parametrize("configured_through", ["env_var", "api_key"]) +def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, configured_through): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still serve the request, since the + bearer token alone signs it.""" + if configured_through == "env_var": + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") + else: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + client = _sync_client_returning_converse_response() + + response = BedrockConverseLLM().completion( + **_completion_kwargs( + optional_params={"maxTokens": 16, "aws_profile_name": "litellm-no-such-aws-profile"}, + litellm_params={}, + client=client, + api_key="bedrock-bearer-token" if configured_through == "api_key" else None, + ) + ) + + assert response.choices[0].message.content == "hi" + assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 08d01127eba..50f8bbcf584 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1033,3 +1033,29 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert credentials.token == "assumed-session-token" assert aws_region_name == "us-east-1" assert "aws_external_id" not in optional_params + + +def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still serve the request, since the + bearer token alone signs it.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(titan_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/amazon.titan-embed-text-v1", + input=test_input, + client=client, + aws_region_name="us-west-2", + aws_profile_name="litellm-no-such-aws-profile", + ) + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py index 7c36b2aa75f..0b11a66c100 100644 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -135,3 +135,24 @@ class TestBedrockImageGeneration: assert response is not None assert len(response.data) > 0 mock_bedrock_image_gen.assert_called_once() + + +def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageGeneration()._prepare_request( + model="amazon.nova-canvas-v1:0", + prompt="A cute baby sea otter", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + api_key=None, + logging_obj=Mock(), + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py index 020b8df1276..58411a9ae18 100644 --- a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py +++ b/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py @@ -3,6 +3,7 @@ import base64 import io from typing import cast +from unittest.mock import Mock, patch import httpx import pytest @@ -655,3 +656,23 @@ def test_transform_response_empty_images_without_error_raises(): raw_response=resp, logging_obj=None, # type: ignore[arg-type] ) + + +def test_prepare_request_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageEdit()._prepare_request( + model="amazon.nova-canvas-v1:0", + image=[io.BytesIO(b"fake-png")], + prompt="make it warmer", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + logging_obj=Mock(), + api_key=None, + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 1d583c16ad7..023e2d8843f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -3098,3 +3098,91 @@ async def test_a_provider_that_keeps_rejecting_is_not_retried_forever_on_the_asy ) assert len(recorder.bodies) == 2 + + +CONTAINER_NOT_FOUND_BODY = { + "error": { + "message": "Container with id 'cntr_gone' not found.", + "type": "invalid_request_error", + "param": None, + "code": None, + } +} + +INVALID_API_KEY_BODY = { + "error": { + "message": "Incorrect API key provided: sk-proj-***. You can find your API key at https://platform.openai.com/account/api-keys.", + "type": "invalid_request_error", + "param": None, + "code": "invalid_api_key", + }, + "status": 401, +} + +CONTAINER_LIST_BODY = { + "object": "list", + "data": [{"id": "cntr_a", "object": "container", "created_at": 1, "status": "running", "name": "a"}], + "first_id": "cntr_a", + "last_id": "cntr_a", + "has_more": True, +} + + +def _container_sync_client(response: httpx.Response) -> HTTPHandler: + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: response)) + return client + + +def _container_async_client(response: httpx.Response) -> AsyncHTTPHandler: + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(lambda _request: response)) + return client + + +def test_container_retrieve_handler_raises_upstream_error_status_and_message(): + from litellm.llms.openai.containers.transformation import OpenAIContainerConfig + + with pytest.raises(BaseLLMException) as exc_info: + BaseLLMHTTPHandler().container_retrieve_handler( + container_id="cntr_gone", + container_provider_config=OpenAIContainerConfig(), + litellm_params=GenericLiteLLMParams(api_key="sk-test"), + logging_obj=Mock(), + client=_container_sync_client(httpx.Response(404, json=CONTAINER_NOT_FOUND_BODY)), + ) + + assert exc_info.value.status_code == 404 + assert exc_info.value.message == "Container with id 'cntr_gone' not found." + + +@pytest.mark.asyncio +async def test_async_container_list_handler_raises_upstream_error_status_and_message(): + from litellm.llms.openai.containers.transformation import OpenAIContainerConfig + + with pytest.raises(BaseLLMException) as exc_info: + await BaseLLMHTTPHandler().async_container_list_handler( + container_provider_config=OpenAIContainerConfig(), + litellm_params=GenericLiteLLMParams(api_key="sk-rejected"), + logging_obj=Mock(), + client=_container_async_client(httpx.Response(401, json=INVALID_API_KEY_BODY)), + ) + + assert exc_info.value.status_code == 401 + assert exc_info.value.message == INVALID_API_KEY_BODY["error"]["message"] + + +@pytest.mark.asyncio +async def test_async_container_list_handler_transforms_success_response(): + from litellm.llms.openai.containers.transformation import OpenAIContainerConfig + + response = await BaseLLMHTTPHandler().async_container_list_handler( + container_provider_config=OpenAIContainerConfig(), + litellm_params=GenericLiteLLMParams(api_key="sk-test"), + logging_obj=Mock(), + limit=1, + client=_container_async_client(httpx.Response(200, json=CONTAINER_LIST_BODY)), + ) + + assert [container.id for container in response.data] == ["cntr_a"] + assert response.has_more is True diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index 8dc4620dd1b..b6281834f24 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -10,11 +10,11 @@ Tests the cost calculation for Dashscope models including: import math import os +from datetime import datetime, timezone import pytest # Add the project root to Python path - import litellm from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, @@ -526,3 +526,139 @@ class TestDashscopeCostCalculator: assert prompt_cost == 0.0 assert math.isclose(completion_cost, 500 * 1.6e-06, rel_tol=1e-10) + + OFF_PEAK_WINDOW = "14:00-00:00" + INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc) + OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc) + + def _register_off_peak_flat_model(self, model_key: str, off_peak_pricing: dict) -> None: + litellm.model_cost[model_key] = { + "litellm_provider": "dashscope", + "mode": "chat", + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 3e-06, + "off_peak_pricing": off_peak_pricing, + } + + def test_dashscope_off_peak_window_swaps_in_the_off_peak_rates(self): + """ + Regression (LIT-6782): a deployment configured with off_peak_pricing kept billing the + standard dashscope rates inside its window, while the same block on a deepseek + deployment billed the off-peak rates. + """ + self._register_off_peak_flat_model( + "dashscope/deepseek-off-peak-test", + { + "hours_utc": self.OFF_PEAK_WINDOW, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1e-07, + }, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=200, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="deepseek-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, (600 * 1.2e-06) + (300 * 1e-07) + (100 * 3e-06), rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="deepseek-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + + assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10) + assert math.isclose(peak_completion_cost, 200 * 4.8e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_window_overrides_the_selected_tier(self): + """An open off-peak window bills the whole request at the flat off-peak rates, whichever tier + the input volume selected.""" + self._register_tiered_model( + "dashscope/qwen-tiered-off-peak-test", + [ + {"range": [0, 1000], "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06}, + {"range": [1000, 2000], "input_cost_per_token": 8e-07, "output_cost_per_token": 3.2e-06}, + ], + ) + litellm.model_cost["dashscope/qwen-tiered-off-peak-test"]["off_peak_pricing"] = { + "hours_utc": self.OFF_PEAK_WINDOW, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + } + usage = Usage(prompt_tokens=1500, completion_tokens=300) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-tiered-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, 1500 * 1e-07, rel_tol=1e-10) + assert math.isclose(completion_cost, 300 * 4e-07, rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="qwen-tiered-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + + assert math.isclose(peak_prompt_cost, 1500 * 8e-07, rel_tol=1e-10) + assert math.isclose(peak_completion_cost, 300 * 3.2e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_rates_left_unset_keep_the_standard_rates(self): + """A block that only overrides the input rate leaves output and cache reads on the standard + rates, and an explicit reasoning rate is never swapped out.""" + self._register_off_peak_flat_model( + "dashscope/qwen-partial-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1.2e-06}, + ) + litellm.model_cost["dashscope/qwen-partial-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06 + usage = Usage( + prompt_tokens=1000, + completion_tokens=200, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-partial-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, (700 * 1.2e-06) + (300 * 2e-07), rel_tol=1e-10) + assert math.isclose(completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_output_rate_covers_reasoning_without_a_dedicated_rate(self): + """Reasoning tokens on a model with no dedicated reasoning rate follow the off-peak output + rate, the same way they follow the standard output rate outside the window.""" + self._register_off_peak_flat_model( + "dashscope/qwen-reasoning-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "output_cost_per_token": 2.4e-06}, + ) + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + _, completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_defaults_to_the_current_time(self): + """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the + default current time.""" + self._register_off_peak_flat_model( + "dashscope/qwen-all-day-off-peak-test", + {"hours_utc": "00:00-00:00", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 2.4e-06}, + ) + usage = Usage(prompt_tokens=1000, completion_tokens=200) + + prompt_cost, completion_cost = dashscope_cost_per_token(model="qwen-all-day-off-peak-test", usage=usage) + + assert math.isclose(prompt_cost, 1000 * 1.2e-06, rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 295121167d6..a453708040e 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1329,6 +1329,523 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +COMPRESSED_MARKER = "[compressed document; retrieve the full text with hash=b573993006976af767214fac]" + + +class StructuredRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites whole messages via structured_messages and leaves + texts untouched, the way message-compressing guardrails do.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class ToolOutputRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses the first tool-result row, the way Headroom does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class DroppingRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first user row and drops the last row, so the + rewrite can only land through the full-conversion fallback.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten[:-1]} + + +def _texts(item: dict) -> list[str]: + content = item.get("content") + if isinstance(content, str): + return [content] + return [part["text"] for part in content] + + +class TestStructuredMessagesWriteBack: + """A guardrail's structured_messages rewrite must land in the Responses request, + not only the per-text mapping the chat handler shares with it.""" + + @pytest.mark.asyncio + async def test_list_input_gets_rewritten_messages_and_keeps_instructions(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + {"role": "user", "content": "memo " * 400}, + {"role": "assistant", "content": "Understood."}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert [_texts(item) for item in user_items] == [[COMPRESSED_MARKER], ["What is the codename?"]] + assert not any(item.get("role") == "system" for item in result["input"]) + assert _texts(next(item for item in result["input"] if item.get("role") == "assistant")) == ["Understood."] + + @pytest.mark.asyncio + async def test_string_input_becomes_rewritten_message_list(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "memo " * 400} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert [_texts(item) for item in result["input"]] == [[COMPRESSED_MARKER]] + assert "instructions" not in result + + @pytest.mark.asyncio + async def test_developer_item_preserved_verbatim_by_row_patch(self): + handler = OpenAIResponsesHandler() + developer_item = {"role": "developer", "content": "Always answer in French."} + data = { + "model": "gpt-5.6", + "input": [ + developer_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is developer_item + assert developer_item["content"] == "Always answer in French." + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_reasoning_and_function_call_items_survive_tool_output_compression(self): + handler = OpenAIResponsesHandler() + reasoning_item = { + "id": "rs_123", + "type": "reasoning", + "summary": [], + "encrypted_content": "gAAAAA-signed-reasoning", + } + function_call_item = { + "id": "fc_123", + "type": "function_call", + "call_id": "call_abc", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + "status": "completed", + } + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + reasoning_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_abc", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + assert result["input"][0] is reasoning_item + assert reasoning_item["encrypted_content"] == "gAAAAA-signed-reasoning" + assert result["input"][1] is function_call_item + assert function_call_item["id"] == "fc_123" + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_abc", + "output": COMPRESSED_MARKER, + } + assert result["input"][3] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_web_search_call_item_preserved_verbatim(self): + handler = OpenAIResponsesHandler() + web_search_item = { + "id": "ws_123", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "codename memo"}, + } + data = { + "model": "gpt-5.6", + "input": [ + web_search_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is web_search_item + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_row_count_change_falls_back_to_full_conversion(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "developer", "content": "Always answer in French."}, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, DroppingRewriteGuardrail()) + + assert len(result["input"]) == 2 + developer = next(item for item in result["input"] if item.get("role") == "developer") + assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}] + assert _texts(next(item for item in result["input"] if item.get("role") == "user")) == [COMPRESSED_MARKER] + + @pytest.mark.asyncio + async def test_same_inputs_object_back_keeps_the_text_mapping(self): + handler = OpenAIResponsesHandler() + original_input = [ + {"role": "user", "content": "Hello"}, + {"role": "user", "content": [{"type": "input_text", "text": "Again"}]}, + ] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, MockGuardrail()) + + assert result["input"] is original_input + assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]] + + +class AllToolOutputsRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses every tool-result row.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if isinstance(m, dict) and m.get("role") == "tool" else m + for m in messages + ] + return {**inputs, "structured_messages": rewritten} + + +class AssistantRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first assistant row's content.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "assistant") + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first else m for i, m in enumerate(messages)] + return {**inputs, "structured_messages": rewritten} + + +class DictStructuredMessagesGuardrail(CustomGuardrail): + """Guardrail that hands back a raw evaluation dict instead of a message list, + the way HiddenLayer v2 does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} + + +def _parallel_tool_call_input() -> list: + return [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"id": "fc_2", "type": "function_call", "call_id": "call_2", "name": "read_b", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo " * 400}, + {"type": "function_call_output", "call_id": "call_2", "output": "note " * 400}, + {"role": "user", "content": "What is the codename?"}, + ] + + +class TestProvenancePatching: + """The O(n) provenance pass must keep patching rewritten rows in place for the + shapes real agent loops produce, and fall back safely everywhere else.""" + + @pytest.mark.asyncio + async def test_parallel_tool_call_outputs_both_patched(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + fc_1, fc_2 = raw_input[0], raw_input[1] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AllToolOutputsRewriteGuardrail()) + + assert result["input"][0] is fc_1 + assert result["input"][1] is fc_2 + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_1", "output": COMPRESSED_MARKER} + assert result["input"][3] == {"type": "function_call_output", "call_id": "call_2", "output": COMPRESSED_MARKER} + assert result["input"][4] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_assistant_turn_with_tool_call_keeps_items_verbatim(self): + handler = OpenAIResponsesHandler() + assistant_item = {"role": "assistant", "content": "Let me read the memo."} + function_call_item = { + "id": "fc_9", + "type": "function_call", + "call_id": "call_9", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + } + data = { + "model": "gpt-5.6", + "input": [ + assistant_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_9", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["input"][0] is assistant_item + assert result["input"][1] is function_call_item + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_9", "output": COMPRESSED_MARKER} + + @pytest.mark.asyncio + async def test_rewrite_of_merged_tool_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + @pytest.mark.asyncio + async def test_rewrite_of_lone_function_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + raw_input = data["input"] + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + def test_provenance_bails_on_non_mapping_item(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance(["not a mapping"], []) is None + + def test_provenance_bails_when_expected_messages_disagree(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance([{"role": "user", "content": "hi"}], [{"role": "user", "content": "bye"}]) is None + + def test_provenance_bails_on_unpredicted_merge(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + ] + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + assert len(expected) == 1 + assert _input_item_provenance(raw_input, expected) is None + + def test_provenance_maps_and_taints_parallel_tool_calls(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = _parallel_tool_call_input() + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + provenance = _input_item_provenance(raw_input, expected) + assert provenance is not None + item_for_message, tainted = provenance + assert tainted == {0} + assert dict(item_for_message) == {1: 2, 2: 3, 3: 4} + + +class TestDictStructuredMessagesGuard: + """A guardrail handing back a non-list structured_messages payload must not + blow up the request; the write-back is skipped instead.""" + + @pytest.mark.asyncio + async def test_list_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + original_input = [{"role": "user", "content": "Hello"}] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] is original_input + assert result["input"] == [{"role": "user", "content": "Hello"}] + + @pytest.mark.asyncio + async def test_string_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "Hello there"} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] == "Hello there" + + +class SystemRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the system row, the way prompt-hardening guardrails do.""" + + def __init__(self, rewritten_content: Any = COMPRESSED_MARKER): + super().__init__() + self.rewritten_content = rewritten_content + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system") + rewritten = [ + {**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class TestPatchEdgeBranches: + @pytest.mark.asyncio + async def test_multimodal_user_item_rewritten_through_conversion(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "memo " * 400}]}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert _texts(result["input"][0]) == [COMPRESSED_MARKER] + assert result["input"][1] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_instructions_rewrite_lands_in_instructions_field(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages(data, SystemRewriteGuardrail()) + + assert result["instructions"] == COMPRESSED_MARKER + assert result["input"][0] is user_item + + @pytest.mark.asyncio + async def test_non_string_instructions_rewrite_falls_back(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages( + data, SystemRewriteGuardrail(rewritten_content=[{"type": "text", "text": COMPRESSED_MARKER}]) + ) + + assert result["input"][0] is not user_item + + @pytest.mark.asyncio + async def test_unpredicted_merge_falls_back_through_patch(self): + handler = OpenAIResponsesHandler() + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "memo " * 400}, + ] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert _texts(user_items[0]) == [COMPRESSED_MARKER] + + def test_item_rewrite_field_ignores_non_string_type(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _item_rewrite_field + + assert _item_rewrite_field({"type": 123, "content": "hello"}) is None + + class ToolEditingGuardrail(CustomGuardrail): """Guardrail that rewrites the flattened chat tools it was handed through ``edit``""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index b0ffd1845fe..4ac072d0ca6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -1718,6 +1718,8 @@ class TestResponsesSurfaceSharesTheEffortRule: ("gpt-5.6-sol", None, False), ("gpt-5.6-terra", "none", True), ("gpt-5.6-terra", "medium", False), + ("gpt-6-astra", None, False), + ("gpt-6-astra", "low", False), ], ) def test_temperature_follows_the_resolved_effort( diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 9c5bd34d59a..c86ce4df2ac 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1505,3 +1505,17 @@ class TestACatalogueOlderThanTheCodeDoesNotStripTemperature: drop_params=True, ) assert "temperature" not in mapped + + +def test_gpt_6_astra_takes_the_reasoning_series_request_shape(): + params = litellm.get_optional_params( + model="gpt-6-astra", + custom_llm_provider="openai", + max_tokens=100, + reasoning_effort="max", + verbosity="low", + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + assert params["reasoning_effort"] == "max" + assert params["verbosity"] == "low" diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py index 1095819c98c..107a1afb2c6 100644 --- a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py +++ b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py @@ -41,6 +41,8 @@ from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config # Models that MUST be classified as GPT-5 (routed through GPT-5 reasoning path) GPT5_MODELS = [ + "gpt-6-astra", + "openai/gpt-6-astra", "gpt-5", "gpt-5.1", "gpt-5.2", @@ -120,6 +122,8 @@ class TestOpenAIGPT5ConfigIsModelGpt5Model: # /v1/responses bridge (when reasoning_effort is set and tools are passed) on # is_model_gpt_5_4_plus_model, so the gpt-5.6 family must land on the True side. GPT5_4_PLUS_MODELS = [ + "gpt-6-astra", + "openai/gpt-6-astra", "gpt-5.4", "gpt-5.5", "gpt-5.5-pro", diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/test_litellm/llms/openai/test_openai.py new file mode 100644 index 00000000000..136b837f191 --- /dev/null +++ b/tests/test_litellm/llms/openai/test_openai.py @@ -0,0 +1,52 @@ +import pytest + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://southcentralus.privatelink.api.openai.com/v1", + "https://eu.api.openai.com/v1", + "https://us.api.openai.com/v1", + "HTTPS://API.OPENAI.COM/v1/", + ], +) +def test_get_stream_options_defaults_include_usage_on_every_openai_backed_host(api_base): + """ + PrivateLink and regional hostnames reach the real OpenAI backend, so a stream with no caller + stream_options must ask for the usage chunk exactly as the default base does. Regression guard + for LIT-6875: spend for those deployments fell back to local token counting. + """ + assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == { + "stream_options": {"include_usage": True} + } + + +@pytest.mark.parametrize( + "api_base", + [ + "https://my-gateway.example/v1", + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_get_stream_options_leaves_foreign_hosts_without_a_usage_default(api_base): + """Only the host decides: an OpenAI-compatible backend elsewhere may not support stream_options at all.""" + assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {} + + +@pytest.mark.parametrize( + "api_base", + ["https://southcentralus.privatelink.api.openai.com/v1", "https://my-gateway.example/v1"], +) +def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api_base): + caller_options = {"include_usage": False} + assert OpenAIChatCompletion().get_stream_options(stream_options=caller_options, api_base=api_base) == { + "stream_options": caller_options + } diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index 3ae29e411e8..d3c21c5bd5a 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -7,7 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.token_counter import token_counter -from litellm.llms.openai.common_utils import BaseOpenAILLM +from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base # Test parameters for different API functions API_FUNCTION_PARAMS = [ @@ -392,3 +392,22 @@ async def test_async_genuine_bad_request_still_raises(provider, stream): with pytest.raises(litellm.BadRequestError): await _call_and_drain() + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ("https://api.openai.com/v1", True), + ("https://api.openai.com:443/v1/", True), + ("https://southcentralus.privatelink.api.openai.com/v1", True), + ("https://eu.api.openai.com/v1", True), + ("HTTPS://API.OPENAI.COM/v1", True), + ("https://my-gateway.example/v1", False), + ("https://api.openai.com.evil.example/v1", False), + ("https://notapi.openai.com/v1", False), + ("https://gateway.example/v1?upstream=api.openai.com", False), + ("not a url", False), + ], +) +def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected): + assert is_openai_backed_api_base(api_base) is expected diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index d1d751989ea..dddc95bf54a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -964,6 +964,48 @@ def test_construct_target_url_with_version_prefix(): assert str(target_url) == expected_url +@pytest.mark.parametrize( + ("requested_route", "expected_url"), + [ + ( + "/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + ), + ( + "/projects/test-project/locations/global/publishers/anthropic/models/count-tokens:rawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/count-tokens:rawPredict", + ), + ( + "/projects/other-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-6:rawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:rawPredict", + ), + ( + "/projects/test-project/locations/global/cachedContents", + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents", + ), + ( + "/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + ), + ( + "/v1beta1/projects/test-project/locations/global/cachedContents", + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents", + ), + ], +) +def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None: + from litellm.llms.vertex_ai.common_utils import construct_target_url + + target_url = construct_target_url( + base_url="https://aiplatform.googleapis.com", + requested_route=requested_route, + vertex_project="test-project", + vertex_location="global", + ) + + assert str(target_url) == expected_url + + def test_fix_enum_types(): """ Test _fix_enum_types function removes enum fields when type is not string. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 588eba4adb7..775f6e5f3b8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -3508,6 +3508,169 @@ def _create_oauth2_server( ) +def _create_id_lookup_oauth2_server(): + return _create_oauth2_server( + server_id="oauth-server-id", + name="oauth-server-name", + server_name="oauth-server-name", + alias="oauth-server-alias", + ) + + +@pytest.mark.asyncio +async def test_authorize_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object(discoverable_endpoints, "encrypt_value_helper", return_value="encrypted-state"), # test-quality-ok: flow seam + ): + response = await discoverable_endpoints.authorize( + request=request, + client_id=server.client_id, + mcp_server_name=server.server_id, + redirect_uri="http://localhost:62646/callback", + state="test_state", + ) + + assert response.status_code == 307 + assert "https://provider.com/oauth/authorize" in response.headers["location"] + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +@pytest.mark.asyncio +async def test_token_endpoint_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"access_token": "token", "token_type": "Bearer"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object(discoverable_endpoints, "get_async_httpx_client", return_value=client), # test-quality-ok: HTTP seam + ): + result = await discoverable_endpoints.token_endpoint( + request=request, + grant_type="authorization_code", + code="test_code", + redirect_uri="http://localhost:62646/callback", + client_id=server.client_id, + mcp_server_name=server.server_id, + client_secret=server.client_secret, + ) + + assert json.loads(result.body)["access_token"] == "token" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +@pytest.mark.asyncio +async def test_register_client_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server().model_copy( + update={"client_id": None, "client_secret": None, "registration_url": "https://provider.com/oauth/register"} + ) + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"client_id": "registered-client", "client_secret": "registered-secret"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object(discoverable_endpoints, "_read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: request seam + patch.object(discoverable_endpoints, "get_async_httpx_client", return_value=client), # test-quality-ok: HTTP seam + ): + result = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_id) + + assert json.loads(result.body)["client_id"] == "registered-client" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +@pytest.mark.asyncio +async def test_protected_resource_metadata_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + ): + result = await discoverable_endpoints._build_oauth_protected_resource_response( + request=request, + mcp_server_name=server.server_id, + use_standard_pattern=True, + ) + + assert result["authorization_servers"] == ["https://llm.example.com/mcp"] + assert result["resource"] == f"https://llm.example.com/mcp/{server.server_id}" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +def test_authorization_server_metadata_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + ): + result = discoverable_endpoints._build_oauth_authorization_server_response( + request=request, + mcp_server_name=server.server_id, + ) + + assert result["scopes_supported"] == server.scopes + assert result["issuer"] == f"https://llm.example.com/{server.server_id}" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + @pytest.mark.asyncio async def test_authorize_root_resolves_single_oauth2_server(): """When /authorize is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 91e870d2d95..3321da83007 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -129,6 +129,32 @@ class TestMCPServerManager: assert added_server.args == ["-m", "server"] assert added_server.env == {"DEBUG": "1", "TEST": "1"} + def test_get_mcp_server_by_id_allows_internal_or_unspecified_client_ip(self): + manager = MCPServerManager() + server = MCPServer( + server_id="private-server", + name="private-server", + transport=MCPTransport.http, + available_on_public_internet=False, + ) + manager.registry[server.server_id] = server + + assert manager.get_mcp_server_by_id(server.server_id) is server + assert manager.get_mcp_server_by_id(server.server_id, client_ip="10.0.0.1") is server + + def test_get_mcp_server_by_id_rejects_private_server_for_public_ip(self): + manager = MCPServerManager() + server = MCPServer( + server_id="private-server", + name="private-server", + transport=MCPTransport.http, + available_on_public_internet=False, + ) + manager.registry[server.server_id] = server + + with patch.object(manager, "_get_general_settings", return_value={}): + assert manager.get_mcp_server_by_id(server.server_id, client_ip="8.8.8.8") is None + async def test_create_mcp_client_stdio(self): """Test creating MCP client for stdio transport""" manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 82528c58ae0..383b72e5c58 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -10,15 +10,42 @@ from unittest.mock import AsyncMock, patch import pytest -from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentAccess, AgentRequestHandler, RestrictedAgentAccess, UnrestrictedAgentAccess, + accessible_agents, ) +def _registry_with(*agent_names: str) -> AgentRegistry: + registry: Final = AgentRegistry() + registry.load_agents_from_config( + [ + { + "agent_name": name, + "agent_card_params": {"name": name, "url": "http://localhost", "version": "1.0.0"}, + } + for name in agent_names + ] + ) + return registry + + +def _agent_id(registry: AgentRegistry, agent_name: str) -> str: + agent: Final = registry.get_agent_by_name(agent_name) + assert agent is not None + return agent.agent_id + + +async def _single_context(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + return [user_api_key_auth] + + @pytest.mark.asyncio class TestAgentRequestHandler: """ @@ -265,6 +292,78 @@ class TestAgentRequestHandler: ) assert result == UnrestrictedAgentAccess() + async def test_accessible_agents_hides_ungranted_agents_from_non_admins(self): + """LIT-6862: a key with no agent grant on itself or its team must list nothing, + while a proxy admin with the same lack of grants still lists every agent.""" + registry: Final = _registry_with("alpha", "beta") + internal_user: Final = UserAPIKeyAuth( + api_key="test-key", user_id="alice", team_id="team-no-perms", user_role=LitellmUserRoles.INTERNAL_USER + ) + proxy_admin: Final = UserAPIKeyAuth( + api_key="admin-key", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + async def no_grant_anywhere(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess: + return UnrestrictedAgentAccess() + + assert ( + await accessible_agents(internal_user, registry.get_agent_list(), no_grant_anywhere, _single_context) == () + ) + assert { + agent.agent_name + for agent in await accessible_agents( + proxy_admin, registry.get_agent_list(), no_grant_anywhere, _single_context + ) + } == {"alpha", "beta"} + + async def test_accessible_agents_lists_only_granted_agents(self): + """A grant for one agent lists that agent and hides the ungranted one.""" + registry: Final = _registry_with("alpha", "beta") + granted_user: Final = UserAPIKeyAuth( + api_key="test-key", user_id="bob", team_id="team-granted", user_role=LitellmUserRoles.INTERNAL_USER + ) + + async def alpha_only(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess: + return RestrictedAgentAccess(frozenset({_agent_id(registry, "alpha")})) + + listed: Final = await accessible_agents(granted_user, registry.get_agent_list(), alpha_only, _single_context) + assert [agent.agent_name for agent in listed] == ["alpha"] + + async def test_accessible_agents_resolves_dashboard_session_through_real_teams_and_user(self): + """LIT-6862: a dashboard session carries the shared litellm-dashboard team id, which holds no + grants. Listing must union the grants of the user's real teams and of the user row instead + of treating the session as ungranted or as unrestricted.""" + registry: Final = _registry_with("alpha", "beta", "gamma") + session: Final = UserAPIKeyAuth( + api_key="session-key", + user_id="alice", + team_id=UI_SESSION_TOKEN_TEAM_ID, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + admitted_user: Final = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + grants: Final = { + "team-granted": RestrictedAgentAccess(frozenset({_agent_id(registry, "alpha")})), + "team-no-perms": UnrestrictedAgentAccess(), + UI_SESSION_TOKEN_TEAM_ID: UnrestrictedAgentAccess(), + } + + async def effective_contexts(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + assert user_api_key_auth is session + return [ + session.model_copy(update={"team_id": "team-granted"}), + session.model_copy(update={"team_id": "team-no-perms"}), + admitted_user, + ] + + async def resolve_access(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess: + if user_api_key_auth is admitted_user: + return RestrictedAgentAccess(frozenset({_agent_id(registry, "beta")})) + assert user_api_key_auth.team_id is not None + return grants[user_api_key_auth.team_id] + + listed: Final = await accessible_agents(session, registry.get_agent_list(), resolve_access, effective_contexts) + assert {agent.agent_name for agent in listed} == {"alpha", "beta"} + async def test_get_allowed_agents_for_key_via_access_group_ids(self): """ Test that _get_allowed_agents_for_key includes agents from key's access_group_ids diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 067f5a9f64c..13d6cd8a68c 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -11,7 +11,6 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( RestrictedAgentAccess, - UnrestrictedAgentAccess, ) from litellm.proxy.agent_endpoints.endpoints import ( _attach_keys_to_agents, @@ -550,9 +549,9 @@ class TestAgentRBACProxyAdminViewOnly: self.allowed_agents_spy.assert_awaited_once() def test_should_still_redact_secrets_for_view_only_admin(self): - """An unrestricted viewer sees the same agents as an admin but with keys + """A viewer granted every agent sees the same agents as an admin but with keys stripped; litellm_params secrets never appear in either response.""" - self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() + self.allowed_agents_spy.return_value = RestrictedAgentAccess(frozenset({"agent-1", "agent-2"})) viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py index 8da365cb587..1db53638070 100644 --- a/tests/test_litellm/proxy/auth/test_litellm_license.py +++ b/tests/test_litellm/proxy/auth/test_litellm_license.py @@ -2,6 +2,8 @@ import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey + from litellm.proxy.auth.litellm_license import LicenseCheck @@ -30,3 +32,70 @@ def test_is_over_limit(): assert license_check.is_over_limit(101) is False assert license_check.is_over_limit(100) is False assert license_check.is_over_limit(99) is False + + +def test_heuristic_v2_router_limit() -> None: + """Only the signed license's auto_router feature lifts the one-router limit; an API-verified + license (no airgapped data) and an airgapped license without the feature keep it.""" + license_check = LicenseCheck() + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]} + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = { + "expiration_date": "2999-01-01", + "allowed_features": ["sso", "auto_router", "audit_logs"], + } + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01"} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = None + assert license_check.heuristic_v2_router_limit() == 1 + + +def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]: + import base64 + + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import padding, rsa + + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + message = json.dumps( + {"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]} + ).encode() + signature = private_key.sign( + message, + padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH), + hashes.SHA256(), + ) + return private_key.public_key(), base64.b64encode(message + b"." + signature).decode() + + +def test_expired_or_unreadable_license_grants_no_features() -> None: + """The verifier stores the signed payload only after the expiry check passes and clears it when a + later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit.""" + license_check = LicenseCheck() + public_key, valid_key = _signed_license("2999-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.heuristic_v2_router_limit() is None + + _, expired_key = _signed_license("2000-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True + assert license_check.airgapped_license_data is None + assert license_check.heuristic_v2_router_limit() == 1 + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True + assert license_check.airgapped_license_data is None + + +def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None: + license_check = LicenseCheck() + public_key, license_key = _signed_license("2999-01-01") + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True + assert license_check.heuristic_v2_router_limit() is None diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index ef560bd1b7d..fcfb9342176 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -1,4 +1,6 @@ +import io import json +from typing import get_type_hints from unittest.mock import AsyncMock, MagicMock, patch import orjson @@ -18,9 +20,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_parsed_body, _safe_get_request_query_params, _safe_set_request_parsed_body, + coerce_numeric_form_fields, get_form_data, get_request_body, get_tags_from_request_body, + numeric_form_fields, populate_request_with_path_params, ) @@ -1029,3 +1033,79 @@ class TestGetRequestBody: mock_request = MagicMock() mock_request.method = "GET" assert await get_request_body(mock_request) == {} + + +class TestNumericFormFields: + def test_image_edit_schema_yields_only_n(self): + from litellm.types.images.main import ImageEditRequestParams + + assert dict(numeric_form_fields(get_type_hints(ImageEditRequestParams))) == {"n": int} + + def test_qualifiers_and_optionality_are_unwrapped(self): + from typing import Optional + + from typing_extensions import Annotated, NotRequired, ReadOnly, Required, TypedDict + + class Schema(TypedDict, total=False): + plain: int + optional: Optional[int] + piped: int | None + read_only: ReadOnly[int | None] + not_required: NotRequired[ReadOnly[int]] + required: Required[ReadOnly[Annotated[float, "meta"]]] + + assert dict(numeric_form_fields(get_type_hints(Schema))) == { + "plain": int, + "optional": int, + "piped": int, + "read_only": int, + "not_required": int, + "required": float, + } + + def test_non_scalar_and_bool_fields_are_skipped(self): + from typing import Any, Literal, Optional, Union + + from typing_extensions import TypedDict + + class Schema(TypedDict, total=False): + flag: bool + optional_flag: Optional[bool] + text: str + choice: Optional[Literal["high", "low"]] + numbers: list[int] + mapping: Optional[dict[str, Any]] + ambiguous: Union[int, str] + + assert dict(numeric_form_fields(get_type_hints(Schema))) == {} + + +class TestCoerceNumericFormFields: + numeric_fields = {"n": int, "temperature": float} + + def test_numeric_strings_are_parsed(self): + assert coerce_numeric_form_fields( + parsed_body={"n": "2", "temperature": "0.5"}, + numeric_fields=self.numeric_fields, + ) == {"n": 2, "temperature": 0.5} + + def test_other_fields_keep_their_string_values(self): + result = coerce_numeric_form_fields( + parsed_body={"size": "1024x1024", "prompt": "2", "quality": "high"}, + numeric_fields=self.numeric_fields, + ) + assert result == {"size": "1024x1024", "prompt": "2", "quality": "high"} + + def test_unparseable_value_is_left_for_the_provider_to_reject(self): + assert coerce_numeric_form_fields( + parsed_body={"n": "two", "temperature": ""}, + numeric_fields=self.numeric_fields, + ) == {"n": "two", "temperature": ""} + + def test_already_typed_and_non_string_values_pass_through(self): + buffer = io.BytesIO(b"png") + result = coerce_numeric_form_fields( + parsed_body={"n": 3, "temperature": None, "image": buffer}, + numeric_fields=self.numeric_fields, + ) + assert result == {"n": 3, "temperature": None, "image": buffer} diff --git a/tests/test_litellm/proxy/container_endpoints/__init__.py b/tests/test_litellm/proxy/container_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/container_endpoints/test_endpoints.py b/tests/test_litellm/proxy/container_endpoints/test_endpoints.py new file mode 100644 index 00000000000..1beff4c82ba --- /dev/null +++ b/tests/test_litellm/proxy/container_endpoints/test_endpoints.py @@ -0,0 +1,134 @@ +import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.container_endpoints import endpoints, ownership +from litellm.types.containers.main import ContainerListResponse, ContainerObject + +PROXY_SERVER_STUB = SimpleNamespace( + general_settings={}, + prisma_client=None, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", +) +ADMIN = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) +NON_ADMIN = UserAPIKeyAuth(user_id="user-1") + + +@pytest.fixture(autouse=True) +def clear_allowed_container_ids_cache(): + ownership._ALLOWED_CONTAINER_IDS_CACHE.cache_dict.clear() + ownership._ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.clear() + yield + ownership._ALLOWED_CONTAINER_IDS_CACHE.cache_dict.clear() + ownership._ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.clear() + + +def _client(auth: UserAPIKeyAuth) -> TestClient: + app = FastAPI() + app.include_router(endpoints.router) + app.dependency_overrides[user_api_key_auth] = lambda: auth + return TestClient(app) + + +def _container(container_id: str) -> ContainerObject: + return ContainerObject(id=container_id, object="container", created_at=1, status="active") + + +def _page(*container_ids: str, has_more: bool) -> ContainerListResponse: + return ContainerListResponse( + object="list", + data=[_container(container_id) for container_id in container_ids], + has_more=has_more, + ) + + +def _upstream_pages(monkeypatch, pages_by_after) -> MagicMock: + processor_cls = MagicMock( + side_effect=lambda data: SimpleNamespace( + base_process_llm_request=AsyncMock(return_value=pages_by_after[data["after"]]) + ) + ) + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", processor_cls) + return processor_cls + + +def _forwarded_pages(processor_cls: MagicMock): + return [(call.kwargs["data"]["after"], call.kwargs["data"]["limit"]) for call in processor_cls.call_args_list] + + +def test_list_containers_forwards_typed_pagination_params_for_admins(monkeypatch): + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) + processor_cls = _upstream_pages(monkeypatch, {"cntr_prev": _page("cntr_next", has_more=True)}) + + response = _client(ADMIN).get( + "/v1/containers", + params={"limit": "1", "order": "desc", "after": "cntr_prev"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [item["id"] for item in response.json()["data"]] == ["cntr_next"] + assert response.json()["has_more"] is True + assert _forwarded_pages(processor_cls) == [("cntr_prev", 1)] + assert processor_cls.call_args.kwargs["data"]["order"] == "desc" + + +def test_list_containers_rejects_a_non_integer_limit(monkeypatch): + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) + processor_cls = _upstream_pages(monkeypatch, {}) + + response = _client(ADMIN).get( + "/v1/containers", + params={"limit": "abc"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 422 + processor_cls.assert_not_called() + + +def test_list_containers_pages_upstream_until_non_admin_keys_see_their_containers(monkeypatch): + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) + table = AsyncMock() + table.find_many.return_value = [SimpleNamespace(model_object_id="container:openai:cntr_owned")] + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))), + ) + processor_cls = _upstream_pages( + monkeypatch, + { + None: _page("cntr_other", has_more=True), + "cntr_other": _page("cntr_owned", has_more=False), + }, + ) + + response = _client(NON_ADMIN).get( + "/v1/containers", + params={"limit": "1"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + body = response.json() + assert [item["id"] for item in body["data"]] == ["cntr_owned"] + assert body["first_id"] == "cntr_owned" + assert body["last_id"] == "cntr_owned" + assert body["has_more"] is False + assert _forwarded_pages(processor_cls) == [(None, 100), ("cntr_other", 100)] diff --git a/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py b/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py new file mode 100644 index 00000000000..a471f915071 --- /dev/null +++ b/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py @@ -0,0 +1,62 @@ +import sys +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.container_endpoints import endpoints, handler_factory + +PROXY_SERVER_STUB = SimpleNamespace( + general_settings={}, + prisma_client=None, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", +) + + +def _client() -> TestClient: + app = FastAPI() + app.include_router(endpoints.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="user-1") + return TestClient(app) + + +def test_list_container_files_forwards_declared_query_params(monkeypatch): + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + AsyncMock(return_value=("cntr_123", "openai")), + ) + processor_cls = MagicMock() + processor_cls.return_value.base_process_llm_request = AsyncMock( + return_value={"object": "list", "data": [], "has_more": True} + ) + monkeypatch.setattr(handler_factory, "ProxyBaseLLMRequestProcessing", processor_cls) + + response = _client().get( + "/v1/containers/cntr_123/files", + params={"limit": "1", "order": "desc", "after": "cfile_prev", "unknown": "x"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json()["has_more"] is True + assert processor_cls.return_value.base_process_llm_request.await_args.kwargs["route_type"] == "alist_container_files" + forwarded = processor_cls.call_args.kwargs["data"] + assert forwarded["container_id"] == "cntr_123" + assert forwarded["limit"] == "1" + assert forwarded["order"] == "desc" + assert forwarded["after"] == "cfile_prev" + assert "unknown" not in forwarded diff --git a/tests/test_litellm/proxy/db/test_replica_identity.py b/tests/test_litellm/proxy/db/test_replica_identity.py index ecfc6433ab1..9738fc9bd98 100644 --- a/tests/test_litellm/proxy/db/test_replica_identity.py +++ b/tests/test_litellm/proxy/db/test_replica_identity.py @@ -29,7 +29,7 @@ def test_hands_the_alter_statement_to_the_prisma_cli(): return subprocess.CompletedProcess(cmd, 0) with patch( - "litellm_proxy_extras.replica_identity.subprocess.run", side_effect=capture + "litellm_proxy_extras.replica_identity.run_prisma", side_effect=capture ): applied = apply_replica_identity_full( schema_path="/somewhere/schema.prisma", @@ -60,7 +60,7 @@ def test_hands_the_alter_statement_to_the_prisma_cli(): ) def test_every_failure_is_reported_instead_of_raised(failure): with patch( - "litellm_proxy_extras.replica_identity.subprocess.run", side_effect=failure + "litellm_proxy_extras.replica_identity.run_prisma", side_effect=failure ): assert ( apply_replica_identity_full( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 953e3de1519..1e44d29b610 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5527,10 +5527,6 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca error frame instead. The finish chunk is withheld while the end-of-stream scan runs, so on a block it is dropped rather than relayed before the frame.""" - from litellm.llms import load_guardrail_translation_mappings - from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( - unified_guardrail as unified_module, - ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -5569,20 +5565,16 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca yield _chunk("the forbidden ") yield _chunk("topic answer", finish_reason="stop") - unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() - try: - with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response) + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response) - out = [] - async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"), - response=_mock_stream(), - request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"}, - ): - out.append(item) - finally: - unified_module.endpoint_guardrail_translation_mappings = None + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"), + response=_mock_stream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"}, + ): + out.append(item) assert len(out) == 2 assert isinstance(out[0], ModelResponseStream) @@ -5792,3 +5784,25 @@ async def test_apply_guardrail_debug_log_masks_signed_request_headers(): assert header_lines, "expected the signed-request debug line to be logged" assert any("X-Amz-Security-Token" in message for message in header_lines) assert all(session_token not in message for message in rendered_messages) + + +@pytest.mark.asyncio +async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The guardrail's AWS profile does not exist, so resolving SigV4 credentials + raises; with a bearer token configured the guardrail must still run, since + the bearer token alone signs the request.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_profile_name="litellm-no-such-aws-profile", + ) + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"action": "NONE", "assessments": []} + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response) as mock_post: + response = await guardrail.make_bedrock_api_request(source="INPUT", messages=[{"role": "user", "content": "hello"}]) + + assert response["action"] == "NONE" + assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index d842a1ee5f9..f4af77d2e40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -833,3 +833,31 @@ async def test_many_blocks_scanned_at_request_level_and_can_block(): sent_texts = [c["text"] for m in body_messages for c in m["content"]] assert sent_texts == [f"b{i}" for i in range(25)] assert all(len(m["content"]) <= 10 for m in body_messages) + + +@pytest.mark.asyncio +async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """Same bearer-token rule as ApplyGuardrail: the guardrail's AWS profile does + not exist, yet the InvokeGuardrailChecks call still goes out on the bearer + token and its verdict is enforced.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + g = BedrockGuardrail( + checks=CONTENT_FILTER_CHECKS, + content_filter_threshold=0.5, + aws_profile_name="litellm-no-such-aws-profile", + ) + payload = {"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.8}]}}} + post = AsyncMock(return_value=_mock_http_response(200, payload)) + + with patch.object(g.async_handler, "post", new=post): + with pytest.raises(HTTPException) as exc: + await g.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hi"}], + request_data={"messages": []}, + ) + + assert exc.value.detail["bedrock_guardrail_checks"] == [ + {"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8} + ] + assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index c04fb7b30ec..400eaf8ab3f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -193,6 +193,33 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( assert "headroom" in _applied_guardrails(request_data) +@pytest.mark.asyncio +async def test_apply_guardrail_leaves_background_requests_uncompressed( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"model": "gpt-4o", "background": True} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ) as post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result is inputs + post.assert_not_awaited() + assert _recorded_guardrail_entries(request_data) == [] + + def _recorded_guardrail_response(request_data: dict) -> dict: entries = request_data["metadata"]["standard_logging_guardrail_information"] assert len(entries) == 1 @@ -954,6 +981,40 @@ async def test_passthrough_handler_does_not_log_headroom_as_run( assert "headroom" not in _applied_guardrails(data) +@pytest.mark.asyncio +async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream( + guardrail: HeadroomGuardrail, +): + """Regression for LIT-6494: on /v1/responses the compressed messages must be + written back into `input`, not only the retrieve tool into `tools`, or the + model keeps reading the full document and never calls headroom_retrieve.""" + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + + data = { + "model": "gpt-5.6", + "instructions": ORIGINAL_MESSAGES[0]["content"], + "input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]], + "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}], + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + ): + result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] + assert [item["content"] for item in result["input"]] == [ + COMPRESSED_MESSAGES_WITH_HASH[0]["content"], + ORIGINAL_MESSAGES[2]["content"], + ORIGINAL_MESSAGES[3]["content"], + ] + assert "A" * 5000 not in json.dumps(result["input"]) + assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME] + + @pytest.mark.asyncio async def test_apply_guardrail_http_error_raises(): guardrail = _make_guardrail() @@ -1950,6 +2011,58 @@ def _openai_text_payload(content: str) -> dict: return _openai_completion_payload({"role": "assistant", "content": content}, "stop") +def _responses_retrieve_tool_definition() -> dict: + return {"type": "function", **_retrieve_tool_definition()["function"]} + + +def _openai_responses_payload(output_item: dict) -> dict: + return { + "id": "resp_ccr", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-4o", + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "text": {"format": {"type": "text"}}, + "truncation": "disabled", + } + + +def _openai_responses_retrieve_call_payload() -> dict: + return _openai_responses_payload( + { + "type": "function_call", + "id": "fc_ccr", + "call_id": "call_ccr", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + "status": "completed", + } + ) + + +def _openai_responses_text_payload(text: str) -> dict: + return _openai_responses_payload( + { + "type": "message", + "id": "msg_ccr", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ) + + @pytest.mark.parametrize( "call_type, stream, tools, expect_conversion", [ @@ -1958,12 +2071,14 @@ def _openai_text_payload(content: str) -> dict: (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), (CallTypes.acompletion, True, None, False), - (CallTypes.aresponses, True, [_retrieve_tool_definition()], False), + (CallTypes.aresponses, True, [_retrieve_tool_definition()], True), + (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True), + (CallTypes.aresponses, False, [_retrieve_tool_definition()], False), (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), ], ) @pytest.mark.asyncio -async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions( +async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses( guardrail: HeadroomGuardrail, call_type: CallTypes, stream: bool, @@ -1986,6 +2101,22 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple assert kwargs["stream"] is True +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail): + kwargs = { + "model": "gpt-4o", + "stream": True, + "background": True, + "tools": [_responses_retrieve_tool_definition()], + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses) + + assert result is kwargs + assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs + assert kwargs["stream"] is True + + @pytest.mark.asyncio async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs( guardrail: HeadroomGuardrail, @@ -2094,6 +2225,117 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( assert not any(key.startswith("_headroom_interception") for key in followup_body) +@pytest.mark.asyncio +async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """Regression test for LIT-6481: streaming /v1/responses must resolve the + retrieve tool call server-side exactly like streaming /chat/completions does, + instead of streaming a headroom_retrieve function_call to the client.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.aresponses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = [event async for event in response] + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any("function_call" in str(getattr(event, "type", "")) for event in events) + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["input"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + + +def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """The synchronous responses() path converts the stream the same way, so it + must hand back a stream iterator with the resolved answer rather than the + completed response object.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = litellm.responses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = list(response) + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert len(upstream.calls) == 2 + assert not json.loads(upstream.calls[1].request.content).get("stream") + + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. # diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index a28a2a71613..a579370ad3c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -75,19 +75,29 @@ class _NoopTranslation(BaseTranslation): return response +def _patch_translation_mappings(monkeypatch, mappings): + """Point the unified guardrail at ``mappings`` for one test, restored by pytest. + + Every override goes through this one seam: competing writers to the same state + are what leaked a stale handler map into unrelated test files (LIT-6834). + """ + monkeypatch.setattr(unified_module, "load_guardrail_translation_mappings", lambda: mappings) + + @pytest.fixture(autouse=True) -def _inject_mcp_handler_mapping(): +def _inject_mcp_handler_mapping(monkeypatch): """Inject MCP handler mapping so the unified guardrail can run inside tests.""" - unified_module.endpoint_guardrail_translation_mappings = { - CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, - CallTypes.anthropic_messages: _NoopTranslation, - CallTypes.ocr: OCRHandler, - CallTypes.aocr: OCRHandler, - CallTypes.responses: OpenAIResponsesHandler, - CallTypes.aresponses: OpenAIResponsesHandler, - } - yield - unified_module.endpoint_guardrail_translation_mappings = None + _patch_translation_mappings( + monkeypatch, + { + CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, + CallTypes.anthropic_messages: _NoopTranslation, + CallTypes.ocr: OCRHandler, + CallTypes.aocr: OCRHandler, + CallTypes.responses: OpenAIResponsesHandler, + CallTypes.aresponses: OpenAIResponsesHandler, + }, + ) class TestUnifiedLLMGuardrails: @@ -396,7 +406,7 @@ class TestUnifiedLLMGuardrails: class TestAsyncPostCallStreamingIteratorHook: @pytest.mark.asyncio - async def test_streaming_content_not_lost_on_sampled_chunks(self): + async def test_streaming_content_not_lost_on_sampled_chunks(self, monkeypatch): """ Verify that every chunk's content is preserved in the output stream. @@ -442,10 +452,7 @@ class TestUnifiedLLMGuardrails: return responses_so_far - # Override the mapping to use our content-clearing translation - unified_module.endpoint_guardrail_translation_mappings = { - CallTypes.acompletion: _ContentClearingTranslation, - } + _patch_translation_mappings(monkeypatch, {CallTypes.acompletion: _ContentClearingTranslation}) handler = UnifiedLLMGuardrails() guardrail = RecordingGuardrail() @@ -885,12 +892,8 @@ class TestStreamingTransform: completions streaming surface.""" @pytest.fixture(autouse=True) - def _use_openai_handler_mapping(self): - unified_module.endpoint_guardrail_translation_mappings = { - CallTypes.acompletion: OpenAIChatCompletionsHandler, - } - yield - unified_module.endpoint_guardrail_translation_mappings = None + def _use_openai_handler_mapping(self, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler}) @pytest.mark.asyncio async def test_block_only_drops_text_rewrites(self): @@ -1719,6 +1722,10 @@ class TestAppliedGuardrailsReflectsExecution: decision and marks itself only when it actually ran (LIT-4650). Ordinary guardrails are still auto-marked by the hook after dispatch.""" + @pytest.fixture(autouse=True) + def _use_texts_only_mapping(self, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.pass_through: _TextsOnlyTranslation}) + @staticmethod def _data(guardrail): return { @@ -1728,7 +1735,6 @@ class TestAppliedGuardrailsReflectsExecution: } async def _run(self, guardrail): - unified_module.endpoint_guardrail_translation_mappings = {CallTypes.pass_through: _TextsOnlyTranslation} data = self._data(guardrail) await UnifiedLLMGuardrails().async_pre_call_hook( user_api_key_dict=None, @@ -1830,10 +1836,8 @@ class TestStreamingHttpErrorFrames: silently truncates the SSE stream (PR #38722 defect 1).""" @pytest.fixture(autouse=True) - def _use_real_mappings(self): - unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() - yield - unified_module.endpoint_guardrail_translation_mappings = None + def _use_real_mappings(self, monkeypatch): + _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) @pytest.mark.asyncio async def test_chat_eos_block_emits_data_error_frame(self): @@ -1938,10 +1942,8 @@ class TestStreamingGuardrailInformationBucket: guardrail_information write was diverted and /spend/logs showed null.""" @pytest.fixture(autouse=True) - def _use_real_mappings(self): - unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() - yield - unified_module.endpoint_guardrail_translation_mappings = None + def _use_real_mappings(self, monkeypatch): + _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) @pytest.mark.asyncio async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self): @@ -2038,11 +2040,7 @@ class TestStreamingScanDedup: @pytest.fixture(autouse=True) def _use_real_mappings(self, monkeypatch): - monkeypatch.setattr( - unified_module, - "endpoint_guardrail_translation_mappings", - load_guardrail_translation_mappings(), - ) + _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) @pytest.mark.asyncio async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): @@ -2239,3 +2237,55 @@ class TestStreamingScanDedup: assert out == chunks assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + +class TestTranslationMappingsAreReadLive: + """The hooks must read the handler map on every call, never memoize it on the module. + + A second module-level cache is what let one test's handler map outlive its own + teardown and decide how unrelated files translated their streams (LIT-6834). + """ + + @staticmethod + def _ocr_request(guardrail): + return { + "guardrail_to_apply": guardrail, + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234", + }, + } + + async def _run_pre_call(self, guardrail): + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + cache=DualCache(), + data=self._ocr_request(guardrail), + call_type=CallTypes.aocr.value, + ) + + @pytest.mark.asyncio + async def test_remapping_between_calls_changes_which_handler_runs(self, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.completion: _NoopTranslation}) + unmapped = RecordingGuardrail() + await self._run_pre_call(unmapped) + assert unmapped.apply_calls == [] + + _patch_translation_mappings(monkeypatch, {CallTypes.aocr: OCRHandler}) + mapped = RecordingGuardrail() + await self._run_pre_call(mapped) + assert [call["input_type"] for call in mapped.apply_calls] == ["request"] + + @pytest.mark.asyncio + async def test_module_exposes_no_second_assignable_handler_map(self, monkeypatch): + _patch_translation_mappings(monkeypatch, {CallTypes.aocr: OCRHandler}) + guardrail = RecordingGuardrail() + await self._run_pre_call(guardrail) + + assert len(guardrail.apply_calls) == 1 + assert not [ + name + for name, value in vars(unified_module).items() + if isinstance(value, dict) and CallTypes.aocr in value + ] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index b99dcb062b4..18aa43f7d1c 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -802,7 +802,7 @@ async def test_bedrock_guardrail_prepare_request_with_api_key(): @pytest.mark.asyncio -async def test_bedrock_guardrail_prepare_request_without_api_key(): +async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch): """Test _prepare_request method falls back to SigV4 when no api_key is provided""" from unittest.mock import Mock, patch @@ -820,18 +820,13 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): # Test data without api_key test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) with ( - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, patch("botocore.awsrequest.AWSRequest") as mock_aws_request, ): - # Mock no AWS_BEARER_TOKEN_BEDROCK - mock_get_secret.return_value = None - # Mock SigV4Auth mock_sigv4_instance = Mock() mock_sigv4_auth.return_value = mock_sigv4_instance @@ -857,7 +852,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): @pytest.mark.asyncio -async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): +async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(monkeypatch): """Test _prepare_request method uses Bearer token from environment when available""" from unittest.mock import Mock, patch @@ -875,15 +870,9 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): # Test data without api_key test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-456") - with ( - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, - patch("botocore.awsrequest.AWSRequest") as mock_aws_request, - ): - - mock_get_secret.return_value = "env-bearer-token-456" + with patch("botocore.awsrequest.AWSRequest") as mock_aws_request: mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 56661b5b843..836668de0c8 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -155,29 +155,103 @@ def test_duplicate_config_guardrail_names_get_distinct_stable_ids(): registry_module.guardrail_initializer_registry.pop("dup_name_test", None) -def test_update_in_memory_guardrail(): - handler = InMemoryGuardrailHandler() - handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail( - guardrail_name="test-guardrail", - default_on=False, - event_hook=GuardrailEventHooks.pre_call, - ) +def _register_mode_following_initializer(guardrail_type: str): + """Registers like the shipped initializers do: construct, then add the instance to litellm's callbacks.""" + import litellm + from litellm.proxy.guardrails import guardrail_registry as registry_module - handler.update_in_memory_guardrail( - "123", - Guardrail( - guardrail_name="test-guardrail", - litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True), - ), - ) - - assert ( - handler.guardrail_id_to_custom_guardrail["123"].should_run_guardrail( - data={}, event_type=GuardrailEventHooks.pre_call + def _initializer(litellm_params, guardrail): + callback = CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + event_hook=GuardrailEventHooks(litellm_params.mode), + default_on=True, ) - is True + litellm.logging_callback_manager.add_litellm_callback(callback) + return callback + + registry_module.guardrail_initializer_registry[guardrail_type] = _initializer + return registry_module + + +def _mode_following_db_row(guardrail_id: str, mode: str, description: str = "") -> Guardrail: + """The raw row GuardrailRegistry.update_guardrail_in_db hands back: litellm_params is a plain dict.""" + return Guardrail( + guardrail_id=guardrail_id, + guardrail_name="mode-following", + litellm_params={"guardrail": "mode_following_test", "mode": mode, "default_on": True}, + guardrail_info={"description": description}, ) - assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call + + +def _live_instances_named(name: str) -> int: + return sum(1 for cb_list in _all_callback_lists() for cb in cb_list if getattr(cb, "guardrail_name", None) == name) + + +def test_update_in_memory_guardrail_raw_db_row_mode_change_gates_at_the_new_stage(): + registry_module = _register_mode_following_initializer("mode_following_test") + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler = InMemoryGuardrailHandler() + handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call"), source="db") + original = handler.guardrail_id_to_custom_guardrail["123"] + + handler.update_in_memory_guardrail("123", _mode_following_db_row("123", "post_call")) + + replacement = handler.guardrail_id_to_custom_guardrail["123"] + assert replacement is not original + assert replacement.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True + assert replacement.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False + assert all(original not in cb_list for cb_list in lists) + assert _live_instances_named("mode-following") == 1 + assert handler.IN_MEMORY_GUARDRAILS["123"]["litellm_params"].mode == "post_call" + assert handler.get_source("123") == "db" + finally: + registry_module.guardrail_initializer_registry.pop("mode_following_test", None) + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_update_in_memory_guardrail_unchanged_params_keep_the_live_instance(): + registry_module = _register_mode_following_initializer("mode_following_test") + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler = InMemoryGuardrailHandler() + handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call", "old"), source="db") + original = handler.guardrail_id_to_custom_guardrail["123"] + + handler.update_in_memory_guardrail("123", _mode_following_db_row("123", "pre_call", "new")) + + assert handler.guardrail_id_to_custom_guardrail["123"] is original + assert handler.IN_MEMORY_GUARDRAILS["123"]["guardrail_info"] == {"description": "new"} + assert _live_instances_named("mode-following") == 1 + finally: + registry_module.guardrail_initializer_registry.pop("mode_following_test", None) + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_update_in_memory_guardrail_invalid_row_keeps_the_previous_instance_enforcing(): + registry_module = _register_mode_following_initializer("mode_following_test") + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler = InMemoryGuardrailHandler() + handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call"), source="db") + + with pytest.raises(ValueError, match="not in the supported event hooks"): + handler.update_in_memory_guardrail("123", _mode_following_db_row("123", "during_call")) + + restored = handler.guardrail_id_to_custom_guardrail["123"] + assert restored.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True + assert handler.IN_MEMORY_GUARDRAILS["123"]["litellm_params"].mode == "pre_call" + assert _live_instances_named("mode-following") == 1 + finally: + registry_module.guardrail_initializer_registry.pop("mode_following_test", None) + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: @@ -557,7 +631,7 @@ def test_presidio_siblings_are_tracked_and_deleted_together(): cb_list[:] = snapshot -def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_stage(): +def test_update_in_memory_guardrail_rebuilds_presidio_siblings_and_keeps_their_stage(): import litellm handler = InMemoryGuardrailHandler() @@ -591,11 +665,15 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st ) handler.update_in_memory_guardrail(guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail=updated) - assert [callback.pii_entities_config for callback in tracked] == [{"EMAIL_ADDRESS": "MASK"}] * 3 + rebuilt = _presidio_callbacks_in(litellm.callbacks) + assert len(rebuilt) == 3 + assert [callback.pii_entities_config for callback in rebuilt] == [{"EMAIL_ADDRESS": "MASK"}] * 3 assert [ - (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked + (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in rebuilt ] == roles_before - assert _presidio_callbacks_in(litellm.callbacks) == tracked + assert not any(previous in rebuilt for previous in tracked) + assert handler.guardrail_id_to_custom_guardrail[PRESIDIO_SIBLINGS_GID] is rebuilt[0] + assert handler.guardrail_id_to_sibling_callbacks[PRESIDIO_SIBLINGS_GID] == tuple(rebuilt[1:]) finally: for cb_list, snapshot in zip(lists, snapshots): cb_list[:] = snapshot diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index 91a011a8234..203391aadad 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -5,10 +5,13 @@ from typing import Any, Dict import orjson import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.image_endpoints import endpoints @@ -115,3 +118,52 @@ async def test_image_generation_prompt_rerouting(monkeypatch): assert captured_route_request_data["prompt"] == "sanitized prompt" assert "messages" not in captured_route_request_data assert response.headers.get("x-callback-test") == "value" + + +def _image_edit_client(monkeypatch, captured: Dict[str, Any]) -> TestClient: + class CaptureProcessing: + def __init__(self, data: Dict[str, Any]) -> None: + captured.update(data) + + async def base_process_llm_request(self, **_: Any) -> Dict[str, Any]: + return {"data": [{"b64_json": "aGk="}]} + + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", CaptureProcessing) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + + app = FastAPI() + app.include_router(endpoints.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app) + + +def test_image_edit_multipart_n_reaches_the_provider_as_an_int(monkeypatch): + """A multipart `n` must not arrive as the string Starlette parsed it into.""" + captured: Dict[str, Any] = {} + + response = _image_edit_client(monkeypatch, captured).post( + "/v1/images/edits", + files={"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")}, + data={"model": "nova-canvas", "prompt": "add a hat", "n": "2", "size": "1024x1024"}, + ) + + assert response.status_code == 200 + assert captured["n"] == 2 + assert isinstance(captured["n"], int) + assert captured["size"] == "1024x1024" + assert captured["prompt"] == "add a hat" + + +def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch): + """An unparseable `n` still reaches the provider, which rejects it as before.""" + captured: Dict[str, Any] = {} + + response = _image_edit_client(monkeypatch, captured).post( + "/v1/images/edits", + files={"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")}, + data={"model": "nova-canvas", "prompt": "add a hat", "n": "two"}, + ) + + assert response.status_code == 200 + assert captured["n"] == "two" diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 022aeff4e20..0d3ea5863a2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2014,6 +2014,67 @@ async def test_get_users_user_id_partial_match(mocker): assert captured_where_conditions["user_id"]["in"] == ["user1", "user2", "user3"] +def test_get_users_search_matches_user_id_or_email(mocker): + """ + `search` ORs a case-insensitive contains match over user_id and user_email on both the rows + query and the count, while the legacy `user_email` param keeps filtering only user_email. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + searched_user_id = "a6f5c02b-0163-45ce-815f-f88d10e95686" + mock_user_row = mocker.MagicMock() + mock_user_row.user_id = searched_user_id + mock_user_row.model_dump.return_value = { + "user_id": searched_user_id, + "user_email": "search@example.com", + "user_role": "internal_user", + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + find_many_wheres = [] + count_wheres = [] + + async def mock_find_many(*args, **kwargs): + find_many_wheres.append(kwargs["where"]) + return [mock_user_row] + + async def mock_count(*args, **kwargs): + count_wheres.append(kwargs["where"]) + return 1 + + async def mock_key_count(*args, **kwargs): + return 0 + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many + mock_prisma_client.db.litellm_usertable.count = mock_count + mock_prisma_client.db.litellm_verificationtoken.count = mock_key_count + mocker.patch( # test-quality-ok: /user/list reads prisma_client off proxy_server at call time + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + search_response = client.get("/user/list", params={"search": "A6F5C02B-0163"}) + assert search_response.status_code == 200, search_response.text + expected_or = ( + {"user_id": {"contains": "A6F5C02B-0163", "mode": "insensitive"}}, + {"user_email": {"contains": "A6F5C02B-0163", "mode": "insensitive"}}, + ) + assert find_many_wheres == [{"OR": expected_or}] + assert count_wheres == [{"OR": expected_or}] + assert [user["user_id"] for user in search_response.json()["users"]] == [searched_user_id] + assert search_response.json()["total"] == 1 + + legacy_response = client.get("/user/list", params={"user_email": "search@example.com"}) + assert legacy_response.status_code == 200, legacy_response.text + assert find_many_wheres[-1] == {"user_email": {"contains": "search@example.com", "mode": "insensitive"}} + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + def test_update_internal_user_params_reset_max_budget_with_none(): """ Test that _update_internal_user_params allows setting max_budget to None. diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5fa59a85c9d..c69f8f20a13 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -28,9 +28,18 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_team_models, ) from litellm.proxy.utils import PrismaClient +from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +async def _passthrough_row(update_data): + return update_data + + +async def _write_empty_row(**kwargs): + return await kwargs["write_row"]({}) + + class MockPrismaClient: def __init__( self, @@ -1191,7 +1200,7 @@ class TestTeamModelSiblingRouting: team_id = "team_no_alias" public_name = "gpt-4.1-mini" - async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client): + async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client, slot=None): return MagicMock(model_id=str(uuid.uuid4())) mock_team_model_add = AsyncMock() @@ -1372,7 +1381,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert result.get("model_name", "").startswith("model_name_test_team_123_") @@ -1435,7 +1445,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1481,7 +1490,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=None, ) @@ -1490,39 +1498,72 @@ class TestTeamModelUpdate: mock_delete.assert_not_called() @pytest.mark.asyncio - async def test_rename_with_prisma_none_clears_patch_model_name(self): - """Rename path must clear patch_data.model_name even when prisma is unavailable (P1).""" + async def test_a_refused_row_write_leaves_the_team_untouched(self): + """The team's model list autocommits, so it is written only after the row write succeeded: a + refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name + whose row never changed.""" + from fastapi import HTTPException + from litellm.proxy.management_endpoints.model_management_endpoints import ( - _update_existing_team_model_assignment, + _update_team_model_in_db, ) from litellm.types.router import ModelInfo db_model = Deployment( - model_name="model_name_team_123_uuid1", + model_name="gpt-4o", litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), - model_info=ModelInfo( - team_id="team_123", team_public_model_name="old-public-name" + model_info=ModelInfo(), + ) + user_api_key_dict = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN) + events: list[str] = [] + written: dict[str, object] = {} + + def patch_data() -> updateDeployment: + return updateDeployment(model_name="team-public", model_info=ModelInfo(team_id="team_123")) + + async def refuse_row(update_data): + events.append("row") + raise HTTPException(status_code=403, detail="slot held") + + async def accept_row(update_data): + events.append("row") + written.update(update_data) + return update_data + + async def team_add(**_): + events.append("team_model_add") + + with ( + patch( # test-quality-ok: the team auth check needs a live DB; the write order is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action", + AsyncMock(return_value=True), ), - ) - patch_data = updateDeployment( - model_name="new-public-name", - model_info=ModelInfo(team_id="team_123"), - ) - user_api_key_dict = UserAPIKeyAuth( - user_id="test_user", - user_role=LitellmUserRoles.PROXY_ADMIN, - ) + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_add, + ), + ): + with pytest.raises(HTTPException): + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=refuse_row, + ) + assert events == ["row"] - await _update_existing_team_model_assignment( - team_id="team_123", - public_model_name="new-public-name", - db_model=db_model, - patch_data=patch_data, - user_api_key_dict=user_api_key_dict, - prisma_client=None, - ) - - assert patch_data.model_name is None + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=accept_row, + ) + assert events == ["row", "row", "team_model_add"] + assert str(written["model_name"]).startswith("model_name_team_123_") + assert "team-public" in str(written["model_info"]) @pytest.mark.asyncio async def test_rename_handles_legacy_string_model_info(self): @@ -1574,7 +1615,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1614,7 +1654,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert "403" in str(exc_info.value) @@ -1900,7 +1941,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) # team ACL must not be touched on a no-op edit @@ -4311,6 +4353,321 @@ class TestStrategyRouterWriteValidation: is None ) + @staticmethod + def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router: + return Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}}, + { + "model_name": "held-v2", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": "held-id"}, + }, + ], + heuristic_v2_router_limit=lambda: limit, + ) + + class _FakeTx: + """Stands in for a prisma transaction: records the raw statements and exposes the model table.""" + + def __init__(self, db_held: int) -> None: + self.db_held = db_held + self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] + self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) + + async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: + self.raw_calls.append((sql, args)) + return [{"held": self.db_held}] if "count(*)" in sql else [] + + async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + class _FakeDb: + """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" + + def __init__(self, db_held: int, existing_row: object = None) -> None: + self.db = self + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held) + self.litellm_proxymodeltable = MagicMock( + create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) + ) + + def tx(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self.tx_obj + + _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} + _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}} + + @pytest.mark.parametrize( + "incoming,existing,expected", + [ + (_V2, None, _V2), + (_V2, _V1, _V2), + (None, _V1, _V1), + (None, None, None), + ("no-config", _V2, _V2), + ], + ) + def test_effective_complexity_router_config( + self, incoming: object, existing: object, expected: object + ) -> None: + """A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _effective_complexity_router_config, + ) + from litellm.types.router import updateLiteLLMParams + + incoming_params = None if incoming is None else updateLiteLLMParams( + complexity_router_config=None if incoming == "no-config" else incoming + ) + existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing) + assert _effective_complexity_router_config(incoming_params, existing_params) == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,effective_config,db_held,config_holds_one,model_id,expected", + [ + (1, _V2, 1, False, None, "refused"), + (1, _V2, 0, True, None, "refused"), + (1, _V2, 0, False, None, "reserved"), + (1, _V2, 0, False, "held-id", "reserved"), + (2, _V2, 1, False, None, "reserved"), + (1, _V1, 5, True, None, "plain"), + (1, None, 5, True, None, "plain"), + (None, _V2, 5, True, None, "plain"), + ], + ) + async def test_heuristic_v2_slot_matrix( + self, + limit: int | None, + effective_config: object, + db_held: int, + config_holds_one: bool, + model_id: str | None, + expected: str, + ) -> None: + """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows + (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL + parameter, and every other write runs on the plain client with no lock.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + HEURISTIC_V2_SLOT_LOCK_KEY, + _heuristic_v2_slot, + ) + + fake = self._FakeDb(db_held) + live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None + with ( + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here + "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", + new=AsyncMock(), + ) as published, + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id): + pass + assert exc_info.value.status_code == 403 + assert "At most 1 auto-router" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables: + handle = tables + if expected == "plain": + await handle.create(data={}) + fake.litellm_proxymodeltable.create.assert_awaited_once_with(data={}) + assert fake.tx_obj.raw_calls == [] + return + assert handle is fake.tx_obj.litellm_proxymodeltable + published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") + (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls + assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql + assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,) + assert count_params == (model_id or "",) + + @pytest.mark.asyncio + async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: + """team_model_add needs a second pool connection, so it must run only after the slot transaction + (and its advisory lock) has closed; a pool-sized burst of team creates would otherwise stall on the + lock holder waiting for a connection the waiters are occupying.""" + from contextlib import asynccontextmanager + + from litellm.proxy.management_endpoints.model_management_endpoints import _add_team_model_to_db + from litellm.types.router import ModelInfo + + events: list[str] = [] + created = MagicMock(model_id="row-1") + + @asynccontextmanager + async def slot(): + events.append("slot-enter") + yield MagicMock(create=AsyncMock(return_value=created)) + events.append("slot-exit") + + async def team_model_add(**_: object) -> None: + events.append("team_model_add") + + deployment = Deployment( + model_name="public-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + model_info=ModelInfo(id="row-1", team_id="team-1"), + ) + with ( + patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_model_add, + ), + ): + result = await _add_team_model_to_db( + model_params=deployment, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + prisma_client=MagicMock(), + slot=slot(), + ) + + assert result is created + assert events == ["slot-enter", "slot-exit", "team_model_add"] + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="second-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "At most 1 auto-router" in str(exc_info.value.message) + fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() + fake.litellm_proxymodeltable.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: the write must be refused before this DB step runs + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=self._db_complexity_router(model_id)), + ), + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: the helper's team bookkeeping needs a live DB; the row writer it is handed is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=_write_empty_row), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)), + user_api_key_dict=admin, + ) + assert exc_info.value.status_code == 403 + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + from litellm.types.router import ModelInfo, updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "model_name": "my-auto-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": model_id}, + } + existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"] + fake = self._FakeDb(db_held=1, existing_row=existing_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._V2), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + @pytest.mark.asyncio async def test_update_model_rejects_prefix_strip(self): from litellm.proxy._types import ProxyException diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index a1c38d26b9d..d35b77f732c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -45,6 +45,10 @@ from litellm.types.router import ( from litellm.types.utils import Usage +async def _passthrough_row(update_data): + return update_data + + def test_model_info_accepts_valid_ptu_fields(): info = ModelInfo( id="x", @@ -385,6 +389,7 @@ class TestTeamModelUpdateValidatesBeforeWriting: patch_data=patch_data, user_api_key_dict=MagicMock(), prisma_client=MagicMock(), + write_row=_passthrough_row, ) return result, touched @@ -914,6 +919,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: patch_data=patch, user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN), prisma_client=MagicMock(), + write_row=_passthrough_row, ) assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 088e82b370f..019ebc9807c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2910,6 +2910,7 @@ async def test_update_team_with_team_member_budget_duration( "metadata": {"team_member_budget_id": "budget_123"}, } mock_existing_team.metadata = {"team_member_budget_id": "budget_123"} + mock_existing_team.members_with_roles = [] mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) @@ -9897,11 +9898,11 @@ class TestResolveTeamAccessGroupResources: assert resolved.access_group_mcp_server_ids == ["mcp-1"] assert resolved.access_group_agent_ids == ["agent-1"] assert [ - (d.access_group_id, d.access_group_name, d.models) + (d.access_group_id, d.access_group_name, d.models, d.mcp_server_ids, d.agent_ids) for d in (resolved.access_group_details or []) ] == [ - ("ag-1", "shared-models", ("gpt-4", "claude-3")), - ("ag-2", "extra-models", ("claude-3", "gemini")), + ("ag-1", "shared-models", ("gpt-4", "claude-3"), ("mcp-1",), ()), + ("ag-2", "extra-models", ("claude-3", "gemini"), (), ("agent-1",)), ] @pytest.mark.asyncio @@ -11290,6 +11291,78 @@ async def test_patch_preserves_required_metadata_key_that_post_would_wipe(): assert patch_meta == {"cost_center": "FINOPS-1", "team_notes": "edited"} # preserved by PATCH +_STORED_METADATA_WITH_BUDGET: Final = { + "team_member_budget_id": "budget-existing-123", + "team_member_key_duration": "30d", + "logging": [{"callback_name": "langfuse", "callback_type": "success"}], + "cost_center": "cc-1234", +} + + +async def _written_metadata_with_budget(kind, body): + """Like ``_written_metadata`` but the team already owns a member budget row.""" + from litellm.proxy._types import LiteLLM_BudgetTable + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( # test-quality-ok: update_team imports update_budget at call time; the module attribute is its only seam + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + AsyncMock(return_value=LiteLLM_BudgetTable(budget_id="budget-existing-123")), + ), + ): + return await _written_metadata(kind, dict(_STORED_METADATA_WITH_BUDGET), body) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +@pytest.mark.parametrize( + "body", + [ + {"team_member_budget": 50.0}, + {"team_member_budget_duration": "1d"}, + {"team_member_tpm_limit": 500}, + {"team_member_rpm_limit": 5}, + ], + ids=lambda body: next(iter(body)), +) +async def test_team_member_budget_only_update_preserves_stored_metadata(kind, body): + """LIT-5150: a budget-only update that omits ``metadata`` must not replace the + stored metadata JSON with just ``{"team_member_budget_id": ...}``.""" + assert await _written_metadata_with_budget(kind, body) == _STORED_METADATA_WITH_BUDGET + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_team_member_key_duration_only_update_preserves_stored_metadata(kind): + """LIT-5150: a metadata-backed field sent alone is merged into the stored + metadata instead of becoming the whole metadata JSON.""" + written = await _written_metadata_with_budget(kind, {"team_member_key_duration": "7d"}) + + assert written == {**_STORED_METADATA_WITH_BUDGET, "team_member_key_duration": "7d"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_explicit_null_metadata_with_budget_field_still_clears_metadata(kind): + """``metadata: null`` is an explicit clear, so only the server-owned budget link survives.""" + written = await _written_metadata_with_budget(kind, {"metadata": None, "team_member_budget": 7.0}) + + assert written == {"team_member_budget_id": "budget-existing-123"} + + +@pytest.mark.asyncio +async def test_metadata_only_update_keeps_team_member_budget_link(): + """LIT-5150: rewriting metadata without any team member field must not drop the + server-owned ``team_member_budget_id``, or the member budget silently resets.""" + body = {"metadata": {"cost_center": "cc-9999"}} + + post_meta = await _written_metadata_with_budget("post", body) + patch_meta = await _written_metadata_with_budget("patch", body) + + assert post_meta == {"cost_center": "cc-9999", "team_member_budget_id": "budget-existing-123"} + assert patch_meta == {**_STORED_METADATA_WITH_BUDGET, "cost_center": "cc-9999"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "body, field, expected", @@ -11318,17 +11391,15 @@ async def test_top_level_fields_identical_post_and_patch(body, field, expected): @pytest.mark.asyncio async def test_patch_strips_system_managed_metadata_key_like_post(): """A caller cannot inject/overwrite server-owned keys via PATCH any more than - via POST: team_member_budget_id is stripped from the write in both.""" + via POST: the stored team_member_budget_id wins over the caller's value in both.""" existing = {"team_member_budget_id": "budget-123", "cost_center": "1234"} body = {"metadata": {"team_member_budget_id": "HACKED", "cost_center": "9999"}} post_meta = await _written_metadata("post", existing, body) patch_meta = await _written_metadata("patch", existing, body) - assert "team_member_budget_id" not in post_meta - assert "team_member_budget_id" not in patch_meta - assert post_meta == {"cost_center": "9999"} - assert patch_meta == {"cost_center": "9999"} + assert post_meta == {"cost_center": "9999", "team_member_budget_id": "budget-123"} + assert patch_meta == {"cost_center": "9999", "team_member_budget_id": "budget-123"} @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 9d1975513a1..c7696079adc 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -292,8 +292,8 @@ class TestUnifiedGuardrailCallTypeResolution: with patch.object( unified_guardrail_module, - "endpoint_guardrail_translation_mappings", - {CallTypes.pass_through: mock_handler_class}, + "load_guardrail_translation_mappings", + lambda: {CallTypes.pass_through: mock_handler_class}, ): result = await unified.async_post_call_success_hook( data=data, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 961479c0393..6735f2a3780 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -5,6 +5,7 @@ import pytest from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _base_vertex_proxy_route, + _upstream_headers_for_vertex_route, ) from litellm.types.router import DeploymentTypedDict @@ -348,6 +349,93 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): assert headers_passed_through is False +VERTEX_ANTHROPIC_MODELS_PREFIX = "v1/projects/test-project/locations/global/publishers/anthropic/models/" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_segment", "expects_anthropic_beta"), + [ + ("count-tokens:rawPredict", False), + ("claude-sonnet-4-6:streamRawPredict", True), + ], +) +async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( + model_segment: str, expects_anthropic_beta: bool +): + with ( + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.proxy_server.llm_router", None + ), + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router" + ) as mock_pt_router, + patch( # test-quality-ok: the route offers no injection point for its header preparation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", + new_callable=AsyncMock, + ) as mock_prep_headers, + patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route, + patch( # test-quality-ok: the route calls auth directly rather than through Depends + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", + new_callable=AsyncMock, + return_value=False, + ), + ): + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ( + { + "anthropic-beta": "tool-search-tool-2025-10-19,web-search-2025-03-05", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + }, + "https://aiplatform.googleapis.com", + False, + "test-project", + "global", + ) + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key") + + await _base_vertex_proxy_route( + endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", + request=MagicMock(), + fastapi_response=MagicMock(), + get_vertex_pass_through_handler=MagicMock(), + ) + + upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] + assert ("anthropic-beta" in upstream_headers) is expects_anthropic_beta + assert upstream_headers["Authorization"] == "Bearer vertex-access-token" + assert upstream_headers["content-type"] == "application/json" + + +def test_upstream_headers_for_vertex_route_filters_anthropic_beta_by_route(): + headers = { + "Anthropic-Beta": "effort-2025-11-24", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + + count_tokens_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}count-tokens:rawPredict", headers + ) + model_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}claude-sonnet-4-6:rawPredict", headers + ) + + assert dict(count_tokens_headers) == { + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + assert dict(model_headers) == headers + + @pytest.mark.asyncio async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): """ diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py index 6ebb10eff76..41c3003af34 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py @@ -104,97 +104,6 @@ class TestPromptVersioning: assert get_base_prompt_id(prompt_id="jack") == "jack" assert get_base_prompt_id(prompt_id="my_prompt.v10") == "my_prompt" - def test_get_latest_version_prompt_id(self): - """ - Test that get_latest_version_prompt_id returns the highest version - """ - from litellm.proxy.prompts.prompt_endpoints import get_latest_version_prompt_id - - # Mock prompt IDs dictionary - all_prompt_ids = { - "jack.v1": {}, - "jack.v2": {}, - "jack.v3": {}, - "jane.v1": {}, - "simple_prompt": {}, - } - - # Test with base prompt ID - should return latest version - assert ( - get_latest_version_prompt_id( - prompt_id="jack", all_prompt_ids=all_prompt_ids - ) - == "jack.v3" - ) - - # Test with versioned prompt ID - should still return latest version - assert ( - get_latest_version_prompt_id( - prompt_id="jack.v1", all_prompt_ids=all_prompt_ids - ) - == "jack.v3" - ) - - # Test with single version - assert ( - get_latest_version_prompt_id( - prompt_id="jane", all_prompt_ids=all_prompt_ids - ) - == "jane.v1" - ) - - # Test with non-versioned prompt - assert ( - get_latest_version_prompt_id( - prompt_id="simple_prompt", all_prompt_ids=all_prompt_ids - ) - == "simple_prompt" - ) - - # Test with non-existent prompt - assert ( - get_latest_version_prompt_id( - prompt_id="nonexistent", all_prompt_ids=all_prompt_ids - ) - == "nonexistent" - ) - - def test_construct_versioned_prompt_id(self): - """ - Test that construct_versioned_prompt_id correctly builds versioned IDs - """ - from litellm.proxy.prompts.prompt_endpoints import construct_versioned_prompt_id - - # Test with base prompt ID and version - assert ( - construct_versioned_prompt_id(prompt_id="jack_success", version=4) - == "jack_success.v4" - ) - - # Test with None version - should return base ID unchanged - assert ( - construct_versioned_prompt_id(prompt_id="jack_success", version=None) - == "jack_success" - ) - - # Test with existing versioned ID - should replace version - assert ( - construct_versioned_prompt_id(prompt_id="jack_success.v2", version=4) - == "jack_success.v4" - ) - - # Test with hyphenated prompt ID - assert ( - construct_versioned_prompt_id(prompt_id="my-prompt", version=1) - == "my-prompt.v1" - ) - - # Test with double-digit version - assert ( - construct_versioned_prompt_id(prompt_id="test_prompt", version=10) - == "test_prompt.v10" - ) - class TestPromptVersionsEndpoint: """ @@ -444,7 +353,7 @@ class TestAdminViewerReadAccess: "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = PromptSpec( + mock_registry.resolve_prompt_spec.return_value = PromptSpec( prompt_id="jack.v2", litellm_params=PromptLiteLLMParams( prompt_id="jack", @@ -453,10 +362,101 @@ class TestAdminViewerReadAccess: ), prompt_info=PromptInfo(prompt_type="db"), ) - mock_registry.IN_MEMORY_PROMPTS = {"jack.v1": {}, "jack.v2": {}} - mock_registry.get_prompt_callback_by_id.return_value = None + mock_registry.get_prompt_callback_for_prompt.return_value = None response = await get_prompt_info(prompt_id="jack", user_api_key_dict=viewer) assert response.prompt_spec.prompt_id == "jack" assert response.prompt_spec.version == 2 + + +class TestConfigPromptInfoWithEnvironment: + """ + Regression: /prompts/{id}/info with an environment param must still resolve + config-file (in-memory) prompts on a DB-backed proxy instead of 400ing. + """ + + def _registry_with_config_prompt(self): + from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry + + registry = InMemoryPromptRegistry() + registry.IN_MEMORY_PROMPTS["envgreet::development"] = PromptSpec( + prompt_id="envgreet", + litellm_params=PromptLiteLLMParams( + prompt_id="envgreet", + prompt_integration="dotprompt", + dotprompt_content="AHOY {{user_message}}", + ), + prompt_info=PromptInfo(prompt_type="config"), + ) + return registry + + def _prisma_client_with_empty_prompt_table(self): + from unittest.mock import AsyncMock + + mock_prisma = MagicMock() + mock_prisma.db.litellm_prompttable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + @pytest.mark.asyncio + async def test_get_prompt_info_with_environment_falls_back_to_registry(self): + from unittest.mock import patch + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.prompts.prompt_endpoints import get_prompt_info + + admin = UserAPIKeyAuth( + api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.proxy_server.prisma_client", + self._prisma_client_with_empty_prompt_table(), + ), + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY", + self._registry_with_config_prompt(), + ), + ): + response = await get_prompt_info( + prompt_id="envgreet", + environment="development", + user_api_key_dict=admin, + ) + + assert response.prompt_spec.prompt_id == "envgreet" + assert response.prompt_spec.litellm_params.dotprompt_content == "AHOY {{user_message}}" + + @pytest.mark.asyncio + async def test_get_prompt_info_with_wrong_environment_still_400s(self): + from unittest.mock import patch + + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.prompts.prompt_endpoints import get_prompt_info + + admin = UserAPIKeyAuth( + api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.proxy_server.prisma_client", + self._prisma_client_with_empty_prompt_table(), + ), + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY", + self._registry_with_config_prompt(), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await get_prompt_info( + prompt_id="envgreet", + environment="production", + user_api_key_dict=admin, + ) + + assert exc_info.value.status_code == 400 + assert "environment production" in exc_info.value.detail diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py index b5792ac7572..387448bb9b3 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -52,8 +52,6 @@ async def test_delete_prompt_success(): with patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - # User passes "test_prompt.v2" - # We simulate that get_prompt_by_id returns the prompt spec for v2 prompt_spec = PromptSpec( prompt_id="test_prompt.v2", litellm_params=PromptLiteLLMParams( @@ -61,7 +59,8 @@ async def test_delete_prompt_success(): ), prompt_info=PromptInfo(prompt_type="db"), ) - mock_registry.get_prompt_by_id.return_value = prompt_spec + mock_registry.resolve_prompt_spec.return_value = prompt_spec + mock_registry.has_config_prompt.return_value = False # Patch the prisma client in the endpoint module with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): @@ -79,7 +78,7 @@ async def test_delete_prompt_success(): # 2. Memory deletion should use base ID mock_registry.delete_prompts_by_base_id.assert_called_once_with( - expected_base_id, environment=None + base_prompt_id=expected_base_id, environment=None ) assert response == { @@ -108,31 +107,14 @@ async def test_delete_prompt_by_base_id_success(): with patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - # User passes "test_prompt" (base ID) - # 1. get_prompt_by_id("test_prompt") -> None (if it's not registered as base) - # 2. It calls get_latest_version_prompt_id -> returns "test_prompt.v3" - # 3. get_prompt_by_id("test_prompt.v3") -> returns Spec - - # Setup mocks behavior - def get_prompt_side_effect(prompt_id): - if prompt_id == "test_prompt": - return None - if prompt_id == "test_prompt.v3": - return PromptSpec( - prompt_id="test_prompt.v3", - litellm_params=PromptLiteLLMParams( - prompt_id="test_prompt", prompt_integration="dotprompt" - ), - prompt_info=PromptInfo(prompt_type="db"), - ) - return None - - mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect - mock_registry.IN_MEMORY_PROMPTS = { - "test_prompt.v1": {}, - "test_prompt.v2": {}, - "test_prompt.v3": {}, - } + mock_registry.resolve_prompt_spec.return_value = PromptSpec( + prompt_id="test_prompt.v3", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + mock_registry.has_config_prompt.return_value = False # Patch the prisma client in the endpoint module with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): @@ -150,7 +132,7 @@ async def test_delete_prompt_by_base_id_success(): # 2. Memory deletion should use base ID mock_registry.delete_prompts_by_base_id.assert_called_once_with( - expected_base_id, environment=None + base_prompt_id=expected_base_id, environment=None ) assert response == { @@ -169,11 +151,12 @@ async def test_delete_prompt_environment_scope_reaches_db_and_registry(): with patch( # test-quality-ok: stubs the collaborator so the test pins what the endpoint deletes "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - mock_registry.get_prompt_by_id.return_value = PromptSpec( + mock_registry.resolve_prompt_spec.return_value = PromptSpec( prompt_id="test_prompt.v2", litellm_params=PromptLiteLLMParams(prompt_id="test_prompt", prompt_integration="dotprompt"), - prompt_info=PromptInfo(prompt_type="db"), + prompt_info=PromptInfo(prompt_type="db", environment="production"), ) + mock_registry.has_config_prompt.return_value = False with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): # test-quality-ok: proxy_server module global is the endpoint's only injection point response = await delete_prompt( @@ -185,7 +168,7 @@ async def test_delete_prompt_environment_scope_reaches_db_and_registry(): mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( where={"prompt_id": "test_prompt", "environment": "production"} ) - mock_registry.delete_prompts_by_base_id.assert_called_once_with("test_prompt", environment="production") + mock_registry.delete_prompts_by_base_id.assert_called_once_with(base_prompt_id="test_prompt", environment="production") assert response == {"message": "Prompt test_prompt deleted successfully from production"} @@ -218,24 +201,8 @@ async def test_get_prompt_info_by_base_id(): prompt_info=PromptInfo(prompt_type="db"), ) - # When get_prompt_by_id is called with "test_prompt", return None (so it searches versions) - # When called with "test_prompt.v3", return the spec - def get_prompt_side_effect(prompt_id): - if prompt_id == "test_prompt": - return None - if prompt_id == "test_prompt.v3": - return prompt_spec_v3 - return None - - mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect - mock_registry.IN_MEMORY_PROMPTS = { - "test_prompt.v1": {}, - "test_prompt.v2": {}, - "test_prompt.v3": {}, - } - - # We also need to mock get_prompt_callback_by_id to avoid content extraction errors/logic - mock_registry.get_prompt_callback_by_id.return_value = None + mock_registry.resolve_prompt_spec.return_value = prompt_spec_v3 + mock_registry.get_prompt_callback_for_prompt.return_value = None response = await get_prompt_info( prompt_id="test_prompt", user_api_key_dict=mock_user_auth @@ -284,7 +251,7 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404(): "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = existing_prompt + mock_registry.has_config_prompt.return_value = False with pytest.raises(HTTPException) as exc_info: await patch_prompt( @@ -325,7 +292,7 @@ async def test_patch_prompt_merges_unsent_fields_from_db_row_not_stale_memory(): "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = stale_in_memory + mock_registry.has_config_prompt.return_value = False mock_registry.reload_prompt.side_effect = lambda prompt: prompt response = await patch_prompt( @@ -487,7 +454,7 @@ async def test_patch_prompt_info_only_keeps_legacy_keyed_row_patchable(): "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = existing_prompt + mock_registry.has_config_prompt.return_value = False await patch_prompt( prompt_id="agent-prompt", diff --git a/tests/test_litellm/proxy/prompts/test_prompt_environment.py b/tests/test_litellm/proxy/prompts/test_prompt_environment.py index ecd89afefbe..3cb647dea13 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_environment.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_environment.py @@ -191,11 +191,7 @@ async def test_update_prompt_stores_environment_and_created_by(): with patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - mock_registry.get_prompt_by_id.return_value = PromptSpec( - prompt_id="my_prompt.v1", - litellm_params=request.litellm_params, - prompt_info=PromptInfo(prompt_type="db"), - ) + mock_registry.has_config_prompt.return_value = False mock_registry.initialize_prompt.return_value = PromptSpec( prompt_id="my_prompt.v2", litellm_params=request.litellm_params, @@ -239,7 +235,8 @@ async def test_delete_prompt_scoped_to_environment(): prompt_info=PromptInfo(prompt_type="db"), environment="staging", ) - mock_registry.get_prompt_by_id.return_value = prompt_spec + mock_registry.resolve_prompt_spec.return_value = prompt_spec + mock_registry.has_config_prompt.return_value = False with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): await delete_prompt( @@ -251,3 +248,6 @@ async def test_delete_prompt_scoped_to_environment(): mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( where={"prompt_id": "test_prompt", "environment": "staging"} ) + mock_registry.delete_prompts_by_base_id.assert_called_once_with( + base_prompt_id="test_prompt", environment="staging" + ) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py index 3008821974e..a0743b65dd7 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_registry.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -1,26 +1,35 @@ import pytest import litellm -from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry +from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry, parse_prompt_version from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec -def _db_prompt_spec(content: str) -> PromptSpec: +def _db_prompt_spec(content: str, environment: str = "development", version: int = 1) -> PromptSpec: return PromptSpec( - prompt_id="greeting.v1", + prompt_id=f"greeting.v{version}", litellm_params=PromptLiteLLMParams( prompt_id="greeting", prompt_integration="dotprompt", prompt_data={"content": content, "metadata": {}}, ), prompt_info=PromptInfo(prompt_type="db"), + version=version, + environment=environment, ) -def _served_content(registry: InMemoryPromptRegistry) -> str: - callback = registry.get_prompt_callback_by_id("greeting.v1") +def _resolved_callback(registry: InMemoryPromptRegistry, environment: str | None = None) -> CustomPromptManagement: + spec = registry.resolve_prompt_spec("greeting", environment=environment) + assert spec is not None + callback = registry.get_prompt_callback_for_prompt(prompt=spec) assert callback is not None - return callback.prompt_manager.get_prompt("greeting").content + return callback + + +def _served_content(registry: InMemoryPromptRegistry, environment: str | None = None) -> str: + return _resolved_callback(registry, environment=environment).prompt_manager.get_prompt("greeting").content @pytest.fixture @@ -32,32 +41,34 @@ def isolated_callbacks(monkeypatch: pytest.MonkeyPatch) -> list: def test_sync_prompt_from_db_reloads_row_edited_elsewhere(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) - stale_callback = registry.get_prompt_callback_by_id("greeting.v1") + stale_callback = _resolved_callback(registry) assert _served_content(registry) == "begin every reply with AHOY" registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY")) assert _served_content(registry) == "begin every reply with HOWDY" - assert registry.get_prompt_by_id("greeting.v1").litellm_params.prompt_data["content"] == "begin every reply with HOWDY" + reloaded_spec = registry.resolve_prompt_spec("greeting", environment="development") + assert reloaded_spec is not None + assert reloaded_spec.litellm_params.prompt_data["content"] == "begin every reply with HOWDY" assert stale_callback not in isolated_callbacks - assert isolated_callbacks == [registry.get_prompt_callback_by_id("greeting.v1")] + assert isolated_callbacks == [_resolved_callback(registry)] def test_sync_prompt_from_db_keeps_unchanged_row_in_place(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) - first_callback = registry.get_prompt_callback_by_id("greeting.v1") + first_callback = _resolved_callback(registry) registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) - assert registry.get_prompt_callback_by_id("greeting.v1") is first_callback + assert _resolved_callback(registry) is first_callback assert isolated_callbacks == [first_callback] def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) - stale_callback = registry.get_prompt_callback_by_id("greeting.v1") + stale_callback = _resolved_callback(registry) reloaded = registry.reload_prompt(prompt=_db_prompt_spec("begin every reply with HOWDY")) @@ -70,7 +81,7 @@ def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_ca def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) - old_callback = registry.get_prompt_callback_by_id("greeting.v1") + old_callback = _resolved_callback(registry) broken = PromptSpec( prompt_id="greeting.v1", @@ -80,63 +91,134 @@ def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolate prompt_data={"content": "begin every reply with HOWDY", "metadata": {}}, ), prompt_info=PromptInfo(prompt_type="db"), + version=1, + environment="development", ) with pytest.raises(ValueError, match="Unsupported prompt"): registry.reload_prompt(prompt=broken) - assert registry.get_prompt_callback_by_id("greeting.v1") is old_callback + assert _resolved_callback(registry) is old_callback assert _served_content(registry) == "begin every reply with AHOY" assert isolated_callbacks == [old_callback] -def _versioned_prompt_spec(version: int, environment: str) -> PromptSpec: - return PromptSpec( - prompt_id=f"greeting.v{version}", +def test_environments_sharing_a_prompt_id_keep_separate_templates(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + assert _served_content(registry, environment="development") == "begin every reply with AHOY" + assert _served_content(registry, environment="production") == "begin every reply with HOWDY" + assert _resolved_callback(registry, environment="development") is not _resolved_callback( + registry, environment="production" + ) + + +def test_default_resolution_prefers_production(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + assert _served_content(registry) == "begin every reply with HOWDY" + + +@pytest.mark.parametrize("environment", ["staging", "qa"]) +def test_default_resolution_serves_the_only_environment_present(isolated_callbacks: list, environment: str) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment=environment)) + + assert _served_content(registry) == "begin every reply with AHOY" + + +def test_resolution_picks_exact_version_and_latest_within_an_environment(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development", version=1)) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with YO", environment="development", version=2)) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production", version=1)) + + exact = registry.resolve_prompt_spec("greeting", version=1, environment="development") + assert exact is not None + assert exact.litellm_params.prompt_data["content"] == "begin every reply with AHOY" + + latest = registry.resolve_prompt_spec("greeting", environment="development") + assert latest is not None + assert latest.litellm_params.prompt_data["content"] == "begin every reply with YO" + + assert registry.resolve_prompt_spec("greeting", version=3, environment="development") is None + + +def test_resolution_returns_none_for_unknown_environment_or_prompt(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + + assert registry.resolve_prompt_spec("greeting", environment="production") is None + assert registry.resolve_prompt_spec("no_such_prompt") is None + + +def test_delete_prompts_by_base_id_scoped_to_one_environment(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + production_callback = _resolved_callback(registry, environment="production") + + deleted = registry.delete_prompts_by_base_id(base_prompt_id="greeting", environment="development") + + assert deleted == ["greeting.v1::development"] + assert registry.resolve_prompt_spec("greeting", environment="development") is None + assert _resolved_callback(registry, environment="production") is production_callback + assert _served_content(registry, environment="production") == "begin every reply with HOWDY" + + deleted_rest = registry.delete_prompts_by_base_id(base_prompt_id="greeting") + + assert deleted_rest == ["greeting.v1::production"] + assert registry.resolve_prompt_spec("greeting") is None + + +def test_has_config_prompt_matches_any_version_of_the_base_id(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + config_spec = PromptSpec( + prompt_id="greeting", litellm_params=PromptLiteLLMParams( prompt_id="greeting", prompt_integration="dotprompt", - prompt_data={"content": f"begin every reply with AHOY v{version}", "metadata": {}}, + prompt_data={"content": "begin every reply with AHOY", "metadata": {}}, ), - prompt_info=PromptInfo(prompt_type="db", environment=environment), - version=version, - environment=environment, + prompt_info=PromptInfo(prompt_type="config"), ) + registry.initialize_prompt(prompt=config_spec) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + assert registry.has_config_prompt(base_prompt_id="greeting") is True + assert registry.has_config_prompt(base_prompt_id="other_prompt") is False def test_delete_prompts_by_base_id_removes_the_callbacks_from_litellm_callbacks(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() - registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development")) - registry.initialize_prompt(prompt=_versioned_prompt_spec(2, "development")) + registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY", version=1)) + registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with YO", version=2)) assert len(isolated_callbacks) == 1 - deleted = registry.delete_prompts_by_base_id("greeting") + deleted = registry.delete_prompts_by_base_id(base_prompt_id="greeting") - assert sorted(deleted) == ["greeting.v1", "greeting.v2"] - assert registry.get_prompt_by_id("greeting.v1") is None - assert registry.get_prompt_callback_by_id("greeting.v2") is None + assert sorted(deleted) == ["greeting.v1::development", "greeting.v2::development"] + assert registry.resolve_prompt_spec("greeting") is None assert isolated_callbacks == [] -def test_delete_prompts_by_base_id_environment_scope_keeps_other_environments(isolated_callbacks: list) -> None: +def test_remove_prompt_is_a_no_op_for_an_unknown_registry_key(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() - registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development")) - registry.initialize_prompt(prompt=_versioned_prompt_spec(2, "production")) - production_callback = registry.get_prompt_callback_by_id("greeting.v2") + registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) - deleted = registry.delete_prompts_by_base_id("greeting", environment="development") + registry.remove_prompt(registry_key="not_there.v1::development") - assert deleted == ["greeting.v1"] - assert registry.get_prompt_by_id("greeting.v1") is None - assert registry.get_prompt_by_id("greeting.v2") is not None - assert registry.get_prompt_callback_by_id("greeting.v2") is production_callback - - -def test_remove_prompt_is_a_no_op_for_an_unknown_id(isolated_callbacks: list) -> None: - registry = InMemoryPromptRegistry() - registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development")) - - registry.remove_prompt(prompt_id="not_there.v1") - - assert registry.get_prompt_by_id("greeting.v1") is not None + assert registry.resolve_prompt_spec("greeting") is not None assert len(isolated_callbacks) == 1 + + +@pytest.mark.parametrize( + ("raw_version", "expected"), + [(2, 2), ("2", 2), (None, None), ("v2", None), (True, None), (2.0, None)], +) +def test_parse_prompt_version_accepts_integers_and_json_strings(raw_version: object, expected: int | None) -> None: + assert parse_prompt_version(raw_version) == expected diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 1ab18639fff..dcfad8f6815 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import ( resolve_routing_plugins, validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, + validate_heuristic_v2_router_limit, ) from .conftest import normalize @@ -193,6 +194,120 @@ def test_validate_deployment_complexity_router_placement_leaves_valid_deployment assert model["litellm_params"] == litellm_params +def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + } + + +def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None: + """Same reason as the two validators above: the proxy router swallows registration errors, so + an over-limit config.yaml must fail here instead of booting with a silently missing router.""" + with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info: + validate_heuristic_v2_router_limit( + [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1 + ) + assert "'auto_router' feature lifts the limit" in str(exc_info.value) + + +@pytest.mark.parametrize( + "model_list,limit", + [ + ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None), + ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1), + ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1), + ], +) +def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone( + model_list: list[dict[str, object]], limit: int | None +) -> None: + assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None + + +_TWO_HEURISTIC_V2_ROUTERS_YAML = ( + "model_list:\n" + " - model_name: gpt-4o-mini\n" + " litellm_params:\n" + " model: openai/gpt-4o-mini\n" + " api_key: k\n" + " - model_name: v2-a\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + " - model_name: v2-b\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + "router_settings:\n" + " heuristic_v2_router_limit: 99\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("license_limit", [1, None]) +async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only( + tmp_path, monkeypatch, license_limit: int | None +) -> None: + """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator + cannot grant the entitlement by editing the config, and a licensed proxy boots both routers.""" + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit + ) + + if license_limit is None: + router, _model_list, _general_settings = await ProxyConfig().load_config( + router=None, config_file_path=str(f) + ) + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + return + + with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license( + tmp_path, monkeypatch +) -> None: + """config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB + is refused at registration because the router was built with the license's ceiling.""" + from litellm.types.router import Deployment + + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML.replace(" - model_name: v2-b\n", " - model_name: v1-b\n", 1).replace( + "classifier_type: heuristic_v2\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + "classifier_type: heuristic\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + )) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1) + + router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() == 1 + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"}) + assert router.upsert_deployment(db_row) is None + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + + def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key(): model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}} diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 0085b6ebd36..a176e91eaa4 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -421,6 +421,76 @@ def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_u assert forwarded_config["aws_region_name"] == "eu-west-1" +def test_rag_query_forwards_managed_store_credentials_to_search(client_internal_user): + """ + Regression for LIT-6773: the registry store's api_key / api_base and its + provider extras (Milvus outputFields, milvus_text_field) must reach the + vector store search the way the direct /v1/vector_stores/{id}/search + endpoint forwards them. Pre-fix the RAG path allowlisted them away and a + managed Milvus store 500'd with "MILVUS_API_KEY is not set". + """ + import litellm + from litellm import Router + from litellm.types.vector_stores import VectorStoreSearchResponse + + mock_vector_store = { + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "litellm_params": { + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "api_base": "http://127.0.0.1:19530", + "api_key": "root:Milvus", + "litellm_embedding_model": "multilingual-e5-large", + "milvus_text_field": "book_intro_text", + "outputFields": ["book_intro_text"], + }, + } + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + ) + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + + with ( + patch("litellm.vector_stores.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch.object(litellm, "vector_store_registry", mock_registry), # test-quality-ok: seeds the store under test + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: mock-response router for completion + patch( # test-quality-ok: store access is not under test, so the request reaches the search boundary + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "which database is built for similarity search?"}], + "retrieval_config": {"vector_store_id": "customer_kb", "custom_llm_provider": "milvus", "top_k": 2}, + }, + ) + + assert response.status_code == 200, response.json() + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["vector_store_id"] == "customer_kb" + assert search_kwargs["custom_llm_provider"] == "milvus" + assert search_kwargs["max_num_results"] == 2 + assert search_kwargs["api_base"] == "http://127.0.0.1:19530" + assert search_kwargs["api_key"] == "root:Milvus" + assert search_kwargs["litellm_embedding_model"] == "multilingual-e5-large" + assert search_kwargs["milvus_text_field"] == "book_intro_text" + assert search_kwargs["outputFields"] == ["book_intro_text"] + + @pytest.mark.parametrize( "blocked_key", ["embedding_model", "litellm_embedding_model", "litellm_embedding_config", "litellm_credential_name"], diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 60a32102946..f4cd8814bc1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -4234,6 +4234,91 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): assert call_args[2] == [api_key] +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens(): + """ + Regression test for LIT-4929: the logs table showed the summed session cost but + only the last call's token usage. Every row of a multi-round session must carry + the session-wide prompt, completion and total token sums from the aggregate + query, while rows outside a session carry none of them. + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + session_id = "sess-multi-round-tokens" + api_key = "hashed-key-xyz" + dict_rows = [ + { + "request_id": "req-1", + "session_id": session_id, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 10, + "prompt_tokens": 7, + "completion_tokens": 3, + }, + { + "request_id": "req-2", + "session_id": session_id, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 50, + "prompt_tokens": 35, + "completion_tokens": 15, + }, + { + "request_id": "req-3", + "session_id": None, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 5, + "prompt_tokens": 4, + "completion_tokens": 1, + }, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, + "session_total_spend": 0.06, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_total_prompt_tokens": 42, + "session_total_completion_tokens": 18, + "session_total_tokens": 60, + } + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=3, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + rows = result["data"] + session_rows = rows[:2] + assert [row["session_total_tokens"] for row in session_rows] == [60, 60] + assert [row["session_total_prompt_tokens"] for row in session_rows] == [42, 42] + assert [row["session_total_completion_tokens"] for row in session_rows] == [18, 18] + assert [(row["total_tokens"], row["prompt_tokens"], row["completion_tokens"]) for row in session_rows] == [ + (10, 7, 3), + (50, 35, 15), + ] + + token_keys = ("session_total_tokens", "session_total_prompt_tokens", "session_total_completion_tokens") + assert all(key not in rows[2] for key in token_keys) + + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_session_cache_hit_count(): """ diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 323930eee60..6a6db5ab7fe 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3399,6 +3399,138 @@ def test_get_spend_logs_id_prefers_the_response_id_over_the_standard_logging_id( ) +@pytest.mark.asyncio +async def test_spend_log_request_id_is_the_message_id_a_bridged_streaming_caller_was_streamed(): + """A streaming /v1/messages call against a non-Anthropic model is served a msg_ id the + adapter mints itself, and it is the only request id that call ever shows the caller, so + GET /spend/logs?request_id=msg_... has to land on the row.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( + AnthropicResponsesStreamWrapper, + ) + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, + ) + + logging_obj = Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(timezone.utc), + litellm_call_id="6825cafe-0000-4000-8000-000000000001", + function_id="1234", + ) + logging_obj.optional_params = {} + + completed_response = ResponsesAPIResponse( + id="resp_01Lit6825Bridged", + object="response", + created_at=1767225600, + model="gpt-5.6", + status="completed", + output=[ + { + "id": "msg_bridged_output", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "epsilon", "annotations": []}], + } + ], + usage=ResponseAPIUsage(input_tokens=12, output_tokens=5, total_tokens=17), + ) + + async def _responses_stream(): + yield {"type": "response.created"} + yield {"type": "response.output_text.delta", "item_id": "msg_bridged_output", "delta": "epsilon"} + yield ResponseCompletedEvent(type="response.completed", response=completed_response) + + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_responses_stream(), + model="gpt-5.6", + litellm_logging_obj=logging_obj, + ) + sse_frames = [frame.decode() async for frame in wrapper.async_anthropic_sse_wrapper()] + + message_start_frames = [f for f in sse_frames if f.startswith("event: message_start\n")] + assert len(message_start_frames) == 1 + streamed_message_id = json.loads(message_start_frames[0].split("data: ", 1)[1])["message"]["id"] + assert streamed_message_id.startswith("msg_") + + _, _, logged_response = logging_obj._success_handler_helper_fn( + result=ResponseCompletedEvent(type="response.completed", response=completed_response), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert logged_response.id == streamed_message_id + payload = get_logging_payload( + kwargs={ + "call_type": "anthropic_messages", + "model": "gpt-5.6", + "litellm_call_id": "6825cafe-0000-4000-8000-000000000001", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=logged_response, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["request_id"] == streamed_message_id + + +@pytest.mark.asyncio +async def test_spend_log_request_id_is_untouched_when_no_message_id_was_streamed(): + """Only the bridged streaming adapter mints a msg_ id of its own, so every other + /v1/messages call must keep the id its own response carried.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, + ) + + logging_obj = Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(timezone.utc), + litellm_call_id="6825cafe-0000-4000-8000-000000000002", + function_id="1234", + ) + logging_obj.optional_params = {} + + completed_response = ResponsesAPIResponse( + id="resp_01Lit6825Unbridged", + object="response", + created_at=1767225600, + model="gpt-5.6", + status="completed", + output=[ + { + "id": "msg_unbridged_output", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "epsilon", "annotations": []}], + } + ], + usage=ResponseAPIUsage(input_tokens=12, output_tokens=5, total_tokens=17), + ) + + _, _, logged_response = logging_obj._success_handler_helper_fn( + result=ResponseCompletedEvent(type="response.completed", response=completed_response), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert logged_response.id + assert not logged_response.id.startswith("msg_") + + def test_batch_cost_row_does_not_collide_with_the_batch_creation_row(): """Creating a batch writes a row keyed by the batch's own id, so keying the cost row the same way makes the insert a duplicate of it. request_id is the primary key and the @@ -3956,3 +4088,207 @@ def test_caller_forged_router_metadata_is_discarded(bucket): ) metadata = json.loads(payload["metadata"]) assert metadata["router_metadata"] is None + + +ANTHROPIC_MESSAGES_RESPONSE: Final = { + "id": "msg_01Lit6806NonStreaming", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "epsilon"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 14, "output_tokens": 4}, +} + +ANTHROPIC_MESSAGES_SSE_CHUNKS: Final = ( + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_01Lit6806Streaming",' + '"type":"message","role":"assistant","model":"claude-haiku-4-5","content":[],' + '"usage":{"input_tokens":14,"output_tokens":1}}}\n\n', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,' + '"content_block":{"type":"text","text":""}}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + '"delta":{"type":"text_delta","text":"epsilon"}}\n\n', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + '"usage":{"output_tokens":4}}\n\n', + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", +) + + +def _anthropic_messages_logging_obj(*, stream: bool) -> Any: + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + stream=stream, + call_type="anthropic_messages", + start_time=datetime.datetime.now(timezone.utc), + litellm_call_id="6806cafe-0000-4000-8000-000000000001", + function_id="1234", + ) + logging_obj.optional_params = {} + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + return logging_obj + + +def _spend_log_request_id(response_obj: Any, kwargs: dict) -> str: + payload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + return payload["request_id"] + + +def test_spend_log_request_id_is_the_message_id_a_non_streaming_messages_caller_received(): + """ + POST /v1/messages hands the caller `id: msg_...`, the only request id they ever see, so + GET /spend/logs?request_id=msg_... has to find the row. + """ + logging_obj = _anthropic_messages_logging_obj(stream=False) + + logged_response = logging_obj._handle_anthropic_messages_response_logging( + result=ANTHROPIC_MESSAGES_RESPONSE + ) + + assert logged_response.id == "msg_01Lit6806NonStreaming" + assert ( + _spend_log_request_id( + response_obj=logged_response, + kwargs={ + "call_type": "anthropic_messages", + "model": "claude-haiku-4-5", + "litellm_call_id": "6806cafe-0000-4000-8000-000000000001", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + ) + == "msg_01Lit6806NonStreaming" + ) + + +def test_spend_log_request_id_is_the_message_id_a_streaming_messages_caller_received(): + """ + The streaming leg of /v1/messages logs through the Anthropic passthrough handler, which used + to stamp litellm_call_id over the msg_ id carried by the message_start event. + """ + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( + AnthropicPassthroughLoggingHandler, + ) + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + logging_obj = _anthropic_messages_logging_obj(stream=True) + logging_obj.model_call_details["stream"] = True + + logged = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude-haiku-4-5"}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.datetime.now(timezone.utc), + all_chunks=list(ANTHROPIC_MESSAGES_SSE_CHUNKS), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert logged["result"].id == "msg_01Lit6806Streaming" + assert ( + _spend_log_request_id( + response_obj=logged["result"], + kwargs={ + **logged["kwargs"], + "call_type": "anthropic_messages", + "litellm_call_id": "6806cafe-0000-4000-8000-000000000001", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + ) + == "msg_01Lit6806Streaming" + ) + + +def test_spend_log_request_id_still_falls_back_to_litellm_call_id_without_a_provider_id(): + """ + Anthropic-compatible upstreams that omit `id` must keep landing on litellm_call_id rather + than on a fresh chatcmpl- uuid nobody can look up. + """ + logging_obj = _anthropic_messages_logging_obj(stream=True) + logging_obj.model_call_details["stream"] = True + + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( + AnthropicPassthroughLoggingHandler, + ) + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=litellm.ModelResponse(id="chatcmpl-generated"), + model="claude-haiku-4-5", + kwargs={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + logging_obj=logging_obj, + ) + assert logging_obj.model_call_details["complete_streaming_response"].id == ( + "6806cafe-0000-4000-8000-000000000001" + ) + + +def test_spend_log_request_id_for_chat_completions_is_untouched(): + """ + /v1/chat/completions callers look their rows up by the chatcmpl- id in the response body. + """ + assert ( + _spend_log_request_id( + response_obj=litellm.ModelResponse(id="chatcmpl-EJvWIw3DAhuKYuwp3jJI4Pnhp2vjv", choices=[]), + kwargs={ + "call_type": "acompletion", + "model": "gpt-5.6", + "litellm_call_id": "6806cafe-0000-4000-8000-000000000002", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + ) + == "chatcmpl-EJvWIw3DAhuKYuwp3jJI4Pnhp2vjv" + ) + + +def test_spend_log_request_id_is_the_response_id_a_bridged_messages_caller_received(): + """ + /v1/messages against a non-Anthropic model answers with the Responses id the caller then + looks their row up by, so the row must not fall back to a fresh chatcmpl- uuid. + """ + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + logging_obj = _anthropic_messages_logging_obj(stream=False) + bridged_response = ResponsesAPIResponse( + id="resp_01Lit6806Bridged", + object="response", + created_at=1767225600, + model="gpt-5.6", + status="completed", + output=[ + { + "id": "msg_bridged_output", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "delta", "annotations": []}], + } + ], + usage=ResponseAPIUsage(input_tokens=13, output_tokens=5, total_tokens=18), + ) + + logged_response = logging_obj._handle_anthropic_messages_response_logging(result=bridged_response) + + assert logged_response.id == "resp_01Lit6806Bridged" + assert ( + _spend_log_request_id( + response_obj=logged_response, + kwargs={ + "call_type": "anthropic_messages", + "model": "gpt-5.6", + "litellm_call_id": "6806cafe-0000-4000-8000-000000000003", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + ) + == "resp_01Lit6806Bridged" + ) diff --git a/tests/test_litellm/proxy/test_blocked_response_usage.py b/tests/test_litellm/proxy/test_blocked_response_usage.py index 37aea8fe3aa..4f20f35e94b 100644 --- a/tests/test_litellm/proxy/test_blocked_response_usage.py +++ b/tests/test_litellm/proxy/test_blocked_response_usage.py @@ -68,12 +68,10 @@ async def test_success_hook_attaches_original_response_on_block(): user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") data = {"guardrail_to_apply": guardrail, "model": "gpt-4o"} - # Inject our translation for the inferred call type (the module global is - # cached across tests, so patch it directly rather than the loader). with patch.object( ug, - "endpoint_guardrail_translation_mappings", - { + "load_guardrail_translation_mappings", + lambda: { CallTypes.acompletion: lambda: translation, CallTypes.completion: lambda: translation, }, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 6d6aad22ca3..ea665b60b19 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1540,8 +1540,8 @@ class TestCommonRequestProcessingHelpers: expected_error_data = { "error": { "message": "Error processing stream start", - "type": "None", - "param": "None", + "type": "internal_server_error", + "param": None, "code": str(status.HTTP_500_INTERNAL_SERVER_ERROR), } } @@ -1569,8 +1569,8 @@ class TestCommonRequestProcessingHelpers: expected_error_data = { "error": { "message": "Content blocked by guardrail", - "type": "None", - "param": "None", + "type": "invalid_request_error", + "param": None, "code": "400", } } @@ -1934,6 +1934,104 @@ class TestCommonRequestProcessingHelpers: assert mock_tracer.trace.call_count == 0 +def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]: + if isinstance(node, dict): + return tuple( + found + for key, value in node.items() + for found in _stringified_none_paths(value, f"{path}.{key}") + ) + if isinstance(node, (list, tuple)): + return tuple( + found + for index, value in enumerate(node) + for found in _stringified_none_paths(value, f"{path}[{index}]") + ) + return (path,) if node == "None" else () + + +def _blocked_guardrail_exception() -> HTTPException: + return HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": {"action": "GUARDRAIL_INTERVENED"}, + "guardrailIdentifier": "gf3sc1mzinjw", + "guardrailVersion": "DRAFT", + }, + ) + + +class TestGuardrailBlockErrorPayloadNeverStringifiesNone: + """Regression for LIT-6808: a blocked-guardrail error body carried the literal string + "None" for type and param instead of a real error type and JSON null.""" + + def test_non_streaming_block_payload_carries_a_real_type_and_null_param(self): + from litellm.proxy.common_request_processing import ( + proxy_exception_from_http_exception, + ) + + payload = json.loads( + json.dumps(proxy_exception_from_http_exception(_blocked_guardrail_exception(), {}).to_dict()) + ) + + assert _stringified_none_paths(payload) == () + assert payload["type"] == "invalid_request_error" + assert payload["param"] is None + assert payload["code"] == "400" + assert payload["message"] == "Violated guardrail policy" + + def test_streaming_block_frame_carries_a_real_type_and_null_param(self): + from litellm.proxy.common_request_processing import sse_error_payload + + error_status, error_obj = sse_error_payload(_blocked_guardrail_exception()) + frame = json.loads(json.dumps({"error": dict(error_obj)})) + + assert error_status == 400 + assert _stringified_none_paths(frame["error"]) == () + assert frame["error"]["type"] == "invalid_request_error" + assert frame["error"]["param"] is None + assert frame["error"]["code"] == "400" + + @pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (401, "authentication_error"), + (403, "permission_error"), + (404, "invalid_request_error"), + (429, "rate_limit_error"), + (500, "internal_server_error"), + (503, "internal_server_error"), + ], + ) + def test_status_code_decides_the_type_when_the_exception_carries_none(self, status_code, expected_type): + from litellm.proxy.common_request_processing import ( + proxy_exception_from_http_exception, + ) + + payload = proxy_exception_from_http_exception( + HTTPException(status_code=status_code, detail="blocked"), {} + ).to_dict() + + assert payload["type"] == expected_type + assert payload["param"] is None + + def test_a_type_and_param_the_exception_carries_win_over_the_fallback(self): + from litellm.proxy.common_request_processing import ( + proxy_exception_from_http_exception, + ) + + exc = HTTPException(status_code=400, detail="unknown model") + exc.type = "authentication_error" + exc.param = "model" + + payload = proxy_exception_from_http_exception(exc, {}).to_dict() + + assert payload["type"] == "authentication_error" + assert payload["param"] == "model" + + class TestExtractErrorFromSSEChunk: """Tests for _extract_error_from_sse_chunk function""" @@ -2999,6 +3097,25 @@ class TestHandleLLMApiExceptionDictDetail: assert proxy_exc.message == "Content blocked by guardrail" assert proxy_exc.provider_specific_fields is None + async def test_blocked_guardrail_error_body_never_carries_the_string_none(self): + """Regression for LIT-6808: the error body a blocked request returns must carry a real + error type and JSON null rather than the literal string "None".""" + proxy_exc = await self._invoke(_blocked_guardrail_exception()) + payload = json.loads(json.dumps(proxy_exc.to_dict())) + + assert _stringified_none_paths(payload) == () + assert payload["type"] == "invalid_request_error" + assert payload["param"] is None + + async def test_unclassified_exception_error_body_never_carries_the_string_none(self): + """The same holds on the generic fallback, where nothing carries a type at all.""" + proxy_exc = await self._invoke(ValueError("Something broke")) + payload = json.loads(json.dumps(proxy_exc.to_dict())) + + assert _stringified_none_paths(payload) == () + assert payload["type"] == "internal_server_error" + assert payload["param"] is None + async def test_not_found_error_preserves_404(self): """NotFoundError with status_code=404 should map to ProxyException code=404.""" from litellm.exceptions import NotFoundError diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 4ed6a468371..aef045b4709 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12065,10 +12065,15 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp } return row - def served_content() -> str: - callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1") + def served_callback(): + spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_sync") + assert spec is not None + callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=spec) assert callback is not None - return callback.prompt_manager.get_prompt("greeting_sync").content + return callback + + def served_content() -> str: + return served_callback().prompt_manager.get_prompt("greeting_sync").content prisma_client = MagicMock() try: @@ -12080,7 +12085,7 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) assert served_content() == "Begin every reply with HOWDY" - assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")] + assert litellm.callbacks == [served_callback()] finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_sync") @@ -12119,22 +12124,25 @@ async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkey ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("broken_sync.v1") is None - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1") is not None - assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1")] + assert IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("broken_sync") is None + healthy_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("healthy_sync") + assert healthy_spec is not None + healthy_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=healthy_spec) + assert healthy_callback is not None + assert litellm.callbacks == [healthy_callback] finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("healthy_sync") IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("broken_sync") @pytest.mark.asyncio -async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collide_on_a_versioned_id(monkeypatch): +async def test_init_prompts_in_db_syncs_every_environment_sharing_a_versioned_id(monkeypatch): from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.proxy.proxy_server import ProxyConfig monkeypatch.setattr(litellm, "callbacks", []) - def db_row(environment: str, content: str, updated_at: datetime) -> MagicMock: + def db_row(environment: str, content: str) -> MagicMock: row = MagicMock() row.model_dump.return_value = { "prompt_id": "greeting_env", @@ -12150,30 +12158,41 @@ async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collid ), "prompt_info": json.dumps({"prompt_type": "db"}), "created_at": None, - "updated_at": updated_at, + "updated_at": None, } return row - freshly_patched = db_row( - "production", "Begin every reply with HOWDY", datetime(2026, 8, 26, 12, 0, tzinfo=timezone.utc) - ) - stale_sibling = db_row( - "development", "Begin every reply with AHOY", datetime(2026, 8, 26, 11, 0, tzinfo=timezone.utc) - ) + def served_content(environment: str | None) -> str: + spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_env", environment=environment) + assert spec is not None + callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=spec) + assert callback is not None + return callback.prompt_manager.get_prompt("greeting_env").content prisma_client = MagicMock() try: - prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[freshly_patched, stale_sibling]) + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[ + db_row("development", "Begin every reply with AHOY"), + db_row("production", "Begin every reply with HOWDY"), + ] + ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - first_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") - assert first_callback is not None - assert first_callback.prompt_manager.get_prompt("greeting_env").content == "Begin every reply with HOWDY" + assert served_content("development") == "Begin every reply with AHOY" + assert served_content("production") == "Begin every reply with HOWDY" + assert served_content(None) == "Begin every reply with HOWDY" + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[ + db_row("development", "Begin every reply with YO"), + db_row("production", "Begin every reply with HOWDY"), + ] + ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") is first_callback - assert litellm.callbacks == [first_callback] + assert served_content("development") == "Begin every reply with YO" + assert served_content("production") == "Begin every reply with HOWDY" finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env") @@ -12216,13 +12235,14 @@ async def test_init_prompts_in_db_unloads_rows_deleted_on_another_worker(monkeyp return_value=[_prompt_db_row("greeting_del", _dotprompt_params("greeting_del"))] ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_del.v1") is not None + loaded_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_del") + assert loaded_spec is not None + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=loaded_spec) is not None prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[]) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("greeting_del.v1") is None - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_del.v1") is None + assert IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_del") is None assert litellm.callbacks == [] finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_del") @@ -12253,10 +12273,12 @@ async def test_init_prompts_in_db_keeps_config_prompts_when_their_id_has_no_db_r await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_cfg") is not None + surviving_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_cfg") + assert surviving_spec is not None + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=surviving_spec) is not None assert len(litellm.callbacks) == 1 finally: - IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id="greeting_cfg") + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_cfg") @pytest.mark.asyncio @@ -12272,7 +12294,9 @@ async def test_init_prompts_in_db_keeps_the_in_memory_copy_when_a_row_fails_to_p return_value=[_prompt_db_row("greeting_broken", _dotprompt_params("greeting_broken"))] ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - loaded_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_broken.v1") + loaded_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_broken") + assert loaded_spec is not None + loaded_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=loaded_spec) assert loaded_callback is not None prisma_client.db.litellm_prompttable.find_many = AsyncMock( @@ -12280,7 +12304,9 @@ async def test_init_prompts_in_db_keeps_the_in_memory_copy_when_a_row_fails_to_p ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_broken.v1") is loaded_callback + kept_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_broken") + assert kept_spec is not None + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=kept_spec) is loaded_callback assert litellm.callbacks == [loaded_callback] finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_broken") @@ -12314,8 +12340,9 @@ async def test_init_prompts_in_db_keeps_a_prompt_created_while_the_sync_was_read prisma_client.db.litellm_prompttable.find_many = AsyncMock(side_effect=create_prompt_behind_the_select) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - surviving_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_race.v1") - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("greeting_race.v1") is not None + surviving_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_race") + assert surviving_spec is not None + surviving_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=surviving_spec) assert surviving_callback is not None assert litellm.callbacks == [surviving_callback] finally: diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcaad968663..f16c6c937d0 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -942,6 +942,47 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map( assert response["max_output_tokens"] == 8000 +def test_create_model_info_response_uses_deployment_mode_for_auto_router(): + router = litellm.Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + }, + { + "model_name": "claude-auto", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": "claude-sonnet", + "MEDIUM": "claude-sonnet", + "COMPLEX": "claude-sonnet", + } + }, + "complexity_router_default_model": "claude-sonnet", + }, + "model_info": { + "mode": "chat", + "max_input_tokens": 1_000_000, + "max_output_tokens": 128_000, + }, + }, + ] + ) + + response = create_model_info_response( + model_id="claude-auto", + provider="openai", + llm_router=router, + get_model_info=_raise_unmapped, + ) + + assert response["mode"] == "chat" + assert response["max_input_tokens"] == 1_000_000 + assert response["max_output_tokens"] == 128_000 + + def test_create_model_info_response_deployment_limits_override_cost_map(): router = MagicMock() router.get_configured_token_limits.return_value = (200000, None) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index e99e34d65d4..2ba58bd5644 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -726,10 +726,7 @@ async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, from litellm.proxy.prompts import prompt_registry monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_by_id", lambda *a, **kw: None - ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: None + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None ) data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} await proxy_logging._process_prompt_template( @@ -752,11 +749,11 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging, monkeypatch.setattr( prompt_registry.IN_MEMORY_PROMPT_REGISTRY, - "get_prompt_callback_by_id", + "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: prompt_spec + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec ) logging_obj = MagicMock() @@ -802,11 +799,11 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi prompt_spec.litellm_params = MagicMock(prompt_id="x") monkeypatch.setattr( prompt_registry.IN_MEMORY_PROMPT_REGISTRY, - "get_prompt_callback_by_id", + "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: prompt_spec + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec ) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt")) @@ -820,6 +817,82 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi ) +@pytest.mark.asyncio +async def test_process_prompt_template_resolves_the_requested_environment(proxy_logging, monkeypatch): + from litellm.proxy.prompts import prompt_registry + + prompt_spec = MagicMock() + prompt_spec.litellm_params = MagicMock(prompt_id="greeting") + resolve_calls: list[dict] = [] + + def fake_resolve(prompt_id, version=None, environment=None): + resolve_calls.append({"prompt_id": prompt_id, "version": version, "environment": environment}) + return prompt_spec + + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", fake_resolve) + monkeypatch.setattr( + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_for_prompt", lambda *a, **kw: MagicMock() + ) + logging_obj = MagicMock() + logging_obj.async_get_chat_completion_prompt = AsyncMock( + return_value=("m", [{"role": "user", "content": "rendered"}], {}) + ) + data: Dict[str, Any] = { + "messages": [{"role": "user", "content": "orig"}], + "model": "m", + "prompt_id": "greeting", + "prompt_version": 1, + "prompt_environment": "development", + } + await proxy_logging._process_prompt_template( + data=data, + litellm_logging_obj=logging_obj, + prompt_id="greeting", + prompt_version=1, + call_type="completion", + ) + + assert resolve_calls == [{"prompt_id": "greeting", "version": 1, "environment": "development"}] + assert "prompt_environment" not in data + assert "prompt_id" not in data + assert data["messages"] == [{"role": "user", "content": "rendered"}] + + +@pytest.mark.asyncio +async def test_pre_call_hook_matches_a_prompt_version_sent_as_a_json_string(proxy_logging, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.prompts import prompt_registry + + prompt_spec = MagicMock() + prompt_spec.litellm_params = MagicMock(prompt_id="greeting") + resolve_calls: list[dict] = [] + + def fake_resolve(prompt_id, version=None, environment=None): + resolve_calls.append({"prompt_id": prompt_id, "version": version, "environment": environment}) + return prompt_spec + + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", fake_resolve) + monkeypatch.setattr( + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_for_prompt", lambda *a, **kw: MagicMock() + ) + logging_obj = MagicMock() + logging_obj.async_get_chat_completion_prompt = AsyncMock( + return_value=("m", [{"role": "user", "content": "rendered"}], {}) + ) + data: Dict[str, Any] = { + "messages": [{"role": "user", "content": "orig"}], + "model": "m", + "prompt_id": "greeting", + "prompt_version": "2", + "litellm_logging_obj": logging_obj, + } + + result = await proxy_logging.pre_call_hook(user_api_key_dict=UserAPIKeyAuth(), data=data, call_type="completion") + + assert resolve_calls == [{"prompt_id": "greeting", "version": 2, "environment": None}] + assert result["messages"] == [{"role": "user", "content": "rendered"}] + + @pytest.mark.asyncio async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(proxy_logging, monkeypatch): from litellm.proxy.prompts import prompt_registry @@ -829,11 +902,11 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p prompt_spec.litellm_params = MagicMock(prompt_id="resolved-id") monkeypatch.setattr( prompt_registry.IN_MEMORY_PROMPT_REGISTRY, - "get_prompt_callback_by_id", + "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: prompt_spec + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec ) logging_obj = MagicMock() diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 7d72121456a..93049b21460 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -120,9 +120,7 @@ async def test_delete_vector_store_checks_access(): "team_id": "team_456", } ) - mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=mock_vector_store - ) + mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=mock_vector_store) # User from different team should get 403 user_api_key_dict = UserAPIKeyAuth(team_id="team_789") @@ -134,9 +132,115 @@ async def test_delete_vector_store_checks_access(): ): with patch("litellm.vector_store_registry", None): with pytest.raises(HTTPException) as exc_info: - await delete_vector_store( - data=request, user_api_key_dict=user_api_key_dict - ) + await delete_vector_store(data=request, user_api_key_dict=user_api_key_dict) assert exc_info.value.status_code == 403 assert "Access denied" in exc_info.value.detail + + +_UNSCOPED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_unscoped", + "custom_llm_provider": "openai", + "team_id": None, +} +_TEAM_A_OWNED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_team_a", + "custom_llm_provider": "openai", + "team_id": "team_a", +} +_UI_CREATED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_ui_created", + "custom_llm_provider": "openai", + "team_id": "litellm-dashboard", +} + + +async def _listed_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + list_vector_stores, + ) + + with patch( # test-quality-ok: the list route reads rows through this module-level DB helper, no injection seam + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[_UNSCOPED, _TEAM_A_OWNED, _UI_CREATED]), + ): + response = await list_vector_stores(user_api_key_dict=user_api_key_dict) + return sorted(vs["vector_store_id"] for vs in response["data"]) + + +@pytest.mark.asyncio +async def test_list_vector_stores_hides_ungranted_stores_from_non_admin_keys(): + """A store with no team_id and no allowlist entry is not listed for a key it was never granted to; + only team ownership or an explicit object_permission grant makes a store visible.""" + assert await _listed_ids(UserAPIKeyAuth()) == [] + assert await _listed_ids(UserAPIKeyAuth(team_id="team_a")) == ["vs_team_a"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", vector_stores=["vs_unscoped"]), + ) + ) == ["vs_unscoped"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + team_object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-2", vector_stores=["vs_unscoped"] + ), + ) + ) == ["vs_unscoped"] + assert await _listed_ids(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) == [ + "vs_team_a", + "vs_ui_created", + "vs_unscoped", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_team_ids", "session_key_grants", "expected"), + [ + ([], None, []), + ([], ["vs_unscoped"], ["vs_unscoped"]), + (["team_a"], None, ["vs_team_a"]), + (["team_a", "team_granted"], None, ["vs_team_a", "vs_unscoped"]), + ], +) +async def test_list_vector_stores_dashboard_session_resolves_real_teams( + user_team_ids: list[str], session_key_grants: list[str] | None, expected: list[str] +): + """A dashboard session lists through the user's real teams plus the session key's own grants: stores created + from the dashboard (team_id litellm-dashboard) are not visible just because every session shares that team id, + while stores owned by or granted to one of the user's teams, or granted to the session key itself, are.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + + alice = UserAPIKeyAuth( + team_id="litellm-dashboard", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=( + LiteLLM_ObjectPermissionTable(object_permission_id="op-4", vector_stores=session_key_grants) + if session_key_grants is not None + else None + ), + ) + teams = { + "team_a": LiteLLM_TeamTableCachedObj(team_id="team_a"), + "team_granted": LiteLLM_TeamTableCachedObj( + team_id="team_granted", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-3", vector_stores=["vs_unscoped"]), + ), + } + + async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj: + return teams[team_id] + + with ( + patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam + "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object + ), + patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam + "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids", + new=AsyncMock(return_value=user_team_ids), + ), + ): + assert await _listed_ids(alice) == expected diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index 51d03544910..264bcd6fb75 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -388,6 +388,72 @@ async def test_aquery_does_not_forward_connection_override_keys_to_search(): assert not (blocked & set(search_kwargs.keys())) +@pytest.mark.asyncio +async def test_aquery_forwards_vector_store_params_to_search_but_not_completion(): + """ + Regression for LIT-6773: the server-trusted vector_store_params (a managed + store's litellm_params) must reach the search call wholesale, including the + connection keys the caller allowlist blocks, while the caller's own + retrieval_config overrides stay blocked, the caller's top-level api_key and + api_base stay on the completion only, and the completion never inherits the + store's connection params. + """ + from unittest.mock import AsyncMock + + from litellm.types.vector_stores import VectorStoreSearchResponse + + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + ) + fake_completion = AsyncMock( + return_value=ModelResponse( + id="chatcmpl-test", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="gpt-4o-mini", + ) + ) + with ( + patch("litellm.vector_stores.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch("litellm.acompletion", new=fake_completion), # test-quality-ok: the completion boundary under test + ): + await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + api_key="sk-llm-key", + api_base="https://llm.example.com", + retrieval_config={ + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "api_base": "https://attacker.example.com", + "api_key": "attacker-key", + }, + vector_store_params={ + "vector_store_id": "customer_kb", + "custom_llm_provider": "milvus", + "api_base": "http://127.0.0.1:19530", + "api_key": "root:Milvus", + "milvus_text_field": "book_intro_text", + "outputFields": ["book_intro_text"], + }, + ) + + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["vector_store_id"] == "customer_kb" + assert search_kwargs["custom_llm_provider"] == "milvus" + assert search_kwargs["api_base"] == "http://127.0.0.1:19530" + assert search_kwargs["api_key"] == "root:Milvus" + assert search_kwargs["milvus_text_field"] == "book_intro_text" + assert search_kwargs["outputFields"] == ["book_intro_text"] + fake_completion.assert_awaited_once() + completion_kwargs = fake_completion.await_args.kwargs + assert completion_kwargs["api_key"] == "sk-llm-key" + assert completion_kwargs["api_base"] == "https://llm.example.com" + assert not ({"milvus_text_field", "outputFields"} & set(completion_kwargs)) + + def test_rag_call_types_are_registered(): """ query/aquery/ingest/aingest are @client-decorated entry points, so their diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index a3dd5688ad1..761e87ac764 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -1,6 +1,7 @@ import asyncio import time from types import TracebackType +from typing import Final from unittest.mock import MagicMock, patch @@ -294,3 +295,39 @@ async def test_azure_health_check_honors_deployment_realtime_protocol(): model_params={"realtime_protocol": "GA"}, ) assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + + +class _ConnectThatStopsAfterCapturingTheUrl: + url: str | None = None + + def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl": + self.url = url + return self + + async def __aenter__(self) -> None: + raise RuntimeError("backend url captured, nothing to bridge") + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + +@pytest.mark.asyncio +async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route(): + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure_ai/gpt-realtime-mini", + websocket=MagicMock(), + api_base="https://my-project.services.ai.azure.com", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + ) + assert connect.url == ( + "wss://my-project.services.ai.azure.com/openai/realtime" + "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini" + ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 4a03913f55a..719d51c11e3 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -25,6 +25,7 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, StreamingChoices, + Usage, ) CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" @@ -527,6 +528,20 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids[0].startswith("resp_") +def test_completed_event_restores_usage_hidden_by_stream_options_none(): + final_chunk = _chunk("", finish_reason="stop") + final_chunk._hidden_params = {"usage": Usage(prompt_tokens=117, completion_tokens=5, total_tokens=122)} + iterator = _build_iterator([_chunk("the document says hello"), final_chunk]) + + events = list(iterator) + + completed = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert completed.response.usage.input_tokens == 117 + assert completed.response.usage.output_tokens == 5 + + def test_object_tool_call_arguments_stream_as_valid_json(): """A provider that sends decoded object arguments must still stream valid JSON. diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index 36199b45847..123ada83ca4 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -330,36 +330,46 @@ ROUTER_CONFIG: Final = json.dumps( ) -class FailingRouteLayer: - """Route layer whose embedding call fails, as it does when the prompt exceeds the encoder's window.""" - - def __call__(self, text: str) -> Any: - raise ValueError( - "Internal_litellm_router API call failed. Error: litellm.InternalServerError: " - "input is too large to process. increase the physical batch size" - ) - - class FixedRouteLayer: - """Route layer that returns whatever the test tells it to, recording the text it was asked about.""" + """Route layer that returns whatever the test tells it to for the query vector it is handed.""" def __init__(self, route_choice: Any) -> None: self.route_choice = route_choice - self.seen_text: str | None = None - def __call__(self, text: str) -> Any: - self.seen_text = text + async def acall(self, vector: Any) -> Any: return self.route_choice +def _embedding_response(input: List[str]) -> Any: + import litellm + + return litellm.EmbeddingResponse( + data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))] + ) + + class StubEmbeddingRouter: - """Stands in for the LiteLLM Router when the route index has to be built for real.""" + """Stands in for the LiteLLM Router, recording the text and kwargs each query embedding was made with.""" + + def __init__(self) -> None: + self.seen_text: str | None = None + self.aembedding_kwargs: Dict[str, Any] | None = None def embedding(self, input: List[str], model: str, **kwargs: Any) -> Any: - import litellm + return _embedding_response(input) - return litellm.EmbeddingResponse( - data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))] + async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any: + self.seen_text = input[0] + self.aembedding_kwargs = kwargs + return _embedding_response(input) + + +class FailingEmbeddingRouter(StubEmbeddingRouter): + """Router whose query embedding fails, as it does when the prompt exceeds the encoder's window.""" + + async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any: + raise ValueError( + "litellm.InternalServerError: input is too large to process. increase the physical batch size" ) @@ -369,7 +379,7 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs: auto_router_config=ROUTER_CONFIG, default_model="fallback-model", embedding_model="text-embedding-3-small", - litellm_router_instance=litellm_router_instance or MagicMock(), + litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(), **kwargs, ) auto_router.routelayer = routelayer @@ -381,7 +391,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: @pytest.mark.asyncio async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self): - auto_router: Final = _auto_router(FailingRouteLayer()) + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter()) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -440,8 +450,8 @@ class TestAutoRouterAlwaysResolvesARoutableModel: async def test_should_still_route_to_the_matched_route_when_one_matches(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -451,7 +461,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: assert result is not None assert result.model == "code-model" - assert layer.seen_text == "fix this stack trace" + assert router.seen_text == "fix this stack trace" class TestAutoRouterEmbeddingInputCap: @@ -483,8 +493,8 @@ class TestAutoRouterRoutesResponsesApiInput: async def test_should_route_a_string_input_when_messages_is_none(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -498,14 +508,14 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "code-model" assert result.messages is None - assert layer.seen_text == "fix this stack trace" + assert router.seen_text == "fix this stack trace" @pytest.mark.asyncio async def test_should_route_a_list_input_with_instructions_when_messages_is_none(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -525,13 +535,13 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "code-model" - assert layer.seen_text is not None - assert "fix this stack trace" in layer.seen_text + assert router.seen_text is not None + assert "fix this stack trace" in router.seen_text @pytest.mark.asyncio async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self): - layer: Final = FixedRouteLayer(None) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -540,12 +550,12 @@ class TestAutoRouterRoutesResponsesApiInput: ) assert result is None - assert layer.seen_text is None + assert router.seen_text is None @pytest.mark.asyncio async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self): - layer: Final = FixedRouteLayer(None) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -555,4 +565,42 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "fallback-model" - assert layer.seen_text == "" + assert router.seen_text == "" + + +class TestAutoRouterAttributesItsEmbeddingSpend: + """The query embedding is billed to the key that sent the request, like any other call it made.""" + + @pytest.mark.asyncio + async def test_should_forward_the_callers_identity_to_the_query_embedding_minus_its_budget_reservation(self): + from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY + + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=router) + request_kwargs: Final = { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_team_id": "team-1", + "user_api_key_budget_reservation": {"reservation_id": "r-1"}, + }, + "litellm_session_id": "session-1", + } + + result: Final = await auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + + assert result is not None + assert router.seen_text == "fix this stack trace" + assert router.aembedding_kwargs is not None + forwarded: Final = router.aembedding_kwargs["metadata"] + assert forwarded["user_api_key"] == "hashed-key" + assert forwarded["user_api_key_team_id"] == "team-1" + assert forwarded[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "autorouter_classifier" + assert "user_api_key_budget_reservation" not in forwarded + assert router.aembedding_kwargs["litellm_session_id"] == "session-1" + assert router.aembedding_kwargs["proxy_server_request"] == { + "body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]} + } diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index aa1b51afe10..da3791da39a 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -14,6 +14,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY @@ -1085,6 +1086,165 @@ class TestRouterComplexityDeploymentMethods: router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + @staticmethod + def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": classifier_type, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + "model_info": {"id": model_id}, + } + + _POOL: dict[str, object] = { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}, + } + + def test_heuristic_v2_ceiling_keeps_the_first_router_and_drops_the_rest(self) -> None: + """The proxy runs with ignore_invalid_deployments, so the second heuristic_v2 router is dropped + at registration while a heuristic (v1) sibling and the first v2 router stay routable.""" + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + self._router_row("v1-c", "id-c", "heuristic"), + ], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert sorted(router.complexity_routers) == ["v1-c", "v2-a"] + assert router.get_deployment(model_id="id-b") is None + + def test_heuristic_v2_ceiling_raises_without_ignore_invalid_deployments(self) -> None: + with pytest.raises(ValueError, match="At most 1 auto-router"): + Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: 1, + ) + + def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None: + """The Router never caches the limit: when the resolver's answer moves (the proxy re-verified + its license), the next registration and the next limit query see the new value.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + limits["value"] = 1 + assert router.heuristic_v2_router_limit_violation() is not None + assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + + def test_heuristic_v2_ceiling_tightening_refuses_the_edit_and_keeps_the_live_router(self) -> None: + """Two heuristic_v2 routers registered under an unlimited ceiling, then the ceiling drops to one: + an edit to either must be refused before its live row is popped, or the failed re-add and + the failed restore would drop a serving router while the write reports success.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + assert router.upsert_deployment(Deployment(**self._router_row("v2-a-renamed", "id-a", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.get_deployment(model_id="id-a") is not None + + assert router.upsert_deployment(Deployment(**self._router_row("v1-a", "id-a", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-a", "v2-b"] + + def test_config_deployments_excludes_db_rows(self) -> None: + """The proxy counts config.yaml routers from here and DB rows from the database, so a DB-loaded + row (``model_info.db_model``) must not show up twice.""" + router = Router(model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")]) + db_row = self._router_row("v2-db", "id-db", "heuristic_v2") + db_row["model_info"] = {"id": "id-db", "db_model": True} + assert router.upsert_deployment(Deployment(**db_row)) is not None + + assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"] + assert count_heuristic_v2_routers(router.config_deployments()) == 1 + + def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None: + """A rollback after a failed upsert re-admits state that was already serving, so it must not be + judged by a ceiling that tightened since: converting one of two live heuristic_v2 routers to a + config whose registration fails must leave it serving its previous v2 configuration.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + broken = self._router_row("v1-a", "id-a", "heuristic") + broken["litellm_params"]["complexity_router_config"]["tiers"] = {} + assert router.upsert_deployment(Deployment(**broken)) is None + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + live = router.get_deployment(model_id="id-a") + assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2" + assert router.heuristic_v2_router_limit_violation() is not None + + def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None: + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ] + ) + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None: + """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot + while a different deployment switching to heuristic_v2 is refused.""" + router = Router( + model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert router.heuristic_v2_router_limit_violation() is not None + + edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2") + assert router.upsert_deployment(Deployment(**edited)) is not None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + + assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"] + def test_hybrid_initialization_waits_for_later_pool_deployments(self): router = Router( model_list=[ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 60433921de6..b5651062098 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1,11 +1,12 @@ import asyncio +import itertools +import json +from collections.abc import Sequence +from typing import Final from unittest.mock import AsyncMock, patch import pytest - -import json - import litellm from litellm.caching.dual_cache import DualCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( @@ -102,11 +103,10 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): # Deterministic routing: first selection uses seq[0], second selection attempts seq[1] # unless the list has been filtered to length=1 by deployment affinity. - choice_calls = {"count": 0} + choice_calls: Final = itertools.count(1) - def deterministic_choice(seq): - choice_calls["count"] += 1 - if choice_calls["count"] == 1: + def deterministic_choice(seq: Sequence[dict[str, object]]) -> dict[str, object]: + if next(choice_calls) == 1: return seq[0] return seq[1] if len(seq) > 1 else seq[0] @@ -998,3 +998,212 @@ async def test_model_group_affinity_config_overrides_global(): ) # All deployments returned (user-key affinity disabled for this group) assert len(filtered) == 2 + + +def _jwt_metadata(user_id: str) -> dict[str, str | None]: + return {"user_api_key_hash": None, "user_api_key_user_id": user_id} + + +def _two_deployments(model_group: str) -> list[dict]: + return [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "openai-deployment-a"}, + }, + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "openai-deployment-b"}, + }, + ] + + +@pytest.mark.asyncio +async def test_async_jwt_user_affinity_routes_to_same_deployment(): + """ + JWT-authenticated proxy requests carry no `user_api_key_hash`, only `user_api_key_user_id`. + They must still pin to one deployment per user. + """ + model_group = "gpt-5.4-mini" + router = litellm.Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-a"}, + "model_info": {"id": "openai-deployment-a"}, + }, + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-b"}, + "model_info": {"id": "openai-deployment-b"}, + }, + ], + optional_pre_call_checks=["deployment_affinity"], + ) + + choice_calls = {"count": 0} + + def deterministic_choice(seq): + choice_calls["count"] += 1 + if choice_calls["count"] == 1: + return seq[0] + return seq[1] if len(seq) > 1 else seq[0] + + with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=deterministic_choice, + ): + first_response = await router.acompletion( + model=model_group, + messages=[{"role": "user", "content": "Reply with the single word ok"}], + mock_response="ok", + metadata=_jwt_metadata("jwt-user-alice"), + ) + second_response = await router.acompletion( + model=model_group, + messages=[{"role": "user", "content": "Reply with the single word ok"}], + mock_response="ok", + metadata=_jwt_metadata("jwt-user-alice"), + ) + + first_model_id = first_response._hidden_params["model_id"] + assert first_model_id in ("openai-deployment-a", "openai-deployment-b") + assert second_response._hidden_params["model_id"] == first_model_id + + +@pytest.mark.asyncio +async def test_proxy_jwt_auth_metadata_pins_per_user(): + """ + The metadata the proxy stamps for a JWT caller (`UserAPIKeyAuth(api_key=None, user_id=)`) + must claim a pin and be read back by the filter, and another JWT user must not inherit it. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + def proxy_request(user_id: str) -> dict[str, object]: + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id), + _metadata_variable_name="metadata", + ) + + alice_request = proxy_request("jwt-user-alice") + alice_metadata = alice_request["metadata"] + assert isinstance(alice_metadata, dict) + assert alice_metadata["user_api_key_hash"] is None + + await callback.async_pre_call_deployment_hook( + kwargs={ + **alice_request, + "metadata": {**alice_metadata, "deployment_model_name": model_group}, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + alice_pinned = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=alice_request, + parent_otel_span=None, + ) + assert [deployment["model_info"]["id"] for deployment in alice_pinned] == ["openai-deployment-b"] + + bob_filtered = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=proxy_request("jwt-user-bob"), + parent_otel_span=None, + ) + assert bob_filtered == healthy_deployments + + +@pytest.mark.asyncio +async def test_jwt_user_id_never_reads_a_virtual_key_pin(): + """ + A JWT user id that happens to equal a virtual key's 64-hex hash must not read that key's pin. + """ + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + key_hash = "a" * 64 + + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {"user_api_key_hash": key_hash, "deployment_model_name": model_group}, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + key_pinned = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"user_api_key_hash": key_hash}}, + parent_otel_span=None, + ) + assert [deployment["model_info"]["id"] for deployment in key_pinned] == ["openai-deployment-b"] + + lookalike_jwt_user = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": _jwt_metadata(key_hash)}, + parent_otel_span=None, + ) + assert lookalike_jwt_user == healthy_deployments + + +@pytest.mark.asyncio +async def test_virtual_key_hash_wins_over_user_id_for_affinity(): + """ + A virtual-key caller with a user id pins on the key hash, so two keys owned by one user + keep independent pins. + """ + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": { + "user_api_key_hash": "key-one", + "user_api_key_user_id": "shared-user", + "deployment_model_name": model_group, + }, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + other_key_same_user = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"user_api_key_hash": "key-two", "user_api_key_user_id": "shared-user"}}, + parent_otel_span=None, + ) + assert other_key_same_user == healthy_deployments diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 0007f09896a..238d0546518 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,8 +1,13 @@ +from collections.abc import Mapping + import pytest from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + is_heuristic_v2_router, strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, @@ -369,3 +374,55 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie flat param on an s3_vectors vector store, so an unscoped gate would reject a valid deployment. Either complexity field names one on its own, which is what the load itself requires.""" assert carries_complexity_router_settings(model, present_fields) is scoped + + +@pytest.mark.parametrize( + "litellm_params,expected", + [ + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False), + ({"model": "auto_router/complexity_router"}, False), + ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False), + ({}, False), + ], +) +def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None: + """Only a complexity router whose config selects heuristic_v2 counts toward the license limit.""" + assert is_heuristic_v2_router(litellm_params) is expected + + +def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None: + v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}} + rows: list[Mapping[str, object]] = [ + {"model_name": "a", "litellm_params": v2}, + {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "c", "litellm_params": v2}, + {"model_name": "d"}, + {"model_name": "e", "litellm_params": "not a mapping"}, + ] + assert count_heuristic_v2_routers(rows) == 2 + assert count_heuristic_v2_routers(()) == 0 + + +@pytest.mark.parametrize( + "held,limit,violates", + [ + (1, 1, False), + (2, 1, True), + (0, 1, False), + (5, None, False), + (3, 3, False), + (4, 3, True), + ], +) +def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None: + violation = heuristic_v2_limit_violation(held=held, limit=limit) + assert (violation is not None) is violates + if violation is not None: + assert f"At most {limit} auto-router" in violation + assert f"would make {held}" in violation + assert "license" not in violation diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index a0dbf3b6637..7b2e45ab3ed 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -371,3 +371,20 @@ class TestKimiK3AdvertisesItsDocumentedLevels: "low", "high", ) + + +class TestGpt6AstraAdvertisesItsDocumentedLevels: + def test_the_entry_advertises_low_through_max_without_none(self, local_model_cost_map): + """OpenAI documents low, medium, high, xhigh and max for gpt-6-astra. Unlike gpt-5.6-sol it + does not take none, so a group must not offer none and must offer max.""" + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="gpt-6-astra", custom_llm_provider="openai")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "low", + "medium", + "high", + "xhigh", + "max", + ) diff --git a/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py b/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py new file mode 100644 index 00000000000..92af1b1dba4 --- /dev/null +++ b/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py @@ -0,0 +1,55 @@ +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm import cost_per_token, get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + +REPO_ROOT: Final = Path(__file__).parents[2] +MODEL: Final = "azure_ai/grok-4.6" +SOURCE: Final = ( + "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/" + "grok-4-6-comes-to-microsoft-foundry-models-built-for-long-horizon-reasoning-and-/4547578" +) +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) + + +def _cost_map_entry(path: Path) -> dict[str, object]: + return COST_MAP_ADAPTER.validate_json(path.read_bytes())[MODEL] + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_azure_ai_grok_4_6_is_priced_and_routed() -> None: + routed_model, provider, _, _ = get_llm_provider(model=MODEL) + assert (routed_model, provider) == ("grok-4.6", "azure_ai") + + info = get_model_info(model=routed_model, custom_llm_provider=provider) + assert info["litellm_provider"] == "azure_ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == 2e-06 + assert info["output_cost_per_token"] == 6e-06 + assert info["cache_read_input_token_cost"] == 5e-07 + assert info["max_input_tokens"] == 200000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_web_search"] is True + + prompt_cost, completion_cost = cost_per_token(model=MODEL, prompt_tokens=1_000_000, completion_tokens=1_000_000) + assert prompt_cost == pytest.approx(2.0) + assert completion_cost == pytest.approx(6.0) + + +def test_azure_ai_grok_4_6_entry_source_and_backup_match() -> None: + main_entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json") + backup_entry = _cost_map_entry(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json") + + assert main_entry["source"] == SOURCE + assert backup_entry == main_entry diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7c2174018e8..6dc3b2790c9 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4473,3 +4473,17 @@ def test_explicit_pricing_precedes_private_provider_response_model( ) assert selected == expected + + +def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map): + """gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens.""" + from litellm.cost_calculator import batch_cost_calculator + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, model="gpt-6-astra", custom_llm_provider="openai" + ) + + assert prompt_cost == pytest.approx(1000 * 5e-6) + assert completion_cost == pytest.approx(500 * 2.5e-5) diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py index 1c12a48ed9d..a60fa9466e6 100644 --- a/tests/test_litellm/test_gpt_5_5_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py @@ -47,8 +47,7 @@ def test_azure_ai_gpt_5_5_model_info(model): routed_model, provider, _, _ = get_llm_provider(model=model) assert routed_model == model.split("/", 1)[1] - # azure_ai/* models resolve under the azure provider in get_llm_provider - assert provider == "azure" + assert provider == "azure_ai" def test_azure_ai_gpt_5_5_backup_matches_main(): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3c8bf142835..9a703635ab6 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -784,6 +784,19 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo assert model_info.get("mode") == "responses" +def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-6-astra", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + ) + + assert model == "gpt-6-astra" + assert model_info.get("mode") == "responses" + + def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" from litellm.main import responses_api_bridge_check @@ -1133,6 +1146,82 @@ def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_st assert model_info.get("mode") != "responses" +@pytest.mark.parametrize( + "api_base", + [ + "https://southcentralus.privatelink.api.openai.com/v1", + "https://privatelink.corp.api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://api.openai.com/v1/", + "HTTPS://API.OPENAI.COM/v1", + ], +) +def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base): + """ + A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a + port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects + function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as + it does for the literal default URL. Regression guard for GH #39353. + """ + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "api_base", + [ + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base): + """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend.""" + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch): + """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges.""" + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1") + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes(): """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base.""" from litellm.main import responses_api_bridge_check @@ -3259,3 +3348,43 @@ def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeyp assert getattr(response.usage, "cost", None) == pytest.approx(0.42) assert response._hidden_params.get("response_cost") is None assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) + + +FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" + + +def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" + ).mock(return_value=httpx.Response(200, json={"text": "hello"})) + + response: Final = litellm.transcription( + model="azure_ai/whisper-1", + file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.text == "hello" + + +def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" + ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) + + response: Final = litellm.speech( + model="azure_ai/tts-1", + input="hello", + voice="alloy", + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.content == b"mp3-bytes" diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py index c0860a5b55f..bdb2dc26813 100644 --- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -52,6 +52,12 @@ PRIORITY_LONG_CONTEXT = { "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, }, + "gpt-6-astra": { + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + }, } EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT} @@ -63,6 +69,7 @@ NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5") def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.add_known_models() @lru_cache(maxsize=2) @@ -114,6 +121,7 @@ TIERED_COST_CASES = [ ("gpt-5.6-sol", "priority", 1.6e-05, 6e-05), ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05), ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06), + ("gpt-6-astra", "priority", 4e-05, 0.00015), ] diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/test_litellm/test_responses_id_security.py index 763ee4dac00..d35b9563888 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/test_litellm/test_responses_id_security.py @@ -14,7 +14,9 @@ from litellm.proxy.hooks.responses_id_security import ( _is_responses_api_create_route, ) from litellm.types.llms.openai import ( + GenericEvent, ResponseCompletedEvent, + ResponseCreatedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, ) @@ -691,6 +693,126 @@ class TestAsyncPostCallStreamingIteratorHook: assert not responses_id_security._is_encrypted_response_id(streamed_id) +class TestStreamedGenericEventIdEncryption: + """A background stream carries event types with no typed model, which arrive as + GenericEvent holding a plain dict. Those used to skip encryption while their typed + siblings were encrypted, so one stream advertised two ids and the unencrypted one + skipped the ownership check. Asserts the property rather than one event type: every + id a client can see is the same encrypted id, and the raw one appears in no frame.""" + + RAW_ID = "resp_rawprovider123" + + @staticmethod + async def _agen(chunks): + for chunk in chunks: + yield chunk + + @classmethod + def _typed_event(cls, event_type): + return { + ResponsesAPIStreamEvents.RESPONSE_CREATED: ResponseCreatedEvent, + ResponsesAPIStreamEvents.RESPONSE_COMPLETED: ResponseCompletedEvent, + }[event_type]( + type=event_type, + response=ResponsesAPIResponse( + id=cls.RAW_ID, + created_at=0, + model="gpt-5.1", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + @classmethod + def _background_stream(cls): + return [ + cls._typed_event(ResponsesAPIStreamEvents.RESPONSE_CREATED), + GenericEvent( + type="response.queued", + response={"id": cls.RAW_ID, "status": "queued"}, + ), + GenericEvent(type="keepalive"), + GenericEvent( + type="response.some_event_openai_adds_later", + response={"id": cls.RAW_ID, "status": "in_progress"}, + ), + cls._typed_event(ResponsesAPIStreamEvents.RESPONSE_COMPLETED), + ] + + @staticmethod + def _advertised_ids(events): + nested = (getattr(event, "response", None) for event in events) + return [ + payload["id"] if isinstance(payload, dict) else payload.id + for payload in nested + if payload is not None + ] + [ + event.id for event in events if isinstance(getattr(event, "id", None), str) + ] + + async def _drain(self, responses_id_security, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij") + + mock_auth = MagicMock() + mock_auth.user_id = "user-a" + mock_auth.team_id = "team-a" + mock_auth.request_route = "/v1/responses" + + return [ + out + async for out in responses_id_security.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_auth, + response=self._agen(self._background_stream()), + request_data={}, + ) + ] + + @pytest.mark.asyncio + async def test_every_event_advertises_the_same_encrypted_id( + self, responses_id_security, monkeypatch + ): + events = await self._drain(responses_id_security, monkeypatch) + advertised = self._advertised_ids(events) + + assert len(advertised) == 4 + assert len(set(advertised)) == 1 + + streamed_id = advertised[0] + assert streamed_id != self.RAW_ID + assert responses_id_security._is_encrypted_response_id(streamed_id) + assert responses_id_security._decrypt_response_id(streamed_id) == ( + self.RAW_ID, + "user-a", + "team-a", + ) + + @pytest.mark.asyncio + async def test_raw_provider_id_never_reaches_the_client( + self, responses_id_security, monkeypatch + ): + events = await self._drain(responses_id_security, monkeypatch) + + assert [self.RAW_ID in event.model_dump_json() for event in events] == [ + False + ] * len(events) + + @pytest.mark.asyncio + async def test_sibling_fields_survive_the_rewrite( + self, responses_id_security, monkeypatch + ): + _, queued, keepalive, later, _ = await self._drain( + responses_id_security, monkeypatch + ) + + assert queued.response["status"] == "queued" + assert later.response["status"] == "in_progress" + assert keepalive.type == "keepalive" + assert getattr(keepalive, "response", None) is None + + class TestAsyncPostCallSuccessHook: """Test async_post_call_success_hook function""" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3b4c80b6b4f..e4236afc586 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7358,6 +7358,71 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +def test_get_configured_mode_reads_deployment_model_info(): + router = litellm.Router( + model_list=[ + { + "model_name": "chat-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"mode": "chat"}, + } + ] + ) + + assert router.get_configured_mode("chat-model") == "chat" + + +def test_get_configured_mode_returns_none_for_unset_or_unknown(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-mode-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_configured_mode("no-mode-model") is None + assert router.get_configured_mode("not-a-real-model") is None + + +def test_get_configured_mode_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"mode": "chat"}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) + + +def test_get_configured_mode_treats_malformed_values_as_absent(): + malformed = ["", " ", 12345, ["chat"], {"mode": "chat"}, True] + router = litellm.Router( + model_list=[ + { + "model_name": f"bad-mode-{i}", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"mode": bad}, + } + for i, bad in enumerate(malformed) + ] + ) + + for i in range(len(malformed)): + assert router.get_configured_mode(f"bad-mode-{i}") is None + + def test_get_configured_display_name_reads_deployment_model_info(): router = litellm.Router( model_list=[ @@ -9275,9 +9340,11 @@ class _FallbackAttemptRecorder(CustomLogger): def __init__(self): super().__init__() self.failed_targets = [] + self.breadcrumbs_per_target = [] async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception): self.failed_targets.append(kwargs.get("model")) + self.breadcrumbs_per_target.append(kwargs.get("metadata", {}).get("previous_models", ())) def _cyclic_fallback_router(num_retries=0): @@ -9348,14 +9415,16 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): A retry has to be configured for the walk state to reach log_retry at all.""" router = _cyclic_fallback_router(num_retries=1) capture = _LogCapture(logging.ERROR) + recorder = _FallbackAttemptRecorder() - await _drive_cyclic_fallback(router, capture) + await _drive_cyclic_fallback(router, capture, recorder) - assert router.previous_models, "no retry breadcrumbs were recorded" + breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] + assert breadcrumbs, "no retry breadcrumbs were recorded" assert any( - "fallback_depth" in breadcrumb for breadcrumb in router.previous_models + "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs ), "no breadcrumb carried router walk state, so this test cannot see the leak" - for breadcrumb in router.previous_models: + for breadcrumb in breadcrumbs: assert "attempted_targets" not in breadcrumb @@ -9393,15 +9462,94 @@ async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_ke container still reaches the breadcrumb, but the raw secret never does, whatever key holds it.""" router = _cyclic_fallback_router(num_retries=1) capture = _LogCapture(logging.ERROR) + metadata = {} - await _drive_cyclic_fallback(router, capture, **request_kwargs) + await _drive_cyclic_fallback(router, capture, metadata=metadata, **request_kwargs) - assert router.previous_models, "no retry breadcrumbs were recorded" - dumped = json.dumps(router.previous_models, default=str) + breadcrumbs = metadata["previous_models"] + assert breadcrumbs, "no retry breadcrumbs were recorded" + dumped = json.dumps(breadcrumbs, default=str) assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped +def _always_failing_router(num_retries): + return litellm.Router( + model_list=[ + { + "model_name": "broken-group", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "mock_response": "litellm.InternalServerError", + }, + } + ], + num_retries=num_retries, + ) + + +async def _fail_one_proxy_shaped_request(router, request_marker): + """The proxy hands the router a metadata dict and a proxy_server_request whose body is a + shallow copy of the request, so body["metadata"] is the very same dict the router later + stamps previous_models onto.""" + metadata = {"request_marker": request_marker} + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model="broken-group", + messages=[{"role": "user", "content": "hi"}], + metadata=metadata, + proxy_server_request={ + "url": "http://localhost:4000/v1/chat/completions", + "method": "POST", + "headers": {}, + "body": {"model": "broken-group", "metadata": metadata}, + }, + ) + return metadata["previous_models"] + + +def _nested_breadcrumb_lists(node): + if isinstance(node, dict): + return [v for k, v in node.items() if k == "previous_models"] + [ + found for v in node.values() for found in _nested_breadcrumb_lists(v) + ] + if isinstance(node, (list, tuple)): + return [found for item in node for found in _nested_breadcrumb_lists(item)] + return [] + + +@pytest.mark.asyncio +async def test_retry_breadcrumbs_stay_per_request_and_flat_across_failing_requests(): + """Every failed attempt appends a breadcrumb to metadata["previous_models"], and the proxy's + request snapshot aliases that same metadata dict. Kept on the Router and copied wholesale, + each breadcrumb embedded every earlier one from every earlier request, so the breadcrumb + tree, and with it the debug repr of the kwargs, roughly doubled on each failed attempt until + a single-worker proxy spent minutes in the redaction regex and stopped answering.""" + router = _always_failing_router(num_retries=2) + + breadcrumbs_per_request = [ + await _fail_one_proxy_shaped_request(router, f"request-{request_number}") for request_number in range(1, 7) + ] + + for request_number, breadcrumbs in enumerate(breadcrumbs_per_request, start=1): + assert len(breadcrumbs) == 3, "one initial attempt plus two retries failed, each leaving one breadcrumb" + assert {breadcrumb["metadata"]["request_marker"] for breadcrumb in breadcrumbs} == {f"request-{request_number}"} + for breadcrumb in breadcrumbs: + assert _nested_breadcrumb_lists(breadcrumb) == [] + assert len({len(repr(breadcrumbs)) for breadcrumbs in breadcrumbs_per_request}) == 1 + + +@pytest.mark.asyncio +async def test_retry_breadcrumbs_keep_only_the_last_four_attempts(): + router = _always_failing_router(num_retries=6) + + breadcrumbs = await _fail_one_proxy_shaped_request(router, "request-1") + + assert len(breadcrumbs) == 4 + assert [breadcrumb["metadata"]["attempted_retries"] for breadcrumb in breadcrumbs] == [3, 4, 5, 6] + + @pytest.mark.asyncio async def test_fallback_traceback_stays_available_at_debug_level(): """Dropping the stack from the ERROR line is only safe because the fallback path still @@ -12553,3 +12701,33 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo bucket = captured.get("litellm_metadata") or captured["metadata"] assert captured["model_info"]["id"] == "provisional-dep" assert bucket["litellm_gateway_injected_cache"] == "" + + +def test_get_configured_mode_reads_deployment_model_info(): + router = Router( + model_list=[ + { + "model_name": "my-tts", + "litellm_params": {"model": "openai/some-unmapped-mode-model"}, + "model_info": {"mode": "audio_speech"}, + } + ] + ) + + assert router.get_configured_mode("my-tts") == "audio_speech" + + +@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": " "}, {"mode": 123}]) +def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info): + router = Router( + model_list=[ + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/some-unmapped-mode-model"}, + "model_info": model_info, + } + ] + ) + + assert router.get_configured_mode("plain-model") is None + assert router.get_configured_mode("unknown-model") is None diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 99c69d6bc31..106a1bca5f2 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -40,6 +40,7 @@ from litellm.utils import ( _is_streaming_request, _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, + async_post_call_success_deployment_hook, client, get_llm_provider, get_non_default_completion_params, @@ -5808,6 +5809,90 @@ class TestHuggingFaceConfigFetch: assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS +@pytest.mark.asyncio +async def test_success_deployment_hook_chains_past_callback_returning_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-5863): the dispatcher must run every callback, chaining each non-None + result into the next call, instead of returning at the first callback answering non-None. + A guardrail answering with the unmodified response used to starve every callback after it.""" + from litellm.types.utils import ModelResponse + + original = ModelResponse() + replacement = ModelResponse() + + class PassthroughLogger(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + return response + + class ReplacingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: list = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.seen.append(response) + return replacement + + class ObservingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: list = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.seen.append(response) + return None + + replacer = ReplacingLogger() + observer = ObservingLogger() + monkeypatch.setattr(litellm, "callbacks", [PassthroughLogger(), replacer, observer]) + + result = await async_post_call_success_deployment_hook( + request_data={}, response=original, call_type=CallTypes.acompletion + ) + + assert replacer.seen == [original] + assert observer.seen == [replacement] + assert result is replacement + + +@pytest.mark.asyncio +async def test_registered_guardrail_does_not_starve_vector_store_search_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-5863): with any guardrail registered ahead of the lazily-appended + VectorStorePreCallHook, /v1/chat/completions responses lost + provider_specific_fields["search_results"] because the guardrail answered the unmodified + response and the dispatcher stopped there.""" + from types import SimpleNamespace + + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + VectorStorePreCallHook, + ) + from litellm.types.utils import ModelResponse + + search_results: Final = [{"search_query": "coolant", "data": [{"content": [{"text": "Cryoline-9", "type": "text"}]}]}] + logging_obj = SimpleNamespace(model_call_details={"search_results": search_results}) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Cryoline-9"}}]) + + monkeypatch.setattr( + litellm, + "callbacks", + [CustomGuardrail(guardrail_name="dummy-guardrail"), VectorStorePreCallHook()], + ) + + result = await async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.acompletion, + ) + + provider_fields = result.choices[0].message.provider_specific_fields + assert provider_fields is not None + assert provider_fields["search_results"] == search_results + + class TestIsVisionExplicitlyDisabled: """github_copilot and chatgpt run an OAuth device flow inside get_llm_provider; the explicit-disable lookup must adopt the declared prefix instead of resolving it, exactly @@ -5836,3 +5921,135 @@ class TestIsVisionExplicitlyDisabled: is_vision_explicitly_disabled("fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731") is True ) assert is_vision_explicitly_disabled("anthropic/claude-sonnet-4-5") is False + + +class TestVerboseRequestLineRedaction: + """`litellm.set_verbose = True` echoes the caller's kwargs back as a `litellm.completion(...)` + line on stdout, so a credential kwarg lands in whatever collects stdout: a terminal, a + container log drain, a CI job log. Credential-named kwargs must not survive that echo, + at any nesting depth, while ordinary params still must, or the line stops telling the + developer what they called.""" + + FAKE_API_KEY: Final = "sk-fake-lit6823-0000000000000000" + + def _verbose_request_line(self, capsys, monkeypatch, **kwargs) -> str: + monkeypatch.setattr(litellm, "set_verbose", True) + monkeypatch.setattr("litellm._logging.set_verbose", True) + capsys.readouterr() + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + mock_response="hi", + **kwargs, + ) + captured: Final = capsys.readouterr() + return "\n".join(line for line in (captured.out + captured.err).splitlines() if "litellm.completion(" in line) + + def test_api_key_never_reaches_the_request_line(self, capsys, monkeypatch): + printed: Final = self._verbose_request_line(capsys, monkeypatch, api_key=self.FAKE_API_KEY) + + assert "litellm.completion(" in printed + assert self.FAKE_API_KEY not in printed + assert "api_key='REDACTED'" in printed + + def test_credential_headers_never_reach_the_request_line(self, capsys, monkeypatch): + printed: Final = self._verbose_request_line( + capsys, + monkeypatch, + api_key=self.FAKE_API_KEY, + extra_headers={"Authorization": "Bearer fake-lit6823-header", "x-request-id": "abc123"}, + ) + + assert "fake-lit6823-header" not in printed + assert "'Authorization': 'REDACTED'" in printed + assert "'x-request-id': 'abc123'" in printed + + def test_credentials_nested_in_a_list_never_reach_the_request_line(self, capsys, monkeypatch): + printed: Final = self._verbose_request_line( + capsys, + monkeypatch, + api_key=self.FAKE_API_KEY, + extra_body={"providers": [{"name": "openai", "api_key": "sk-fake-lit6823-nested"}]}, + ) + + assert "sk-fake-lit6823-nested" not in printed + assert "'name': 'openai'" in printed + + def test_ordinary_params_still_printed(self, capsys, monkeypatch): + printed: Final = self._verbose_request_line( + capsys, monkeypatch, api_key=self.FAKE_API_KEY, max_tokens=17, temperature=0.25 + ) + + assert "model='gpt-3.5-turbo'" in printed + assert "max_tokens=17" in printed + assert "temperature=0.25" in printed + + +class TestFinalOptionalParamsLineRedaction: + """A verbose run echoes the fully built optional params too, and `extra_body` carries whatever the + caller nested inside it straight onto that line, so a credential tucked in there lands in a terminal + or a log drain in plaintext. It has to be redacted on both surfaces `print_verbose` writes to, and the + line has to keep printing on both, because `litellm.set_verbose` and the DEBUG logger are independent + switches and neither implies the other.""" + + FAKE_NESTED_KEY: Final = "sk-fake-lit6835-nested-0000000000" + + def _complete(self, **kwargs) -> None: + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + mock_response="hi", + **kwargs, + ) + + def _printed_line(self, capsys) -> str: + captured: Final = capsys.readouterr() + return "\n".join( + line for line in (captured.out + captured.err).splitlines() if "Final returned optional params" in line + ) + + def test_nested_credential_is_redacted_when_only_set_verbose_is_on(self, capsys, caplog, monkeypatch): + monkeypatch.setattr(litellm, "set_verbose", True) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + capsys.readouterr() + self._complete(extra_body={"providers": [{"name": "openai", "api_key": self.FAKE_NESTED_KEY}]}) + printed: Final = self._printed_line(capsys) + + assert printed + assert self.FAKE_NESTED_KEY not in printed + assert "'api_key': 'REDACTED'" in printed + assert "'name': 'openai'" in printed + + def test_line_still_reaches_the_logger_when_only_the_debug_logger_is_on(self, capsys, caplog, monkeypatch): + monkeypatch.setattr(litellm, "set_verbose", False) + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + self._complete(extra_body={"providers": [{"name": "openai", "api_key": self.FAKE_NESTED_KEY}]}) + logged: Final = "\n".join( + record.getMessage() + for record in caplog.records + if "Final returned optional params" in record.getMessage() + ) + + assert logged + assert self.FAKE_NESTED_KEY not in logged + assert "'name': 'openai'" in logged + + def test_nothing_is_emitted_when_neither_verbose_switch_is_on(self, capsys, caplog, monkeypatch): + monkeypatch.setattr(litellm, "set_verbose", False) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + capsys.readouterr() + self._complete(extra_body={"providers": [{"name": "openai", "api_key": self.FAKE_NESTED_KEY}]}) + captured: Final = capsys.readouterr() + + assert "Final returned optional params" not in captured.out + captured.err + assert self.FAKE_NESTED_KEY not in captured.out + captured.err + + def test_ordinary_optional_params_still_reach_the_line(self, capsys, caplog, monkeypatch): + monkeypatch.setattr(litellm, "set_verbose", True) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + capsys.readouterr() + self._complete(max_tokens=17, temperature=0.25) + printed: Final = self._printed_line(capsys) + + assert "'max_tokens': 17" in printed + assert "'temperature': 0.25" in printed diff --git a/type-discipline-budget.json b/type-discipline-budget.json index cbcb5dca443..2f85128b4b6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22328 }, "LIT002": { - "limit": 26763 + "limit": 26760 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16477 + "limit": 16470 }, "LIT011": { - "limit": 5519 + "limit": 5516 }, "LIT012": { "limit": 4489 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx index 52efd6407f8..edbb897fb2f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx @@ -65,11 +65,13 @@ describe("PromptTable", () => { expect(within(rows[1]).getByText("prompt-older")).toBeInTheDocument(); }); - it("should call onPromptClick when the prompt ID is clicked", async () => { + it("should call onPromptClick with the row's environment, defaulting to development", async () => { const user = userEvent.setup(); render(); await user.click(screen.getByRole("button", { name: "prompt-newer" })); - expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-newer"); + expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-newer", "production"); + await user.click(screen.getByRole("button", { name: "prompt-older" })); + expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-older", "development"); }); it("should label the environment and default missing environments to development", () => { @@ -83,7 +85,7 @@ describe("PromptTable", () => { render(); await user.click(screen.getByTestId("prompt-actions-prompt-newer")); await user.click(await screen.findByTestId("prompt-action-delete")); - expect(mockOnDeleteClick).toHaveBeenCalledWith("prompt-newer", "prompt-newer"); + expect(mockOnDeleteClick).toHaveBeenCalledWith("prompt-newer", "prompt-newer", "production"); }); it("should copy the prompt ID through the actions menu", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx index 47d4f64f254..c766042ac44 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx @@ -13,8 +13,8 @@ import { ModelGroupInfo } from "./prompt_utils"; interface PromptTableProps { promptsList: PromptSpec[]; isLoading: boolean; - onPromptClick?: (id: string) => void; - onDeleteClick?: (id: string, name: string) => void; + onPromptClick?: (id: string, environment: string) => void; + onDeleteClick?: (id: string, name: string, environment: string) => void; accessToken: string | null; isAdmin: boolean; } @@ -74,7 +74,9 @@ const PromptTable: React.FC = ({ prompt.prompt_id || String(index)} + getRowId={(prompt, index) => + prompt.prompt_id ? `${prompt.prompt_id}::${prompt.environment || "development"}` : String(index) + } sortingMode="client" sorting={sorting} onSortingChange={setSorting} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx index ae584ef6df6..f927a6d1486 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx @@ -64,7 +64,7 @@ function PromptModelCell({ prompt, modelHubData }: { prompt: PromptSpec; modelHu interface PromptRowActionsProps { prompt: PromptSpec; isAdmin: boolean; - onDeleteClick?: (id: string, name: string) => void; + onDeleteClick?: (id: string, name: string, environment: string) => void; } function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsProps) { @@ -91,7 +91,13 @@ function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsPr onDeleteClick?.(prompt.prompt_id, prompt.prompt_id || "Unknown Prompt")} + onClick={() => + onDeleteClick?.( + prompt.prompt_id, + prompt.prompt_id || "Unknown Prompt", + prompt.environment || "development", + ) + } > Delete @@ -106,8 +112,8 @@ function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsPr interface PromptTableColumnsDeps { modelHubData: Map; isAdmin: boolean; - onPromptClick?: (id: string) => void; - onDeleteClick?: (id: string, name: string) => void; + onPromptClick?: (id: string, environment: string) => void; + onDeleteClick?: (id: string, name: string, environment: string) => void; } export const getPromptTableColumns = ({ @@ -128,7 +134,11 @@ export const getPromptTableColumns = ({ title={row.original.prompt_id} titleClassName="font-mono text-xs font-normal" className="max-w-60" - onClick={onPromptClick ? () => onPromptClick(row.original.prompt_id) : undefined} + onClick={ + onPromptClick + ? () => onPromptClick(row.original.prompt_id, row.original.environment || "development") + : undefined + } /> ), }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx index 3b0ee3a3ce2..a1c7280422f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx @@ -16,21 +16,31 @@ vi.mock("./PromptTable", () => ({ __esModule: true, default: ({ isLoading, + onPromptClick, onDeleteClick, }: { isLoading: boolean; - onDeleteClick: (id: string, name: string) => void; + onPromptClick: (id: string, environment: string) => void; + onDeleteClick: (id: string, name: string, environment: string) => void; }) => (
{isLoading ? "table-loading" : "table-loaded"} - +
), })); -vi.mock("./prompt_info", () => ({ __esModule: true, default: () =>
prompt-info-view
})); +vi.mock("./prompt_info", () => ({ + __esModule: true, + default: ({ initialEnvironment }: { initialEnvironment?: string }) => ( +
prompt-info-view:{initialEnvironment ?? "none"}
+ ), +})); vi.mock("./add_prompt_form", () => ({ __esModule: true, default: ({ visible }: { visible: boolean }) => (visible ?
add-prompt-form
: null), @@ -141,6 +151,22 @@ describe("PromptsPanel toolbar", () => { }); }); +describe("PromptsPanel row navigation", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockGetPromptsList.mockResolvedValue({ prompts: [] } as never); + }); + + it("should open the info view preselected to the clicked row's environment", async () => { + const user = userEvent.setup(); + renderPanel("Admin"); + + await user.click(await screen.findByRole("button", { name: "row-open" })); + + expect(screen.getByText("prompt-info-view:staging")).toBeInTheDocument(); + }); +}); + describe("PromptsPanel delete confirmation", () => { beforeEach(() => { vi.clearAllMocks(); @@ -154,13 +180,13 @@ describe("PromptsPanel delete confirmation", () => { await user.click(await screen.findByRole("button", { name: "row-delete" })); - expect(await screen.findByText(/delete prompt: my-prompt/i)).toBeInTheDocument(); + expect(await screen.findByText(/the staging copy of prompt: my-prompt/i)).toBeInTheDocument(); expect(screen.getByText(/cannot be undone/i)).toBeInTheDocument(); expect(mockDeletePromptCall).not.toHaveBeenCalled(); await user.click(screen.getByRole("button", { name: /^delete$/i })); - await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1")); + await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1", "staging")); }); it("should abandon the delete when the confirmation is dismissed", async () => { @@ -168,11 +194,11 @@ describe("PromptsPanel delete confirmation", () => { renderPanel("Admin"); await user.click(await screen.findByRole("button", { name: "row-delete" })); - await screen.findByText(/delete prompt: my-prompt/i); + await screen.findByText(/the staging copy of prompt: my-prompt/i); await user.click(screen.getByRole("button", { name: /cancel/i })); - await waitFor(() => expect(screen.queryByText(/delete prompt: my-prompt/i)).not.toBeInTheDocument()); + await waitFor(() => expect(screen.queryByText(/the staging copy of prompt: my-prompt/i)).not.toBeInTheDocument()); expect(mockDeletePromptCall).not.toHaveBeenCalled(); }); @@ -187,14 +213,14 @@ describe("PromptsPanel delete confirmation", () => { renderPanel("Admin"); await user.click(await screen.findByRole("button", { name: "row-delete" })); - await screen.findByText(/delete prompt: my-prompt/i); + await screen.findByText(/the staging copy of prompt: my-prompt/i); await user.click(screen.getByRole("button", { name: /^delete$/i })); - await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1")); + await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1", "staging")); await user.keyboard("{Escape}"); - expect(screen.getByText(/delete prompt: my-prompt/i)).toBeInTheDocument(); + expect(screen.getByText(/the staging copy of prompt: my-prompt/i)).toBeInTheDocument(); finishDelete(); - await waitFor(() => expect(screen.queryByText(/delete prompt: my-prompt/i)).not.toBeInTheDocument()); + await waitFor(() => expect(screen.queryByText(/the staging copy of prompt: my-prompt/i)).not.toBeInTheDocument()); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx index 3594e933704..2d8d905c480 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx @@ -41,11 +41,12 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { const [isLoading, setIsLoading] = useState(true); const [selectedEnvironment, setSelectedEnvironment] = useState(undefined); const [selectedPromptId, setSelectedPromptId] = useState(null); + const [selectedPromptEnvironment, setSelectedPromptEnvironment] = useState(undefined); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [showEditorView, setShowEditorView] = useState(false); const [editPromptData, setEditPromptData] = useState(null); const [isDeleting, setIsDeleting] = useState(false); - const [promptToDelete, setPromptToDelete] = useState<{ id: string; name: string } | null>(null); + const [promptToDelete, setPromptToDelete] = useState<{ id: string; name: string; environment: string } | null>(null); // Admin Viewer follows the read-parity rule: see prompts, no writes. const canModify = userRole ? isProxyAdminRole(userRole) : false; @@ -71,8 +72,9 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { fetchPrompts(); }, [accessToken, selectedEnvironment]); - const handlePromptClick = (promptId: string) => { + const handlePromptClick = (promptId: string, environment: string) => { setSelectedPromptId(promptId); + setSelectedPromptEnvironment(environment); }; const handleAddPrompt = () => { @@ -111,8 +113,8 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { setSelectedPromptId(null); }; - const handleDeleteClick = (promptId: string, promptName: string) => { - setPromptToDelete({ id: promptId, name: promptName }); + const handleDeleteClick = (promptId: string, promptName: string, environment: string) => { + setPromptToDelete({ id: promptId, name: promptName, environment }); }; const handleDeleteConfirm = async () => { @@ -120,8 +122,8 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { setIsDeleting(true); try { - await deletePromptCall(accessToken, promptToDelete.id); - toast.success(`Prompt "${promptToDelete.name}" deleted successfully`); + await deletePromptCall(accessToken, promptToDelete.id, promptToDelete.environment); + toast.success(`Prompt "${promptToDelete.name}" deleted successfully from ${promptToDelete.environment}`); fetchPrompts(); // Refresh the list } catch (error) { console.error("Error deleting prompt:", error); @@ -148,6 +150,7 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { ) : selectedPromptId ? ( setSelectedPromptId(null)} accessToken={accessToken} isAdmin={canModify} @@ -219,7 +222,8 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { Delete Prompt - Are you sure you want to delete prompt: {promptToDelete.name} ? This action cannot be undone. + Are you sure you want to delete the {promptToDelete.environment} copy of prompt: {promptToDelete.name}? + This action cannot be undone. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx index a1b4ad52634..7fa44a4dfe5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx @@ -44,4 +44,28 @@ describe("PromptCodeSnippets", () => { expect(screen.getByRole("combobox", { name: "Language" })).toHaveTextContent("Python (OpenAI SDK)"); }); + + it("includes the viewed environment in every generated request", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + render( + , + ); + await user.click(screen.getByRole("button", { name: /get code/i })); + await screen.findByText("Generated Code"); + + await user.click(screen.getByRole("button", { name: /copy to clipboard/i })); + expect(await navigator.clipboard.readText()).toContain('"prompt_environment": "development"'); + + await user.click(screen.getByRole("tab", { name: "With Version" })); + await user.click(screen.getByRole("button", { name: /copy to clipboard/i })); + const versionSnippet = await navigator.clipboard.readText(); + expect(versionSnippet).toContain('"prompt_environment": "development"'); + expect(versionSnippet).toContain('"prompt_version": 2'); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx index af7d6421265..a6adc160674 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx @@ -22,6 +22,7 @@ interface PromptCodeSnippetsProps { promptVariables?: Record; accessToken: string | null; version?: string; + environment?: string; proxySettings?: { PROXY_BASE_URL?: string; LITELLM_UI_API_DOC_BASE_URL?: string | null; @@ -34,6 +35,7 @@ const PromptCodeSnippets: React.FC = ({ promptVariables = {}, accessToken, version = "1", + environment, proxySettings, }) => { const syntaxTheme = useSyntaxTheme(coy); @@ -64,6 +66,9 @@ const PromptCodeSnippets: React.FC = ({ // Generate code based on selected language and tab const generateCode = () => { const hasVariables = Object.keys(promptVariables).length > 0; + const curlEnvironment = environment ? `,\n "prompt_environment": "${environment}"` : ""; + const pythonEnvironment = environment ? `,\n "prompt_environment": "${environment}"` : ""; + const jsEnvironment = environment ? `,\n prompt_environment: "${environment}"` : ""; if (selectedLanguage === "curl") { if (selectedTab === "basic") { @@ -72,7 +77,7 @@ const PromptCodeSnippets: React.FC = ({ -H 'Authorization: Bearer ${effectiveApiKey}' \\ -d '{ "model": "${model}", - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${curlEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 6).replace(/\n/g, "\n ")}` @@ -85,7 +90,7 @@ const PromptCodeSnippets: React.FC = ({ -H 'Authorization: Bearer ${effectiveApiKey}' \\ -d '{ "model": "${model}", - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${curlEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 6).replace(/\n/g, "\n ")}` @@ -104,7 +109,7 @@ const PromptCodeSnippets: React.FC = ({ -H 'Authorization: Bearer ${effectiveApiKey}' \\ -d '{ "model": "${model}", - "prompt_id": "${promptId}", + "prompt_id": "${promptId}"${curlEnvironment}, "prompt_version": ${version}, "messages": [ { @@ -127,7 +132,7 @@ client = openai.OpenAI( response = client.chat.completions.create( model="${model}", extra_body={ - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${pythonEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` @@ -145,7 +150,7 @@ response = client.chat.completions.create( {"role": "user", "content": "hi"} ], extra_body={ - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${pythonEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` @@ -163,7 +168,7 @@ response = client.chat.completions.create( {"role": "user", "content": "Who are u"} ], extra_body={ - "prompt_id": "${promptId}", + "prompt_id": "${promptId}"${pythonEnvironment}, "prompt_version": ${version} } ) @@ -186,9 +191,9 @@ async function main() { model: "${model}", ${ hasVariables - ? `prompt_id: "${promptId}", + ? `prompt_id: "${promptId}"${jsEnvironment}, prompt_variables: ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` - : `prompt_id: "${promptId}"` + : `prompt_id: "${promptId}"${jsEnvironment}` } }); @@ -206,9 +211,9 @@ async function main() { ], ${ hasVariables - ? `prompt_id: "${promptId}", + ? `prompt_id: "${promptId}"${jsEnvironment}, prompt_variables: ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` - : `prompt_id: "${promptId}"` + : `prompt_id: "${promptId}"${jsEnvironment}` } }); @@ -224,7 +229,7 @@ async function main() { messages: [ { role: "user", content: "Who are u" } ], - prompt_id: "${promptId}", + prompt_id: "${promptId}"${jsEnvironment}, prompt_version: ${version} }); @@ -241,7 +246,7 @@ main();`; if (isModalVisible) { setGeneratedCode(generateCode()); } - }, [isModalVisible, selectedLanguage, selectedTab, promptId, model, promptVariables]); + }, [isModalVisible, selectedLanguage, selectedTab, promptId, model, promptVariables, version, environment]); return ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx index 5194cdd4e63..3afc00e37a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx @@ -2,7 +2,9 @@ import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PromptEditorHeader from "./PromptEditorHeader"; -vi.mock("./PromptCodeSnippets", () => ({ default: () => })); +vi.mock("./PromptCodeSnippets", () => ({ + default: ({ environment }: { environment?: string }) => , +})); describe("PromptEditorHeader", () => { it("preserves navigation, naming, and save actions", () => { @@ -48,5 +50,6 @@ describe("PromptEditorHeader", () => { ); expect(screen.getByRole("combobox", { name: "Environment" })).toHaveTextContent(label); + expect(screen.getByRole("button", { name: "Get Code" })).toHaveAttribute("data-environment", environment); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx index eea9755054f..04cac01365a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx @@ -89,6 +89,7 @@ const PromptEditorHeader: React.FC = ({ promptVariables={promptVariables} accessToken={accessToken} version={version?.replace("v", "") || "1"} + environment={environment} proxySettings={proxySettings} /> {editMode && onShowHistory && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx index bb29f12ac42..b14d5c3d91f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx @@ -12,7 +12,9 @@ vi.mock("@/components/networking", () => ({ })); vi.mock("./prompt_editor_view/PromptCodeSnippets", () => ({ - default: () =>
, + default: ({ environment }: { environment?: string }) => ( +
+ ), })); const promptWithoutTemplate = { @@ -29,6 +31,67 @@ const promptWithoutTemplate = { environments: [], }; +describe("PromptInfoView environment scoping", () => { + beforeEach(() => { + vi.mocked(networking.getPromptInfo).mockReset().mockResolvedValue(promptWithoutTemplate); + vi.mocked(networking.getPromptVersions).mockReset().mockResolvedValue({ prompts: [] }); + }); + + it("fetches the initial environment it was opened with", async () => { + render( + , + ); + + await screen.findByRole("tab", { name: "Raw JSON" }); + expect(networking.getPromptInfo).toHaveBeenCalledWith("sk-test", "support-reply", "staging"); + }); + + it("fetches the serve default when opened without an environment", async () => { + render(); + + await screen.findByRole("tab", { name: "Raw JSON" }); + expect(networking.getPromptInfo).toHaveBeenCalledWith("sk-test", "support-reply", undefined); + }); +}); + +describe("PromptInfoView code snippets", () => { + beforeEach(() => { + vi.mocked(networking.getPromptVersions).mockReset().mockResolvedValue({ prompts: [] }); + }); + + it.each([ + ["a prompt with several environments", "staging", ["development", "staging"]], + ["a config prompt with no environment list", "development", []], + ])("hands the viewed environment of %s to the code snippets", async (_label, environment, environments) => { + vi.mocked(networking.getPromptInfo) + .mockReset() + .mockResolvedValue({ + ...promptWithoutTemplate, + prompt_spec: { ...promptWithoutTemplate.prompt_spec, environment }, + environments, + }); + + render( + , + ); + + await screen.findByRole("tab", { name: "Raw JSON" }); + expect(screen.getByTestId("prompt-code-snippets")).toHaveAttribute("data-environment", environment); + }); +}); + describe("PromptInfoView tabs", () => { beforeEach(() => { vi.mocked(networking.getPromptInfo).mockReset().mockResolvedValue(promptWithoutTemplate); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx index c6a6f09fcaa..062e1d84a0a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx @@ -20,6 +20,7 @@ import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from " export interface PromptInfoProps { promptId: string; + initialEnvironment?: string; onClose: () => void; accessToken: string | null; isAdmin: boolean; @@ -27,7 +28,15 @@ export interface PromptInfoProps { onEdit?: (promptData: any) => void; } -const PromptInfoView: React.FC = ({ promptId, onClose, accessToken, isAdmin, onDelete, onEdit }) => { +const PromptInfoView: React.FC = ({ + promptId, + initialEnvironment, + onClose, + accessToken, + isAdmin, + onDelete, + onEdit, +}) => { const [promptData, setPromptData] = useState(null); const [promptTemplate, setPromptTemplate] = useState(null); const [rawApiResponse, setRawApiResponse] = useState(null); @@ -43,7 +52,7 @@ const PromptInfoView: React.FC = ({ promptId, onClose, accessTo const [selectedVersion, setSelectedVersion] = useState(null); const [loadingVersions, setLoadingVersions] = useState(false); - // Initial fetch — no environment filter, gets default + all environments list + // Fetches the requested environment (or the serve-time default when omitted) plus the environments list const fetchPromptInfo = async (environment?: string) => { try { setLoading(true); @@ -89,7 +98,7 @@ const PromptInfoView: React.FC = ({ promptId, onClose, accessTo setSelectedEnv(null); setEnvironments([]); setVersionHistory([]); - fetchPromptInfo(); + fetchPromptInfo(initialEnvironment); }, [promptId, accessToken]); // When environment changes (user clicks tab), re-fetch — skip initial mount @@ -212,6 +221,7 @@ const PromptInfoView: React.FC = ({ promptId, onClose, accessTo promptVariables={extractTemplateVariables(promptTemplate?.content)} accessToken={accessToken} version={currentVersion} + environment={selectedEnv ?? promptData.environment} />