diff --git a/.github/e2e-stack/down.sh b/.github/e2e-stack/down.sh index 9f72f2d6e64..740d626beea 100755 --- a/.github/e2e-stack/down.sh +++ b/.github/e2e-stack/down.sh @@ -10,7 +10,7 @@ for pid_file in "${STACK_DIR}"/pids/*.pid; do rm -f "${pid_file}" done -for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do +for container in e2e-nginx e2e-keycloak e2e-valkey e2e-jaeger e2e-postgres; do docker rm -f "${container}" >/dev/null 2>&1 done diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index a62358f81ff..238818a0d36 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile( ) HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" + r"|^tests/e2e/idp_realm\.json$" r"|^tests/e2e/gateway/" r"|^\.github/e2e-stack/" r"|^\.github/workflows/test-e2e-changed\.yml$" diff --git a/.github/e2e-stack/start-idp.sh b/.github/e2e-stack/start-idp.sh new file mode 100644 index 00000000000..e59a7ade34c --- /dev/null +++ b/.github/e2e-stack/start-idp.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +KEYCLOAK_IMAGE="${E2E_KEYCLOAK_IMAGE:-quay.io/keycloak/keycloak@sha256:ff4257d0d64efbe99ed1ddfaf07765cc3c36dc7518bf8324d41961327f441c54}" +KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}" +POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}" +: "${DATABASE_HOST:?}" "${DATABASE_PORT:?}" "${DATABASE_USER:?}" "${DATABASE_PASSWORD:?}" "${DATABASE_NAME:?}" + +DB_HOST="${DATABASE_HOST}" +DB_NETWORK_ARGS=(--network bridge) +IDP_NETWORK_ARGS=(-p "127.0.0.1:${KEYCLOAK_PORT}:${KEYCLOAK_PORT}") +if [[ "$(uname)" == "Linux" ]]; then + DB_NETWORK_ARGS=(--network host) + IDP_NETWORK_ARGS=(--network host) +elif [[ "${DB_HOST}" == "127.0.0.1" || "${DB_HOST}" == "localhost" ]]; then + DB_HOST=host.docker.internal +fi + +docker run --rm "${DB_NETWORK_ARGS[@]}" -e "PGPASSWORD=${DATABASE_PASSWORD}" \ + "${POSTGRES_IMAGE}" psql -h "${DB_HOST}" -p "${DATABASE_PORT}" \ + -U "${DATABASE_USER}" -d "${DATABASE_NAME}" -v ON_ERROR_STOP=1 \ + -c 'CREATE SCHEMA IF NOT EXISTS keycloak' >/dev/null + +docker rm -f e2e-keycloak >/dev/null 2>&1 || true +docker run -d --name e2e-keycloak "${IDP_NETWORK_ARGS[@]}" --memory 1536m \ + -v "${REPO_ROOT}/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \ + -e KC_DB=postgres -e "KC_DB_URL_HOST=${DB_HOST}" -e "KC_DB_URL_PORT=${DATABASE_PORT}" \ + -e "KC_DB_URL_DATABASE=${DATABASE_NAME}" -e KC_DB_SCHEMA=keycloak \ + -e "KC_DB_USERNAME=${DATABASE_USER}" -e "KC_DB_PASSWORD=${DATABASE_PASSWORD}" \ + -e KC_DB_POOL_INITIAL_SIZE=2 -e KC_DB_POOL_MIN_SIZE=2 -e KC_DB_POOL_MAX_SIZE=10 \ + -e "KC_HTTP_PORT=${KEYCLOAK_PORT}" -e KC_BOOTSTRAP_ADMIN_USERNAME=admin \ + -e KC_BOOTSTRAP_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret \ + "${KEYCLOAK_IMAGE}" start-dev --import-realm >/dev/null + +deadline=$((SECONDS + ${E2E_KEYCLOAK_STARTUP_TIMEOUT:-300})) +until curl -fsS --connect-timeout 2 --max-time 3 \ + "http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/.well-known/openid-configuration" >/dev/null 2>&1; do + if ((SECONDS >= deadline)); then + echo 'e2e-stack: timed out waiting for the Keycloak realm' >&2 + exit 1 + fi + sleep 2 +done +echo 'e2e-stack: Keycloak realm is up' diff --git a/.github/e2e-stack/up.sh b/.github/e2e-stack/up.sh index 2f2e6c6f9a8..a789a570483 100755 --- a/.github/e2e-stack/up.sh +++ b/.github/e2e-stack/up.sh @@ -25,6 +25,7 @@ DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}" DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}" JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}" JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}" +KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}" MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}" @@ -124,6 +125,9 @@ SERVER_ENV=( "OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}" "SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem" "PYTHONPATH=${REPO_ROOT}" + "JWT_PUBLIC_KEY_URL=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/protocol/openid-connect/certs" + "JWT_ISSUER=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e" + "JWT_AUDIENCE=litellm-e2e" ) if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json" @@ -132,6 +136,8 @@ fi cd "${REPO_ROOT}" +env "${SERVER_ENV[@]}" "E2E_KEYCLOAK_PORT=${KEYCLOAK_PORT}" bash .github/e2e-stack/start-idp.sh + log "running migrations" env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1 @@ -200,6 +206,9 @@ LITELLM_MASTER_KEY=${MASTER_KEY} REDIS_HOST=127.0.0.1 REDIS_PORT=${REDIS_PORT} E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT} +E2E_KEYCLOAK_URL=http://127.0.0.1:${KEYCLOAK_PORT} +E2E_KEYCLOAK_ADMIN_USER=admin +E2E_KEYCLOAK_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME} EOF diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index e85a397cbd2..1a2c81d1f92 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -127,6 +127,8 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac - Low: anything else worth noting: naming, cleanup, an edge case nobody hits Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a human reader + If you assumed something instead of testing it, e.g. "only reproduces with X on" or "no + user-observable behavior difference", list it here too with what breaks if it is wrong Leave this section empty if there are none --> ## QA runbook diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index e6d2264fbf0..9c7e0db7065 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -81,7 +81,7 @@ jobs: run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py - name: test_e2e_changed_gate - run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py - name: router_code_coverage run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 23ab6dfcfe4..1db597ff673 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -27,6 +27,8 @@ jobs: sparse-checkout: | .github/e2e-stack tests/e2e/access_control + tests/e2e/management/test_jwt_management_e2e.py + tests/e2e/other/test_jwt_auth_e2e.py persist-credentials: false ref: ${{ github.sha }} @@ -45,7 +47,8 @@ jobs: --jq '.[] | select(.status != "removed") | .filename')" gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}" tests="$(printf '%s\n' "${files}" \ - | python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)" + | python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py \ + tests/e2e/management/test_jwt_management_e2e.py tests/e2e/other/test_jwt_auth_e2e.py)" echo "tests=${tests}" >> "${GITHUB_OUTPUT}" if [ -n "${tests}" ]; then echo "any=true" >> "${GITHUB_OUTPUT}" diff --git a/.gitignore b/.gitignore index deb0acae56e..7da917ce450 100644 --- a/.gitignore +++ b/.gitignore @@ -147,3 +147,6 @@ crash.*.log ui/litellm-dashboard/out/ litellm.log + +.coverage-rust +coverage-rust.xml diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 92b73867e67..3733072a948 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -3,7 +3,8 @@ The gateway exposes the LLM data-plane surface: chat/completions, embeddings, audio, batches, files, fine-tuning, rerank, ocr, rag, video, search, image, responses, vector stores, passthrough providers, realtime websockets, MCP -tool-call endpoints, and operational endpoints (/health, /metrics). +tool-call endpoints, and operational endpoints (/health, /metrics, and the +/debug/memory/summary read of the serving worker's RSS). Any path not listed here is dropped from the gateway process so management/UI endpoints don't ride on the same pods. @@ -121,6 +122,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset( "/docs/oauth2-redirect", "/redoc", "/test", + "/debug/memory/summary", } ) diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 692a799e783..20fd1a722dc 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -257,6 +257,14 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. - name: DATABASE_SCHEMA value: {{ .schema | quote }} {{- end }} +{{- if .sslMode }} +- name: DATABASE_SSLMODE + value: {{ .sslMode | quote }} +{{- end }} +{{- if .sslRootCert }} +- name: DATABASE_SSLROOTCERT + value: {{ .sslRootCert | quote }} +{{- end }} {{- if and .useIAMAuth .useAzureEntraAuth }} {{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }} {{- end }} diff --git a/helm/litellm/templates/ingress.yaml b/helm/litellm/templates/ingress.yaml index 732564b280f..d42558b9396 100644 --- a/helm/litellm/templates/ingress.yaml +++ b/helm/litellm/templates/ingress.yaml @@ -89,7 +89,7 @@ at "/" Prefix would swallow the whole backend management API) instead of adding to it. */}} -{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}} +{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}} apiVersion: networking.k8s.io/v1 kind: Ingress metadata: @@ -129,6 +129,8 @@ spec: # --- Gateway data plane --- # Exact /test only (see the $gatewayPrefixes comment above); # /test/* MCP management endpoints fall to the backend catch-all. + # Exact /debug/memory/summary reads a serving worker's RSS (the e2e memory + # gate); the rest of /debug/* stays on the backend. - path: /test pathType: Exact backend: @@ -136,6 +138,13 @@ spec: name: {{ $gatewayName }} port: number: {{ $gatewayPort }} + - path: /debug/memory/summary + pathType: Exact + backend: + service: + name: {{ $gatewayName }} + port: + number: {{ $gatewayPort }} {{- range $gatewayPrefixes }} {{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" . "pathType" "Prefix") }} {{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" . $pathType) }} diff --git a/helm/litellm/tests/database_auth_tests.yaml b/helm/litellm/tests/database_auth_tests.yaml index adbe14c59c2..add531f68bb 100644 --- a/helm/litellm/tests/database_auth_tests.yaml +++ b/helm/litellm/tests/database_auth_tests.yaml @@ -4,6 +4,7 @@ templates: - gateway/configmap.yaml - backend/deployment.yaml - backend/configmap.yaml + - migrations-job.yaml values: - ./values/required.yaml tests: @@ -67,6 +68,82 @@ tests: value: "true" any: true + - it: emits no TLS env by default + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLMODE + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLROOTCERT + any: true + + - it: writer sslMode and sslRootCert reach gateway and backend as DATABASE_SSLMODE and DATABASE_SSLROOTCERT + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + set: + database.writer.useIAMAuth: true + database.writer.sslMode: verify-full + database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLMODE + value: verify-full + any: true + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLROOTCERT + value: /etc/ssl/certs/ca-certificates.crt + any: true + + - it: writer sslMode and sslRootCert reach the collector sidecar and the migrations job, which dial Postgres themselves + set: + gateway.collector.enabled: true + database.connectionPool.enabled: true + database.writer.sslMode: verify-full + database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: collector + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_SSLMODE + value: verify-full + any: true + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: DATABASE_SSLROOTCERT + value: /etc/ssl/certs/ca-certificates.crt + any: true + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLMODE + value: verify-full + any: true + template: migrations-job.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLROOTCERT + value: /etc/ssl/certs/ca-certificates.crt + any: true + template: migrations-job.yaml + - it: writer rejects both token sources at once template: gateway/deployment.yaml set: diff --git a/helm/litellm/tests/ingress_controller_tests.yaml b/helm/litellm/tests/ingress_controller_tests.yaml index 40790ba674a..aa30db3c9c1 100644 --- a/helm/litellm/tests/ingress_controller_tests.yaml +++ b/helm/litellm/tests/ingress_controller_tests.yaml @@ -97,6 +97,16 @@ tests: name: RELEASE-NAME-litellm-gateway port: number: 4000 + - contains: + path: spec.rules[0].http.paths + content: + path: /debug/memory/summary + pathType: Exact + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 - equal: path: spec.rules[0].http.paths[-1] value: diff --git a/helm/litellm/tests/ingress_extra_paths_tests.yaml b/helm/litellm/tests/ingress_extra_paths_tests.yaml index fc7d5943278..1305af15ae2 100644 --- a/helm/litellm/tests/ingress_extra_paths_tests.yaml +++ b/helm/litellm/tests/ingress_extra_paths_tests.yaml @@ -288,6 +288,17 @@ tests: - failedTemplate: errorMessage: "ingress.extraPaths[0]: path /test with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it" + - it: rejects an entry that would take over the exact /debug/memory/summary route + set: + ingress.enabled: true + ingress.extraPaths: + - path: /debug/memory/summary + pathType: Exact + service: backend + asserts: + - failedTemplate: + errorMessage: "ingress.extraPaths[0]: path /debug/memory/summary with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it" + - it: allows a built-in path under a different pathType, which is a distinct rule set: ingress.enabled: true diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 1873219d1ea..4ca54131d6a 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -208,6 +208,11 @@ database: name: litellm-writer-secret usernameKey: username passwordKey: password + # libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the + # in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS. + # sslRootCert on its own implies sslMode verify-full + sslMode: "" + sslRootCert: "" # Optional read-replica routing. When `reader.host` is set, the proxy routes # reads (find_*, count, group_by, query_raw/_first) to this endpoint while diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md deleted file mode 100644 index ae8ae5a6870..00000000000 --- a/litellm-rust/ADDING_A_PROVIDER.md +++ /dev/null @@ -1,29 +0,0 @@ -# Adding a provider / route to litellm-rust - -Everything for a route lives in `crates/core/src//`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint. - -1. **Entrypoint** — `mod.rs`: `pub async fn (request) -> CoreResult`, the Rust equivalent of `litellm.()`, plus a `_stream` variant when the route streams. It is the only thing a host touches. -2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`. -3. **Provider config** — `crates/core/src/providers///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. -4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response. - -## Coding standards - -Before writing new logic, look for an existing base to extend. When a change is -“the same behavior for one more provider/endpoint/integration”, the codebase -almost always already has a shared abstraction for it (for example, provider -`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared -helpers in `litellm_core_utils/`, typed request/response models, or factory -functions). Find it first with a search, then add the new variant by inheriting -from or composing that base, overriding only what genuinely differs (model -name, parameter mapping, or auth). - -Never copy an existing implementation and edit it in place, and never hand-roll -a parallel version of logic a base already provides. If you catch yourself -writing a second copy of a pattern that exists twice already, stop and extract a -base instead: put the shared shape in one place and make both call sites thin -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 the commands under "Checks" in [CLAUDE.md](CLAUDE.md). diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md deleted file mode 100644 index 17856218e60..00000000000 --- a/litellm-rust/AGENTS.md +++ /dev/null @@ -1,45 +0,0 @@ -# AGENTS.md - -litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. - -## Crates - -| Crate | Role | -|-------|------| -| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | -| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. | -| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. | -| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | -| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | -| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | - -Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate. - -## Where a route lives - -A top-level LiteLLM call is a module under `crates/core/src//`, shaped like `messages`: - -``` -core/src/messages/ - mod.rs # pub async fn messages(..) -> CoreResult<..> (+ messages_stream for SSE) - types.rs # request/response types, MessagesRequest - transformation.rs # the provider template trait - prepare.rs # provider resolution, auth headers, URL - handler.rs # the provider call - client.rs # the shared reqwest client -``` - -Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched. - -Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. - -Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional. - -## Style - -All Rust in `litellm-rust/` follows the official Rust Style Guide: -https://doc.rust-lang.org/style-guide/ - -`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style. - -Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md deleted file mode 100644 index dfacf37b6cd..00000000000 --- a/litellm-rust/CLAUDE.md +++ /dev/null @@ -1,189 +0,0 @@ -# CLAUDE.md - -This file defines the rules for Rust work in LiteLLM. - -## Provider Coding Standards - -Before writing new logic, look for an existing base to extend. When a change is -“the same behavior for one more provider/endpoint/integration”, the codebase -almost always already has a shared abstraction for it (for example, provider -`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared -helpers in `litellm_core_utils/`, typed request/response models, or factory -functions). Find it first with a search, then add the new variant by inheriting -from or composing that base, overriding only what genuinely differs (model -name, parameter mapping, or auth). - -Never copy an existing implementation and edit it in place, and never hand-roll -a parallel version of logic a base already provides. If you catch yourself -writing a second copy of a pattern that exists twice already, stop and extract a -base instead: put the shared shape in one place and make both call sites thin -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. - -## Crates (see AGENTS.md) - -`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. -`litellm-config` is the config-loading boundary and returns resolved core types. -`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and -`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop` -holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate -is a layer or shared foundation, not a route; add modules, not crates. - -## Core Boundary - -`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()` -is `litellm_core::messages::messages(request).await`: you call it, it does the -provider call, and you get a typed non-streaming response back. - -Route-level Rust structure mirrors LiteLLM's Python responsibilities: -- `core/src//` owns the route end to end: the public entrypoint fn named - after the route in `mod.rs`, the request/response types (`types.rs`), the - provider template trait (`transformation.rs`), the provider/auth/URL - resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that - performs the call (`handler.rs`). `core/src/messages` is the reference. -- `core/src/providers///transformation.rs` owns the - provider-specific transform. For Anthropic Messages, this means - `core/src/providers/anthropic/messages/transformation.rs`. -- Handlers live in `core`, never in a host. `ai-gateway` must not contain a - route handler that talks to a provider; its axum route reads the HTTP request, - picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals - Python objects and calls the same entrypoint. - -Streaming keeps the same shape: the route entrypoint has a `_stream` -variant in `core` that returns the upstream response so a host can splice it to -its own caller; the host still owns no provider logic. - -Call-hook and lifecycle instrumentation, including phase timing, usage -accumulation, and callback payload construction, always lives in `core`. -Hosts feed observed events into core and dispatch the completed payloads through -their I/O logger; hosts must not own callback orchestration. - -Allowed in `core`: -- The public entrypoint for a top-level LiteLLM call -- Request/response transforms and stream chunk normalization -- Provider resolution, auth header construction, and URL building -- The provider HTTP call itself, through a shared reused client with connect and - request timeouts -- Shared data types and validation errors -- Deterministic token/cost helper logic - -Not allowed in `core`: -- Serving HTTP: axum routes, extractors, and transport concerns stay in the host -- Filesystem access -- Database access -- Config file reading and rollout state -- Logging callbacks, spend writes, or custom callbacks -- Global mutable runtime state - -Env reads in `core` are limited to credential fallback inside a route's -`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when -no key is passed. Everything else config-shaped is resolved by the host and -passed in. - -Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`) -predate this rule and are being moved into `core` route modules; do not add new -ones there, and prefer moving one when you touch it. - -Python owns rollout state and fallback while Rust is being introduced. Rust -paths must be off by default until parity tests prove equivalence with Python. -A new provider/route may instead be implemented rust-only with no Python -reference; then the Python interface is a thin dispatch that calls Rust with no -fallback, and you state the rust-only choice explicitly in the PR. Either way -the Python side stays minimal (it only marshals inputs and calls the Rust -interface), never add a per-route feature flag, and never push provider -dispatch into `litellm/main.py`; put it in a thin dispatch class under -`litellm/llms///`. - -## Production Bar - -Rust code in this workspace is held to a strict parity and robustness bar from -the first PR: - -- Correctness parity is proven with tests. Do not rely on README claims or - manual inspection for a port that mirrors Python behavior. -- Every provider transform must have unit tests for supported-parameter - filtering, request body shape, response normalization, missing/null fields, - and bad-input errors. -- When Rust is exposed through Python, add Python tests that prove disabled, - enabled, and unavailable-bridge fallback behavior. -- Avoid panics on user/provider input. Return typed errors and let the host map - them to Python exceptions or HTTP responses. -- OCR handles documents that often contain personal data. Do not log document - contents, base64 payloads, provider response bodies, or secrets. -- Error messages must be useful but data-minimized. Truncate or sanitize any - upstream body before it crosses a host boundary. -- Treat empty or whitespace-only credentials, URLs, and config values as absent - at the host/config resolution layer. -- Preserve Python output shape intentionally. If a field is always serialized as - `null` for Python parity, leave a short comment explaining that parity choice. - -## Network I/O Rules - -These rules apply to every module that executes network I/O, whether it is a -`core` route handler or a host such as `ai-gateway`: - -- Set connect and full-request timeouts. No unbounded waits. -- Reuse HTTP clients; do not construct clients per request. -- Prefer rustls TLS for portable Python wheels and Linux images unless there is - a documented reason not to. -- Add request IDs and structured tracing at the host layer, without logging OCR - document contents or secrets. -- Do not echo raw upstream response bodies to callers. Sanitize and bound them. -- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is - impossible by construction and documented. - -## Rust Style Guide - -All Rust in `litellm-rust/` follows the official Rust Style Guide: -https://doc.rust-lang.org/style-guide/ - -`rustfmt` implements the guide's formatting rules by default, so the mechanical -side is enforced for you: run `cargo fmt` before committing and CI gates every -PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add -a `rustfmt.toml` that diverges from the default style; the default style *is* the -guide. - -The guide also covers conventions rustfmt cannot auto-apply; follow these too: -- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for - types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and - statics; acronyms count as one word (`HttpClient`, not `HTTPClient`). -- Ordering and grouping the guide prescribes: imports grouped std / external / - crate-local, derives before other attributes, and consistent item order. -- Idioms the guide recommends over the formatter fighting you (e.g. prefer - restructuring an over-long expression rather than forcing an awkward wrap). - -## Constants - -Magic numbers and fixed strings go in a crate-level `constants.rs`, never -hardcoded inline — the Rust mirror of Python's `litellm/constants.py`. - -- Each crate that needs them has `src/constants.rs` (declared `mod constants;`); - import from it (`use crate::constants::...`). Don't scatter `const` values at - the top of feature modules. -- An env-overridable tunable still lives in `constants.rs` as its `DEFAULT_*` - value; the env read (with fallback to that default) happens at the host/config - resolution layer, not in `core`/`providers`. -- Exception: a value that is purely local to one function and has no meaning - elsewhere may stay inline, but prefer `constants.rs` when in doubt. - -## Checks - -Run these before pushing Rust changes. The same checks run in GitHub Actions -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 --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 -the existing Python output with the Rust-backed output. diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7b0b593b70f..7e3d25e9c5d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1948,12 +1948,17 @@ dependencies = [ "azure_core", "azure_identity", "base64 0.22.1", + "bytes", "data-url", + "futures-util", "gcp_auth", + "mime_guess", "moka", "rand 0.8.7", "reqwest 0.12.28", "rstest", + "rustls 0.23.42", + "rustls-native-certs", "serde", "serde_json", "serde_path_to_error", @@ -1962,6 +1967,7 @@ dependencies = [ "subtle", "thiserror 2.0.19", "tokio", + "tokio-tungstenite", "tracing", "tracing-subscriber", "url", @@ -1974,12 +1980,12 @@ version = "0.1.0" dependencies = [ "criterion", "futures-util", - "litellm-ai-gateway", "litellm-core", "litellm-python-interop", "litellm-token-counter", "pyo3", "pyo3-async-runtimes", + "rstest", "serde", "serde_json", "tokio", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 5f25e69a1f8..5c72c86d6ef 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -16,6 +16,7 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +bytes = "1" tracing = "0.1" tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } diff --git a/litellm-rust/README.md b/litellm-rust/README.md deleted file mode 100644 index 650d38753e7..00000000000 --- a/litellm-rust/README.md +++ /dev/null @@ -1,56 +0,0 @@ -# LiteLLM Rust - -This workspace contains the staged Rust implementation for LiteLLM. - -`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call -that makes the LLM call and hands back a typed response, the same shape as -`litellm.messages()` in Python. - -```rust -let response = litellm_core::messages::messages(MessagesRequest { - model: "claude-sonnet-4-5", - body, - api_key: Some(key), - .. -}) -.await?; -``` - -Python continues to own configuration, retries, routing policy, logging, -callbacks, spend tracking, and customer plugins until each Rust path has parity -coverage and production evidence. - -## Crates - -| Crate | Role | -|-------|------| -| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | -| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | -| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | -| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | - -Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. - -## Layout - -```text -crates/ - core/ The SDK: route modules + provider transforms. - src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client - src/providers/anthropic/messages/transformation.rs - config/ Config loading and resolved deployments. - ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. - python-interop/ Domain-neutral PyO3 conversion and GIL primitives. - python-bridge/ PyO3 API adapter for Python LiteLLM. -``` - -The folder shape follows the Python provider tree: -`core/src/providers///transformation.rs`. The bridge exposes one -function per top-level route, mirroring the core entrypoints. - -## Checks - -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 deleted file mode 100644 index 952bbc38b43..00000000000 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ /dev/null @@ -1,53 +0,0 @@ -# Provider coding standards (litellm-rust) - -Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response. - -## Provider resolution - -1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string. -2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers. - -## Transforms and the base config - -3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src//transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`). -4. Each provider implements that trait as a `const __CONFIG` in `core/src/providers///transformation.rs`, mirroring the Python provider tree. -5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it. -6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers. - -## Boundaries - -7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request. -8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers///`; a route is a module, never a new crate. -9. Route entry point stays thin: `core::::()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them. -10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`. - -## Types and errors - -11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string. -12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input. -13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating. -14. Early returns over deep nesting; small focused files over god modules. -15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test. - -## Safety and data minimization - -16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary. -17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer. -18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS. - -## Tests and rollout - -19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity. -20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping. -21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR. - -## Python bridge (SDK side) - -22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust. -23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms///` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method. -24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_`. - -## Checks before push - -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/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md deleted file mode 100644 index b2fd583316b..00000000000 --- a/litellm-rust/crates/ai-gateway/AGENTS.md +++ /dev/null @@ -1,54 +0,0 @@ -# ai-gateway — folder architecture - -The Axum server that fronts the Rust gateway. It owns transport + config + auth -only; deployment selection lives in `core::router`, and the LLM call itself -(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint -such as `litellm_core::messages::messages`. No provider handler lives here. - -``` -src/ - main.rs # entrypoint: build AppState (router + master key), bind, serve - state.rs # AppState — shared Arc + master_key - auth/ # authentication as an axum extractor — added to handler args - mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY) - routes/ # one module per route, all matching the same template - AGENTS.md # ← the route template (read this before adding a route) - mod.rs # app(): merges every module's router() - health.rs # simple route (one file): router() + liveness/readiness - realtime/ # route with logic → axum surface + a no-axum service: - mod.rs # router() + handler + WS<->events adapter (the axum surface) - service.rs # business logic (select deployment, call provider) — no axum, testable -``` - -## Rules - -- **Routes follow one template.** Each route module exposes - `pub fn router() -> Router`; `routes/mod.rs` only merges them. Simple - routes are one file; non-trivial routes are a folder (`handler`/`service`/ - `transport`). See `routes/AGENTS.md`. -- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's - args; it runs during extraction. Never re-implement the check per route. -- **Handlers are thin.** A handler validates and delegates to its `service`. No - business logic, no provider calls, no transforms in handlers. -- **Services call `core`, they don't reimplement it.** A `service` picks the - deployment and calls the `core` route entrypoint. Provider resolution, auth - headers, URL building, and the HTTP call are `core`'s job; a service that - builds a provider request itself is a bug (`routes/messages/service.rs` is - the reference). -- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in - `state.rs`; read env/config only in `main.rs` when building state. - -## Auth (interim) - -A single **master key** (`LITELLM_MASTER_KEY`), enforced by the -`auth::RequireMasterKey` extractor: any caller presenting it as -`Authorization: Bearer ` may invoke the gateway. Fails closed (500) when -unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to -override). Full per-key auth + budgets/rate-limits are delegated to the Python -proxy in a later phase. Health routes don't add the extractor (unauthenticated). - -## Python interop - -Python-backed loading lives in `litellm-config` and is **load-time only**. The -gateway's `python-config` feature forwards to that crate. The realtime data path -never takes the GIL. diff --git a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md deleted file mode 100644 index 6d090cf4c8e..00000000000 --- a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md +++ /dev/null @@ -1,14 +0,0 @@ -# ai-gateway architecture - -The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an -API callback: it POSTs each finished session to the LiteLLM proxy, which records -spend and runs the usual callbacks. - -```mermaid -flowchart LR - C[client] <--> G[Rust ai-gateway
LLM inference] - G <--> O[OpenAI realtime] - G -. spend tracking callback .-> P[litellm proxy] - F[litellm-config
load-time only] --> G - F -. Python backend .-> P -``` diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 74cf66e88a2..dfa61226d4e 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -13,6 +13,11 @@ name = "litellm-ai-gateway" path = "src/main.rs" required-features = ["server"] +[[bin]] +name = "trace-parity-gateway" +path = "src/bin/trace_parity_gateway.rs" +required-features = ["trace-parity"] + [dependencies] tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } diff --git a/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md b/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md deleted file mode 100644 index 84e926af243..00000000000 --- a/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md +++ /dev/null @@ -1,55 +0,0 @@ -# Realtime gateway benchmark — pool on/off - -Measures what the gateway adds over talking to OpenAI's realtime WebSocket -directly, and what the pre-warmed connection pool removes. See -`../../src/routes/realtime/README.md` for how the pool works. - -## Results - -5000 calls / 500 concurrency, gateway at 10 instances, pool ON -(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice. -Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade, -**session** = upgrade → `session.created` (the phase the pool removes), -**1st-audio** = `response.create` → first audio delta (OpenAI inference), -**total** = full wall-clock. - -| metric | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI | -| ------------------ | ------------- | ----------------- | ------------- | ---------- | -| success rate (%) | 99.8 | 99.8 | — | — | -| dial p50 (ms) | 276 | 158 | −118 | **faster** | -| session p50 (ms) | 7 | 0 | −7 | **faster** | -| 1st-audio p50 (ms) | 440 | 664 | +224 | slower¹ | -| total p50 (ms) | 816 | 1010 | +194 | slower¹ | -| total p95 (ms) | 2152 | 1970 | −182 | **faster** | -| total p99 (ms) | 2692 | 2610 | −82 | **faster** | - -The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the -**session phase sub-millisecond** at the median — ~76% of connects hit the pool, -~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead: -`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran -slower during the gateway legs and drags `total p50` with it. - -**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the -fresh-dial overhead the pool removes. - -## Reproduce - -The load generator lives in a separate repo: -**https://github.com/ishaan-berri/litellm-realtime-bench** - -```bash -git clone https://github.com/ishaan-berri/litellm-realtime-bench -cd litellm-realtime-bench && go build -o wsbench . - -# Direct to OpenAI (baseline) -./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60 - -# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0 -./wsbench -host -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60 -``` - -Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`, -`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At -500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was -used here for 10 instances). The bench repo's README covers running 500-concurrency -legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.** diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 6f48f38c9f6..b17f17de11f 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -277,7 +277,7 @@ fn core_error_kind(error: &Error) -> &'static str { Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", - Error::MissingField(_) => "MissingField", + Error::MissingField(_) | Error::MissingDocumentUrl => "MissingField", Error::Http { .. } => "HttpError", Error::InvalidResponse(_) => "InvalidResponse", Error::Network(_) => "NetworkError", diff --git a/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs new file mode 100644 index 00000000000..9036deb9871 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs @@ -0,0 +1,40 @@ +use std::io::Read; + +use serde::Deserialize; +use serde_json::Value; + +#[derive(Deserialize)] +struct Input { + model_alias: String, + provider_model: String, + api_base: String, + body: Value, +} + +#[tokio::main] +async fn main() { + let mut input = String::new(); + if let Err(error) = std::io::stdin().read_to_string(&mut input) { + fail(error); + } + let input: Input = match serde_json::from_str(&input) { + Ok(input) => input, + Err(error) => fail(error), + }; + let result = litellm_ai_gateway::trace_parity::traced_messages_request( + input.model_alias, + input.provider_model, + input.api_base, + input.body, + ) + .await; + match serde_json::to_string(&result) { + Ok(result) => println!("{result}"), + Err(error) => fail(error), + } +} + +fn fail(error: impl std::fmt::Display) -> ! { + eprintln!("{error}"); + std::process::exit(1) +} diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 7f3b6b0650f..f86dd778424 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; -use std::sync::Arc; use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; @@ -10,106 +8,21 @@ use litellm_core::auth::error::MissingCredential; use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; use litellm_core::responses::types::ResponsesWsEvent; use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; -use tokio::net::TcpStream; -use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; -use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use crate::io::tls::connect_upstream; +use litellm_core::responses::websocket::{ResponsesUpstreamWs, connect_upstream}; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, }; const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -pub type ResponsesUpstreamWs = WebSocketStream>; type UpstreamTx = SplitSink; type UpstreamRx = SplitStream; -#[derive(Clone)] -pub struct ResponsesWebSocketConnection { - socket: Arc>>, -} - -impl ResponsesWebSocketConnection { - pub async fn connect_url( - url: &str, - headers: &HashMap, - timeout: Option, - ) -> Result { - let mut request = url - .into_client_request() - .map_err(|error| Error::Network(error.to_string()))?; - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|error| Error::InvalidRequest(error.to_string()))?; - let header_value = HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string()))?; - request.headers_mut().insert(header_name, header_value); - } - let connect = connect_upstream(request); - let result = match timeout { - Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - Error::Network("Responses WebSocket connection timed out".to_string()) - })?, - None => connect.await, - }; - let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { - status: response.status().as_u16(), - body: String::new(), - }, - other => Error::Network(other.to_string()), - })?; - Ok(Self { - socket: Arc::new(Mutex::new(Some(socket))), - }) - } - - pub async fn send_text(&self, text: String) -> Result<(), Error> { - let mut socket = self.socket.lock().await; - let Some(socket) = socket.as_mut() else { - return Err(Error::Network("Responses WebSocket is closed".to_string())); - }; - socket - .send(Message::Text(text)) - .await - .map_err(|error| Error::Network(error.to_string())) - } - - pub async fn recv_text(&self) -> Result, Error> { - let mut socket_guard = self.socket.lock().await; - let Some(socket) = socket_guard.as_mut() else { - return Ok(None); - }; - match socket.next().await { - Some(Ok(Message::Text(text))) => Ok(Some(text)), - Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) - .map(Some) - .map_err(|error| Error::InvalidResponse(error.to_string())), - Some(Ok(Message::Close(_))) | None => Ok(None), - Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Network(error.to_string())), - } - } - - pub async fn close(&self) -> Result<(), Error> { - let mut socket = self.socket.lock().await; - if let Some(socket) = socket.as_mut() { - socket - .close(None) - .await - .map_err(|error| Error::Network(error.to_string()))?; - } - *socket = None; - Ok(()) - } -} - pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { api_key .map(str::trim) diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 39465e28e84..3334053a0a4 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -118,7 +118,8 @@ impl IntoResponse for MessagesRouteError { | Error::Connect(_) | Error::InvalidResponse(_) | Error::InvalidType { .. } - | Error::MissingField(_) => ( + | Error::MissingField(_) + | Error::MissingDocumentUrl => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/ai-gateway/src/trace_parity.rs index 21123df3f1c..00c9b53e691 100644 --- a/litellm-rust/crates/ai-gateway/src/trace_parity.rs +++ b/litellm-rust/crates/ai-gateway/src/trace_parity.rs @@ -10,6 +10,7 @@ use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; use serde::Serialize; use serde_json::Value; use tower::ServiceExt; +use tracing::instrument::WithSubscriber; use crate::io::realtime_pool::RealtimePool; use crate::routes; @@ -21,6 +22,38 @@ pub struct GatewayResponse { pub body: Value, } +#[derive(Debug, Serialize)] +pub struct TracedGatewayResponse { + pub response: Option, + pub error: Option, + pub trace: Vec, +} + +pub async fn traced_messages_request( + model_alias: String, + provider_model: String, + api_base: String, + body: Value, +) -> TracedGatewayResponse { + let trace = litellm_core::observability::FunctionTrace::default(); + let result = messages_request(model_alias, provider_model, api_base, body) + .with_subscriber(trace.dispatcher()) + .await; + let events = trace.events(); + match result { + Ok(response) => TracedGatewayResponse { + response: Some(response), + error: None, + trace: events, + }, + Err(error) => TracedGatewayResponse { + response: None, + error: Some(error.to_string()), + trace: events, + }, + } +} + pub async fn messages_request( model_alias: String, provider_model: String, diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs index 05f7d9610d5..ac37440d682 100644 --- a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -2,10 +2,10 @@ //! API has to resolve its own crypto provider, in a test binary where nothing //! has installed a process-wide one, and has to leave it uninstalled. -use std::collections::HashMap; use std::time::Duration; -use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; +use futures_util::{sink, stream}; +use litellm_ai_gateway::io::responses_ws::async_responses_websocket; use tokio::net::TcpListener; async fn dead_tls_server() -> u16 { @@ -30,10 +30,15 @@ async fn dead_tls_server() -> u16 { async fn dialing_wss_returns_an_error_instead_of_panicking() { let port = dead_tls_server().await; - let result = ResponsesWebSocketConnection::connect_url( - &format!("wss://127.0.0.1:{port}/"), - &HashMap::new(), + let result = async_responses_websocket( + "gpt-5", + Some("test-key"), + Some(&format!("wss://127.0.0.1:{port}/")), + None, Some(Duration::from_secs(10)), + |_| {}, + stream::empty(), + sink::drain(), ) .await; diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index aee8b4937ef..9ba7bfb5323 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -2,6 +2,6 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate. -Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback dispatch. Env reads are limited to credential fallback in a route's `prepare.rs`. +Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`. Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates. diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md deleted file mode 100644 index 5d36305ded5..00000000000 --- a/litellm-rust/crates/core/CLAUDE.md +++ /dev/null @@ -1,66 +0,0 @@ -# CLAUDE.md - -Rules for `litellm-rust/crates/core`. - -## Responsibility - -`core` is the LiteLLM SDK in Rust: it makes the LLM call. Every top-level -LiteLLM call has a public entrypoint here, named after the route -(`messages::messages()` is the Rust equivalent of `litellm.messages()`), and -calling it returns a typed non-streaming response. - -Allowed: -- The public entrypoint for a route, plus its `_stream` variant when the - route supports streaming. -- Provider resolution, auth header construction, URL building, and the provider - HTTP call (shared reused client, connect + request timeouts). -- Shared request/response structs. -- Typed errors with stable, non-sensitive messages. -- Deterministic validation helpers. -- Serialization helpers that intentionally mirror Python output shape. -- Route templates that match Python base config responsibilities, such as - `messages::transformation::AnthropicMessagesProviderConfig`. - -Not allowed: -- Serving HTTP: axum routers, extractors, and other transport concerns. -- Filesystem, database, or cache access. -- Config file reading or rollout state; the host resolves those and passes them - in. Env reads are limited to credential fallback in a route's `prepare.rs`. -- Logging callbacks, tracing spans, spend writes, or customer callbacks. -- Provider-specific branching that belongs in `providers`. -- Panics for user/provider-controlled input. - -## Typed Contracts (core rule) - -Trait and function boundaries MUST be strongly typed. No stringly-typed JSON -(`&str` / `String` / `Vec` / bare `serde_json::Value`) as a transform -input or output. Parse wire bytes into typed structs/enums at the host edge; -`core` and `providers` operate only on those types (e.g. `RealtimeEvent`, -`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a -typed field on a struct, not a raw string threaded through the API. - -## Structure - -Use route names directly under `src/`: `messages`, `ocr`, future -`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not -invent broad names like `engine` for route contracts. - -`src/messages` is the reference shape for a route module: - -``` -mod.rs pub async fn messages(..) (+ messages_stream) -types.rs request/response types -transformation.rs the provider template trait -prepare.rs provider resolution, auth headers, URL -handler.rs the provider call -client.rs the shared reqwest client -``` - -## Parity Rules - -- Every shared type used by a provider transform needs unit tests for - serialization shape. -- If Python parity requires always emitting a `null` field instead of omitting - it, document that in code and pin it with a test. -- Error enums should preserve enough detail for Python/HTTP hosts to map errors - consistently without exposing document contents or upstream bodies. diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index a2433435e34..09c526f73cf 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -6,25 +6,27 @@ license.workspace = true repository.workspace = true autotests = false -[[test]] -name = "workspace_crate_allowlist" -path = "tests/workspace_crate_allowlist.rs" - [dependencies] +bytes.workspace = true +futures-util.workspace = true base64.workspace = true azure_core.workspace = true azure_identity.workspace = true data-url = "0.3.2" gcp_auth.workspace = true moka.workspace = true +mime_guess = "2.0.5" rand.workspace = true reqwest.workspace = true +rustls.workspace = true +rustls-native-certs.workspace = true serde.workspace = true serde_json.workspace = true serde_path_to_error = "0.1" strum.workspace = true subtle.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["sync"] } +tokio-tungstenite.workspace = true thiserror.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, optional = true } diff --git a/litellm-rust/crates/core/src/auth/credential.rs b/litellm-rust/crates/core/src/auth/credential.rs index b5235b6780c..c64d331b877 100644 --- a/litellm-rust/crates/core/src/auth/credential.rs +++ b/litellm-rust/crates/core/src/auth/credential.rs @@ -9,6 +9,21 @@ use crate::AuthError; use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; +pub fn credential_index(requested: &str, names: &[String]) -> Option { + names.iter().position(|name| name == requested) +} + +pub fn credential_default_fields<'a>( + supplied: &[String], + credential_fields: &'a [String], +) -> Vec<&'a str> { + credential_fields + .iter() + .filter(|name| !supplied.contains(name)) + .map(String::as_str) + .collect() +} + #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { Path(PathBuf), diff --git a/litellm-rust/crates/core/src/auth/mod.rs b/litellm-rust/crates/core/src/auth/mod.rs index 35d9c676f65..2940a983fb9 100644 --- a/litellm-rust/crates/core/src/auth/mod.rs +++ b/litellm-rust/crates/core/src/auth/mod.rs @@ -49,6 +49,7 @@ impl Sourced { pub use credential::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, + credential_default_fields, credential_index, }; pub use http::{CredentialPlacement, RequestAuth}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; diff --git a/litellm-rust/crates/core/src/call_lifecycle/README.md b/litellm-rust/crates/core/src/call_lifecycle/README.md deleted file mode 100644 index 692e249ef27..00000000000 --- a/litellm-rust/crates/core/src/call_lifecycle/README.md +++ /dev/null @@ -1,167 +0,0 @@ -# Call lifecycle - -`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call -types migrated to Rust. It owns lifecycle ordering, phase timing, and trace -observer calls. It must not know about OCR, chat, messages, responses, -completions, provider auth, request transforms, or response normalization. - -Call-type modules own their domain behavior. For example, OCR owns document -payloads, OCR provider transforms, safe document fetch, guardrail payload shape, -callback payload shape, and provider HTTP execution. - -## Runtime order - -Every wrapped call runs in this order: - -1. `async_pre_call_hook` -2. `async_during_call_hook` -3. provider call -4. `async_log_success_event` or `async_log_failure_event` - -`async_pre_call_hook` receives the initial LiteLLM request shape. It is where -pre-call custom guardrails run. - -`async_during_call_hook` converts the initial request into the provider-ready -request. It is where provider config selection, parameter mapping, auth/header -resolution, request transforms, and during-call guardrails belong. - -The provider call receives only the provider-ready request. It should execute -I/O and call the provider response transform. - -Success and failure callbacks receive `CallLifecycleTiming`. Callback failures -must not replace the original provider or guardrail result. - -## Trace contract - -The lifecycle runner records: - -- full call start and end time -- `pre_call` phase timing -- `during_call` phase timing -- `provider_call` phase timing -- `success_callback` phase timing -- `failure_callback` phase timing - -`CallLifecycleObserver` receives phase start and end events. The default -observer is a no-op. Future OTEL support should implement this observer instead -of editing OCR, chat, messages, responses, completions, or provider modules. - -## Required shape - -Each migrated call type should use this folder shape: - -```text -litellm-rust/crates/ai-gateway/src// - mod.rs # thin public entrypoint - types.rs # public request, prepared request, provider request, response types - prepare.rs # model/provider/callback/guardrail setup - hooks.rs # CallLifecycleHooks implementation - handler.rs # provider I/O and response normalization - tests.rs # call-type lifecycle and handler tests -``` - -Provider transforms can live in `litellm-rust/crates/core/src/providers/...`. -Shared call-type helpers can live beside the call type, but generic lifecycle -code stays in this folder. - -## Core API - -The prepared request implements `CallLifecycleRequest`: - -```rust -impl CallLifecycleRequest for PreparedMessagesRequest { - fn lifecycle_context(&self) -> CallLifecycleContext { - CallLifecycleContext::new( - "messages", - self.model.clone(), - self.custom_llm_provider.clone(), - self.litellm_call_id.clone(), - ) - } -} -``` - -The call-type hooks implement `CallLifecycleHooks`: - -```rust -impl CallLifecycleHooks< - PreparedMessagesRequest, - ProviderMessagesRequest, - MessagesResponse, -> for MessagesLifecycleHooks { - fn async_pre_call_hook(...) { - // run pre-call custom guardrails against the LiteLLM request shape - } - - fn async_during_call_hook(...) { - // map params, validate env, transform request, run during-call guardrails - } - - fn async_log_success_event(...) { - // call async_log_success_event on configured custom loggers - } - - fn async_log_failure_event(...) { - // call async_log_failure_event without swallowing the original error - } -} -``` - -The public entrypoint stays thin: - -```rust -pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { - let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?; - - CallLifecycle::default() - .run_request(request, &hooks, execute_messages_provider_call) - .await -} -``` - -Use `run_request` for new call types. Keep `run` available only for specialized -tests or existing code that already has a `CallLifecycleContext`. - -## Adding a new call type - -1. Add `/types.rs` - -Define the public request accepted by the bridge, the prepared request used by -the lifecycle runner, and the provider request consumed by the handler. - -2. Implement `CallLifecycleRequest` - -Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`. -Do not put provider-specific logic here. - -3. Add `/prepare.rs` - -Resolve model/provider once, generate or preserve `litellm_call_id`, construct -callback and guardrail runners, and return `PreparedCall`. - -4. Add `/hooks.rs` - -Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction, -provider config selection, param mapping, request transform, during-call -guardrail payload construction, and callback payload construction here. - -5. Add `/handler.rs` - -Execute the provider request and normalize the provider response. Do not repeat -provider-specific transforms here; call the provider config. - -6. Add tests - -Cover hook order, success callback payload, failure callback payload, pre-call -guardrail blocking before provider I/O, during-call body mutation, and provider -error mapping. - -## Review checklist - -- Core lifecycle has no call-type or provider-specific branches -- Public call-type entrypoint only prepares and calls `run_request` -- Provider behavior lives behind provider config/transformation code -- Hook method names map to the Python custom logger and guardrail concepts -- Phase timing is recorded once in lifecycle, not separately per call type -- Callback failures never hide the original provider or guardrail error -- Tests prove the provider socket is not touched when pre-call guardrails block diff --git a/litellm-rust/crates/core/src/call_lifecycle/host.rs b/litellm-rust/crates/core/src/call_lifecycle/host.rs new file mode 100644 index 00000000000..ac6ddf99b9e --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/host.rs @@ -0,0 +1,121 @@ +use std::future::Future; +use std::pin::Pin; + +pub enum HostCallStep { + Host(O), + Complete(C), +} + +pub type HostCallFuture<'a, O, C> = + Pin, crate::Error>> + Send + 'a>>; + +pub trait HostCall: Send + Sync { + type Operation: Send + 'static; + type Result: Send + 'static; + type Complete: Send + 'static; + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete>; + + fn interrupt( + &mut self, + failure: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete>; +} + +pub enum HostStep { + Ready(V), + Suspend(S), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum HostPhase { + Setup, + DeploymentPreCall, + Prepare, + Execute, + ConstructResponse, + DeploymentPostCall, + Finalize, + Success, + MapFailure, + DeploymentFailure, + Failure, + AsyncFailure, + Complete, +} + +#[derive(Clone, Debug)] +pub enum HostFailure { + Error(crate::Error), + Cancelled(crate::Error), +} + +pub struct HostLifecycle { + phase: HostPhase, + asynchronous: bool, +} + +impl HostLifecycle { + pub fn new(asynchronous: bool) -> Self { + Self { + phase: HostPhase::Setup, + asynchronous, + } + } + + pub fn phase(&self) -> HostPhase { + self.phase + } + + pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option { + if let Err(failure) = result { + if self.phase == HostPhase::DeploymentFailure { + self.phase = HostPhase::Failure; + return None; + } + let error = match failure { + HostFailure::Cancelled(error) => { + self.phase = HostPhase::Complete; + return Some(error); + } + HostFailure::Error(error) => error, + }; + match self.phase { + HostPhase::Failure | HostPhase::AsyncFailure => { + self.advance(); + return None; + } + HostPhase::Success => self.phase = HostPhase::Complete, + HostPhase::Execute | HostPhase::ConstructResponse => { + self.phase = HostPhase::MapFailure; + } + _ => self.phase = HostPhase::Failure, + } + return Some(error); + } + self.advance(); + None + } + + fn advance(&mut self) { + self.phase = match self.phase { + HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall, + HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare, + HostPhase::Prepare => HostPhase::Execute, + HostPhase::Execute => HostPhase::ConstructResponse, + HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall, + HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize, + HostPhase::Finalize => HostPhase::Success, + HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure, + HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure, + HostPhase::Failure if self.asynchronous => HostPhase::AsyncFailure, + HostPhase::Failure + | HostPhase::AsyncFailure + | HostPhase::Success + | HostPhase::Complete => HostPhase::Complete, + }; + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 637c156e192..5c752a73899 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -3,6 +3,10 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH}; use crate::Error; +pub mod host; +#[cfg(test)] +#[path = "../../tests/host_lifecycle.rs"] +mod host_tests; pub mod types; pub use types::{ diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index 9469d379462..1babb0078b8 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -46,9 +46,10 @@ pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10; +pub(crate) const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600; pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; -pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; +pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024; pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10; pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120; @@ -63,3 +64,6 @@ pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY"; pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://"; pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr"; pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1"; + +pub(crate) const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com"; +pub(crate) const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY"; diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index fa4a9d36e03..359ad56c336 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -1,6 +1,6 @@ use thiserror::Error as ThisError; -#[derive(Debug, ThisError, PartialEq, Eq)] +#[derive(Clone, Debug, ThisError, PartialEq, Eq)] pub enum Error { #[error("expected {expected}, got {actual}")] InvalidType { @@ -9,6 +9,8 @@ pub enum Error { }, #[error("missing required field: {0}")] MissingField(&'static str), + #[error("Document URL is required")] + MissingDocumentUrl, #[error("invalid response: {0}")] InvalidResponse(String), #[error("invalid provider: {0}")] @@ -52,6 +54,17 @@ pub enum Error { Unsupported(&'static str), } +impl Error { + pub const fn http_status_code(&self) -> Option { + match self { + Self::InvalidRequest(_) => Some(400), + Self::MissingDocumentUrl => Some(500), + Self::Http { status, .. } => Some(*status), + _ => None, + } + } +} + #[derive(Debug, ThisError)] pub(crate) enum MediaError { #[error("media URL rejected by network policy")] @@ -106,6 +119,7 @@ impl From for Error { fn from(error: crate::ocr::error::OcrRequestError) -> Self { match error { crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field), + crate::ocr::error::OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl, error => Self::InvalidRequest(error.to_string()), } } diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs new file mode 100644 index 00000000000..4c8455a171c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs @@ -0,0 +1,131 @@ +use super::super::OcrAdapter; +use crate::Error; +use crate::ocr::OcrClient; +use crate::ocr::codecs::cohere::{ + CohereParams, CohereResponse, transform_request, transform_response, validate_document, +}; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{credential_env, transform_request_body}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use crate::providers::azure_ai::auth::AzureAuthInputs; +use crate::url_utils::ApiUrl; + +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; + +pub(crate) struct AzureCohereAdapter; + +impl OcrAdapter for AzureCohereAdapter { + type ProviderResponse = CohereResponse; + const PROVIDER: OcrProvider = OcrProvider::AzureAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let params = super::super::super::wire::decode_request_value::( + serde_json::Value::Object(request.optional_params.clone()), + "optional_params", + )?; + let mut config = AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); + let base = request + .connection + .api_base + .clone() + .or_else(|| credential_env(AZURE_AI_API_BASE_ENV)) + .filter(|base| !base.trim().is_empty()) + .ok_or_else(|| { + Error::Auth( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(), + ) + })?; + let headers = + super::validate_ai_environment(&request.connection, &config, &credential_env).await?; + validate_document(&request.document)?; + let remote = request.document.source().starts_with("http://") + || request.document.source().starts_with("https://"); + let document = inline_remote_document( + client.document_fetcher(), + request.document.clone(), + &request.connection, + ) + .await?; + let body = transform_request(&request.model, document, params)?; + transform_request_body( + client, + request, + &complete_url(&base)?, + &headers, + !remote, + body, + |body| { + validate_document(&body.document)?; + validate_inline_document(&body.document) + }, + ) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + transform_response(&request.model, response) + } +} + +fn complete_url(base: &str) -> Result { + let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + if !matches!(url.scheme(), "http" | "https") { + return Err(invalid_api_base().into()); + } + let path = url.path().trim_end_matches('/').to_string(); + if path.ends_with("/v2/parse") { + url.set_path(&path); + return Ok(url.into()); + } + url.set_path(path.strip_suffix("/models").unwrap_or(&path)); + ApiUrl::parse(url.as_str()) + .and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base().into()) +} + +fn invalid_api_base() -> OcrRequestError { + OcrRequestError::RequestField { + path: "api_base".into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() { + for suffix in [ + "", + "/models", + "/providers/cohere/v2", + "/providers/cohere/v2/parse", + ] { + assert_eq!( + complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(), + "https://example.com/providers/cohere/v2/parse?tenant=a" + ); + } + assert_eq!( + complete_url("https://example.com/v2/parse?tenant=a").unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + assert!(complete_url("relative/path").is_err()); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs index 71ca69ddc58..e90c27ba59d 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs @@ -10,7 +10,6 @@ use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; use crate::ocr::prepare::{credential_env, transform_request_body}; use crate::ocr::registry::OcrProvider; use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat}; -use crate::ocr::wire::DecodedOcrResponse; use crate::providers::azure_ai::auth::AzureAuthInputs; use crate::url_utils::ApiUrl; @@ -32,18 +31,19 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { client: &OcrClient, ) -> Result { let params = map_ocr_params(request)?; - let config = AzureAuthInputs::from_sourced_optional_params( + let mut config = AzureAuthInputs::from_sourced_optional_params( &request.optional_params, &request.input_sources, ) .map_err(Error::from)?; + config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); let headers = validate_environment(&request.connection, &config, &credential_env).await?; let endpoint = nonblank(request.connection.api_base.clone()) .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) .ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?; let url = get_complete_url(&endpoint, &request.model, ¶ms)?; let body = document_intelligence::transform_ocr_request(request.document.clone())?; - transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await + transform_request_body(client, request, &url, &headers, false, body, |_| Ok(())).await } fn transform_ocr_response( @@ -61,7 +61,7 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { url: &str, headers: &[(String, String)], request: &LiteLLMOcrRequest, - ) -> Result, OcrError> { + ) -> Result, OcrError> { polling::read_operation_response( client.polling_http(), response, @@ -69,6 +69,7 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { headers, &request.connection, request.response_format()? == OcrResponseFormat::Native, + &request.hooks, ) .await } diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs index 1bddea0da4f..6ed1e4441d4 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Duration; use reqwest::Url; @@ -9,6 +10,7 @@ use crate::ocr::codecs::document_intelligence::{ AzureDocumentIntelligenceOperation, OperationStatus, }; use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError}; +use crate::ocr::hooks::OcrHooks; use crate::ocr::types::OcrConnection; use crate::ocr::wire::DecodedOcrResponse; @@ -19,24 +21,33 @@ pub(super) async fn read_operation_response( headers: &[(String, String)], connection: &OcrConnection, native: bool, + hooks: &Arc, ) -> Result, OcrError> { if response.status() != reqwest::StatusCode::ACCEPTED { - return read_json_response(response, native).await; + let bytes = + crate::ocr::client::read_response_bytes(response, connection.max_response_bytes) + .await?; + crate::ocr::handler::post_call(hooks, &bytes).await?; + return Ok(crate::ocr::wire::decode_response(&bytes, native)?); } let location = response .headers() .get("operation-location") .and_then(|value| value.to_str().ok()) - .ok_or(OcrPollingError::PollLocation)?; + .ok_or(OcrPollingError::PollLocation)? + .to_string(); let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?; - let operation = Url::parse(location).map_err(|_| OcrPollingError::PollOrigin)?; + let operation = Url::parse(&location).map_err(|_| OcrPollingError::PollOrigin)?; if original.origin() != operation.origin() || !operation.username().is_empty() || operation.password().is_some() { return Err(OcrPollingError::PollOrigin.into()); } - poll_operation(http_client, operation, headers, connection, native).await + let bytes = + crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?; + crate::ocr::handler::post_call(hooks, &bytes).await?; + poll_operation(http_client, operation, headers, connection, native, hooks).await } async fn poll_operation( @@ -45,6 +56,7 @@ async fn poll_operation( headers: &[(String, String)], connection: &OcrConnection, native: bool, + hooks: &Arc, ) -> Result, OcrError> { let deadline = Instant::now() .checked_add(connection.poll_timeout) @@ -75,12 +87,19 @@ async fn poll_operation( .max(1); let decoded = tokio::time::timeout_at( deadline, - read_json_response::(response, native), + read_json_response::( + response, + native, + connection.max_response_bytes, + ), ) .await .map_err(|_| OcrPollingError::PollTimeout)??; match &decoded.data.status { - Some(OperationStatus::Succeeded) => return Ok(decoded), + Some(OperationStatus::Succeeded) => { + crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?; + return Ok(decoded); + } Some(OperationStatus::Running | OperationStatus::NotStarted) => { tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry))) .await diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs index 3107494d39e..8639590b05c 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs @@ -33,13 +33,16 @@ impl OcrAdapter for AzureMistralAdapter { known: params, extra_params: _extra_params, } = _prepare_ocr_request::(request)?; - let config = AzureAuthInputs::from_sourced_optional_params( + let mut config = AzureAuthInputs::from_sourced_optional_params( &request.optional_params, &request.input_sources, ) .map_err(Error::from)?; - let headers = validate_environment(&request.connection, &config, &credential_env).await?; + config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; + let headers = validate_environment(&request.connection, &config, &credential_env).await?; + let retains_document = !request.document.source().starts_with("http://") + && !request.document.source().starts_with("https://"); let document = inline_remote_document( client.document_fetcher(), request.document.clone(), @@ -47,9 +50,15 @@ impl OcrAdapter for AzureMistralAdapter { ) .await?; let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; - transform_request_body(client, request, &url, &headers, body, |body| { - validate_inline_document(&body.document) - }) + transform_request_body( + client, + request, + &url, + &headers, + retains_document, + body, + |body| validate_inline_document(&body.document), + ) .await } @@ -83,12 +92,15 @@ fn get_complete_url( }) } -async fn validate_environment( +pub(in crate::ocr::adapters) async fn validate_environment( connection: &OcrConnection, config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, OcrError> { if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + if config.azure_ad_token_provider.is_some() { + super::resolve_entra(config, env_lookup).await?; + } super::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs index 9c02a7471c9..3d30ae6d6bd 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs @@ -1,3 +1,4 @@ +mod cohere; mod document_intelligence; mod mistral; @@ -10,8 +11,10 @@ use crate::ocr::error::OcrError; use crate::ocr::types::OcrConnection; use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService}; +pub(crate) use cohere::AzureCohereAdapter; pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter; pub(crate) use mistral::AzureMistralAdapter; +pub(super) use mistral::validate_environment as validate_ai_environment; async fn resolve_entra( config: &AzureAuthInputs, @@ -22,6 +25,10 @@ async fn resolve_entra( .get_or_init(AzureAuthService::default) .get_azure_ad_token(config, env_lookup) .await + .or_else(|error| match error { + crate::AuthError::EmptyAzureToken => Ok(None), + other => Err(other), + }) .map(|credential| { credential.map(|credential| { let source = credential.source(); diff --git a/litellm-rust/crates/core/src/ocr/adapters/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs new file mode 100644 index 00000000000..933ead7f7f7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs @@ -0,0 +1,123 @@ +use super::OcrAdapter; +use crate::Error; +use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE}; +use crate::ocr::OcrClient; +use crate::ocr::codecs::cohere::{ + CohereParams, CohereResponse, transform_request, transform_response, validate_document, +}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{credential_env, transform_request_body}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; +use crate::url_utils::ApiUrl; + +pub(crate) struct CohereAdapter; + +impl OcrAdapter for CohereAdapter { + type ProviderResponse = CohereResponse; + const PROVIDER: OcrProvider = OcrProvider::Cohere; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let params = super::super::wire::decode_request_value::( + serde_json::Value::Object(request.optional_params.clone()), + "optional_params", + )?; + let headers = validate_environment(&request.connection, &credential_env)?; + let url = complete_url( + request + .connection + .api_base + .as_deref() + .unwrap_or(COHERE_PARSE_API_BASE), + )?; + let body = transform_request(&request.model, request.document.clone(), params)?; + transform_request_body(client, request, &url, &headers, true, body, |body| { + validate_document(&body.document) + }) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + transform_response(&request.model, response) + } +} + +fn complete_url(base: &str) -> Result { + let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(invalid_api_base().into()); + } + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base().into()) +} + +fn invalid_api_base() -> OcrRequestError { + OcrRequestError::RequestField { + path: "api_base".into(), + } +} + +fn validate_environment( + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(COHERE_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| { + Error::Auth("Missing COHERE_API_KEY - set it in the environment or pass api_key".into()) + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() { + for suffix in ["", "/v2", "/v2/parse"] { + assert_eq!( + complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + } + } + + #[test] + fn rejects_invalid_urls_and_blank_keys() { + assert!(complete_url("relative/path").is_err()); + assert!(complete_url("ftp://example.com").is_err()); + assert!(matches!( + validate_environment( + &OcrConnection { + api_key: Some(" ".into()), + ..Default::default() + }, + &|_| None, + ), + Err(OcrError::Public(Error::Auth(_))) + )); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs index ea569ffb34f..cdbc2c3effc 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs @@ -33,7 +33,7 @@ impl OcrAdapter for MistralAdapter { let url = get_complete_url(request.connection.api_base.as_deref())?; let body = mistral::transform_ocr_request(&request.model, request.document.clone(), ¶ms)?; - transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await + transform_request_body(client, request, &url, &headers, true, body, |_| Ok(())).await } fn transform_ocr_response( diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs index 9171d11836c..d473fcad280 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -5,15 +5,16 @@ use serde::de::DeserializeOwned; use super::OcrClient; use super::error::{OcrError, OcrResponseError}; use super::registry::OcrProvider; -use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat}; -use super::wire::DecodedOcrResponse; +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; mod azure; +mod cohere; mod mistral; mod reducto; mod vertex; -pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; +pub(crate) use azure::{AzureCohereAdapter, AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; +pub(crate) use cohere::CohereAdapter; pub(crate) use mistral::MistralAdapter; pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter}; pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter}; @@ -55,18 +56,27 @@ pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static { _url: &str, _headers: &[(String, String)], request: &LiteLLMOcrRequest, - ) -> impl Future, OcrError>> + Send - { - let retain_native = request - .response_format() - .map(|format| format == OcrResponseFormat::Native); - async move { super::client::read_json_response(response, retain_native?).await } + ) -> impl Future< + Output = Result, OcrError>, + > + Send { + async move { + let bytes = + super::client::read_response_bytes(response, request.connection.max_response_bytes) + .await?; + super::handler::post_call(&request.hooks, &bytes).await?; + Ok(super::wire::decode_response( + &bytes, + request.response_format()? == super::types::OcrResponseFormat::Native, + )?) + } } } macro_rules! for_each_ocr_adapter { ($callback:ident) => { $callback! { + Cohere, $crate::ocr::adapters::CohereAdapter, $crate::ocr::adapters::CohereAdapter, Cohere; + AzureCohere, $crate::ocr::adapters::AzureCohereAdapter, $crate::ocr::adapters::AzureCohereAdapter, AzureAi; Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral; AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi; AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs index 062a0071a34..8889bcd1b45 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs @@ -27,7 +27,7 @@ impl OcrAdapter for ReductoLegacyAdapter { } = _prepare_ocr_request::(request)?; let headers = super::validate_environment(&request.connection, &credential_env)?; let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; - let document = guardrail_document(request, &url).await?; + let (document, headers) = guardrail_document(request, &url, &headers).await?; let document = super::prepare_document(client, document, &request.connection, &headers).await?; let body = reducto::transform_legacy_ocr_request(&request.model, document, ¶ms)?; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs index 7621d0d326a..2dafe291674 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs @@ -93,7 +93,7 @@ pub(super) async fn prepare_document( .map_err(crate::error::TransportError::from)?; let uploaded = crate::ocr::client::read_json_response::< crate::ocr::codecs::reducto::ReductoUploadResponse, - >(response, false) + >(response, false, connection.max_response_bytes) .await? .data; let file_id = uploaded diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs index a49f8105e26..c272d31b67e 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs @@ -27,7 +27,7 @@ impl OcrAdapter for ReductoV3Adapter { } = _prepare_ocr_request::(request)?; let headers = super::validate_environment(&request.connection, &credential_env)?; let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; - let document = guardrail_document(request, &url).await?; + let (document, headers) = guardrail_document(request, &url, &headers).await?; let document = super::prepare_document(client, document, &request.connection, &headers).await?; let body = reducto::transform_v3_ocr_request(&request.model, document, ¶ms)?; diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs index ef188f8b9ac..d16b3e7f386 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs @@ -57,9 +57,15 @@ impl OcrAdapter for VertexDeepSeekAdapter { let document = request.document.clone(); let body = deepseek::transform_ocr_request(&provider_model(&request.model), document, ¶ms)?; - transform_request_body(client, request, &url, &authentication.headers, body, |_| { - Ok(()) - }) + transform_request_body( + client, + request, + &url, + &authentication.headers, + false, + body, + |_| Ok(()), + ) .await } diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs index f3335bf497c..88c61725cee 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs @@ -54,6 +54,8 @@ impl OcrAdapter for VertexMistralAdapter { &location, &request.model, )?; + let retains_document = !request.document.source().starts_with("http://") + && !request.document.source().starts_with("https://"); let document = inline_remote_document( client.document_fetcher(), request.document.clone(), @@ -66,6 +68,7 @@ impl OcrAdapter for VertexMistralAdapter { request, &url, &authentication.headers, + retains_document, body, |body| validate_inline_document(&body.document), ) diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index ab2d098d0bb..394ca778d2f 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,10 +1,10 @@ use std::sync::OnceLock; use std::time::Duration; +use bytes::{Bytes, BytesMut}; use serde::de::DeserializeOwned; -use super::error::OcrError; -use super::handler::perform_ocr_request; +use super::error::{OcrError, OcrResponseError}; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; use super::wire::{DecodedOcrResponse, decode_response}; use crate::Error; @@ -32,6 +32,10 @@ impl OcrClient { }) } + pub fn shared() -> Result { + shared_client() + } + #[tracing::instrument( name = "ocr", target = "litellm::function_trace", @@ -39,7 +43,34 @@ impl OcrClient { skip_all )] pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result { - perform_ocr_request(self, request).await + use super::{ + NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost, + OcrHostOperation, OcrHostResult, + }; + + let host = OcrHookHost::new(request.hooks.clone()); + let mut request = Some(request); + let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all()) + else { + return Err(Error::InvalidRequest( + "native OCR host admission declined".into(), + )); + }; + let mut result = None; + loop { + match call.resume(result.take()).await? { + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().ok_or_else(|| { + Error::InvalidRequest("OCR request was already projected".into()) + })?), + false, + )))) + } + OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await), + OcrCallStep::Complete(response) => return Ok(response), + } + } } pub(crate) fn provider_http(&self) -> &reqwest::Client { @@ -77,7 +108,7 @@ fn no_redirect_http() -> Result { .map_err(TransportError::from) } -pub async fn ocr(request: LiteLLMOcrRequest) -> Result { +pub(crate) fn shared_client() -> Result { static CLIENT: OnceLock> = OnceLock::new(); let client = CLIENT .get_or_init(|| { @@ -88,18 +119,50 @@ pub async fn ocr(request: LiteLLMOcrRequest) -> Result Result { + shared_client()?.perform(request).await } pub async fn read_json_response( response: reqwest::Response, native: bool, + max_response_bytes: usize, ) -> Result, OcrError> { + let bytes = read_response_bytes(response, max_response_bytes).await?; + Ok(decode_response(&bytes, native)?) +} + +pub(crate) async fn read_response_bytes( + mut response: reqwest::Response, + max_response_bytes: usize, +) -> Result { let status = response.status(); - let bytes = response - .bytes() - .await - .map_err(crate::error::TransportError::from)?; + let limit = if status.is_success() { + max_response_bytes + } else { + max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)) + }; + if status.is_success() + && response + .content_length() + .is_some_and(|length| length > limit as u64) + { + return Err(OcrResponseError::TooLarge { limit }.into()); + } + let mut bytes = BytesMut::new(); + while let Some(chunk) = response.chunk().await.map_err(transport_error)? { + let remaining = limit.saturating_sub(bytes.len()); + if status.is_success() && chunk.len() > remaining { + return Err(OcrResponseError::TooLarge { limit }.into()); + } + bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); + if !status.is_success() && bytes.len() == limit { + break; + } + } if !status.is_success() { return Err(crate::error::TransportError::Http { status: status.as_u16(), @@ -107,5 +170,41 @@ pub async fn read_json_response( } .into()); } - Ok(decode_response(&bytes, native)?) + Ok(bytes.freeze()) +} + +pub(crate) fn transport_error(error: reqwest::Error) -> Error { + if error.is_timeout() { + return Error::Http { + status: 408, + body: "OCR request timed out".into(), + }; + } + crate::error::TransportError::from(error).into() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn request_timeout_has_an_http_408_status() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let _connection = listener.accept().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + }); + let error = reqwest::Client::new() + .get(format!("http://{address}")) + .timeout(Duration::from_millis(10)) + .send() + .await + .unwrap_err(); + assert!(matches!( + transport_error(error), + Error::Http { status: 408, .. } + )); + server.abort(); + } } diff --git a/litellm-rust/crates/core/src/ocr/codecs/cohere.rs b/litellm-rust/crates/core/src/ocr/codecs/cohere.rs new file mode 100644 index 00000000000..649432f39d3 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/cohere.rs @@ -0,0 +1,254 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; + +use crate::ocr::document::InlineDocument; +use crate::ocr::error::{OcrRequestError, OcrResponseError}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum OutputFormat { + #[default] + Markdown, + Blocks, +} + +#[derive(Deserialize)] +pub(crate) struct CohereParams { + #[serde(default)] + pub output_format: OutputFormat, +} + +#[derive(Deserialize, Serialize)] +pub(crate) struct CohereRequest { + pub model: String, + pub document: OcrDocument, + pub output_format: OutputFormat, +} + +pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), OcrRequestError> { + let OcrDocument::ImageUrl { image_url, .. } = document else { + return Err(OcrRequestError::CohereImageOnly); + }; + if image_url.is_empty() { + return Err(OcrRequestError::CohereImageOnly); + } + if let Some(inline) = InlineDocument::parse(image_url)? { + if !inline.mime_type().type_.eq_ignore_ascii_case("image") { + return Err(OcrRequestError::CohereImageOnly); + } + inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + } + Ok(()) +} + +#[derive(Deserialize)] +pub(crate) struct CohereResponse { + #[serde(default)] + pages: Vec, + meta: Option, +} + +#[derive(Deserialize)] +struct CoherePage { + index: Option, + markdown: Option, + blocks: Option>>, +} + +#[derive(Deserialize)] +struct CohereMarkdown { + #[serde(default)] + content: String, + images: Option>>, +} + +#[derive(Deserialize)] +struct CohereMeta { + billed_units: Option, +} + +#[derive(Deserialize)] +struct CohereBilledUnits { + pages: Option, +} + +pub(crate) fn transform_response( + model: &str, + response: CohereResponse, +) -> Result { + let pages_processed = response + .meta + .and_then(|meta| meta.billed_units) + .and_then(|units| units.pages) + .map(Ok) + .unwrap_or_else(|| { + i64::try_from(response.pages.len()).map_err(|_| OcrResponseError::NumericRange("pages")) + })?; + let pages = response + .pages + .into_iter() + .enumerate() + .map(|(position, page)| { + let index = page.index.map(Ok).unwrap_or_else(|| { + i64::try_from(position).map_err(|_| OcrResponseError::NumericRange("page index")) + })?; + let (content, images) = page + .markdown + .map(|markdown| { + let images = + markdown + .images + .filter(|images| !images.is_empty()) + .map(|images| { + images + .into_iter() + .map(|mut image| { + if let Some(Value::Object(bbox)) = + image.get("bounding_box").cloned() + { + image.insert("bbox".into(), Value::Object(bbox)); + } + Value::Object(image) + }) + .collect::>() + }); + (markdown.content, images) + }) + .unwrap_or_default(); + let mut normalized = json!({"index": index, "markdown": content, "images": images}); + if let Some(blocks) = page.blocks { + normalized["blocks"] = json!(blocks); + } + Ok(normalized) + }) + .collect::, OcrResponseError>>()?; + Ok(LiteLLMOcrResponse { + pages, + model: model.into(), + document_annotation: None, + usage_info: Some(json!({"pages_processed": pages_processed})), + object: "ocr".into(), + extra_fields: Map::new(), + provider_native_response: None, + }) +} + +pub(crate) fn transform_request( + model: &str, + document: OcrDocument, + params: CohereParams, +) -> Result { + validate_document(&document)?; + Ok(CohereRequest { + model: model.into(), + document, + output_format: params.output_format, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn response_normalizes_markdown_images_blocks_and_billed_pages() { + let response = serde_json::from_value(json!({ + "pages": [ + { + "type":"markdown", + "index":4, + "markdown":{ + "content":"receipt", + "images":[{ + "id":"image", + "bounding_box":{"top_left_x":1,"bottom_right_x":48}, + "bounding_box_normalized":{"top_left_x":0.04,"bottom_right_x":0.15}, + "description":"scan", + "category":"logo" + }] + } + }, + {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]} + ], + "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}} + })) + .unwrap(); + let normalized = transform_response("parse-v5.0", response).unwrap(); + assert_eq!(normalized.pages[0]["index"], 4); + assert_eq!(normalized.pages[0]["markdown"], "receipt"); + assert_eq!(normalized.pages[0]["images"][0]["bbox"]["top_left_x"], 1); + assert_eq!( + normalized.pages[0]["images"][0]["bounding_box_normalized"]["bottom_right_x"], + 0.15 + ); + assert_eq!(normalized.pages[0]["images"][0]["description"], "scan"); + assert_eq!(normalized.pages[0]["images"][0]["category"], "logo"); + assert_eq!(normalized.pages[1]["index"], 1); + assert_eq!(normalized.pages[1]["markdown"], ""); + assert_eq!(normalized.pages[1]["blocks"][0]["text"]["content"], "total"); + assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 3); + } + + #[test] + fn response_defaults_and_invalid_fields() { + for value in [ + json!({}), + json!({"meta":null}), + json!({"pages":[],"meta":{"billed_units":null}}), + ] { + let normalized = + transform_response("parse", serde_json::from_value(value).unwrap()).unwrap(); + assert!(normalized.pages.is_empty()); + assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 0); + } + for value in [ + json!({"pages":null}), + json!({"pages":[{"markdown":"text"}]}), + json!({"pages":[{"index":"bad"}]}), + ] { + assert!(serde_json::from_value::(value).is_err()); + } + let normalized = transform_response( + "parse", + serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(), + ) + .unwrap(); + assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 1); + assert!(normalized.pages[0]["images"].is_null()); + } + + #[test] + fn request_requires_image_and_supported_output_format() { + for value in [ + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + json!({"type":"image_url","image_url":""}), + json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}), + ] { + assert_eq!( + validate_document(&serde_json::from_value(value).unwrap()), + Err(OcrRequestError::CohereImageOnly) + ); + } + assert!(serde_json::from_value::(json!({"output_format":"html"})).is_err()); + for format in ["markdown", "blocks"] { + assert!( + serde_json::from_value::(json!({"output_format":format})).is_ok() + ); + } + let request = transform_request( + "parse-v5.0", + serde_json::from_value(json!({ + "type":"image_url", + "image_url":"https://example.com/image.png" + })) + .unwrap(), + serde_json::from_value(json!({})).unwrap(), + ) + .unwrap(); + assert_eq!( + serde_json::to_value(request).unwrap()["output_format"], + "markdown" + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs index 98cfc0db78d..7e8ce63b379 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs @@ -12,13 +12,17 @@ pub(crate) fn transform_ocr_request( params: &DeepSeekOcrParams, ) -> Result { if document.source().is_empty() { - return Err(OcrRequestError::MissingField("document URL")); + return Err(OcrRequestError::MissingDocumentUrl); } + let content = OcrDocument::ImageUrl { + image_url: document.source().to_string(), + extra_fields: serde_json::Map::new(), + }; Ok(DeepSeekOcrRequest { model: provider_model.to_string(), messages: vec![DeepSeekOcrMessage { role: UserRole::User, - content: vec![document], + content: vec![content], }], params: params.clone(), }) diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs index 85d1dafa542..9389f93b8e3 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs @@ -163,6 +163,30 @@ mod tests { ); } + #[rstest] + #[case(json!([0, 1, 2]), Some("1,2,3"))] + #[case(json!([2, 0, 0, 1]), Some("1,2,3"))] + #[case(json!([]), None)] + #[case(json!("3-9"), Some("3-9"))] + #[case(json!("1-3, 5"), Some("1-3,5"))] + #[case(json!(["1", "3-5"]), Some("1,3-5"))] + fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) { + assert_eq!( + map(json!({"pages": input})).unwrap().pages.as_deref(), + expected + ); + } + + #[rstest] + #[case(json!("a,b"))] + #[case(json!([-1]))] + #[case(json!([true, false]))] + #[case(json!([1, "2"]))] + #[case(json!(5))] + fn invalid_page_mapping_matches_python(#[case] input: Value) { + assert!(map(json!({"pages": input})).is_err()); + } + #[rstest] #[case(json!(["keyValuePairs"]), "keyValuePairs")] #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")] diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs index 2b848fcfb7a..f76a7c2b232 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs @@ -13,7 +13,7 @@ pub(crate) fn transform_ocr_request( ) -> Result { let source = document.source(); if source.is_empty() { - return Err(OcrRequestError::MissingField("document URL")); + return Err(OcrRequestError::MissingDocumentUrl); } Ok(if let Some(document) = InlineDocument::parse(source)? { DocumentIntelligenceRequest::Base64Source( @@ -46,10 +46,7 @@ pub(crate) fn transform_ocr_response( let mut extra_fields = Map::new(); extra_fields.insert("content".into(), option_value(result.content)); extra_fields.insert("tables".into(), option_value(result.tables)); - extra_fields.insert( - "key_value_pairs".into(), - option_value(result.key_value_pairs), - ); + extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs)); Ok(LiteLLMOcrResponse { pages, model: model.into(), diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs index 5bd7e555a1e..e60f1f5d3d6 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs @@ -114,6 +114,7 @@ mod tests { #[rstest] #[case("table_format", json!("html"))] #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("block"))] #[case("document_annotation_prompt", json!("extract"))] #[case("include_blocks", json!(true))] #[case("id", json!("req-123"))] @@ -133,6 +134,7 @@ mod tests { #[rstest] #[case("pages", json!([0, 2]))] + #[case("pages", json!("0,2-4"))] #[case("include_image_base64", json!(true))] #[case("image_limit", json!(2))] #[case("image_min_size", json!(100))] @@ -196,8 +198,16 @@ mod tests { #[rstest] fn transform_ocr_response_preserves_blocks_and_confidence_scores() { let response: MistralOcrResponse = serde_json::from_value(json!({ - "pages":[{"index":0,"markdown":"hello","blocks":[{"type":"title"}],"confidence_scores":{"mean":0.99}}], + "pages":[{ + "index":0, + "markdown":"hello", + "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], + "dimensions":{"width":612,"height":792,"dpi":72}, + "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], + "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} + }], "model":"returned-model", + "document_annotation":"{\"language\":\"en\"}", "usage_info":{"pages_processed":1} })) .unwrap(); @@ -205,7 +215,20 @@ mod tests { .unwrap() .into_json(); assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); - assert_eq!(result["pages"][0]["confidence_scores"]["mean"], 0.99); + assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1); + assert_eq!( + result["pages"][0]["blocks"][0]["confidence_scores"]["mean"], + 0.98 + ); + assert_eq!( + result["pages"][0]["confidence_scores"]["average_page_confidence_score"], + 0.99 + ); + assert_eq!(result["pages"][0]["images"][0]["id"], "img-0"); + assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72); + assert_eq!(result["model"], "returned-model"); + assert_eq!(result["document_annotation"], "{\"language\":\"en\"}"); + assert_eq!(result["usage_info"]["pages_processed"], 1); } #[rstest] diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs index 0e601cd8319..e0bc8a267d2 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs @@ -3,10 +3,17 @@ use serde_json::{Map, Value}; use crate::ocr::types::OcrDocument; +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum MistralOcrPages { + Range(String), + Indices(Vec), +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub(crate) struct MistralOcrParams { #[serde(skip_serializing_if = "Option::is_none")] - pub pages: Option>, + pub pages: Option, #[serde(skip_serializing_if = "Option::is_none")] pub include_image_base64: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/litellm-rust/crates/core/src/ocr/codecs/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mod.rs index 7c752749901..639b985b9ae 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mod.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod cohere; pub(crate) mod deepseek; pub(crate) mod document_intelligence; pub(crate) mod mistral; diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index e89b1c5c569..82a32ac1ab5 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -2,13 +2,90 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::mime::Mime; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; use reqwest::Url; +use serde_json::Map; use super::error::{OcrError, OcrRequestError, OcrResponseError}; use super::types::{OcrConnection, OcrDocument}; -use crate::constants::OCR_MAX_FETCH_REDIRECTS; +use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS}; use crate::error::{MediaError, TransportError}; use crate::media::{DownloadPolicy, MediaFetcher}; +pub fn encode_file_document( + bytes: &[u8], + file_name: Option<&str>, + mime_type: Option<&str>, +) -> Result { + if bytes.is_empty() { + return Err(OcrRequestError::EmptyFile); + } + if bytes.len() > OCR_INLINE_MAX_BYTES { + return Err(OcrRequestError::InlineDocumentTooLarge); + } + if let Some(value) = mime_type + && !valid_mime_type(value) + { + return Err(OcrRequestError::InvalidMimeType(value.into())); + } + let mime_type = mime_type + .map(str::to_string) + .or_else(|| file_name.map(|name| mime_type_for_name(name).to_string())) + .unwrap_or_else(|| "application/octet-stream".into()); + let source = format!("data:{mime_type};base64,{}", STANDARD.encode(bytes)); + Ok(if mime_type.starts_with("image/") { + OcrDocument::ImageUrl { + image_url: source, + extra_fields: Map::new(), + } + } else { + OcrDocument::DocumentUrl { + document_url: source, + extra_fields: Map::new(), + } + }) +} + +fn valid_mime_type(value: &str) -> bool { + let Some((kind, subtype)) = value.split_once('/') else { + return false; + }; + !kind.is_empty() + && !subtype.is_empty() + && kind.chars().chain(subtype.chars()).all(|character| { + character.is_alphanumeric() || matches!(character, '.' | '+' | '-' | '_') + }) +} + +pub fn mime_type_for_name(name: &str) -> &'static str { + let extension = std::path::Path::new(name) + .extension() + .and_then(|value| value.to_str()) + .unwrap_or_default(); + match extension.to_ascii_lowercase().as_str() { + "pdf" => "application/pdf", + "png" => "image/png", + "jpg" | "jpeg" => "image/jpeg", + "gif" => "image/gif", + "webp" => "image/webp", + "tiff" | "tif" => "image/tiff", + "bmp" => "image/bmp", + _ => mime_guess::from_path(name) + .first_raw() + .unwrap_or("application/octet-stream"), + } +} + +pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str { + match content_type + .and_then(|value| value.split(';').next()) + .map(str::trim) + { + Some(value) if !value.is_empty() && value != "application/octet-stream" => value, + _ => file_name + .map(mime_type_for_name) + .unwrap_or("application/octet-stream"), + } +} + pub(crate) struct InlineDocument<'a>(DataUrl<'a>); impl<'a> InlineDocument<'a> { @@ -95,9 +172,11 @@ fn map_media_error(error: MediaError) -> OcrError { body: "OCR document download failed".into(), } .into(), - MediaError::Timeout => { - TransportError::Network("OCR document download timed out".into()).into() + MediaError::Timeout => TransportError::Http { + status: 408, + body: "OCR document download timed out".into(), } + .into(), MediaError::Transport(error) => error.into(), } } @@ -114,6 +193,90 @@ mod tests { } } + #[test] + fn file_bytes_are_encoded_with_core_owned_mime_policy() { + assert_eq!( + encode_file_document(b"abc", Some("scan.png"), None).unwrap(), + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,YWJj".into(), + extra_fields: Map::new(), + } + ); + assert_eq!( + encode_file_document(b"abc", None, Some("application/pdf")).unwrap(), + document("data:application/pdf;base64,YWJj") + ); + } + + #[test] + fn file_name_mime_mapping_matches_python() { + for (name, expected) in [ + ("document.pdf", "application/pdf"), + ("image.png", "image/png"), + ("photo.jpg", "image/jpeg"), + ("photo.jpeg", "image/jpeg"), + ("animation.gif", "image/gif"), + ("image.webp", "image/webp"), + ("scan.tiff", "image/tiff"), + ("scan.tif", "image/tiff"), + ("bitmap.bmp", "image/bmp"), + ("DOCUMENT.PDF", "application/pdf"), + ("IMAGE.PNG", "image/png"), + ("file.unknown-extension", "application/octet-stream"), + ] { + assert_eq!(mime_type_for_name(name), expected); + } + } + + #[test] + fn upload_mime_mapping_matches_python() { + assert_eq!( + upload_mime_type(Some("report.pdf"), Some("application/octet-stream")), + "application/pdf" + ); + assert_eq!(upload_mime_type(Some("image.png"), None), "image/png"); + assert_eq!(upload_mime_type(None, None), "application/octet-stream"); + assert_eq!( + upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")), + "application/pdf" + ); + assert_eq!( + upload_mime_type( + Some("img.png"), + Some("image/png; charset=utf-8; boundary=something") + ), + "image/png" + ); + } + + #[test] + fn file_encoding_enforces_decoded_size_limit() { + let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1]; + assert_eq!( + encode_file_document(&bytes, None, None), + Err(OcrRequestError::InlineDocumentTooLarge) + ); + let document = encode_file_document(&bytes[..OCR_INLINE_MAX_BYTES], None, None).unwrap(); + let inline = InlineDocument::parse(document.source()).unwrap().unwrap(); + assert_eq!( + inline.decode(OCR_INLINE_MAX_BYTES).unwrap(), + bytes[..OCR_INLINE_MAX_BYTES] + ); + } + + #[test] + fn file_encoding_rejects_empty_bytes_and_invalid_explicit_mime() { + assert!(encode_file_document(b"", None, None).is_err()); + for mime in [ + "text/plain;bad", + "text/plain/extra", + " text/plain", + "text/plain\n", + ] { + assert!(encode_file_document(b"abc", None, Some(mime)).is_err()); + } + } + #[test] fn decodes_data_urls_and_limits_decoded_size() { for (source, expected) in [ diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 522d059ec48..55ea2cbcdae 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -4,15 +4,27 @@ use crate::error::TransportError; #[derive(Debug, Clone, PartialEq, Eq, Error)] pub enum OcrRequestError { + #[error("File is empty or could not be read")] + EmptyFile, + #[error("Invalid MIME type: {0}")] + InvalidMimeType(String), + #[error( + "Cohere Parse only accepts `image_url` documents; document_url and PDF inputs are not supported" + )] + CohereImageOnly, #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")] RequestFormat, #[error("invalid OCR request field: {path}")] RequestField { path: String }, #[error("missing required field: {0}")] MissingField(&'static str), + #[error("Document URL is required")] + MissingDocumentUrl, #[error("invalid OCR document data URI")] InvalidDataUri, - #[error("Reducto requires a reducto:// id or a data URI")] + #[error( + "Reducto requires a reducto:// id or a data URI; plain HTTP URLs are not supported, upload the file first" + )] ReductoSource, #[error("inline OCR document exceeds the size limit")] InlineDocumentTooLarge, @@ -34,6 +46,8 @@ pub enum OcrRequestError { #[derive(Debug, Clone, PartialEq, Eq, Error)] pub enum OcrResponseError { + #[error("OCR response exceeds the size limit of {limit} bytes")] + TooLarge { limit: usize }, #[error("invalid OCR response field: {path}")] ResponseField { path: String }, #[error("OCR response is missing non-empty content")] diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 0b04319d966..cd1d538aaa8 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,15 +1,17 @@ use super::OcrClient; use super::adapters::OcrAdapter; -use super::hooks::OcrLifecycleHooks; +use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest}; use super::registry::OcrAdapterKind; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; use crate::Error; use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext}; +use std::sync::Arc; pub(crate) async fn perform_ocr_request( client: &OcrClient, request: LiteLLMOcrRequest, ) -> Result { + request.response_format()?; let context = CallLifecycleContext::new( "ocr", request.model.clone(), @@ -23,27 +25,71 @@ pub(crate) async fn perform_ocr_request( hooks: request.hooks.clone(), provider_name: context.custom_llm_provider.clone(), }; - CallLifecycle::default().run(context, request, &hooks, |request| async move { - macro_rules! execute_selected_adapter { + CallLifecycle::default() + .run(context, request, &hooks, |request| async move { + PreparedOcrCall::prepare(client.clone(), request) + .await? + .execute() + .await? + .normalize() + }) + .await +} + +pub(crate) struct PreparedOcrCall { + client: OcrClient, + request: LiteLLMOcrRequest, + http: reqwest::Request, +} + +impl PreparedOcrCall { + pub(crate) async fn prepare( + client: OcrClient, + request: LiteLLMOcrRequest, + ) -> Result { + macro_rules! prepare_adapter { ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { match request.adapter { - $( OcrAdapterKind::$variant => execute_ocr_provider_call(client, &$instance, request).await, )+ + $( OcrAdapterKind::$variant => $instance.prepare_request(&request, &client).await?, )+ } }; } - super::adapters::for_each_ocr_adapter!(execute_selected_adapter) - }).await + let http = super::adapters::for_each_ocr_adapter!(prepare_adapter); + Ok(Self { + client, + request, + http, + }) + } + + pub(crate) async fn execute(self) -> Result { + let url = self.http.url().to_string(); + let headers = request_headers(&self.http)?; + let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts( + self.client.provider_http().clone(), + self.http, + )) + .await + .map_err(super::client::transport_error)?; + macro_rules! read_adapter { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + match self.request.adapter { + $( OcrAdapterKind::$variant => { + let decoded = $instance.read_response(&self.client, response, &url, &headers, &self.request).await?; + Ok(OcrProviderResponse { + request: self.request, + data: OcrProviderData::$variant(decoded), + }) + }, )+ + } + }; + } + super::adapters::for_each_ocr_adapter!(read_adapter) + } } -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -async fn execute_ocr_provider_call( - client: &OcrClient, - adapter: &A, - request: LiteLLMOcrRequest, -) -> Result { - let provider_request = adapter.prepare_request(&request, client).await?; - let url = provider_request.url().to_string(); - let headers = provider_request +fn request_headers(request: &reqwest::Request) -> Result, Error> { + request .headers() .iter() .map(|(name, value)| { @@ -53,20 +99,41 @@ async fn execute_ocr_provider_call( .map_err(|_| super::error::OcrRequestError::RequestField { path: "headers".into(), }) + .map_err(Error::from) }) - .collect::, _>>()?; - let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts( - client.provider_http().clone(), - provider_request, - )) - .await - .map_err(crate::error::TransportError::from)?; - let decoded = adapter - .read_response(client, response, &url, &headers, &request) - .await?; - let response = adapter.transform_ocr_response(&request, decoded.data)?; - Ok(LiteLLMOcrResponse { - provider_native_response: decoded.native, - ..response - }) + .collect() } + +macro_rules! provider_data { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + enum OcrProviderData { + $( $variant(super::wire::DecodedOcrResponse<<$adapter as OcrAdapter>::ProviderResponse>), )+ + } + + impl OcrProviderResponse { + pub(crate) fn normalize(self) -> Result { + match self.data { + $( OcrProviderData::$variant(decoded) => { + let response = $instance.transform_ocr_response(&self.request, decoded.data)?; + Ok(LiteLLMOcrResponse { provider_native_response: decoded.native, ..response }) + }, )+ + } + } + } + }; +} + +pub(crate) struct OcrProviderResponse { + request: LiteLLMOcrRequest, + data: OcrProviderData, +} + +pub(crate) async fn post_call(hooks: &Arc, bytes: &[u8]) -> Result<(), Error> { + let original_response = serde_json::Value::String(String::from_utf8_lossy(bytes).into_owned()); + hooks + .post_call(OcrPostCallRequest { original_response }) + .await?; + Ok(()) +} + +super::adapters::for_each_ocr_adapter!(provider_data); diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 7dd3c6bf8b2..3e7507e9ed5 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -24,11 +24,19 @@ pub struct OcrDuringCallRequest { pub model: String, pub custom_llm_provider: String, pub url: String, + pub headers: Vec<(String, String)>, pub body: Value, + #[serde(skip)] + pub retained_fields: Vec, +} + +#[derive(Clone, Debug, Serialize)] +pub struct OcrPostCallRequest { + pub original_response: Value, } pub trait OcrHooks: Send + Sync { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { false } fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { @@ -40,6 +48,9 @@ pub trait OcrHooks: Send + Sync { ) -> OcrHookFuture<'_, OcrDuringCallRequest> { Box::pin(async move { Ok(request) }) } + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { Ok(request) }) + } fn success<'a>( &'a self, _context: &'a CallLifecycleContext, @@ -80,7 +91,7 @@ impl CallLifecycleHooks Self::PreCallFuture<'a> { Box::pin(async move { - if !self.hooks.has_guardrails() { + if !self.hooks.intercepts_requests() { return Ok(request); } let changed = self diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs new file mode 100644 index 00000000000..92c9d4b717c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -0,0 +1,640 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use tokio::sync::{mpsc, oneshot}; + +use super::handler::perform_ocr_request; +use super::hooks::{ + OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, + OcrPreCallRequest, +}; +use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; +use crate::AuthError; +use crate::Error; +use crate::auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use crate::call_lifecycle::host::{ + HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase, +}; +use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; + +pub type NativeResult = Result, Error>; + +#[derive(Debug, PartialEq, Eq)] +pub enum NativeOutcome { + Completed(T), + Declined(OcrDecline), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrDecline { + ProviderWorkflow, + HostOperations, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OcrAdmission { + pub provider_workflow: bool, + pub host_operations: bool, + pub asynchronous: bool, +} + +impl OcrAdmission { + pub const fn all() -> Self { + Self { + provider_workflow: true, + host_operations: true, + asynchronous: false, + } + } +} + +#[derive(Clone, Debug)] +pub enum OcrHostOperation { + ProjectRequest, + Lifecycle(HostPhase), + ConstructResponse(Arc), + MapFailure(Error), + Success { + context: CallLifecycleContext, + response: Arc, + timing: CallLifecycleTiming, + }, + Failure { + context: CallLifecycleContext, + error: Error, + timing: CallLifecycleTiming, + }, + AcquireAzureAdToken, + PreCall(OcrPreCallRequest), + DuringCall(OcrDuringCallRequest), + PostCall(OcrPostCallRequest), +} + +impl OcrHostOperation { + pub const fn phase(&self) -> Option { + match self { + Self::Lifecycle(phase) => Some(*phase), + Self::Success { .. } => Some(HostPhase::Success), + Self::Failure { .. } => Some(HostPhase::Failure), + _ => None, + } + } +} + +pub enum OcrHostResult { + Request(Result<(Box, bool), Error>), + Lifecycle(Result<(), HostFailure>), + AzureAdToken(Result), + PreCall(Result), + DuringCall(Result), + PostCall(Result), +} + +pub type OcrCallStep = HostCallStep; + +pub struct OcrCall { + lifecycle: HostLifecycle, + execution: OcrExecution, + response: Option>, + error: Option, + pending: bool, + completed: bool, + projecting: bool, +} + +impl OcrCall { + pub fn admit(client: OcrClient, admission: OcrAdmission) -> NativeOutcome { + if !admission.provider_workflow { + return NativeOutcome::Declined(OcrDecline::ProviderWorkflow); + } + if !admission.host_operations { + return NativeOutcome::Declined(OcrDecline::HostOperations); + } + NativeOutcome::Completed(Self { + lifecycle: HostLifecycle::new(admission.asynchronous), + execution: OcrExecution::new(client), + response: None, + error: None, + pending: false, + completed: false, + projecting: false, + }) + } + + pub async fn resume(&mut self, result: Option) -> Result { + if self.completed { + return Err(Error::InvalidRequest( + "OCR call cannot be resumed after completion".into(), + )); + } + if self.pending != result.is_some() { + return Err(Error::InvalidRequest( + "OCR host operation result does not match pending state".into(), + )); + } + match &result { + Some(OcrHostResult::Lifecycle(Ok(()))) + if self.lifecycle.phase() == HostPhase::Execute => + { + return Err(Error::InvalidRequest( + "OCR provider operation requires a typed result".into(), + )); + } + Some(result) + if !matches!(result, OcrHostResult::Lifecycle(_)) + && self.lifecycle.phase() != HostPhase::Execute => + { + return Err(Error::InvalidRequest( + "unexpected OCR provider operation result".into(), + )); + } + _ => {} + } + self.pending = false; + let provider_result = match result { + Some(OcrHostResult::Request(result)) if self.projecting => { + self.projecting = false; + match result { + Ok((request, azure_ad_token_provider)) => { + self.execution.request = Some(*request); + self.execution.azure_ad_token_provider = azure_ad_token_provider; + } + Err(error) => self.accept(Err(HostFailure::Error(error))), + } + None + } + Some(OcrHostResult::Request(_)) => { + return Err(Error::InvalidRequest( + "unexpected OCR request projection".into(), + )); + } + Some(OcrHostResult::Lifecycle(result)) => { + self.accept(result); + None + } + result => result, + }; + if self.lifecycle.phase() == HostPhase::Execute { + if self.execution.request.is_none() + && self.execution.execution.is_none() + && !self.execution.completed + { + self.projecting = true; + return Ok(self.host_step(OcrHostOperation::ProjectRequest)); + } + match self.execution.resume(provider_result).await { + Ok(OcrCallStep::Host(operation)) => return Ok(self.host_step(operation)), + Ok(OcrCallStep::Complete(response)) => { + self.response = Some(Arc::new(response)); + self.accept(Ok(())); + } + Err(error) => self.accept(Err(HostFailure::Error(error))), + } + } + if self.error.is_some() { + self.execution.stop().await; + } + let operation = match self.lifecycle.phase() { + HostPhase::Complete => { + self.completed = true; + return match self.error.take() { + Some(error) => Err(error), + None => self + .response + .take() + .map(Arc::unwrap_or_clone) + .map(OcrCallStep::Complete) + .ok_or_else(|| { + Error::InvalidRequest("OCR completed without a response".into()) + }), + }; + } + HostPhase::ConstructResponse => OcrHostOperation::ConstructResponse( + self.response + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))? + .clone(), + ), + HostPhase::MapFailure => OcrHostOperation::MapFailure( + self.error + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))? + .clone(), + ), + HostPhase::Success | HostPhase::Failure => { + let snapshot = self + .execution + .terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone(); + match (self.lifecycle.phase(), snapshot) { + (HostPhase::Success, Some((context, timing))) => OcrHostOperation::Success { + context, + response: self + .response + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))? + .clone(), + timing, + }, + (HostPhase::Failure, Some((context, timing))) => OcrHostOperation::Failure { + context, + error: self + .error + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))? + .clone(), + timing, + }, + (phase, _) => OcrHostOperation::Lifecycle(phase), + } + } + phase => OcrHostOperation::Lifecycle(phase), + }; + Ok(self.host_step(operation)) + } + + fn accept(&mut self, result: Result<(), HostFailure>) { + let cancelled = matches!(&result, Err(HostFailure::Cancelled(_))); + if let Some(error) = self.lifecycle.accept(result) { + if cancelled { + self.error = Some(error); + } else { + self.error.get_or_insert(error); + } + self.execution.cancel(); + } + } + + pub async fn interrupt(&mut self, failure: HostFailure) -> Result { + if self.completed { + return Err(Error::InvalidRequest( + "OCR call cannot be interrupted after completion".into(), + )); + } + self.pending = false; + self.accept(Err(failure)); + self.resume(None).await + } + + fn host_step(&mut self, operation: OcrHostOperation) -> OcrCallStep { + self.pending = true; + OcrCallStep::Host(operation) + } +} + +impl HostCall for OcrCall { + type Operation = OcrHostOperation; + type Result = OcrHostResult; + type Complete = LiteLLMOcrResponse; + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(OcrCall::resume(self, result)) + } + + fn interrupt( + &mut self, + failure: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(OcrCall::interrupt(self, failure)) + } +} + +struct PendingOperation { + operation: OcrHostOperation, + result: oneshot::Sender, +} + +struct OcrExecution { + client: Option, + request: Option, + operations_tx: mpsc::UnboundedSender, + operations_rx: mpsc::UnboundedReceiver, + pending_result: Option>, + execution: Option>>, + completed: bool, + azure_ad_token_provider: bool, + terminal: Arc>>, +} + +impl OcrExecution { + fn new(client: OcrClient) -> Self { + let (operations_tx, operations_rx) = mpsc::unbounded_channel(); + Self { + client: Some(client), + request: None, + operations_tx, + operations_rx, + pending_result: None, + execution: None, + completed: false, + azure_ad_token_provider: false, + terminal: Arc::default(), + } + } + + pub async fn resume(&mut self, result: Option) -> Result { + if self.completed { + return Err(Error::InvalidRequest( + "OCR call cannot be resumed after completion".into(), + )); + } + match (self.pending_result.take(), result) { + (Some(sender), Some(result)) => sender + .send(result) + .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))?, + (None, None) if self.execution.is_none() => self.start(), + (Some(sender), None) => { + self.pending_result = Some(sender); + return Err(Error::InvalidRequest( + "OCR host operation result is required".into(), + )); + } + (None, Some(_)) => { + return Err(Error::InvalidRequest( + "unexpected OCR host operation result".into(), + )); + } + (None, None) => {} + } + + let execution = self.execution.as_mut().ok_or_else(|| { + Error::InvalidRequest("OCR call cannot be resumed after completion".into()) + })?; + tokio::select! { + operation = self.operations_rx.recv() => { + let operation = operation.ok_or_else(|| Error::InvalidRequest("OCR operation channel closed".into()))?; + self.pending_result = Some(operation.result); + Ok(OcrCallStep::Host(operation.operation)) + } + result = execution => { + self.execution = None; + self.completed = true; + result + .map_err(|error| Error::Network(format!("OCR execution task failed: {error}")))? + .map(OcrCallStep::Complete) + } + } + } + + fn start(&mut self) { + let client = self.client.take().expect("admitted OCR call has a client"); + let mut request = self + .request + .take() + .expect("admitted OCR call has a request"); + let intercepts_requests = request.hooks.intercepts_requests(); + if self.azure_ad_token_provider { + request.azure_ad_token_provider = Some(TokenProviderHandle::new(Arc::new( + OcrAzureAdTokenProvider { + operations: self.operations_tx.clone(), + }, + ))); + } + request.hooks = Arc::new(ProtocolHooks { + operations: self.operations_tx.clone(), + intercepts_requests, + terminal: self.terminal.clone(), + }); + self.execution = Some(tokio::spawn(async move { + perform_ocr_request(&client, request).await + })); + } + + fn cancel(&mut self) { + self.pending_result = None; + if let Some(execution) = &self.execution { + execution.abort(); + } + } + + async fn stop(&mut self) { + self.cancel(); + if let Some(execution) = self.execution.as_mut() { + let _ = execution.await; + } + self.execution = None; + } +} + +impl Drop for OcrExecution { + fn drop(&mut self) { + if let Some(execution) = &self.execution { + execution.abort(); + } + } +} + +struct ProtocolHooks { + operations: mpsc::UnboundedSender, + intercepts_requests: bool, + terminal: Arc>>, +} + +#[derive(Debug)] +struct OcrAzureAdTokenProvider { + operations: mpsc::UnboundedSender, +} + +impl TokenProvider for OcrAzureAdTokenProvider { + fn acquire(&self) -> TokenFuture<'_> { + Box::pin(async move { + let (result, receiver) = oneshot::channel(); + self.operations + .send(PendingOperation { + operation: OcrHostOperation::AcquireAzureAdToken, + result, + }) + .map_err(|_| { + AuthError::AzureTokenAcquisition("OCR host driver was abandoned".into()) + })?; + match receiver.await.map_err(|_| { + AuthError::AzureTokenAcquisition( + "OCR token provider operation was abandoned".into(), + ) + })? { + OcrHostResult::AzureAdToken(result) => result, + _ => Err(AuthError::AzureTokenAcquisition( + "invalid OCR token provider host result".into(), + )), + } + }) + } +} + +impl ProtocolHooks { + async fn invoke(&self, operation: OcrHostOperation) -> Result { + let (result, receiver) = oneshot::channel(); + self.operations + .send(PendingOperation { operation, result }) + .map_err(|_| Error::InvalidRequest("OCR host driver was abandoned".into()))?; + receiver + .await + .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into())) + } +} + +impl OcrHooks for ProtocolHooks { + fn intercepts_requests(&self) -> bool { + self.intercepts_requests + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { + match self.invoke(OcrHostOperation::PreCall(request)).await? { + OcrHostResult::PreCall(result) => result, + _ => Err(Error::InvalidRequest( + "invalid OCR pre-call host result".into(), + )), + } + }) + } + + fn during_call( + &self, + request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + match self.invoke(OcrHostOperation::DuringCall(request)).await? { + OcrHostResult::DuringCall(result) => result, + _ => Err(Error::InvalidRequest( + "invalid OCR during-call host result".into(), + )), + } + }) + } + + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + match self.invoke(OcrHostOperation::PostCall(request)).await? { + OcrHostResult::PostCall(result) => result, + _ => Err(Error::InvalidRequest( + "invalid OCR post-call host result".into(), + )), + } + }) + } + + fn success<'a>( + &'a self, + context: &'a CallLifecycleContext, + _response: &'a LiteLLMOcrResponse, + timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + *self + .terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some((context.clone(), timing.clone())); + }) + } + + fn failure<'a>( + &'a self, + context: &'a CallLifecycleContext, + _error: &'a Error, + timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + *self + .terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some((context.clone(), timing.clone())); + }) + } +} + +pub type OcrHostFuture<'a> = Pin + Send + 'a>>; + +pub trait OcrHost: Send + Sync { + fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_>; +} + +pub struct NoopOcrHost; + +impl OcrHost for NoopOcrHost { + fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> { + Box::pin(async move { + match operation { + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( + Error::InvalidRequest("OCR host has no request projection".into()), + )), + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => OcrHostResult::Lifecycle(Ok(())), + OcrHostOperation::AcquireAzureAdToken => { + OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition( + "OCR host has no Azure AD token provider".into(), + ))) + } + OcrHostOperation::PreCall(request) => OcrHostResult::PreCall(Ok(request)), + OcrHostOperation::DuringCall(request) => OcrHostResult::DuringCall(Ok(request)), + OcrHostOperation::PostCall(request) => OcrHostResult::PostCall(Ok(request)), + } + }) + } +} + +pub struct OcrHookHost { + hooks: Arc, +} + +impl OcrHookHost { + pub fn new(hooks: Arc) -> Self { + Self { hooks } + } +} + +impl OcrHost for OcrHookHost { + fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> { + Box::pin(async move { + match operation { + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( + Error::InvalidRequest("OCR hook host has no request projection".into()), + )), + OcrHostOperation::Success { + context, + response, + timing, + } => { + self.hooks.success(&context, &response, &timing).await; + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Failure { + context, + error, + timing, + } => { + self.hooks.failure(&context, &error, &timing).await; + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) => OcrHostResult::Lifecycle(Ok(())), + OcrHostOperation::AcquireAzureAdToken => { + OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition( + "OCR hook host has no Azure AD token provider".into(), + ))) + } + OcrHostOperation::PreCall(request) => { + OcrHostResult::PreCall(self.hooks.pre_call(request).await) + } + OcrHostOperation::DuringCall(request) => { + OcrHostResult::DuringCall(self.hooks.during_call(request).await) + } + OcrHostOperation::PostCall(request) => { + OcrHostResult::PostCall(self.hooks.post_call(request).await) + } + } + }) + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 1e975c3f521..e29fd6ac572 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -5,12 +5,18 @@ mod document; pub mod error; mod handler; pub mod hooks; +mod lifecycle; mod prepare; mod registry; pub mod types; pub mod wire; pub use client::{OcrClient, ocr}; +pub use document::{encode_file_document, mime_type_for_name, upload_mime_type}; +pub use lifecycle::{ + NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, + OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, +}; pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; #[cfg(test)] diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index bf6f924088c..9934a1d9a14 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -62,34 +62,48 @@ pub(crate) async fn transform_request_body( request: &LiteLLMOcrRequest, url: &str, headers: &[(String, String)], + retains_document: bool, body: B, validate: impl FnOnce(&B) -> Result<(), OcrRequestError>, ) -> Result where B: Serialize + DeserializeOwned, { - let body = if request.hooks.has_guardrails() { + let (body, headers) = if request.hooks.intercepts_requests() { + let body = serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { + path: "body".into(), + })?; + let retained_fields = request + .optional_params + .keys() + .filter(|name| body.get(*name).is_some()) + .cloned() + .chain(retains_document.then(|| "document".to_string())) + .collect(); let changed = request .hooks .during_call(OcrDuringCallRequest { model: request.model.clone(), custom_llm_provider: request.adapter.provider().as_str().into(), url: url.into(), - body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { - path: "body".into(), - })?, + headers: headers.to_vec(), + body, + retained_fields, }) .await?; let body = OcrWireBody::::decode(changed.body)?; validate(&body.body)?; - body + (body, changed.headers) } else { - OcrWireBody { - body, - extra: Map::new(), - } + ( + OcrWireBody { + body, + extra: Map::new(), + }, + headers.to_vec(), + ) }; - build_http_request(client, request, url, headers, &body) + build_http_request(client, request, url, &headers, &body) } pub(crate) fn build_http_request( @@ -113,9 +127,10 @@ pub(crate) fn build_http_request( pub(crate) async fn guardrail_document( request: &LiteLLMOcrRequest, url: &str, -) -> Result { - if !request.hooks.has_guardrails() { - return Ok(request.document.clone()); + headers: &[(String, String)], +) -> Result<(OcrDocument, Vec<(String, String)>), OcrError> { + if !request.hooks.intercepts_requests() { + return Ok((request.document.clone(), headers.to_vec())); } let changed = request .hooks @@ -123,14 +138,17 @@ pub(crate) async fn guardrail_document( model: request.model.clone(), custom_llm_provider: request.adapter.provider().as_str().into(), url: url.into(), + headers: headers.to_vec(), body: serde_json::to_value(&request.document).map_err(|_| { OcrRequestError::RequestField { path: "document".into(), } })?, + retained_fields: Vec::new(), }) .await?; - super::wire::decode_request_value(changed.body, "guardrail.document").map_err(OcrError::from) + let document = super::wire::decode_request_value(changed.body, "guardrail.document")?; + Ok((document, changed.headers)) } #[derive(Serialize)] diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs index 1b20a91143b..ed7d4fd5cf2 100644 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -23,6 +23,7 @@ super::adapters::for_each_ocr_adapter!(define_adapter_types); #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum OcrProvider { + Cohere, Mistral, AzureAi, Reducto, @@ -32,6 +33,7 @@ pub(crate) enum OcrProvider { impl OcrProvider { pub(crate) const fn as_str(self) -> &'static str { match self { + Self::Cohere => "cohere", Self::Mistral => "mistral", Self::AzureAi => "azure_ai", Self::Reducto => "reducto", @@ -50,6 +52,7 @@ pub(crate) fn resolve_wire_adapter( custom_llm_provider: OcrProvider::Mistral.as_str(), }); let typed_provider = match provider.custom_llm_provider { + "cohere" => OcrProvider::Cohere, "mistral" => OcrProvider::Mistral, "azure_ai" => OcrProvider::AzureAi, "reducto" => OcrProvider::Reducto, @@ -57,10 +60,17 @@ pub(crate) fn resolve_wire_adapter( value => return Err(Error::InvalidProvider(value.to_string())), }; let adapter = match typed_provider { + OcrProvider::Cohere => OcrAdapterKind::Cohere, OcrProvider::Mistral => OcrAdapterKind::Mistral, OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { OcrAdapterKind::AzureDocumentIntelligence } + OcrProvider::AzureAi + if provider.model.to_ascii_lowercase().contains("cohere") + && provider.model.to_ascii_lowercase().contains("parse") => + { + OcrAdapterKind::AzureCohere + } OcrProvider::AzureAi => OcrAdapterKind::AzureMistral, OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => { OcrAdapterKind::ReductoLegacy @@ -68,12 +78,7 @@ pub(crate) fn resolve_wire_adapter( OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-v3") => { OcrAdapterKind::ReductoV3 } - OcrProvider::Reducto => { - return Err(Error::InvalidRequest(format!( - "unsupported Reducto OCR model: {}", - provider.model - ))); - } + OcrProvider::Reducto => OcrAdapterKind::ReductoV3, OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { OcrAdapterKind::VertexDeepSeek } @@ -107,11 +112,10 @@ mod tests { } #[test] - fn unknown_reducto_models_are_rejected() { - assert!(matches!( - resolve_wire_adapter("reducto/future-parse-model", None), - Err(Error::InvalidRequest(_)) - )); + fn unknown_reducto_models_use_the_current_protocol() { + let (model, adapter) = resolve_wire_adapter("reducto/future-parse-model", None).unwrap(); + assert_eq!(model, "future-parse-model"); + assert_eq!(adapter, OcrAdapterKind::ReductoV3); } #[test] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 06519f86c91..76df8b42806 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -8,7 +8,7 @@ use serde_json::{Map, Value}; use super::hooks::{NoopOcrHooks, OcrHooks}; use super::registry::{OcrAdapterKind, resolve_wire_adapter}; use crate::Error; -use crate::auth::InputSource; +use crate::auth::{InputSource, TokenProviderHandle}; use crate::constants::OCR_HTTP_TIMEOUT_SECS; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -68,6 +68,7 @@ pub struct OcrConnection { pub extra_headers_source: InputSource, pub timeout: Duration, pub max_download_bytes: u64, + pub max_response_bytes: usize, pub poll_timeout: Duration, } @@ -82,6 +83,7 @@ impl Default for OcrConnection { extra_headers_source: InputSource::Deployment, timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, + max_response_bytes: crate::constants::OCR_RESPONSE_MAX_BYTES, poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS), } } @@ -95,6 +97,7 @@ pub struct LiteLLMOcrRequest { pub litellm_call_id: Option, pub optional_params: Map, pub input_sources: BTreeMap, + pub azure_ad_token_provider: Option, pub(crate) adapter: OcrAdapterKind, } @@ -115,6 +118,7 @@ impl LiteLLMOcrRequest { litellm_call_id: None, optional_params, input_sources: BTreeMap::new(), + azure_ad_token_provider: None, adapter: adapter_kind, }) } @@ -132,6 +136,10 @@ impl LiteLLMOcrRequest { .map(|format| format.unwrap_or_default()) } + pub fn provider_name(&self) -> &'static str { + self.adapter.provider().as_str() + } + pub fn with_host_hooks( self, hooks: Arc, @@ -169,6 +177,47 @@ mod tests { use super::*; use serde_json::json; + #[test] + fn document_variants_preserve_provider_fields_when_rewriting_sources() { + for (value, original, replacement, expected) in [ + ( + json!({ + "type":"document_url", + "document_url":"https://example.com/input.pdf", + "document_name":"input.pdf" + }), + "https://example.com/input.pdf", + "data:application/pdf;base64,AA==", + json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,AA==", + "document_name":"input.pdf" + }), + ), + ( + json!({ + "type":"image_url", + "image_url":"https://example.com/input.png", + "detail":"high" + }), + "https://example.com/input.png", + "data:image/png;base64,AA==", + json!({ + "type":"image_url", + "image_url":"data:image/png;base64,AA==", + "detail":"high" + }), + ), + ] { + let document: OcrDocument = serde_json::from_value(value).unwrap(); + assert_eq!(document.source(), original); + assert_eq!( + serde_json::to_value(document.with_source(replacement.into())).unwrap(), + expected + ); + } + } + #[test] fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() { let response = LiteLLMOcrResponse { diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 34d0a7d7b86..6dc6b34b73d 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -3,7 +3,6 @@ use crate::ocr::error::OcrResponseError; use std::collections::BTreeMap; use std::time::Duration; -use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest}; use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument}; use crate::Error; use crate::auth::InputSource; @@ -13,10 +12,58 @@ use serde::{ }; use serde_json::{Map, Value}; +const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"]; +const MISTRAL_OPTION_FIELDS: &[&str] = &[ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "include_blocks", + "id", +]; +const DEEPSEEK_OPTION_FIELDS: &[&str] = + &["stream", "temperature", "max_tokens", "top_p", "n", "stop"]; +const DOCUMENT_INTELLIGENCE_OPTION_FIELDS: &[&str] = &["pages", "features"]; +const REDUCTO_V3_OPTION_FIELDS: &[&str] = &["formatting", "retrieval", "settings"]; +const REDUCTO_LEGACY_OPTION_FIELDS: &[&str] = &["enhance"]; +const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_scope", + "azure_authority_host", + "azure_credential", + "azure_federated_token_file", + "enable_azure_ad_token_refresh", +]; +const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[ + "vertex_credentials", + "vertex_ai_credentials", + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", +]; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OptionalParamSpec { + pub name: &'static str, + pub secret: bool, +} + #[derive(Debug)] pub struct DecodedOcrResponse { pub data: T, pub native: Option, + pub text: String, } #[derive(Deserialize)] @@ -39,11 +86,65 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok() } +pub fn consumed_optional_param_names( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, Error> { + use super::registry::OcrAdapterKind; + + let (_, adapter) = super::registry::resolve_wire_adapter(model, custom_llm_provider)?; + let provider_fields: &[&str] = match adapter { + OcrAdapterKind::Cohere | OcrAdapterKind::AzureCohere => &["output_format"], + OcrAdapterKind::Mistral | OcrAdapterKind::AzureMistral | OcrAdapterKind::VertexMistral => { + MISTRAL_OPTION_FIELDS + } + OcrAdapterKind::AzureDocumentIntelligence => DOCUMENT_INTELLIGENCE_OPTION_FIELDS, + OcrAdapterKind::ReductoV3 => REDUCTO_V3_OPTION_FIELDS, + OcrAdapterKind::ReductoLegacy => REDUCTO_LEGACY_OPTION_FIELDS, + OcrAdapterKind::VertexDeepSeek => DEEPSEEK_OPTION_FIELDS, + }; + let auth_fields: &[&str] = match adapter { + OcrAdapterKind::AzureMistral + | OcrAdapterKind::AzureDocumentIntelligence + | OcrAdapterKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS, + OcrAdapterKind::VertexMistral | OcrAdapterKind::VertexDeepSeek => VERTEX_AUTH_OPTION_FIELDS, + _ => &[], + }; + Ok(COMMON_OPTION_FIELDS + .iter() + .chain(provider_fields) + .chain(auth_fields) + .copied() + .collect()) +} + +pub fn consumed_optional_params( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, Error> { + consumed_optional_param_names(model, custom_llm_provider).map(|names| { + names + .into_iter() + .map(|name| OptionalParamSpec { + name, + secret: matches!( + name, + "azure_ad_token" + | "client_secret" + | "azure_federated_token_file" + | "vertex_credentials" + | "vertex_ai_credentials" + ), + }) + .collect() + }) +} + pub fn decode_request(wire: OcrWireRequest) -> Result { let api_key_source = source_for(&wire.input_sources, "api_key"); let api_base_source = source_for(&wire.input_sources, "api_base"); let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); - let document = decode_request_value(wire.document, "document")?; + let document = decode_document(wire.document)?; let headers = wire .extra_headers .unwrap_or_default() @@ -66,11 +167,28 @@ pub fn decode_request(wire: OcrWireRequest) -> Result }) .transpose()?; let defaults = OcrConnection::default(); + let max_response_bytes = wire + .optional_params + .get("max_response_bytes") + .map(|value| { + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0 && *value <= defaults.max_response_bytes) + .ok_or_else(|| OcrRequestError::RequestField { + path: "max_response_bytes".into(), + }) + }) + .transpose()? + .unwrap_or(defaults.max_response_bytes); let request = LiteLLMOcrRequest::new( wire.model, document, wire.custom_llm_provider.as_deref(), - wire.optional_params, + wire.optional_params + .into_iter() + .filter(|(name, _)| name != "max_response_bytes") + .collect(), )?; let connection = OcrConnection { api_key: nonblank(wire.api_key), @@ -81,6 +199,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result extra_headers_source, timeout: timeout.unwrap_or(defaults.timeout), max_download_bytes: defaults.max_download_bytes, + max_response_bytes, poll_timeout: defaults.poll_timeout, }; Ok(LiteLLMOcrRequest { @@ -90,6 +209,16 @@ pub fn decode_request(wire: OcrWireRequest) -> Result }) } +fn decode_document(value: Value) -> Result { + let kind = value.get("type").and_then(Value::as_str); + let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none() + || matches!(kind, Some("image_url")) && value.get("image_url").is_none(); + if missing_url { + return Err(OcrRequestError::MissingDocumentUrl); + } + decode_request_value(value, "document") +} + fn source_for(sources: &BTreeMap, name: &str) -> InputSource { sources.get(name).copied().unwrap_or_default() } @@ -134,38 +263,81 @@ pub fn decode_response( } else { None }; - Ok(DecodedOcrResponse { data, native }) -} - -pub fn decode_pre_call_result( - original: OcrPreCallRequest, - value: Value, -) -> Result { - #[derive(Deserialize)] - struct Changed { - document: OcrDocument, - #[serde(default)] - optional_params: Map, - } - let changed: Changed = decode_request_value(value, "guardrail")?; - Ok(OcrPreCallRequest { - document: changed.document, - optional_params: Value::Object(changed.optional_params), - ..original + Ok(DecodedOcrResponse { + data, + native, + text: String::from_utf8_lossy(bytes).into_owned(), }) } -pub fn decode_during_call_result( - original: OcrDuringCallRequest, - value: Value, -) -> Result { - #[derive(Deserialize)] - struct Changed { - body: Value, +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn option_projection_is_provider_specific_and_excludes_opaque_fields() { + let mistral = consumed_optional_param_names("mistral/model", None).unwrap(); + assert!(mistral.contains(&"pages")); + assert!(mistral.contains(&"req_format")); + assert!(!mistral.contains(&"vertex_project")); + assert!(!mistral.contains(&"opaque_extension")); + + let vertex = consumed_optional_param_names("vertex_ai/deepseek-ocr", None).unwrap(); + assert!(vertex.contains(&"temperature")); + assert!(vertex.contains(&"vertex_credentials")); + assert!(!vertex.contains(&"pages")); + } + + #[test] + fn optional_param_metadata_marks_only_credentials_as_secret() { + let azure = consumed_optional_params("model", Some("azure_ai")).unwrap(); + assert!( + azure + .iter() + .any(|spec| spec.name == "client_secret" && spec.secret) + ); + assert!( + azure + .iter() + .any(|spec| spec.name == "tenant_id" && !spec.secret) + ); + let vertex = consumed_optional_params("deepseek-ocr", Some("vertex_ai")).unwrap(); + assert!( + vertex + .iter() + .any(|spec| spec.name == "vertex_credentials" && spec.secret) + ); + assert!( + vertex + .iter() + .any(|spec| spec.name == "vertex_project" && !spec.secret) + ); + } + + #[test] + fn activation_includes_migrated_providers() { + assert!(is_supported_request("model", Some("mistral"))); + assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); + assert!(is_supported_request( + "documentintelligence/prebuilt-read", + Some("azure_ai") + )); + assert!(is_supported_request("parse-v3", Some("reducto"))); + assert!(is_supported_request("parse-legacy", Some("reducto"))); + assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); + } + + #[test] + fn missing_document_source_has_a_typed_public_error() { + for document in [ + serde_json::json!({"type": "document_url"}), + serde_json::json!({"type": "image_url"}), + ] { + assert_eq!( + decode_document(document), + Err(OcrRequestError::MissingDocumentUrl) + ); + } } - let changed: Changed = decode_request_value(value, "guardrail")?; - Ok(OcrDuringCallRequest { - body: changed.body, - ..original - }) } diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 5d037e9cf1b..34213e5f6c4 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -1,3 +1,21 @@ +use std::collections::HashMap; +use std::io; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use futures_util::{SinkExt, StreamExt}; +use rustls::{ClientConfig, RootCertStore}; +use tokio::net::TcpStream; +use tokio::sync::Mutex; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::TlsError; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue}; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, +}; + use crate::Error; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; @@ -125,6 +143,137 @@ pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool { ) } +pub type ResponsesUpstreamWs = WebSocketStream>; + +static TLS_CONFIG: OnceLock> = OnceLock::new(); + +fn build_tls_config() -> Result> { + let native = rustls_native_certs::load_native_certs(); + let mut store = RootCertStore::empty(); + let (added, _ignored) = store.add_parsable_certificates(native.certs); + if added == 0 { + return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io( + io::Error::other(format!( + "no usable native root certificates: {:?}", + native.errors + )), + ))); + } + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map(|builder| builder.with_root_certificates(store).with_no_client_auth()) + .map_err(|error| { + Box::new(tokio_tungstenite::tungstenite::Error::Tls( + TlsError::Rustls(error), + )) + }) +} + +fn tls_config() -> Result, Box> { + if let Some(config) = TLS_CONFIG.get() { + return Ok(Arc::clone(config)); + } + let built = Arc::new(build_tls_config()?); + Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) +} + +pub async fn connect_upstream( + request: R, +) -> Result<(ResponsesUpstreamWs, Response), Box> +where + R: IntoClientRequest + Unpin, +{ + let request = request.into_client_request().map_err(Box::new)?; + let connector = match request.uri().scheme_str() { + Some("wss") => Some(Connector::Rustls(tls_config()?)), + _ => None, + }; + connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(Box::new) +} + +#[derive(Clone)] +pub struct ResponsesWebSocketConnection { + socket: Arc>>, +} + +impl ResponsesWebSocketConnection { + pub async fn connect_url( + url: &str, + headers: &HashMap, + timeout: Option, + ) -> Result { + let mut request = url + .into_client_request() + .map_err(|error| Error::Network(error.to_string()))?; + for (name, value) in headers { + let header_name = name + .parse::() + .map_err(|error| Error::InvalidRequest(error.to_string()))?; + let header_value = HeaderValue::from_str(value) + .map_err(|error| Error::InvalidRequest(error.to_string()))?; + request.headers_mut().insert(header_name, header_value); + } + let connect = connect_upstream(request); + let result = match timeout { + Some(timeout) => tokio::time::timeout(timeout, connect) + .await + .map_err(|_| Error::Network("Responses WebSocket connection timed out".into()))?, + None => connect.await, + }; + let (socket, _) = result.map_err(|error| match *error { + tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { + status: response.status().as_u16(), + body: String::new(), + }, + other => Error::Network(other.to_string()), + })?; + Ok(Self { + socket: Arc::new(Mutex::new(Some(socket))), + }) + } + + pub async fn send_text(&self, text: String) -> Result<(), Error> { + let mut socket = self.socket.lock().await; + let Some(socket) = socket.as_mut() else { + return Err(Error::Network("Responses WebSocket is closed".into())); + }; + socket + .send(Message::Text(text)) + .await + .map_err(|error| Error::Network(error.to_string())) + } + + pub async fn recv_text(&self) -> Result, Error> { + let mut socket = self.socket.lock().await; + let Some(socket) = socket.as_mut() else { + return Ok(None); + }; + match socket.next().await { + Some(Ok(Message::Text(text))) => Ok(Some(text)), + Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) + .map(Some) + .map_err(|error| Error::InvalidResponse(error.to_string())), + Some(Ok(Message::Close(_))) | None => Ok(None), + Some(Ok(_)) => Ok(None), + Some(Err(error)) => Err(Error::Network(error.to_string())), + } + } + + pub async fn close(&self) -> Result<(), Error> { + let mut socket = self.socket.lock().await; + if let Some(socket) = socket.as_mut() { + socket + .close(None) + .await + .map_err(|error| Error::Network(error.to_string()))?; + } + *socket = None; + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs index d7d532cfef1..b6dc8d90b93 100644 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -70,7 +70,7 @@ async fn facade_acquires_supplied_entra_token_for_final_request() { struct ReplaceBodyDocument; impl OcrHooks for ReplaceBodyDocument { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index e4c81dea5a7..3fca59033cc 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,5 @@ use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use super::wire::{OcrWireRequest, decode_request}; @@ -124,6 +125,14 @@ async fn immediate_response_normalizes_pages_and_preserves_native() { json!({"width":816,"height":1056,"dpi":96}) ); assert_eq!(result.usage_info, Some(json!({"pages_processed":1}))); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); assert_eq!(result.provider_native_response, Some(operation)); } @@ -169,6 +178,55 @@ async fn accepted_response_polls_to_success_with_only_credentials() { } } +struct SubmissionBoundary { + request_count: Arc>>, +} + +impl super::hooks::OcrHooks for SubmissionBoundary { + fn post_call( + &self, + request: super::hooks::OcrPostCallRequest, + ) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> { + Box::pin(async move { + match self.request_count.lock().unwrap().len() { + 1 => assert_eq!(request.original_response, json!(r#"{"submitted":true}"#)), + 2 => assert!( + request + .original_response + .as_str() + .unwrap() + .contains("succeeded") + ), + count => panic!("unexpected callback after {count} requests"), + } + Ok(request) + }) + } +} + +#[tokio::test] +async fn accepted_response_runs_post_call_before_polling() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(SubmissionBoundary { + request_count: seen.clone(), + }), + ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); +} + #[tokio::test] async fn polling_forwards_bearer_credentials() { let (base, seen, server) = mock_server(vec![ @@ -361,7 +419,7 @@ async fn pre_call_guardrail_receives_caller_pages_before_mapping() { struct RewritePages; impl OcrHooks for RewritePages { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs index 875fc9e3dc6..4ba39561dcd 100644 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/deepseek_ocr.rs @@ -34,6 +34,28 @@ fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { assert!(result.get("ignored").is_none()); } +#[rstest] +#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] +#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] +fn request_maps_both_document_types_to_image_content(#[case] document: Value) { + let source = document + .get("image_url") + .or_else(|| document.get("document_url")) + .unwrap() + .clone(); + let request = transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + serde_json::from_value(document).unwrap(), + &DeepSeekOcrParams::default(), + ) + .unwrap(); + let result = serde_json::to_value(request).unwrap(); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":source}) + ); +} + #[rstest] #[case(json!("# hello"), "# hello")] #[case(json!("{broken"), "{broken")] diff --git a/litellm-rust/crates/core/tests/host_lifecycle.rs b/litellm-rust/crates/core/tests/host_lifecycle.rs new file mode 100644 index 00000000000..19fb946afde --- /dev/null +++ b/litellm-rust/crates/core/tests/host_lifecycle.rs @@ -0,0 +1,116 @@ +use crate::Error; +use crate::call_lifecycle::host::{HostFailure, HostLifecycle, HostPhase}; + +fn run(fail_at: Option, asynchronous: bool) -> (Vec, Vec) { + let mut lifecycle = HostLifecycle::new(asynchronous); + let mut events = Vec::new(); + let mut failures = Vec::new(); + while lifecycle.phase() != HostPhase::Complete { + let phase = lifecycle.phase(); + events.push(phase); + let result = if Some(phase) == fail_at { + Err(HostFailure::Error(Error::InvalidRequest( + "selected failure".into(), + ))) + } else { + Ok(()) + }; + if let Some(error) = lifecycle.accept(result) { + failures.push(error); + } + } + (events, failures) +} + +#[test] +fn public_outcome_is_finalized_before_a_single_terminal_dispatch() { + for asynchronous in [false, true] { + let (events, failures) = run(None, asynchronous); + assert!(failures.is_empty()); + assert_eq!( + &events[events.len() - 2..], + &[HostPhase::Finalize, HostPhase::Success] + ); + assert_eq!( + events + .iter() + .filter(|phase| **phase == HostPhase::Execute) + .count(), + 1 + ); + assert_eq!( + events.contains(&HostPhase::DeploymentPostCall), + asynchronous + ); + } +} + +#[test] +fn only_provider_and_response_construction_failures_use_provider_mapping() { + for phase in [ + HostPhase::Setup, + HostPhase::DeploymentPreCall, + HostPhase::Prepare, + HostPhase::Execute, + HostPhase::ConstructResponse, + HostPhase::DeploymentPostCall, + HostPhase::Finalize, + ] { + let (events, failures) = run(Some(phase), true); + assert_eq!(failures.len(), 1); + assert!(!events.contains(&HostPhase::Success)); + let mapped = matches!(phase, HostPhase::Execute | HostPhase::ConstructResponse); + assert_eq!(events.contains(&HostPhase::MapFailure), mapped); + assert_eq!(events.contains(&HostPhase::DeploymentFailure), mapped); + assert_eq!( + &events[events.len() - 2..], + &[HostPhase::Failure, HostPhase::AsyncFailure] + ); + assert!( + events + .iter() + .filter(|phase| **phase == HostPhase::Execute) + .count() + <= 1 + ); + } +} + +#[test] +fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() { + let mut lifecycle = HostLifecycle::new(true); + while lifecycle.phase() != HostPhase::Execute { + lifecycle.accept(Ok(())); + } + let selected = Error::InvalidRequest("provider".into()); + assert_eq!( + lifecycle.accept(Err(HostFailure::Error(selected.clone()))), + Some(selected) + ); + lifecycle.accept(Ok(())); + for phase in [ + HostPhase::DeploymentFailure, + HostPhase::Failure, + HostPhase::AsyncFailure, + ] { + assert_eq!(lifecycle.phase(), phase); + assert_eq!( + lifecycle.accept(Err(HostFailure::Error(Error::InvalidRequest( + "callback".into() + )))), + None + ); + } + assert_eq!(lifecycle.phase(), HostPhase::Complete); +} + +#[test] +fn cancellation_skips_terminal_dispatch() { + let mut lifecycle = HostLifecycle::new(true); + let error = Error::InvalidRequest("cancelled".into()); + assert_eq!( + lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))), + Some(error) + ); + assert_eq!(lifecycle.phase(), HostPhase::Complete); +} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index cecd8869741..55f8713d76e 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -3,9 +3,16 @@ use std::sync::{Arc, Mutex}; use serde_json::{Value, json}; use super::OcrClient; -use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest}; +use super::hooks::{ + OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, + OcrPreCallRequest, +}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use super::wire::{OcrWireRequest, decode_request}; +use super::{ + NativeOutcome, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHost, + OcrHostOperation, OcrHostResult, +}; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; #[test] @@ -51,7 +58,7 @@ async fn facade_executes_direct_mistral_once() { let result = perform_ocr(wire_request( "mistral/model", &base, - json!({"extract_header":true,"unknown":"ignored"}), + json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), )) .await .unwrap(); @@ -72,6 +79,7 @@ async fn facade_executes_direct_mistral_once() { json!({ "model":"model", "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", "extract_header":true }) ); @@ -124,7 +132,7 @@ struct RecordingHooks { } impl OcrHooks for RecordingHooks { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } @@ -148,6 +156,13 @@ impl OcrHooks for RecordingHooks { }) } + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("post"); + Ok(request) + }) + } + fn success<'a>( &'a self, _context: &'a CallLifecycleContext, @@ -171,6 +186,38 @@ impl OcrHooks for RecordingHooks { } } +struct HeaderEditHooks; + +impl OcrHooks for HeaderEditHooks { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + request + .headers + .push(("x-core-callback".into(), "edited".into())); + Box::pin(async move { Ok(request) }) + } +} + +#[tokio::test] +async fn lifecycle_sends_headers_returned_by_the_typed_during_call_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(HeaderEditHooks), + ..wire_request("mistral/model", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); +} + #[tokio::test] async fn lifecycle_orders_hooks_and_emits_one_success() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; @@ -185,7 +232,10 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { }; perform_ocr(request).await.unwrap(); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]); + assert_eq!( + *events.lock().unwrap(), + ["pre", "during", "post", "success"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -227,3 +277,562 @@ async fn upstream_failure_emits_one_terminal_failure() { assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]); assert_eq!(seen.lock().unwrap().len(), 1); } + +struct AdmissionSpy { + effects: Arc>, +} + +impl OcrHooks for AdmissionSpy { + fn intercepts_requests(&self) -> bool { + *self.effects.lock().unwrap() += 1; + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + *self.effects.lock().unwrap() += 1; + Box::pin(async move { Ok(request) }) + } +} + +#[test] +fn admission_declines_without_invoking_hooks_or_transport() { + for (admission, expected) in [ + ( + OcrAdmission { + provider_workflow: false, + host_operations: true, + asynchronous: false, + }, + OcrDecline::ProviderWorkflow, + ), + ( + OcrAdmission { + provider_workflow: true, + host_operations: false, + asynchronous: false, + }, + OcrDecline::HostOperations, + ), + ] { + let outcome = OcrCall::admit(super::test_support::ocr_client(), admission); + assert!(matches!(outcome, NativeOutcome::Declined(reason) if reason == expected)); + } +} + +#[tokio::test] +async fn fallible_host_phases_do_not_replay_or_reach_transport() { + for failure_phase in ["pre", "during"] { + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut result = None; + let mut phases = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(operation)) => match operation { + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => { + result = Some(OcrHostResult::Lifecycle(Ok(()))) + } + OcrHostOperation::ProjectRequest => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))) + } + OcrHostOperation::AcquireAzureAdToken => { + panic!("test request has no token provider") + } + OcrHostOperation::PreCall(request) => { + phases.push("pre"); + result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { + Err(crate::Error::InvalidRequest("pre failed".into())) + } else { + Ok(request) + })); + } + OcrHostOperation::DuringCall(request) => { + phases.push("during"); + result = Some(OcrHostResult::DuringCall(if failure_phase == "during" { + Err(crate::Error::InvalidRequest("during failed".into())) + } else { + Ok(request) + })); + } + OcrHostOperation::PostCall(_) => panic!("transport should not be reached"), + }, + Err(error) => break error, + Ok(OcrCallStep::Complete(_)) => panic!("failed call completed"), + } + }; + assert!(matches!(error, crate::Error::InvalidRequest(_))); + assert_eq!( + phases + .iter() + .filter(|phase| **phase == failure_phase) + .count(), + 1 + ); + } +} + +#[tokio::test] +async fn invalid_provider_response_runs_post_call_before_normalization_failure() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; + let mut request = Some(wire_request("mistral/model", &base, json!({}))); + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let host = NoopOcrHost; + let mut result = None; + let mut post_calls = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))); + } + Ok(OcrCallStep::Host(operation)) => { + if let OcrHostOperation::PostCall(request) = &operation { + post_calls.push(request.original_response.clone()); + } + result = Some(host.invoke(operation).await); + } + Err(error) => break error, + Ok(OcrCallStep::Complete(_)) => panic!("invalid provider response completed"), + } + }; + server.await.unwrap(); + assert!(matches!(error, crate::Error::InvalidResponse(_))); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]); +} + +#[tokio::test] +async fn direct_native_host_drives_the_same_state_machine() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"native"}] + }))]) + .await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", &base, json!({})) + }; + let NativeOutcome::Completed(mut call) = OcrCall::admit( + super::test_support::ocr_client(), + OcrAdmission { + asynchronous: true, + ..OcrAdmission::all() + }, + ) else { + panic!("supported call declined") + }; + let mut request = Some(request); + let host = NoopOcrHost; + let mut result = None; + let mut operations = Vec::new(); + let response = loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(operation) => { + operations.push(match &operation { + OcrHostOperation::ProjectRequest => "ProjectRequest".into(), + OcrHostOperation::Lifecycle(phase) => format!("{phase:?}"), + OcrHostOperation::PreCall(_) => "PreCall".into(), + OcrHostOperation::DuringCall(_) => "DuringCall".into(), + OcrHostOperation::PostCall(_) => "PostCall".into(), + OcrHostOperation::ConstructResponse(_) => "ConstructResponse".into(), + OcrHostOperation::Success { response, .. } => { + assert_eq!(response.pages[0]["markdown"], "native"); + "Success".into() + } + _ => panic!("unexpected OCR operation"), + }); + result = Some(match operation { + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } + operation => host.invoke(operation).await, + }); + } + OcrCallStep::Complete(response) => break response, + } + }; + server.await.unwrap(); + assert_eq!(response.pages[0]["markdown"], "native"); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!( + operations, + [ + "Setup", + "DeploymentPreCall", + "Prepare", + "ProjectRequest", + "PreCall", + "DuringCall", + "PostCall", + "ConstructResponse", + "DeploymentPostCall", + "Finalize", + "Success", + ] + ); + assert!(matches!( + call.resume(None).await, + Err(crate::Error::InvalidRequest(_)) + )); +} + +#[tokio::test] +async fn public_finalization_failure_never_dispatches_success_or_replays_provider() { + use crate::call_lifecycle::host::{HostFailure, HostPhase}; + + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = Some(wire_request("mistral/model", &base, json!({}))); + let NativeOutcome::Completed(mut call) = OcrCall::admit( + super::test_support::ocr_client(), + OcrAdmission { + asynchronous: true, + ..OcrAdmission::all() + }, + ) else { + panic!("supported call declined") + }; + let selected = crate::Error::InvalidRequest("public metadata failed".into()); + let host = NoopOcrHost; + let mut result = None; + let mut failures = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(operation)) => { + result = Some(match operation { + OcrHostOperation::Lifecycle(HostPhase::Finalize) => { + OcrHostResult::Lifecycle(Err(HostFailure::Error(selected.clone()))) + } + OcrHostOperation::Failure { error, .. } => { + assert_eq!(error, selected); + failures.push("sync"); + OcrHostResult::Lifecycle(Err(HostFailure::Error( + crate::Error::InvalidRequest("failure callback failed".into()), + ))) + } + OcrHostOperation::Lifecycle(HostPhase::AsyncFailure) => { + failures.push("async"); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Success { .. } + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { + panic!("finalization failure used provider/success dispatch") + } + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } + operation => host.invoke(operation).await, + }); + } + Ok(OcrCallStep::Complete(_)) => panic!("failed call completed successfully"), + Err(error) => break error, + } + }; + server.await.unwrap(); + assert_eq!(error, selected); + assert_eq!(failures, ["sync", "async"]); + assert_eq!(seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption() { + use crate::call_lifecycle::host::HostFailure; + + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let host = NoopOcrHost; + let mut result = None; + loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))) + } + OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await), + OcrCallStep::Complete(_) => panic!("provider executed before pre-call result"), + } + } + let selected = crate::Error::InvalidRequest("cancelled".into()); + assert!(matches!( + call.interrupt(HostFailure::Cancelled(selected.clone())).await, + Err(error) if error == selected + )); + assert!( + call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) + .await + .is_err() + ); +} + +#[tokio::test] +async fn missing_host_result_preserves_pending_operation() { + use crate::call_lifecycle::host::HostPhase; + + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + assert!(matches!( + call.resume(None).await.unwrap(), + OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Setup)) + )); + assert!(call.resume(None).await.is_err()); + assert!(matches!( + call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) + .await + .unwrap(), + OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Prepare)) + )); +} + +async fn read_bounded_response( + response: Vec, + limit: usize, +) -> Result { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(&response).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + super::client::read_response_bytes(response, limit), + ) + .await; + server.abort(); + let _ = server.await; + result.expect("bounded reads must finish without waiting for the rest of an oversized body") +} + +#[tokio::test] +async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { + use super::error::{OcrError, OcrResponseError}; + + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", + ] { + assert_eq!( + read_bounded_response(response.as_bytes().to_vec(), 8) + .await + .unwrap(), + "abcdefgh" + ); + } + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", + ] { + assert!(matches!( + read_bounded_response(response.as_bytes().to_vec(), 8).await, + Err(OcrError::Response(OcrResponseError::TooLarge { limit: 8 })) + )); + } +} + +#[tokio::test] +async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining() { + let prefix = "x".repeat(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)); + for headers in ["Content-Length: 1000000", "Transfer-Encoding: chunked"] { + let body = if headers.starts_with("Transfer") { + format!("{:x}\r\n{prefix}\r\n", prefix.len()) + } else { + prefix.clone() + }; + let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); + let error = read_bounded_response(response.into_bytes(), 4096) + .await + .unwrap_err(); + match error { + super::error::OcrError::Transport(crate::error::TransportError::Http { + status, + body, + }) => { + assert_eq!(status, 429); + assert_eq!( + body, + format!( + "{}... (truncated)", + "x".repeat(crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS) + ) + ); + } + error => panic!("unexpected error: {error}"), + } + } +} + +#[test] +fn response_limit_is_validated_and_not_forwarded_to_the_provider() { + let request = wire_request( + "mistral/model", + "http://localhost", + json!({"max_response_bytes": 123}), + ); + assert_eq!(request.connection.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); + for value in [ + json!(0), + json!(-1), + json!(true), + json!("123"), + json!(1.5), + json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1), + Value::Null, + ] { + let wire = serde_json::from_value(json!({ + "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "optional_params": {"max_response_bytes": value} + })).unwrap(); + let Err(error) = decode_request(wire) else { + panic!("invalid response limit accepted") + }; + assert!(error.to_string().contains("max_response_bytes")); + } +} + +#[derive(Debug)] +struct PendingToken { + entered: Arc, + dropped: Arc, +} + +struct TokenFutureDrop(Arc); + +impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } +} + +impl crate::auth::TokenProvider for PendingToken { + fn acquire(&self) -> crate::auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } +} + +#[tokio::test] +async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_is_cancelled() { + use crate::call_lifecycle::host::HostFailure; + use std::future::Future; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::task::Poll; + + for interrupt_acknowledgement in [false, true] { + let entered = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + let request = super::LiteLLMOcrRequest { + connection: super::OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.connection + }, + azure_ad_token_provider: Some(crate::auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))), + ..request + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = entered.notified() => break, + step = call.resume(result.take()) => { + result = Some(match step.unwrap() { + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))), + OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await, + OcrCallStep::Complete(_) => panic!("pending provider completed"), + }); + } + } + } + }).await.unwrap(); + assert!(!dropped.load(Ordering::SeqCst)); + let selected = crate::Error::InvalidRequest("cancelled".into()); + if interrupt_acknowledgement { + let mut acknowledgement = + Box::pin(call.interrupt(HostFailure::Cancelled(selected.clone()))); + std::future::poll_fn(|cx| { + assert!(acknowledgement.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(acknowledgement); + assert!(!dropped.load(Ordering::SeqCst)); + } + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + call.interrupt(HostFailure::Cancelled(selected.clone())), + ) + .await + .unwrap(); + assert!(matches!(result, Err(error) if error == selected)); + assert!( + dropped.load(Ordering::SeqCst), + "cancellation returned while provider captures were still alive" + ); + } +} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 8e86e4713ef..a15e9cae5b5 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use rstest::rstest; use serde_json::{Value, json}; -use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks}; +use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; fn request_body(request: &str) -> Value { @@ -100,6 +100,42 @@ async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { assert!(requests[1].starts_with("POST /parse ")); } +struct ParseBoundary { + request_count: Arc>>, +} + +impl OcrHooks for ParseBoundary { + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + assert_eq!(self.request_count.lock().unwrap().len(), 2); + assert_eq!( + request.original_response, + json!(r#"{"result":{"chunks":[]}}"#) + ); + Ok(request) + }) + } +} + +#[tokio::test] +async fn post_call_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(ParseBoundary { + request_count: seen.clone(), + }), + ..wire_request("reducto/parse-v3", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); +} + #[rstest] #[case(json!({"file_id":""}))] #[case(json!({}))] @@ -148,9 +184,16 @@ async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { fn response_normalization_groups_blocks_and_distinguishes_null_result() { use crate::ocr::codecs::reducto::{ReductoResponse, transform_ocr_response}; - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"chunks":[ - {"blocks":[{"content":"B","bbox":{"page":2},"kind":"table"}]}, - {"blocks":[{"content":"A","bbox":{"page":1},"kind":"text"},{"content":"C","bbox":{"page":1}}]} + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ + {"blocks":[{ + "type":"Table", + "content":"B", + "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, + "confidence":"high", + "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, + "image_url":null + }]}, + {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} ]}}); let response: ReductoResponse = serde_json::from_value(raw).unwrap(); let normalized = transform_ocr_response("parse-v3", response) @@ -158,7 +201,17 @@ fn response_normalization_groups_blocks_and_distinguishes_null_result() { .into_json(); assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["kind"], "table"); + assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["bbox"], + json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) + ); + assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], + 0.95 + ); + assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); assert_eq!(normalized["usage_info"]["pages_processed"], 2); assert_eq!(normalized["usage_info"]["credits"], 3.0); @@ -195,7 +248,7 @@ async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { struct RewriteDocument; impl OcrHooks for RewriteDocument { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs index 6d3061d8f5d..676799eb2fe 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -49,7 +49,7 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { assert!(body.get("extra_body").is_none()); assert_eq!( body["messages"][0]["content"][0], - json!({"type":"document_url","document_url":"gs://bucket/document.pdf"}) + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) ); } diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs deleted file mode 100644 index e7739fe7312..00000000000 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ /dev/null @@ -1,115 +0,0 @@ -//! Enforcement: the litellm-rust workspace has exactly six crates. -//! -//! `core` (the Rust SDK), `token-counter` (standalone input token counting), -//! `config` (the config-loading boundary), -//! `ai-gateway` (the HTTP/WebSocket host), -//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the -//! PyO3 cdylib). Adding or removing a crate must be a -//! deliberate act: this test fails until the allowlist here is updated, forcing -//! whoever changes the crate set to justify the new crate per the rule that a -//! crate is a layer needing independent compilation / its own deps / a separate -//! artifact — and to keep `litellm-rust/AGENTS.md` in sync. -//! -//! Std-only (no toml crate): we scan the workspace manifest's `members = [...]` -//! block and the `crates/` directory directly. - -use std::collections::BTreeSet; -use std::fs; -use std::path::{Path, PathBuf}; - -/// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the -/// workspace legitimately gains or loses a crate. -const EXPECTED_MEMBERS: &[&str] = &[ - "crates/core", - "crates/token-counter", - "crates/config", - "crates/ai-gateway", - "crates/python-interop", - "crates/python-bridge", -]; - -/// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &[ - "core", - "token-counter", - "config", - "ai-gateway", - "python-interop", - "python-bridge", -]; - -const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact)."; - -/// Absolute path to the workspace root (`litellm-rust/`). -fn workspace_root() -> PathBuf { - // CARGO_MANIFEST_DIR is `.../litellm-rust/crates/core`; the workspace root is - // two levels up. - Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../..")) - .canonicalize() - .expect("workspace root should resolve") -} - -/// Parse the `members = [ ... ]` array out of the workspace `[workspace]` table. -/// -/// Minimal hand-rolled scan: find `members`, then collect every double-quoted -/// string up to the closing `]`. Good enough for our fixed manifest shape and -/// keeps this test dependency-free. -fn parse_members(manifest: &str) -> BTreeSet { - let after_members = manifest - .split_once("members") - .map(|(_, rest)| rest) - .expect("workspace manifest should declare members"); - let open = after_members.find('[').expect("members should be an array"); - let close = after_members[open..] - .find(']') - .map(|offset| open + offset) - .expect("members array should be closed"); - let body = &after_members[open + 1..close]; - - let mut members = BTreeSet::new(); - let mut rest = body; - while let Some(start) = rest.find('"') { - let after_quote = &rest[start + 1..]; - let end = after_quote - .find('"') - .expect("opening quote should be matched"); - members.insert(after_quote[..end].to_string()); - rest = &after_quote[end + 1..]; - } - members -} - -/// The crate subdirectory names under `crates/`. -/// -/// A directory counts as a crate only when it holds a `Cargo.toml`; non-crate -/// directories (e.g. docs like `CODING_STANDARDS/`) are ignored so they can live -/// under `crates/` without tripping the crate-proliferation guard. -fn crate_dirs(root: &Path) -> BTreeSet { - fs::read_dir(root.join("crates")) - .expect("crates/ directory should exist") - .filter_map(Result::ok) - .filter(|entry| entry.file_type().map(|ty| ty.is_dir()).unwrap_or(false)) - .filter(|entry| entry.path().join("Cargo.toml").is_file()) - .map(|entry| entry.file_name().to_string_lossy().into_owned()) - .collect() -} - -#[test] -fn workspace_members_match_allowlist() { - let root = workspace_root(); - let manifest = fs::read_to_string(root.join("Cargo.toml")) - .expect("workspace Cargo.toml should be readable"); - - let actual = parse_members(&manifest); - let expected: BTreeSet = EXPECTED_MEMBERS.iter().map(|s| s.to_string()).collect(); - assert_eq!(actual, expected, "{MISMATCH}"); -} - -#[test] -fn crates_directory_matches_allowlist() { - let root = workspace_root(); - - let actual = crate_dirs(&root); - let expected: BTreeSet = EXPECTED_CRATE_DIRS.iter().map(|s| s.to_string()).collect(); - assert_eq!(actual, expected, "{MISMATCH}"); -} diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 42282ca4da4..9262617156b 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,42 @@ -litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop. - -Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. +- Target invariants, not completion claims; these supersede older conflicting bridge guidance +- Keep this crate the product-specific PyO3 consumer of `litellm-python-interop` + - Own registration, input projection, retained Python state, callback invocation, public response/error construction and host scheduling + - Keep value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment in `execution.rs`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` + - Core owns typed native state, admission, lifecycle sequencing, provider preparation/I/O, normalization and terminal-outcome/dispatch decisions + - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers + - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points +- Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work + - Free-threading requires separate runtime/concurrency validation; omitting the attribute does not opt out on PyO3 0.28+ +- Preserve public argument binding and Python object provenance + - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view + - Retain independently captured body/header roots; in-place mutation and logging-envelope field replacement have different effects + - Project only consumed fields at reference read points; no eager whole-graph serialization or equality-based alias reconstruction + - Preserve provider-specific upload/submission/poll observation and encoding boundaries; signed/build-captured bytes must not be silently reserialized +- Only core's typed, effect-free admission may return `Declined`; conversion errors and all post-admission failures are terminal + - Admission cannot invoke hooks, acquire credentials, consume files/iterators, prepare requests or perform I/O + - Disabled/unavailable native execution or an admission decline may select legacy once; callback exceptions never authorize fallback or replay +- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle in `src/lifecycle.rs` + - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`; explicitly tagged `Await`/`Complete` preserve awaitable final values + - Validate Created/Running/Suspended/Closed protocol states; core alone chooses lifecycle phases and result/error policy + - Defer effectful setup/context reads/timestamps until start; unstarted-handle destruction releases inputs independently of Python `finally` + - Catch only the selected await's errors; start/resume errors propagate, `GeneratorExit` closes without further awaits + - Inline hooks preserve caller task/thread/loop and context writes; `into_future` creates a separate task and cannot satisfy this contract + - Delivery follows the binding, not callable type; keep direct, awaited, worker, background and deferred behavior distinct +- Finalize fallible public response/error construction, replacements and metadata under core control before terminal dispatch + - Success/failure handler entry receives the exact selected public response/exception; logging projections/redaction/snapshots retain their own copy contracts + - Ordinary failure-callback errors cannot suppress later eligible sync/async callbacks or replace the mapped provider error; control-flow exceptions have phase-specific policy + - Dispatch errors never replay provider work/accepted dispatch or trigger the opposite outcome; proxy acceptance/rejection releases core-owned deferred success at most once +- Make ownership safe across suspension, re-entry, cancellation and GC + - Keep native provider state typed in core; do not shuttle it through opaque Python transport/response classes + - Prefer one retained `Py` via `PyErr::into_value(py)`; reconstruct transient `PyErr`s, preserving identity, traceback, cause and context + - Traverse every owned Python edge, including duplicate references; traversal cannot call Python + - Take state out and mark Running under a short borrow, release borrows/locks before Python invocation, publish terminal state before finalizer-capable drops + - Close/GC/deferred release are idempotent and re-entry-safe, including during Rust unwinding; release only owned references, never clear caller containers or mask the selected error + - Cancellation signaling is not termination; retain captures until work actually finishes and use a Rust-selected awaited acknowledgement where required, never synchronous close/GC +- Verify behavior through a fresh, provenance-checked installed extension and positive native execution evidence before replacing the custom coroutine + - Cover admitted provider workflows, binding/read-point/identity behavior, failure continuation, finalization, no replay, deferred gates, re-entry, GC and cancellation termination + - Measure real conversion/copy costs before optimizing; preserve input contracts and capture lifetimes with `PyBackedBytes`, and lookup timing when interning names + - Ship accurate `_native.pyi` declarations and typing markers; distinguish Future-returning bindings from coroutine-returning bindings +- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [exception transfer](https://docs.rs/pyo3/0.29.2/pyo3/struct.PyErr.html#method.into_value), [re-entry](https://pyo3.rs/v0.29.2/class/call.html) + - [GIL policy](https://pyo3.rs/v0.29.2/free-threading.html), [experimental async limits](https://pyo3.rs/v0.29.2/async-await.html), [task conversion](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/fn.into_future_with_locals.html), [native cancellation/delivery](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/tokio/fn.future_into_py.html) + - [performance](https://pyo3.rs/v0.29.2/performance.html), [PyBackedBytes](https://docs.rs/pyo3/0.29.2/pyo3/pybacked/struct.PyBackedBytes.html), [typing](https://pyo3.rs/v0.29.2/python-typing-hints.html) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 337a1e8e5ac..42fad740870 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -17,7 +17,6 @@ panic-test = [] trace-parity = [ "dep:tracing", "litellm-core/observability", - "litellm-ai-gateway/trace-parity", ] [dependencies] @@ -25,7 +24,6 @@ futures-util.workspace = true tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-token-counter.workspace = true -litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true @@ -35,6 +33,7 @@ tokio = { workspace = true, features = ["sync"] } [dev-dependencies] criterion.workspace = true +rstest.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/auth.rs b/litellm-rust/crates/python-bridge/src/auth.rs new file mode 100644 index 00000000000..8dc0b7aabf0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/auth.rs @@ -0,0 +1,194 @@ +use litellm_core::auth::{ResolvedCredential, SecretValue}; +use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::PyString; + +#[derive(Clone, Copy)] +pub(crate) struct TokenProviderContract { + callable_error: &'static str, + token_type_error: &'static str, + callback_error: &'static str, +} + +pub(crate) const AZURE_AD_TOKEN_PROVIDER: TokenProviderContract = TokenProviderContract { + callable_error: "Azure AD token provider must be callable", + token_type_error: "Azure AD token must be a string, got {}", + callback_error: "Failed to get Azure AD token: {}", +}; + +pub(crate) struct PythonTokenProvider { + callback: Py, + contract: TokenProviderContract, +} + +impl PythonTokenProvider { + pub(crate) fn select( + provider: Bound<'_, PyAny>, + contract: TokenProviderContract, + ) -> Option { + (provider.is_callable() && provider.is_truthy().unwrap_or(false)).then(|| Self { + callback: provider.unbind(), + contract, + }) + } + + pub(crate) fn acquire(&self, py: Python<'_>) -> PyResult { + let provider = self.callback.bind(py); + if !provider.is_callable() { + return Err(PyTypeError::new_err(self.contract.callable_error)); + } + let token = (|| { + let token = provider.call0()?; + if !token.is_instance_of::() { + let message = PyString::new(py, self.contract.token_type_error) + .call_method1("format", (token.get_type(),))?; + return Err(PyTypeError::new_err(message.unbind())); + } + Ok(token) + })() + .map_err(|error| { + if error.is_instance_of::(py) || !error.is_instance_of::(py) { + return error; + } + match PyString::new(py, self.contract.callback_error) + .call_method1("format", (error.value(py),)) + { + Ok(message) => { + let wrapped = PyRuntimeError::new_err(message.unbind()); + wrapped.set_context(py, Some(error.clone_ref(py))); + wrapped.set_cause(py, Some(error)); + wrapped + } + Err(format_error) => { + format_error.set_context(py, Some(error)); + format_error + } + } + })?; + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new(token.extract::()?), + expires_on: None, + }) + } + + pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.callback) + } +} + +#[cfg(test)] +mod tests { + use pyo3::exceptions::PyRuntimeError; + use pyo3::types::PyDict; + + use super::*; + + #[test] + fn token_callback_preserves_exception_identity_and_explicit_chaining() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +class ProviderError(Exception): + def __format__(self, specification): + return 'unavailable' +ordinary = ProviderError('must use __format__') +type_error = TypeError('signature') +abort = KeyboardInterrupt('cancelled') +def provider(error): + def acquire(): + raise error + return acquire +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + for name in ["ordinary", "type_error", "abort"] { + let original = locals.get_item(name).unwrap().unwrap(); + let callback = locals + .get_item("provider") + .unwrap() + .unwrap() + .call1((&original,)) + .unwrap(); + let provider = + PythonTokenProvider::select(callback, AZURE_AD_TOKEN_PROVIDER).unwrap(); + let error = provider.acquire(py).unwrap_err(); + if name == "ordinary" { + assert!(error.is_instance_of::(py)); + assert!(error.cause(py).unwrap().value(py).is(&original)); + assert!( + error + .value(py) + .getattr("__context__") + .unwrap() + .is(&original) + ); + assert_eq!( + error.value(py).str().unwrap().to_str().unwrap(), + "Failed to get Azure AD token: unavailable" + ); + } else { + assert!(error.value(py).is(&original)); + } + } + }); + } + + #[test] + fn invalid_token_type_formatting_preserves_python_failure_semantics() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +failure = ValueError('formatting failed') +class TokenType(type): + def __format__(cls, specification): + raise failure +class Token(metaclass=TokenType): + pass +def provider(): + return Token() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let provider = PythonTokenProvider::select( + locals.get_item("provider").unwrap().unwrap(), + AZURE_AD_TOKEN_PROVIDER, + ) + .unwrap(); + let error = provider.acquire(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!( + error + .cause(py) + .unwrap() + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn token_string_extraction_errors_are_not_wrapped_as_callback_failures() { + Python::initialize(); + Python::attach(|py| { + let callback = py + .eval(pyo3::ffi::c_str!("lambda: '\\ud800'"), None, None) + .unwrap(); + let provider = PythonTokenProvider::select(callback, AZURE_AD_TOKEN_PROVIDER).unwrap(); + let error = provider.acquire(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index e1f458ea0bc..701c6abb68c 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -22,7 +22,8 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { Error::InvalidProvider(_) | Error::InvalidRequest(_) | Error::InvalidType { .. } - | Error::MissingField(_) => PyValueError::new_err(err.to_string()), + | Error::MissingField(_) + | Error::MissingDocumentUrl => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } @@ -41,6 +42,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidRequest(_) | Error::InvalidType { .. } | Error::MissingField(_) + | Error::MissingDocumentUrl | Error::MissingApiKey { .. } | Error::MissingAzureAiCredentials | Error::MissingAzureDocumentIntelligenceCredentials @@ -49,9 +51,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - Error::Http { status, body } => { - RustUpstreamError::new_err((status, format!("{status}: {body}"))) - } + Error::Http { status, body } => RustUpstreamError::new_err((status, body)), Error::Network(message) | Error::InvalidResponse(message) => { RustUpstreamError::new_err((0u16, message)) } @@ -63,41 +63,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } - -pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::MissingField("document_url" | "image_url") => { - PyValueError::new_err("Document URL is required") - } - Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - other => core_error_to_pyerr(other), - } -} - -#[cfg(test)] -mod ocr_error_tests { - use super::*; - - #[test] - fn ocr_errors_preserve_python_validation_and_provider_details() { - Python::initialize(); - Python::attach(|py| { - for field in ["document_url", "image_url"] { - let mapped = ocr_error_to_pyerr(Error::MissingField(field)); - assert!(mapped.is_instance_of::(py)); - assert_eq!(mapped.value(py).to_string(), "Document URL is required"); - } - let mapped = ocr_error_to_pyerr(Error::Http { - status: 429, - body: r#"{"message":"rate limited"}"#.to_string(), - }); - assert!(mapped.is_instance_of::(py)); - let args: (u16, String) = mapped - .value(py) - .getattr("args") - .and_then(|args| args.extract()) - .expect("OCR failures retain status and unprefixed provider message"); - assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index b57197b9ddf..d8dda10068d 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -1,5 +1,7 @@ use std::future::Future; use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::task::{Context, Poll, Waker}; use std::time::Duration; use futures_util::FutureExt; @@ -28,6 +30,27 @@ where ) } +pub(crate) fn run_sync_value(py: Python<'_>, future: F) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future) +} + +fn run_sync_value_on(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + if Handle::try_current().is_ok() { + return Err(PyRuntimeError::new_err( + "synchronous native routes cannot run from a Tokio context; use the async route", + )); + } + release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))? +} + fn run_sync_on( py: Python<'_>, runtime: &Runtime, @@ -67,6 +90,32 @@ where }) } +pub(crate) fn run_async_value(py: Python<'_>, future: F) -> PyResult> +where + T: for<'py> IntoPyObject<'py> + Send + 'static, + F: Future> + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? }) +} + +pub(crate) fn poll_async_value(py: Python<'_>, future: Pin<&mut F>) -> PyResult> +where + T: Send, + F: Future> + Send, +{ + let result = release_gil(py, || { + let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter(); + std::panic::catch_unwind(AssertUnwindSafe(|| { + future.poll(&mut Context::from_waker(Waker::noop())) + })) + .map_err(panic_to_pyerr) + })?; + match result { + Poll::Ready(result) => result.map(Poll::Ready), + Poll::Pending => Ok(Poll::Pending), + } +} + fn map_core_result(result: Result, map_error: fn(E) -> PyErr) -> PyResult { match result { Ok(value) => Ok(value), @@ -119,11 +168,30 @@ mod tests { use litellm_core::error::Error; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; + use rstest::{fixture, rstest}; use serde::Serializer; use tokio::runtime::Builder; use super::*; + struct InitializedPython; + + impl InitializedPython { + fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } + } + + #[fixture] + #[once] + fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython + } + fn runtime_error(error: Error) -> PyErr { PyRuntimeError::new_err(error.to_string()) } @@ -194,10 +262,84 @@ mod tests { .expect("result should convert") } - #[test] - fn sync_runner_polls_future_on_the_caller_thread() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn inline_poll_releases_gil_and_enters_runtime( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let (sender, receiver) = mpsc::sync_channel(1); + let worker = thread::spawn(move || Python::attach(|_| sender.send(()).unwrap())); + let mut future = Box::pin(async move { + receiver.recv_timeout(Duration::from_secs(2)).unwrap(); + Ok(Handle::try_current().is_ok()) + }); + assert_eq!( + poll_async_value(py, future.as_mut()).unwrap(), + Poll::Ready(true) + ); + worker.join().unwrap(); + }); + } + + #[rstest] + fn inline_poll_contains_panics_and_preserves_python_errors( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let mut panicking = Box::pin(poll_fn(|_| -> Poll> { + panic!("inline native panic") + })); + let error = poll_async_value(py, panicking.as_mut()).unwrap_err(); + assert!(error.is_instance_of::(py)); + let original = PyRuntimeError::new_err("inline failure"); + let identity = original.value(py).clone().unbind(); + let mut failing = Box::pin(async move { Err::<(), _>(original) }); + let error = poll_async_value(py, failing.as_mut()).unwrap_err(); + assert!(error.value(py).is(identity.bind(py))); + }); + } + + #[pyfunction] + fn pending_after_inline_poll(py: Python<'_>) -> PyResult> { + let starts = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&starts); + let mut future = Box::pin(async move { + starts.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(5)).await; + Ok(starts.load(Ordering::SeqCst)) + }); + assert!(poll_async_value(py, future.as_mut())?.is_pending()); + assert_eq!(observed.load(Ordering::SeqCst), 1); + run_async_value(py, future) + } + + #[rstest] + fn inline_pending_future_resumes_on_tokio_without_restarting( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "pending", + wrap_pyfunction!(pending_after_inline_poll, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + "import asyncio\nasync def exercise():\n assert await asyncio.wait_for(pending(), 2) == 1\nasyncio.run(exercise())" + ), + Some(&locals), + Some(&locals), + ).unwrap(); + }); + } + + #[rstest] + fn sync_runner_polls_future_on_the_caller_thread( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let caller_thread = std::thread::current().id(); let result = run_sync( py, @@ -209,10 +351,11 @@ mod tests { }); } - #[test] - fn sync_runner_releases_gil_while_waiting() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_releases_gil_while_waiting( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let result = run_sync( py, async { @@ -230,16 +373,17 @@ mod tests { }); } - #[test] - fn sync_runner_rejects_calls_from_a_tokio_context() { - Python::initialize(); + #[rstest] + fn sync_runner_rejects_calls_from_a_tokio_context( + #[from(initialized_python)] python: &InitializedPython, + ) { let runtime = Builder::new_current_thread() .enable_all() .build() .expect("runtime should build"); let error = runtime.block_on(async { - Python::attach(|py| { + python.attach(|py| { run_sync::(py, async { Ok(true) }, runtime_error) .expect_err("sync route should reject a nested Tokio runtime") }) @@ -251,14 +395,15 @@ mod tests { ); } - #[test] - fn sync_runner_can_drive_a_current_thread_runtime() { - Python::initialize(); + #[rstest] + fn sync_runner_can_drive_a_current_thread_runtime( + #[from(initialized_python)] python: &InitializedPython, + ) { let runtime = Builder::new_current_thread() .enable_all() .build() .expect("runtime should build"); - Python::attach(|py| { + python.attach(|py| { let result = run_sync_on( py, &runtime, @@ -272,10 +417,9 @@ mod tests { }); } - #[test] - fn sync_runner_maps_a_panicked_future() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_maps_a_panicked_future(#[from(initialized_python)] python: &InitializedPython) { + python.attach(|py| { let error = run_sync::( py, poll_fn(|_| -> Poll> { panic!("route future panicked") }), @@ -288,10 +432,11 @@ mod tests { }); } - #[test] - fn sync_runner_maps_a_panicked_error_mapper() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_maps_a_panicked_error_mapper( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let error = run_sync::( py, async { Err(Error::InvalidRequest("invalid".to_string())) }, @@ -304,10 +449,11 @@ mod tests { }); } - #[test] - fn sync_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_surfaces_serializer_panics( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) .expect_err("serializer panic should become a Python exception"); @@ -316,9 +462,10 @@ mod tests { }); } - #[test] - fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { - Python::initialize(); + #[rstest] + fn sync_runner_supports_concurrent_callers_on_the_shared_runtime( + #[from(initialized_python)] _python: &InitializedPython, + ) { let barrier = Arc::new(tokio::sync::Barrier::new(2)); let callers: Vec<_> = (0..2) .map(|_| { @@ -349,10 +496,11 @@ mod tests { assert_eq!(results, vec![true, true]); } - #[test] - fn async_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn async_runner_surfaces_serializer_panics( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let module = PyModule::new(py, "runtime").expect("module should be created"); module .add_function( @@ -386,11 +534,12 @@ asyncio.run(exercise()) }); } - #[test] - fn async_result_delivery_does_not_stall_tokio_workers() { - Python::initialize(); + #[rstest] + fn async_result_delivery_does_not_stall_tokio_workers( + #[from(initialized_python)] python: &InitializedPython, + ) { ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); - Python::attach(|py| { + python.attach(|py| { let module = PyModule::new(py, "runtime").expect("module should be created"); for function in [ wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"), diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index cf0450a1b30..12bc57a8931 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,14 +1,16 @@ +mod auth; mod constants; mod diagnostics; mod errors; mod execution; #[cfg(feature = "trace-parity")] mod function_trace; +mod lifecycle; mod marshal; mod routes; mod token_counter; -use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; +use litellm_core::responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use pyo3::types::PyAny; use serde_json::Value; @@ -64,7 +66,7 @@ impl ResponsesWebSocketConnection { } } -#[pymodule(gil_used = false)] +#[pymodule(gil_used = true)] mod _native { use pyo3::prelude::*; @@ -152,7 +154,6 @@ mod tests { "amessages", "chat_completions", "achat_completions", - "gateway_messages", ] ); } diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs new file mode 100644 index 00000000000..06b32b67fd5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -0,0 +1,391 @@ +use pyo3::exceptions::PyBaseException; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +#[derive(FromPyObject)] +pub(crate) struct PythonLogger(Py); + +impl PythonLogger { + pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> { + self.0.bind(py) + } + + pub(crate) fn clone_ref(&self, py: Python<'_>) -> Self { + Self(self.0.clone_ref(py)) + } + + pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + + pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { + if !self + .object(py) + .getattr("_native_callback_fast_path") + .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + { + return Ok(true); + } + py.import("litellm.rust_bridge.lifecycle")? + .getattr("callbacks_needed")? + .call1((self.object(py), phase))? + .extract() + } + + pub(super) fn success_bookkeeping( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + asynchronous: bool, + ) -> PyResult<()> { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("success_bookkeeping")? + .call1((self.object(py), response, start, end, asynchronous))?; + Ok(()) + } + + pub(super) fn defers_async_logging(&self, py: Python<'_>) -> bool { + self.object(py) + .getattr("_defer_async_logging") + .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + } + + pub(super) fn defer_success( + &self, + py: Python<'_>, + pending: Py, + ) -> PyResult<()> { + self.object(py).setattr("_native_pending_logging", pending) + } + + pub(super) fn sync_success_for_async_call( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + if !self.callbacks_needed(py, "sync_success_async")? { + return Ok(()); + } + self.object(py).call_method1( + "handle_sync_success_callbacks_for_async_calls", + (response, start, end), + )?; + Ok(()) + } + + pub(super) fn failure( + &self, + py: Python<'_>, + error: &Py, + start: &Py, + end: &Option>, + asynchronous: bool, + ) -> PyResult>> { + if !self.callbacks_needed( + py, + if asynchronous { + "async_failure" + } else { + "sync_failure" + }, + )? { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("failure_bookkeeping")? + .call1((self.object(py), error, start, end, asynchronous))?; + return Ok(None); + } + let trace = py + .import("traceback")? + .getattr("format_exception")? + .call1((error,))?; + let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?; + let value = self.object(py).call_method1( + if asynchronous { + "async_failure_handler" + } else { + "failure_handler" + }, + (error, trace, start, end), + )?; + Ok(asynchronous.then(|| value.unbind())) + } + + pub(super) fn restore_context(&self, py: Python<'_>) -> PyResult<()> { + py.import("litellm.utils")? + .getattr("_restore_correlation_context_if_supported")? + .call1((self.object(py),))?; + Ok(()) + } + + pub(super) fn submit_success( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + if !self.callbacks_needed(py, "sync_success")? { + return self.success_bookkeeping(py, response, start, end, false); + } + let context = py.import("contextvars")?.call_method0("copy_context")?; + py.import("litellm.litellm_core_utils.litellm_logging")? + .getattr("executor")? + .call_method1( + "submit", + ( + context.getattr("run")?, + self.object(py).getattr("success_handler")?, + response, + start, + end, + ), + )?; + Ok(()) + } + + pub(super) fn enqueue_success( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + if !self.callbacks_needed(py, "async_success")? { + return self.success_bookkeeping(py, response, start, end, true); + } + let context = py.import("contextvars")?.call_method0("copy_context")?; + let worker = py + .import("litellm.litellm_core_utils.logging_worker")? + .getattr("GLOBAL_LOGGING_WORKER")? + .getattr("ensure_initialized_and_enqueue")?; + let coroutine = self + .object(py) + .call_method1("async_success_handler", (response, start, end))?; + let enqueue = context.call_method1("run", (worker, &coroutine)); + if enqueue.is_err() + && let Err(error) = coroutine.call_method0("close") + { + error.write_unraisable(py, Some(&coroutine)); + } + enqueue.map(|_| ()) + } +} + +pub(super) struct SetupResult<'py>(Bound<'py, PyAny>); + +impl SetupResult<'_> { + pub(super) fn logger(&self) -> PyResult { + self.0.getattr("logger")?.extract() + } + + pub(super) fn kwargs(&self) -> PyResult> { + Ok(self.0.getattr("kwargs")?.extract()?) + } +} + +pub(super) fn setup<'py>( + py: Python<'py>, + call_type: &str, + args: &Py, + kwargs: &Py, + start: &Py, + asynchronous: bool, +) -> PyResult> { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("setup")? + .call1((call_type, args, kwargs, start, asynchronous)) + .map(SetupResult) +} + +pub(super) fn finalize( + py: Python<'_>, + response: &Option>, + logger: &PythonLogger, + kwargs: &Py, + start: &Py, + end: &Option>, +) -> PyResult<()> { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("finalize")? + .call1((response, logger.object(py), kwargs, start, end))?; + Ok(()) +} + +pub(super) fn is_internal_call(py: Python<'_>) -> PyResult { + py.import("litellm._internal_context")? + .getattr("is_internal_call")? + .call_method0("get")? + .extract() +} + +pub(super) struct DeploymentHooks; + +impl DeploymentHooks { + pub(super) fn needed(py: Python<'_>) -> PyResult { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("deployment_callbacks_needed")? + .call0()? + .extract() + } + + pub(super) fn before_call( + py: Python<'_>, + kwargs: &Py, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_pre_call_deployment_hook")? + .call1((kwargs, call_type)) + .map(Bound::unbind) + } + + pub(super) fn after_success( + py: Python<'_>, + kwargs: &Py, + response: &Option>, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_post_call_success_deployment_hook")? + .call1((kwargs, response, call_type)) + .map(Bound::unbind) + } + + pub(super) fn after_failure( + py: Python<'_>, + kwargs: &Py, + error: &Py, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_post_call_failure_deployment_hook")? + .call1((kwargs, error, call_type)) + .map(Bound::unbind) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyTypeError; + + #[test] + fn setup_fields_are_checked_in_order_without_eager_logger_method_reads() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +reads = [] +class Logger: + def __getattribute__(self, name): + reads.append(name) + raise AssertionError('logger methods must remain lazy') +logger = Logger() +class Setup: + @property + def logger(self): + reads.append('logger') + return logger + @property + def kwargs(self): + reads.append('kwargs') + return [] +result = Setup() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let result = SetupResult(locals.get_item("result").unwrap().unwrap()); + let logger = result.logger().unwrap(); + assert!( + logger + .object(py) + .is(locals.get_item("logger").unwrap().unwrap()) + ); + assert_eq!( + locals + .get_item("reads") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + ["logger"] + ); + assert!( + result + .kwargs() + .unwrap_err() + .is_instance_of::(py) + ); + assert_eq!( + locals + .get_item("reads") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + ["logger", "kwargs"] + ); + }); + } + + #[test] + fn logger_resolves_each_callback_at_invocation_and_preserves_arguments() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +calls = [] +response, start, end = object(), object(), object() +class Logger: + @property + def handle_sync_success_callbacks_for_async_calls(self): + generation = len(calls) + def callback(*args): + assert args == (response, start, end) + calls.append(generation) + return callback +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let logger: PythonLogger = locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(); + let response = Some(locals.get_item("response").unwrap().unwrap().unbind()); + let start = locals.get_item("start").unwrap().unwrap().unbind(); + let end = Some(locals.get_item("end").unwrap().unwrap().unbind()); + for _ in 0..2 { + logger + .sync_success_for_async_call(py, &response, &start, &end) + .unwrap(); + } + assert_eq!( + locals + .get_item("calls") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + [0, 1] + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs b/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs new file mode 100644 index 00000000000..17a480a7225 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs @@ -0,0 +1,139 @@ +use std::panic::{AssertUnwindSafe, catch_unwind}; + +use litellm_python_interop::panic_to_pyerr; +use pyo3::exceptions::{PyBaseException, PyRuntimeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; + +pub(super) enum ExecutionStep { + Return(Py), + Await(Py), +} + +pub(super) trait ExecutionBody: Send + Sync { + fn resume(&mut self, result: Option>>) -> PyResult; + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; +} + +enum ExecutionState { + Created(Box), + Running, + Suspended(Box), + Closed, +} + +#[pyclass] +pub(super) struct Execution { + state: ExecutionState, +} + +impl Execution { + pub(super) fn new(body: impl ExecutionBody + 'static) -> Self { + Self { + state: ExecutionState::Created(Box::new(body)), + } + } + + fn advance( + slf: &Bound<'_, Self>, + py: Python<'_>, + result: Option>>, + ) -> PyResult> { + let mut body = { + let mut execution = slf.borrow_mut(); + match (&execution.state, result.is_some()) { + (ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {} + (ExecutionState::Running, _) => { + return Err(PyRuntimeError::new_err("execution is already running")); + } + (ExecutionState::Closed, _) => { + return Err(PyRuntimeError::new_err("execution is closed")); + } + _ => { + return Err(PyRuntimeError::new_err( + "execution requires start before resume and can only start once", + )); + } + } + match std::mem::replace(&mut execution.state, ExecutionState::Running) { + ExecutionState::Created(body) | ExecutionState::Suspended(body) => body, + _ => unreachable!(), + } + }; + let outcome = catch_unwind(AssertUnwindSafe(|| { + let step = body.resume(result)?; + let (tag, value, suspended) = match step { + ExecutionStep::Await(value) => ("Await", value, true), + ExecutionStep::Return(value) => ("Complete", value, false), + }; + let step = py + .import("litellm.rust_bridge.lifecycle")? + .getattr(tag)? + .call1((value,))? + .unbind(); + Ok((step, suspended)) + })) + .map_err(panic_to_pyerr) + .and_then(|result| result); + match outcome { + Ok((step, true)) if matches!(slf.borrow().state, ExecutionState::Running) => { + slf.borrow_mut().state = ExecutionState::Suspended(body); + Ok(step) + } + outcome => { + slf.borrow_mut().state = ExecutionState::Closed; + drop(body); + outcome.and_then(|(step, suspended)| { + if suspended { + Err(PyRuntimeError::new_err( + "execution was closed while running", + )) + } else { + Ok(step) + } + }) + } + } + } +} + +#[pymethods] +impl Execution { + fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult> { + Self::advance(slf, py, None) + } + + fn resume_value( + slf: &Bound<'_, Self>, + py: Python<'_>, + value: Py, + ) -> PyResult> { + Self::advance(slf, py, Some(Ok(value))) + } + + fn resume_error( + slf: &Bound<'_, Self>, + py: Python<'_>, + error: Bound<'_, PyBaseException>, + ) -> PyResult> { + Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any())))) + } + + fn close(slf: &Bound<'_, Self>) { + let state = std::mem::replace(&mut slf.borrow_mut().state, ExecutionState::Closed); + drop(state); + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + match &self.state { + ExecutionState::Created(body) | ExecutionState::Suspended(body) => { + body.traverse(&visit) + } + _ => Ok(()), + } + } + + fn __clear__(slf: &Bound<'_, Self>) { + Self::close(slf); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs new file mode 100644 index 00000000000..014564ae89d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -0,0 +1,1175 @@ +use std::sync::Arc; +use std::task::Poll; + +use futures_util::future::{AbortHandle, Abortable}; +#[cfg(test)] +use litellm_core::call_lifecycle::host::HostCallFuture; +use litellm_core::call_lifecycle::host::{ + HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, +}; +use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; +use tokio::sync::Mutex; + +use crate::execution::{poll_async_value, run_async_value, run_sync_value}; + +mod bindings; +mod handle; +mod preparation; + +use bindings::DeploymentHooks; +pub(crate) use bindings::PythonLogger; +use handle::{Execution, ExecutionBody, ExecutionStep}; + +pub(crate) enum OperationClass { + Phase(HostPhase), + Route, +} + +pub(crate) trait PythonRoute: Send + Sync { + type Call: NativeCall + 'static; + + fn state(&self) -> &PythonCallState; + fn state_mut(&mut self) -> &mut PythonCallState; + fn classify(operation: &::Operation) -> OperationClass; + fn lifecycle_result() -> ::Result; + fn map_error(error: litellm_core::Error) -> PyErr; + fn invoke( + &mut self, + py: Python<'_>, + operation: ::Operation, + ) -> PyResult<::Result>; + fn cleanup(&mut self); + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; +} + +type NativeStep = NativeCallStep<::Operation, ::Complete>; +type NativeResult = Result, litellm_core::Error>; +type HostResumeStep = HostStep::Call>, Py>; + +struct NativeCallState { + call: C, + result: Option>, +} + +enum PendingOperation { + Native, + Host(HostPhase), +} + +struct PythonLifecycle { + route: R, + call: Option>>>, + pending: Option, + native_abort: Option, +} + +pub(crate) fn run_call( + py: Python<'_>, + call: R::Call, + route: R, +) -> PyResult> { + let asynchronous = route.state().asynchronous; + let mut lifecycle = PythonLifecycle { + route, + call: Some(Arc::new(Mutex::new(NativeCallState { call, result: None }))), + pending: None, + native_abort: None, + }; + if asynchronous { + let execution = Py::new(py, Execution::new(lifecycle))?; + return py + .import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) + .map(Bound::unbind); + } + match lifecycle.resume(None)? { + ExecutionStep::Return(value) => Ok(value), + ExecutionStep::Await(_) => Err(pyo3::exceptions::PyRuntimeError::new_err( + "sync call suspended", + )), + } +} + +pub(crate) fn missing_state() -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err("missing native call state") +} + +impl PythonLifecycle { + fn resume_core( + &mut self, + py: Python<'_>, + result: Option::Result, HostFailure>>, + ) -> PyResult> { + let call = Arc::clone(self.call.as_ref().ok_or_else(missing_state)?); + let future = async move { + let mut call = call.lock().await; + let result = match result { + Some(Err(failure)) => call.call.interrupt(failure).await, + Some(Ok(result)) => call.call.resume(Some(result)).await, + None => call.call.resume(None).await, + }; + call.result = Some(result); + Ok(()) + }; + if self.route.state().asynchronous { + let mut future = Box::pin(future); + if let Poll::Ready(()) = poll_async_value(py, future.as_mut())? { + return Ok(HostStep::Ready(self.take_native_result()?)); + } + let (abort, registration) = AbortHandle::new_pair(); + self.native_abort = Some(abort); + self.pending = Some(PendingOperation::Native); + Ok(HostStep::Suspend( + run_async_value(py, async move { + Abortable::new(future, registration) + .await + .map_err(|_| PyRuntimeError::new_err("native execution closed"))? + })? + .unbind(), + )) + } else { + run_sync_value(py, future)?; + Ok(HostStep::Ready(self.take_native_result()?)) + } + } + + fn take_native_result(&self) -> PyResult> { + self.call + .as_ref() + .ok_or_else(missing_state)? + .try_lock() + .map_err(|_| missing_state())? + .result + .take() + .ok_or_else(missing_state)? + .map_err(R::map_error) + } + + fn host_failure( + &mut self, + py: Python<'_>, + error: PyErr, + phase: Option, + ) -> HostFailure { + let native = litellm_core::Error::InvalidRequest(error.to_string()); + let cancelled = !error.is_instance_of::(py); + let failure = if !cancelled { + HostFailure::Error(native) + } else { + HostFailure::Cancelled(native) + }; + let state = self.route.state_mut(); + if state.error.is_none() || (cancelled && phase != Some(HostPhase::DeploymentFailure)) { + state.retain_error(py, error); + } + if state.end.is_none() { + state.end = now(py).ok(); + } + failure + } + + fn drive( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + let mut step = match (self.pending.take(), result) { + (None, None) => self.resume_core(py, None)?, + (Some(PendingOperation::Native), Some(result)) => match result { + Ok(_) => HostStep::Ready(self.take_native_result()?), + Err(error) => { + let failure = self.host_failure(py, error, None); + self.resume_core(py, Some(Err(failure)))? + } + }, + (Some(PendingOperation::Host(phase)), Some(result)) => { + let result = + result.and_then(|value| self.route.state_mut().accept(py, phase, value)); + let result = match result { + Ok(()) => Ok(R::lifecycle_result()), + Err(error) => Err(self.host_failure(py, error, Some(phase))), + }; + self.resume_core(py, Some(result))? + } + _ => return Err(missing_state()), + }; + loop { + let operation = match step { + HostStep::Suspend(awaitable) => return Ok(ExecutionStep::Await(awaitable)), + HostStep::Ready(NativeCallStep::Complete(_)) => { + return self + .route + .state_mut() + .response + .take() + .map(ExecutionStep::Return) + .ok_or_else(missing_state); + } + HostStep::Ready(NativeCallStep::Host(operation)) => operation, + }; + let phase = match R::classify(&operation) { + OperationClass::Phase(phase) => Some(phase), + OperationClass::Route => None, + }; + let result = match phase { + Some(phase) => match self.route.state_mut().invoke(py, phase) { + Ok(HostStep::Suspend(awaitable)) => { + self.pending = Some(PendingOperation::Host(phase)); + return Ok(ExecutionStep::Await(awaitable)); + } + Ok(HostStep::Ready(value)) => self + .route + .state_mut() + .accept(py, phase, value) + .map(|()| R::lifecycle_result()), + Err(error) => Err(error), + }, + None => self.route.invoke(py, operation), + }; + let result = match result { + Ok(result) => Ok(result), + Err(error) => Err(self.host_failure(py, error, phase)), + }; + step = self.resume_core(py, Some(result))?; + } + } +} + +impl ExecutionBody for PythonLifecycle { + fn resume(&mut self, result: Option>>) -> PyResult { + let result = Python::attach(|py| self.drive(py, result)); + match result { + Ok(ExecutionStep::Await(value)) => Ok(ExecutionStep::Await(value)), + result => result.map_err(|error| { + Python::attach(|py| { + self.route + .state_mut() + .error + .take() + .map(|value| PyErr::from_value(value.into_bound(py).into_any())) + .unwrap_or(error) + }) + }), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.route.state().traverse(visit)?; + self.route.traverse(visit) + } +} + +impl PythonLifecycle { + fn clear(&mut self) { + if let Some(abort) = self.native_abort.take() { + abort.abort(); + } + if self.call.take().is_some() { + Python::attach(|py| self.route.state_mut().cleanup(py)); + self.route.cleanup(); + } + } +} + +impl Drop for PythonLifecycle { + fn drop(&mut self) { + self.clear(); + } +} + +pub(crate) struct PythonCallState { + pub args: Py, + pub kwargs: Py, + pub logger: Option, + pub start: Py, + pub end: Option>, + pub response: Option>, + pub error: Option>, + pub asynchronous: bool, + pub internal: bool, + pub call_type: &'static str, +} + +pub(crate) fn now(py: Python<'_>) -> PyResult> { + py.import("datetime")? + .getattr("datetime")? + .call_method0("now") + .map(Bound::unbind) +} + +impl PythonCallState { + fn invoke( + &mut self, + py: Python<'_>, + phase: HostPhase, + ) -> PyResult, Py>> { + match phase { + HostPhase::Setup => self.setup(py)?, + HostPhase::DeploymentPreCall => { + if !DeploymentHooks::needed(py)? { + return Ok(HostStep::Ready(self.kwargs.clone_ref(py).into_any())); + } + return Ok(HostStep::Suspend(DeploymentHooks::before_call( + py, + &self.kwargs, + self.call_type, + )?)); + } + HostPhase::Prepare => self.prepare(py)?, + HostPhase::DeploymentPostCall => { + if !DeploymentHooks::needed(py)? { + return self + .response + .as_ref() + .map(|value| HostStep::Ready(value.clone_ref(py))) + .ok_or_else(missing_state); + } + return Ok(HostStep::Suspend(DeploymentHooks::after_success( + py, + &self.kwargs, + &self.response, + self.call_type, + )?)); + } + HostPhase::Finalize => self.finalize(py)?, + HostPhase::Success => self.dispatch_success(py)?, + HostPhase::DeploymentFailure => { + if let Some(error) = &self.error + && DeploymentHooks::needed(py)? + { + return Ok(HostStep::Suspend(DeploymentHooks::after_failure( + py, + &self.kwargs, + error, + self.call_type, + )?)); + } + } + HostPhase::Failure | HostPhase::AsyncFailure => { + if let Some(awaitable) = + self.dispatch_failure(py, phase == HostPhase::AsyncFailure)? + { + return Ok(HostStep::Suspend(awaitable)); + } + } + HostPhase::Execute + | HostPhase::ConstructResponse + | HostPhase::MapFailure + | HostPhase::Complete => return Err(missing_state()), + } + Ok(HostStep::Ready(py.None())) + } + + fn accept(&mut self, py: Python<'_>, phase: HostPhase, value: Py) -> PyResult<()> { + match phase { + HostPhase::DeploymentPreCall => { + self.kwargs = value.into_bound(py).cast_into::()?.unbind() + } + HostPhase::DeploymentPostCall => self.response = Some(value), + _ => {} + } + Ok(()) + } + + pub fn new( + py: Python<'_>, + args: Py, + kwargs: Py, + asynchronous: bool, + call_type: &'static str, + ) -> PyResult { + Ok(Self { + args, + kwargs, + logger: None, + start: py.None(), + end: None, + response: None, + error: None, + asynchronous, + internal: false, + call_type, + }) + } + + pub fn logger(&self) -> PyResult<&PythonLogger> { + self.logger.as_ref().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized") + }) + } + + pub fn setup(&mut self, py: Python<'_>) -> PyResult<()> { + self.start = now(py)?; + self.internal = bindings::is_internal_call(py)?; + let result = bindings::setup( + py, + self.call_type, + &self.args, + &self.kwargs, + &self.start, + self.asynchronous, + )?; + self.logger = Some(result.logger()?); + self.kwargs = result.kwargs()?; + Ok(()) + } + + pub fn prepare(&mut self, py: Python<'_>) -> PyResult<()> { + self.kwargs = preparation::prepare(py, self.kwargs.bind(py), self.logger()?)?.unbind(); + Ok(()) + } + + pub fn finalize(&self, py: Python<'_>) -> PyResult<()> { + bindings::finalize( + py, + &self.response, + self.logger()?, + &self.kwargs, + &self.start, + &self.end, + ) + } + + pub fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> { + match self.try_dispatch_success(py) { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, self.logger.as_ref().map(|logger| logger.object(py))); + Ok(()) + } + result => result, + } + } + + fn try_dispatch_success(&self, py: Python<'_>) -> PyResult<()> { + let logger = self.logger()?; + let pending = || PendingSuccess { + logger: logger.clone_ref(py), + response: self.response.as_ref().map(|value| value.clone_ref(py)), + start: self.start.clone_ref(py), + end: self.end.as_ref().map(|value| value.clone_ref(py)), + }; + if !self.asynchronous { + if !logger.callbacks_needed(py, "sync_success")? { + return logger.success_bookkeeping( + py, + &self.response, + &self.start, + &self.end, + false, + ); + } + pending().sync(py) + } else { + if !self.internal + && self + .kwargs + .bind(py) + .get_item("fallbacks")? + .is_none_or(|value| value.is_none()) + { + if !logger.callbacks_needed(py, "async_success")? { + logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; + } else if logger.defers_async_logging(py) { + logger.defer_success( + py, + Py::new( + py, + PendingLogging { + pending: Some(pending()), + }, + )?, + )?; + } else { + pending().asynchronous(py)?; + } + } + logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } + } + + pub fn dispatch_failure( + &self, + py: Python<'_>, + asynchronous: bool, + ) -> PyResult>> { + if self.logger.is_none() || (self.asynchronous && self.internal) { + return Ok(None); + } + let Some(error) = &self.error else { + return Ok(None); + }; + self.logger()? + .failure(py, error, &self.start, &self.end, asynchronous) + } + + pub fn cleanup(&mut self, py: Python<'_>) { + if let Some(logger) = self.logger.take() + && let Err(error) = logger.restore_context(py) + { + error.write_unraisable(py, None); + } + } + + pub fn retain_error(&mut self, py: Python<'_>, error: PyErr) { + self.error = Some(error.into_value(py)); + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.args)?; + visit.call(&self.kwargs)?; + if let Some(logger) = &self.logger { + logger.traverse(visit)?; + } + visit.call(&self.start)?; + visit.call(&self.end)?; + visit.call(&self.response)?; + visit.call(&self.error) + } +} + +struct PendingSuccess { + logger: PythonLogger, + response: Option>, + start: Py, + end: Option>, +} + +impl PendingSuccess { + fn sync(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .submit_success(py, &self.response, &self.start, &self.end) + } + + fn asynchronous(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .enqueue_success(py, &self.response, &self.start, &self.end) + } +} + +#[pyclass] +struct PendingLogging { + pending: Option, +} + +#[pymethods] +impl PendingLogging { + fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { + let pending = slf.borrow_mut().pending.take(); + if let Some(pending) = pending + && success + { + match pending.asynchronous(py) { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(pending.logger.object(py))); + } + result => return result, + } + } + Ok(()) + } + + fn __traverse__(&self, visit: pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { + if let Some(pending) = &self.pending { + pending.logger.traverse(&visit)?; + visit.call(&pending.response)?; + visit.call(&pending.start)?; + visit.call(&pending.end)?; + } + Ok(()) + } + + fn __clear__(slf: &Bound<'_, Self>) { + let pending = slf.borrow_mut().pending.take(); + drop(pending); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::types::PyDict; + use std::sync::Mutex; + + static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); + + fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> { + py.import("litellm.litellm_core_utils.logging_worker")? + .setattr("GLOBAL_LOGGING_WORKER", worker) + } + + struct RetainingHost { + retained: Option>, + } + + impl ExecutionBody for RetainingHost { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| Ok(ExecutionStep::Return(py.None()))) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.retained) + } + } + + #[pyfunction] + fn retaining_coroutine(py: Python<'_>, retained: Py) -> PyResult> { + Py::new( + py, + Execution::new(RetainingHost { + retained: Some(retained), + }), + ) + } + + struct AwaitBody(Option>); + + impl ExecutionBody for AwaitBody { + fn resume(&mut self, result: Option>>) -> PyResult { + match self.0.take() { + Some(awaitable) => Ok(ExecutionStep::Await(awaitable)), + None => result + .expect("selected await completed") + .map(ExecutionStep::Return), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + } + + #[pyfunction] + fn await_execution(awaitable: Py) -> Execution { + Execution::new(AwaitBody(Some(awaitable))) + } + + struct CallingBody(Py); + + impl ExecutionBody for CallingBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + } + + #[pyfunction] + fn calling_execution(callback: Py) -> Execution { + Execution::new(CallingBody(callback)) + } + + struct SyntheticCall(bool); + + impl NativeCall for SyntheticCall { + type Operation = (); + type Result = (); + type Complete = (); + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(async move { + match (self.0, result) { + (false, None) => { + self.0 = true; + Ok(NativeCallStep::Host(())) + } + (true, Some(())) => Ok(NativeCallStep::Complete(())), + _ => Err(litellm_core::Error::InvalidRequest( + "invalid synthetic lifecycle state".into(), + )), + } + }) + } + + fn interrupt( + &mut self, + _: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(async { Ok(NativeCallStep::Complete(())) }) + } + } + + struct SyntheticRoute(PythonCallState); + + impl PythonRoute for SyntheticRoute { + type Call = SyntheticCall; + + fn state(&self) -> &PythonCallState { + &self.0 + } + + fn state_mut(&mut self) -> &mut PythonCallState { + &mut self.0 + } + + fn classify(_: &()) -> OperationClass { + OperationClass::Route + } + + fn lifecycle_result() {} + + fn map_error(error: litellm_core::Error) -> PyErr { + crate::errors::core_error_to_pyerr(error) + } + + fn invoke(&mut self, py: Python<'_>, _: ()) -> PyResult<()> { + self.0.response = Some( + pyo3::types::PyString::new(py, "shared lifecycle") + .into_any() + .unbind(), + ); + Ok(()) + } + + fn cleanup(&mut self) {} + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } + } + + #[test] + fn shared_runner_executes_a_non_ocr_adapter() { + Python::initialize(); + Python::attach(|py| { + let route = SyntheticRoute( + PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + false, + "synthetic", + ) + .unwrap(), + ); + let value: String = run_call(py, SyntheticCall(false), route) + .unwrap() + .extract(py) + .unwrap(); + assert_eq!(value, "shared lifecycle"); + }); + } + + #[test] + fn ready_native_lifecycle_completes_without_scheduling() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap(); + let route = SyntheticRoute( + PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + true, + "synthetic", + ) + .unwrap(), + ); + let coroutine = run_call(py, SyntheticCall(false), route).unwrap(); + let completed = coroutine + .call_method1(py, "send", (py.None(),)) + .unwrap_err(); + assert!(completed.is_instance_of::(py)); + assert_eq!( + completed + .value(py) + .getattr("value") + .unwrap() + .extract::() + .unwrap(), + "shared lifecycle", + ); + }); + } + + #[test] + fn python_driver_preserves_inline_await_and_native_ownership() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + py.import("asyncio").unwrap(); + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + let module = PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap(); + let locals = PyDict::new(py); + locals + .set_item("drive", module.getattr("drive").unwrap()) + .unwrap(); + locals + .set_item( + "await_execution", + wrap_pyfunction!(await_execution, py).unwrap(), + ) + .unwrap(); + locals + .set_item( + "calling_execution", + wrap_pyfunction!(calling_execution, py).unwrap(), + ) + .unwrap(); + let probe = std::ffi::CString::new(include_str!("../../tests/lifecycle.py")).unwrap(); + py.run(&probe, Some(&locals), Some(&locals)).unwrap(); + }); + } + + struct ErrorBody(PythonCallState); + + impl ExecutionBody for ErrorBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| { + Err(PyErr::from_value( + self.0.error.take().unwrap().into_bound(py).into_any(), + )) + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.0.traverse(visit) + } + } + + #[pyfunction] + fn error_execution(py: Python<'_>, error: Bound<'_, PyBaseException>) -> Execution { + let mut state = PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + true, + "test", + ) + .unwrap(); + state.retain_error(py, PyErr::from_value(error.into_any())); + Execution::new(ErrorBody(state)) + } + + #[test] + fn retained_exception_frames_and_duplicate_argument_edges_are_collectable() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "error_execution", + wrap_pyfunction!(error_execution, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained: + pass + +def cycle(): + retained = Retained() + try: + raise ValueError('retained traceback') + except ValueError as error: + retained.owner = error_execution(error) + return weakref.ref(retained) + +reference = cycle() +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + fn state( + py: Python<'_>, + logger: Py, + response: Py, + asynchronous: bool, + ) -> PythonCallState { + PythonCallState { + args: PyTuple::empty(py).unbind(), + kwargs: PyDict::new(py).unbind(), + logger: Some(logger.extract(py).unwrap()), + start: py.None(), + end: Some(py.None()), + response: Some(response), + error: None, + asynchronous, + internal: false, + call_type: "test", + } + } + + #[test] + fn success_dispatch_reports_ordinary_failures_without_replacing_response() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +import sys + +response = object() +failure = ValueError('terminal diagnostic') +diagnostics = [] +old_hook = sys.unraisablehook +sys.unraisablehook = lambda event: diagnostics.append(event.exc_value) + +class Logger: + def handle_sync_success_callbacks_for_async_calls(self, *args): + raise failure + +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let response = locals.get_item("response").unwrap().unwrap().unbind(); + let mut lifecycle_state = state( + py, + locals.get_item("logger").unwrap().unwrap().unbind(), + response.clone_ref(py), + true, + ); + lifecycle_state.internal = true; + lifecycle_state.dispatch_success(py).unwrap(); + assert!(lifecycle_state.response.as_ref().unwrap().is(&response)); + py.run( + pyo3::ffi::c_str!( + r#" +assert diagnostics == [failure] +sys.unraisablehook = old_hook +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn retained_failure_preserves_exception_identity() { + Python::initialize(); + Python::attach(|py| { + let logger = PyDict::new(py).into_any().unbind(); + let response = py.None(); + let failure = pyo3::exceptions::PyValueError::new_err("identity"); + let failure_value = failure.value(py).clone().unbind(); + let mut lifecycle_state = state(py, logger, response, false); + lifecycle_state.retain_error(py, failure); + let retained = lifecycle_state.error.take().unwrap(); + assert!(retained.is(&failure_value)); + }); + } + + #[test] + fn deferred_release_uses_release_context_and_allows_reentry_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +import sys +import types +from contextvars import ContextVar + +litellm = types.ModuleType('litellm') +core_utils = types.ModuleType('litellm.litellm_core_utils') +logging_worker = types.ModuleType('litellm.litellm_core_utils.logging_worker') +litellm.litellm_core_utils = core_utils +core_utils.logging_worker = logging_worker +sys.modules['litellm'] = litellm +sys.modules['litellm.litellm_core_utils'] = core_utils +sys.modules['litellm.litellm_core_utils.logging_worker'] = logging_worker + +marker = ContextVar('marker', default='unset') +observed = [] + +class Coroutine: + def close(self): + observed.append('closed') + +class Worker: + def ensure_initialized_and_enqueue(self, coroutine): + observed.append(marker.get()) + pending.release(True) + coroutine.close() + +class Logger: + def async_success_handler(self, *args): + observed.append('created') + return Coroutine() + +worker = Worker() +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + install_logging_worker(py, &locals.get_item("worker").unwrap().unwrap()).unwrap(); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(), + response: Some(py.None()), + start: py.None(), + end: Some(py.None()), + }), + }, + ) + .unwrap(); + locals.set_item("pending", &pending).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['created', 'release', 'closed'] +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn deferred_logging_collects_cycles_through_typed_logger() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!("class Logger: pass\nlogger = Logger()"), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(), + response: None, + start: py.None(), + end: None, + }), + }, + ) + .unwrap(); + locals.set_item("pending", pending).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref +logger.pending = pending +reference = weakref.ref(logger) +del logger, pending +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn coroutine_collects_cycles_retained_by_bridge_host() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "retaining_coroutine", + wrap_pyfunction!(retaining_coroutine, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained: + pass + +def cycle(): + retained = Retained() + coroutine = retaining_coroutine(retained) + retained.coroutine = coroutine + return weakref.ref(retained) + +retained_ref = cycle() +gc.collect() +assert retained_ref() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs b/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs new file mode 100644 index 00000000000..ba4a8bb3739 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs @@ -0,0 +1,314 @@ +use litellm_core::auth::{credential_default_fields, credential_index}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList}; + +struct CredentialEntry<'py>(Bound<'py, PyAny>); + +impl<'py> CredentialEntry<'py> { + fn name(&self) -> PyResult { + self.0.getattr("credential_name")?.extract() + } + + fn values(&self) -> PyResult> { + Ok(self.0.getattr("credential_values")?.cast_into::()?) + } +} + +pub(super) fn prepare<'py>( + py: Python<'py>, + kwargs: &Bound<'py, PyDict>, + logger: &super::PythonLogger, +) -> PyResult> { + let arguments = kwargs.copy()?; + arguments.set_item("litellm_logging_obj", logger.object(py))?; + let litellm = py.import("litellm")?; + inherit_credentials(py, &litellm, &arguments)?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("check_limits")? + .call1((&arguments,))?; + Ok(arguments) +} + +fn inherit_credentials( + py: Python<'_>, + litellm: &Bound<'_, PyModule>, + arguments: &Bound<'_, PyDict>, +) -> PyResult<()> { + let Some(requested) = arguments + .get_item("litellm_credential_name")? + .filter(|value| !value.is_none()) + else { + return Ok(()); + }; + if !requested.is_truthy()? { + return Ok(()); + } + let requested: String = requested.extract()?; + let credentials = litellm.getattr("credential_list")?.cast_into::()?; + let names = credentials + .iter() + .map(|credential| CredentialEntry(credential).name()) + .collect::>>()?; + let Some(index) = credential_index(&requested, &names) else { + py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1( + "warning", + ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()), + )?; + return Ok(()); + }; + let selected = CredentialEntry(credentials.get_item(index)?); + let values = selected.values()?; + let supplied: Vec = arguments.keys().extract()?; + let fields: Vec = values.keys().extract()?; + for name in credential_default_fields(&supplied, &fields) { + if let Some(value) = values.get_item(name)? { + arguments.set_item(name, value)?; + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> { + let litellm = PyModule::new(py, "credential_host")?; + litellm.setattr( + "credential_list", + locals.get_item("credentials").unwrap().unwrap(), + )?; + inherit_credentials( + py, + &litellm, + &locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::()?, + ) + } + + #[test] + fn duplicate_names_select_the_first_entry_without_reading_other_values() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +accesses = [] +class Credential: + def __init__(self, name, values): + self._name = name + self._values = values + @property + def credential_name(self): + accesses.append(('name', self._name)) + return self._name + @property + def credential_values(self): + accesses.append(('values', self._name)) + return self._values +credentials = [ + Credential('ocr-test', {'api_key': 'first'}), + Credential('other', {'api_key': 'unused'}), + Credential('ocr-test', {'api_key': 'later'}), +] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + inherit(py, &locals).unwrap(); + let arguments = locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert_eq!( + arguments + .get_item("api_key") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "first" + ); + let accesses: Vec<(String, String)> = locals + .get_item("accesses") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!( + accesses, + [ + ("name".into(), "ocr-test".into()), + ("name".into(), "other".into()), + ("name".into(), "ocr-test".into()), + ("values".into(), "ocr-test".into()), + ] + ); + }); + } + + #[test] + fn later_invalid_name_still_fails_after_an_earlier_match() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +failure = LookupError('later name') +class Good: + credential_name = 'ocr-test' + credential_values = {'api_key': 'first'} +class Bad: + @property + def credential_name(self): + raise failure +credentials = [Good(), Bad()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + let error = inherit(py, &locals).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn selected_values_must_be_a_dictionary_and_property_errors_keep_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Listed: + credential_name = 'ocr-test' + credential_values = ['not-a-dict'] +credentials = [Listed()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + assert!( + inherit(py, &locals) + .unwrap_err() + .is_instance_of::(py) + ); + + let locals = eval( + py, + c" +failure = RuntimeError('values failed') +class Broken: + credential_name = 'ocr-test' + @property + def credential_values(self): + raise failure +credentials = [Broken()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + let error = inherit(py, &locals).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn explicit_none_is_not_overwritten_and_inherited_objects_keep_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +opaque = object() +class Credential: + credential_name = 'ocr-test' + credential_values = {'api_key': 'credential-key', 'opaque': opaque} +credentials = [Credential()] +arguments = {'litellm_credential_name': 'ocr-test', 'api_key': None} +", + ); + inherit(py, &locals).unwrap(); + let arguments = locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert!(arguments.get_item("api_key").unwrap().unwrap().is_none()); + assert!( + arguments + .get_item("opaque") + .unwrap() + .unwrap() + .is(locals.get_item("opaque").unwrap().unwrap()) + ); + }); + } + + #[test] + fn selection_rereads_the_list_after_name_properties_run() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class First: + @property + def credential_name(self): + credentials[0] = Second() + return 'ocr-test' + credential_values = {'api_key': 'first'} +class Second: + credential_name = 'ocr-test' + credential_values = {'api_key': 'replaced'} +credentials = [First()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + inherit(py, &locals).unwrap(); + let arguments = locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert_eq!( + arguments + .get_item("api_key") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "replaced" + ); + }); + } + + #[test] + fn falsy_credential_names_return_before_loading_credentials() { + Python::initialize(); + Python::attach(|py| { + let litellm = PyModule::new(py, "credential_host").unwrap(); + for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] { + let arguments = PyDict::new(py); + arguments.set_item("litellm_credential_name", name).unwrap(); + inherit_credentials(py, &litellm, &arguments).unwrap(); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index a14e4b55d82..5f7633a64a0 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -1,10 +1,14 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::time::Duration; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::types::PyDict; use serde_json::{Map, Value}; +use litellm_core::auth::InputSource; +use litellm_python_interop::from_py_preserving_errors as from_py; + pub(crate) struct RouteOptions { pub(crate) model: String, pub(crate) api_key: Option, @@ -36,18 +40,18 @@ impl RouteOptions { } } -pub(crate) fn required_value( - name: &'static str, - value: Value, - expected: fn(&Value) -> bool, - expected_name: &'static str, -) -> PyResult { - if expected(&value) { - return Ok(value); +pub(crate) fn required_array(name: &'static str, value: Value) -> PyResult> { + match value { + Value::Array(values) => Ok(values), + _ => Err(PyValueError::new_err(format!("{name} must be a list"))), + } +} + +pub(crate) fn required_object(name: &'static str, value: Value) -> PyResult> { + match value { + Value::Object(values) => Ok(values), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), } - Err(PyValueError::new_err(format!( - "{name} must be a {expected_name}" - ))) } pub(crate) fn object_or_empty( @@ -55,7 +59,7 @@ pub(crate) fn object_or_empty( value: Option, ) -> PyResult> { match value { - Some(value) => object(name, value), + Some(value) => required_object(name, value), None => Ok(Map::new()), } } @@ -64,14 +68,7 @@ fn optional_object( name: &'static str, value: Option, ) -> PyResult>> { - value.map(|value| object(name, value)).transpose() -} - -fn object(name: &'static str, value: Value) -> PyResult> { - match value { - Value::Object(map) => Ok(map), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), - } + value.map(|value| required_object(name, value)).transpose() } pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { @@ -84,6 +81,72 @@ pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option }) } +pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py) -> PyResult> { + py.import("litellm.rust_bridge.timeouts")? + .getattr("timeout_to_seconds")? + .call1((timeout,))? + .extract() +} + +pub(crate) fn project_optional_fields( + kwargs: &Bound<'_, PyDict>, + names: &[&str], +) -> PyResult> { + names + .iter() + .filter_map(|name| match kwargs.get_item(name) { + Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect() +} + +struct RequestFieldSources<'py> { + body: Option>, + credentials: Option>, +} + +impl<'py> RequestFieldSources<'py> { + fn extract(proxy_request: &Bound<'py, PyAny>) -> PyResult { + let proxy_request = proxy_request.cast::()?; + + let body = proxy_request + .get_item("body_fields")? + .or(proxy_request.get_item("body")?); + + let credentials = proxy_request.get_item("credential_fields")?; + + Ok(Self { body, credentials }) + } + + fn contains(&self, name: &str) -> bool { + self.body + .as_ref() + .is_some_and(|fields| fields.contains(name).unwrap_or(false)) + || self + .credentials + .as_ref() + .is_some_and(|fields| fields.contains(name).unwrap_or(false)) + } +} + +pub(crate) fn request_input_sources<'a>( + kwargs: &Bound<'_, PyDict>, + names: impl Iterator, +) -> PyResult> { + let Some(proxy_request) = kwargs.get_item("proxy_server_request")? else { + return Ok(BTreeMap::new()); + }; + + let sources = RequestFieldSources::extract(&proxy_request)?; + + Ok(names + .filter(|name| sources.contains(name)) + .map(|name| (name.to_string(), InputSource::Request)) + .collect()) +} + pub(crate) fn marshal_headers(headers: Option) -> PyResult> { let value = match headers { Some(headers) => headers, @@ -102,3 +165,199 @@ pub(crate) fn marshal_headers(headers: Option) -> PyResult(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn sources( + py: Python<'_>, + proxy: &Bound<'_, PyAny>, + names: &[&str], + ) -> PyResult> { + let kwargs = PyDict::new(py); + kwargs.set_item("proxy_server_request", proxy)?; + request_input_sources(&kwargs, names.iter().copied()) + } + + #[test] + fn required_shapes_preserve_nested_values_and_existing_errors() { + let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]); + assert_eq!( + Value::Array(required_array("messages", nested.clone()).unwrap()), + nested + ); + + let body = json!({"model": "claude", "metadata": {"user": "1"}}); + assert_eq!( + Value::Object(required_object("body", body.clone()).unwrap()), + body + ); + + assert_eq!( + required_array("messages", json!({"role": "user"})) + .unwrap_err() + .to_string(), + "ValueError: messages must be a list" + ); + assert_eq!( + required_object("body", json!([])).unwrap_err().to_string(), + "ValueError: body must be a dict" + ); + } + + #[test] + fn optional_parameters_treat_missing_as_empty() { + assert_eq!( + object_or_empty("optional_params", None).unwrap(), + Map::new() + ); + assert_eq!( + object_or_empty("optional_params", Some(json!({"temperature": 0.2}))).unwrap(), + required_object("optional_params", json!({"temperature": 0.2})).unwrap() + ); + } + + #[test] + fn missing_none_and_empty_proxy_metadata_are_distinct() { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + assert!( + request_input_sources(&kwargs, ["api_key"].into_iter()) + .unwrap() + .is_empty() + ); + + kwargs.set_item("proxy_server_request", py.None()).unwrap(); + assert!( + request_input_sources(&kwargs, ["api_key"].into_iter()) + .unwrap_err() + .is_instance_of::(py) + ); + + kwargs + .set_item("proxy_server_request", PyDict::new(py)) + .unwrap(); + assert!( + request_input_sources(&kwargs, ["api_key"].into_iter()) + .unwrap() + .is_empty() + ); + }); + } + + #[test] + fn body_fields_win_over_body_and_explicit_none_does_not_fall_back() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +proxy = {'body_fields': ['api_key'], 'body': ['api_base']} +none_fields = {'body_fields': None, 'body': ['api_key']} +body_only = {'body': ['api_base']} +", + ); + let named = sources( + py, + &locals.get_item("proxy").unwrap().unwrap(), + &["api_key", "api_base"], + ) + .unwrap(); + assert_eq!(named.get("api_key").copied(), Some(InputSource::Request)); + assert!(!named.contains_key("api_base")); + + assert!( + sources( + py, + &locals.get_item("none_fields").unwrap().unwrap(), + &["api_key"], + ) + .unwrap() + .is_empty() + ); + + let body_only = sources( + py, + &locals.get_item("body_only").unwrap().unwrap(), + &["api_base"], + ) + .unwrap(); + assert_eq!( + body_only.get("api_base").copied(), + Some(InputSource::Request) + ); + }); + } + + #[test] + fn body_and_credential_membership_can_mark_request_fields() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Raising: + def __contains__(self, item): + raise RuntimeError('credential membership') +proxy = { + 'body_fields': ['api_key'], + 'credential_fields': Raising(), +} +credentials_only = {'credential_fields': ['extra_headers']} +erroring = {'body_fields': Raising()} +extra = {'body_fields': ['api_key', 'unused']} +", + ); + let skipped = sources( + py, + &locals.get_item("proxy").unwrap().unwrap(), + &["api_key"], + ) + .unwrap(); + assert_eq!(skipped.get("api_key").copied(), Some(InputSource::Request)); + + let credentials = sources( + py, + &locals.get_item("credentials_only").unwrap().unwrap(), + &["extra_headers"], + ) + .unwrap(); + assert_eq!( + credentials.get("extra_headers").copied(), + Some(InputSource::Request) + ); + + assert!( + sources( + py, + &locals.get_item("erroring").unwrap().unwrap(), + &["api_key"], + ) + .unwrap() + .is_empty() + ); + + let requested = sources( + py, + &locals.get_item("extra").unwrap().unwrap(), + &["api_key"], + ) + .unwrap(); + assert_eq!(requested.len(), 1); + assert_eq!( + requested.get("api_key").copied(), + Some(InputSource::Request) + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs new file mode 100644 index 00000000000..f2997ee278c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs @@ -0,0 +1,12 @@ +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs similarity index 100% rename from litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs rename to litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs new file mode 100644 index 00000000000..f2997ee278c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -0,0 +1,12 @@ +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs similarity index 95% rename from litellm-rust/crates/python-bridge/src/routes/chat_completions.rs rename to litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs index 08ab476005c..e67bfa89cc7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs @@ -9,12 +9,12 @@ use pyo3::prelude::*; use serde_json::Value; use crate::errors::chat_completions_error_to_pyerr; -use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_array}; fn prepare_chat_completions( inputs: ChatCompletionsInputs, ) -> PyResult> + Send + 'static> { - let messages = required_value("messages", inputs.messages, Value::is_array, "list")?; + let messages = required_array("messages", inputs.messages)?; let optional_params = object_or_empty("optional_params", inputs.optional_params)?; let options = RouteOptions::from_python(RouteOptionsInputs { model: inputs.model, @@ -36,7 +36,7 @@ fn prepare_chat_completions( } = options; run_chat_completions(ChatCompletionsRequest { model: &model, - messages, + messages: Value::Array(messages), optional_params, api_key: api_key.as_deref(), api_base: api_base.as_deref(), diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 97313651011..571042062f5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -389,6 +389,82 @@ mod tests { }); } + #[test] + fn missing_and_explicit_none_optional_params_share_the_next_error() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let messages = PyList::empty(py); + let headers = PyList::empty(py); + let omitted = PyDict::new(py); + omitted + .set_item("extra_headers", &headers) + .expect("kwargs should accept extra_headers"); + let explicit = PyDict::new(py); + explicit + .set_item("optional_params", py.None()) + .expect("kwargs should accept optional_params"); + explicit + .set_item("extra_headers", &headers) + .expect("kwargs should accept extra_headers"); + + let omitted_error = module + .getattr("chat_completions") + .and_then(|function| function.call(("model", &messages), Some(&omitted))) + .expect_err("omitted optional_params should reach header validation"); + let explicit_error = module + .getattr("chat_completions") + .and_then(|function| function.call(("model", &messages), Some(&explicit))) + .expect_err("None optional_params should reach header validation"); + assert_eq!( + omitted_error.to_string(), + "ValueError: extra_headers must be a dict" + ); + assert_eq!(explicit_error.to_string(), omitted_error.to_string()); + }); + } + + #[test] + fn chat_completions_decline_keeps_existing_reasons() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let decline = module + .getattr("chat_completions_decline") + .expect("decline helper should be registered"); + let empty = PyList::empty(py); + let unreadable = py + .eval(c"'nope'", None, None) + .expect("string messages should convert"); + + let unknown: Option = decline + .call1(("unknown-model", &empty)) + .and_then(|value| value.extract()) + .expect("unknown providers should decline"); + assert_eq!( + unknown.as_deref(), + Some("provider is not on the rust chat completions path") + ); + + let empty_reason: Option = decline + .call1(("anthropic/claude-sonnet-4-5", &empty)) + .and_then(|value| value.extract()) + .expect("empty lists should decline"); + assert_eq!(empty_reason.as_deref(), Some("empty message list")); + + let unreadable_reason: Option = decline + .call1(("anthropic/claude-sonnet-4-5", unreadable)) + .and_then(|value| value.extract()) + .expect("non-list messages should decline"); + assert_eq!( + unreadable_reason.as_deref(), + Some("unreadable message list") + ); + }); + } + #[test] fn generated_routes_execute_sync_and_async_contracts() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs deleted file mode 100644 index 97ff93f299a..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs +++ /dev/null @@ -1,29 +0,0 @@ -use pyo3::prelude::*; -use serde_json::Value; - -use crate::errors::core_error_to_pyerr; - -#[pyfunction] -fn gateway_messages<'py>( - py: Python<'py>, - model_alias: String, - provider_model: String, - api_base: String, - #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value, -) -> PyResult> { - let future = litellm_ai_gateway::trace_parity::messages_request( - model_alias, - provider_model, - api_base, - body, - ); - crate::execution::run_async( - py, - crate::function_trace::capture(future), - core_error_to_pyerr, - ) -} - -pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { - super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs new file mode 100644 index 00000000000..f2997ee278c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -0,0 +1,12 @@ +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages/value.rs similarity index 94% rename from litellm-rust/crates/python-bridge/src/routes/messages.rs rename to litellm-rust/crates/python-bridge/src/routes/messages/value.rs index f69b5e9251d..b741e54f0ca 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/value.rs @@ -6,12 +6,12 @@ use serde_json::Value; use std::future::Future; use crate::errors::core_error_to_pyerr; -use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, required_object}; fn prepare_messages( inputs: MessagesInputs, ) -> PyResult> + Send + 'static> { - let body = required_value("body", inputs.body, Value::is_object, "dict")?; + let body = required_object("body", inputs.body)?; let options = RouteOptions::from_python(RouteOptionsInputs { model: inputs.model, api_key: inputs.api_key, @@ -32,7 +32,7 @@ fn prepare_messages( } = options; run_messages(MessagesRequest { model: &model, - body, + body: Value::Object(body), api_key: api_key.as_deref(), api_base: api_base.as_deref(), custom_llm_provider: custom_llm_provider.as_deref(), diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 7e81f2ffe9b..97c39a5d6b3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -3,9 +3,6 @@ use pyo3::prelude::*; #[macro_use] mod definition; -#[cfg(feature = "trace-parity")] -mod gateway_messages; - mod audio_transcription; mod chat_completions; mod messages; @@ -16,6 +13,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { audio_transcription::register(module)?; messages::register(module)?; chat_completions::register(module)?; + #[cfg(feature = "trace-parity")] { let trace = PyModule::new(module.py(), "_trace")?; @@ -23,7 +21,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { audio_transcription::register_trace(&trace)?; messages::register_trace(&trace)?; chat_completions::register_trace(&trace)?; - gateway_messages::register_trace(&trace)?; module.add_submodule(&trace)?; } Ok(()) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs new file mode 100644 index 00000000000..1cbe8a179e3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -0,0 +1,161 @@ +use pyo3::exceptions::PyBaseException; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::Value; + +use litellm_core::ocr::LiteLLMOcrResponse; +use litellm_core::ocr::hooks::OcrPreCallRequest; +use litellm_python_interop::to_py_preserving_errors as to_py; + +use crate::lifecycle::PythonLogger; + +pub(super) struct OcrLoggingFields { + model: String, + custom_llm_provider: String, + optional_params: Value, +} + +impl From<&OcrPreCallRequest> for OcrLoggingFields { + fn from(request: &OcrPreCallRequest) -> Self { + Self { + model: request.model.clone(), + custom_llm_provider: request.custom_llm_provider.clone(), + optional_params: request.optional_params.clone(), + } + } +} + +impl PythonLogger { + pub(super) fn update_ocr( + &self, + py: Python<'_>, + kwargs: &Py, + pre_call: &OcrLoggingFields, + secret_fields: &[&str], + url: &str, + ) -> PyResult<()> { + let update = PyDict::new(py); + update.set_item("kwargs", redact(py, kwargs.bind(py), secret_fields)?)?; + update.set_item("model", &pre_call.model)?; + update.set_item( + "optional_params", + redact( + py, + &to_py(py, &pre_call.optional_params)? + .into_bound(py) + .cast_into::()?, + secret_fields, + )?, + )?; + let params = PyDict::new(py); + params.set_item( + "litellm_call_id", + kwargs.bind(py).get_item("litellm_call_id")?, + )?; + params.set_item("api_base", url)?; + for name in ["logger_fn", "litellm_request_debug"] { + if let Some(value) = kwargs.bind(py).get_item(name)? { + params.set_item(name, value)?; + } + } + update.set_item("litellm_params", params)?; + update.set_item("custom_llm_provider", &pre_call.custom_llm_provider)?; + self.object(py) + .call_method("update_from_kwargs", (), Some(&update))?; + Ok(()) + } + + pub(crate) fn pre_ocr( + &self, + py: Python<'_>, + api_key: &Option>, + body: &Bound<'_, PyDict>, + headers: &Bound<'_, PyDict>, + url: &str, + ) -> PyResult<()> { + let additional = PyDict::new(py); + additional.set_item("complete_input_dict", body)?; + additional.set_item("headers", headers)?; + additional.set_item("api_base", url)?; + let kwargs = PyDict::new(py); + kwargs.set_item("input", "OCR document processing")?; + kwargs.set_item("api_key", api_key)?; + kwargs.set_item("additional_args", &additional)?; + if self.callbacks_needed(py, "input")? { + self.object(py).call_method("pre_call", (), Some(&kwargs))?; + } else { + self.object(py) + .call_method("_pre_call", (), Some(&kwargs))?; + self.object(py).call_method0("record_api_call_start_time")?; + } + Ok(()) + } + + pub(crate) fn post_ocr( + &self, + py: Python<'_>, + original_response: &Value, + body: Option<&Py>, + headers: Option<&Py>, + ) -> PyResult<()> { + let additional = PyDict::new(py); + additional.set_item("complete_input_dict", body)?; + additional.set_item("headers", headers)?; + if self.callbacks_needed(py, "input")? { + let kwargs = PyDict::new(py); + kwargs.set_item("original_response", to_py(py, original_response)?)?; + kwargs.set_item("additional_args", &additional)?; + self.object(py) + .call_method("post_call", (), Some(&kwargs))?; + } else { + let response = py + .import("json")? + .call_method1("dumps", (to_py(py, original_response)?,))?; + self.object(py).call_method1( + "record_post_call", + (response, py.None(), py.None(), additional), + )?; + } + Ok(()) + } +} + +fn redact( + py: Python<'_>, + params: &Bound<'_, PyDict>, + secret_fields: &[&str], +) -> PyResult> { + let redacted = PyDict::new(py); + for (name, value) in params { + let name = name.extract::()?; + if name == "proxy_server_request" { + continue; + } + if secret_fields.contains(&name.as_str()) { + redacted.set_item(name, "****")?; + } else { + redacted.set_item(name, value)?; + } + } + Ok(redacted.unbind()) +} + +pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult> { + py.import("litellm.rust_bridge.ocr")? + .getattr("_response")? + .call1((to_py(py, response)?,)) + .map(Bound::unbind) +} + +pub(super) fn map_failure( + py: Python<'_>, + error: &Py, + request: &Bound<'_, PyAny>, + provider: &str, +) -> PyResult> { + Ok(py + .import("litellm.rust_bridge.ocr_lifecycle")? + .getattr("map_failure")? + .call1((error, request, provider))? + .extract()?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs new file mode 100644 index 00000000000..d43c2f88775 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -0,0 +1,264 @@ +use std::io::Read; +use std::path::PathBuf; + +use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedBytes; +#[cfg(test)] +use pyo3::types::PyDict; +use pyo3::types::{PyBytes, PyString}; + +use litellm_core::constants::OCR_INLINE_MAX_BYTES; +use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type}; +use litellm_python_interop::to_py_preserving_errors; + +enum FileBytes { + Python(PyBackedBytes), + Native(Vec), +} + +impl AsRef<[u8]> for FileBytes { + fn as_ref(&self) -> &[u8] { + match self { + Self::Python(bytes) => bytes, + Self::Native(bytes) => bytes, + } + } +} + +fn read_file_input( + py: Python<'_>, + file: &Bound<'_, PyAny>, +) -> PyResult<(FileBytes, Option)> { + if file.is_instance_of::() { + return Err(PyValueError::new_err( + "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", + )); + } + if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { + let path: PathBuf = file.extract()?; + let name = path + .file_name() + .map(|value| value.to_string_lossy().into_owned()); + let bytes = py + .detach(|| { + let mut bytes = Vec::new(); + std::fs::File::open(&path)? + .take(OCR_INLINE_MAX_BYTES as u64 + 1) + .read_to_end(&mut bytes)?; + Ok::<_, std::io::Error>(bytes) + }) + .map_err(|error| { + if error.kind() == std::io::ErrorKind::NotFound { + PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) + } else { + error.into() + } + })?; + return Ok((FileBytes::Native(bytes), name)); + } + if file.is_instance_of::() { + return Ok((FileBytes::Python(file.extract()?), None)); + } + let reader = file + .getattr_opt("read")? + .filter(|value| value.is_callable()); + let Some(reader) = reader else { + return Err(PyValueError::new_err(format!( + "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", + file.get_type(), + ))); + }; + let name = file + .getattr_opt("name")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()?; + let value = reader.call0()?; + let bytes = if value.is_instance_of::() { + FileBytes::Native(value.extract::()?.into_bytes()) + } else if value.is_instance_of::() { + FileBytes::Python(value.extract()?) + } else { + return Err(PyTypeError::new_err(format!( + "OCR file read must return bytes or str, got {}", + value.get_type(), + ))); + }; + Ok((bytes, name)) +} + +pub(super) struct FileDocumentInput { + bytes: FileBytes, + name: Option, + mime_type: Option, +} + +impl FromPyObject<'_, '_> for FileDocumentInput { + type Error = PyErr; + + fn extract(document: Borrowed<'_, '_, PyAny>) -> PyResult { + let py = document.py(); + let mime_type = match document.get_item("mime_type") { + Ok(value) => Some(value.extract::()?), + Err(error) if error.is_instance_of::(py) => None, + Err(error) => return Err(error), + }; + let file = document.get_item("file").map_err(|error| { + if error.is_instance_of::(py) { + PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes") + } else { + error + } + })?; + if file.is_none() { + return Err(PyValueError::new_err( + "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", + )); + } + let (bytes, name) = read_file_input(py, &file)?; + Ok(Self { + bytes, + name, + mime_type, + }) + } +} + +pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult { + py.detach(|| { + encode_file_document( + document.bytes.as_ref(), + document.name.as_deref(), + document.mime_type.as_deref(), + ) + }) + .map_err(|error| PyValueError::new_err(error.to_string())) +} + +#[pyfunction] +fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult> { + to_py_preserving_errors(py, &file_document(py, document.extract()?)?) +} + +#[pyfunction] +fn _ocr_mime_type(file_name: &str) -> String { + mime_type_for_name(file_name).into() +} + +#[pyfunction] +#[pyo3(signature = (file_content, file_name=None, content_type=None))] +fn _ocr_upload_document( + py: Python<'_>, + file_content: &Bound<'_, PyBytes>, + file_name: Option<&str>, + content_type: Option<&str>, +) -> PyResult> { + let bytes: PyBackedBytes = file_content.extract()?; + let document = py + .detach(|| { + encode_file_document( + &bytes, + None, + Some(upload_mime_type(file_name, content_type)), + ) + }) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + to_py_preserving_errors(py, &document) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?; + module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?; + module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?; + module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extraction_validates_required_file_and_optional_mime_type() { + Python::initialize(); + Python::attach(|py| { + for expression in [c"{}", c"{'file': None}"] { + let document = py.eval(expression, None, None).unwrap(); + let error = document.extract::().err().unwrap(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("must include a 'file' field")); + } + for expression in [ + c"{'file': b'abc', 'mime_type': None}", + c"{'file': b'abc', 'mime_type': 7}", + ] { + let document = py.eval(expression, None, None).unwrap(); + let error = document.extract::().err().unwrap(); + assert!(error.is_instance_of::(py)); + } + let document = py.eval(c"{'file': b'abc'}", None, None).unwrap(); + let input: FileDocumentInput = document.extract().unwrap(); + assert_eq!(input.bytes.as_ref(), b"abc"); + assert_eq!(input.name, None); + assert_eq!(input.mime_type, None); + }); + } + + #[test] + fn extraction_validates_mime_type_before_consuming_file() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"class Reader: + def __init__(self): + self.reads = 0 + def read(self): + self.reads += 1 + return b'abc' +reader = Reader() +document = {'file': reader, 'mime_type': 7}", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let document = locals.get_item("document").unwrap().unwrap(); + let error = document.extract::().err().unwrap(); + assert!(error.is_instance_of::(py)); + let reads: usize = locals + .get_item("reader") + .unwrap() + .unwrap() + .getattr("reads") + .unwrap() + .extract() + .unwrap(); + assert_eq!(reads, 0); + }); + } + + #[test] + fn extraction_preserves_reader_key_error_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"failure = KeyError('reader failed') +class Reader: + def read(self): + raise failure +document = {'file': Reader()}", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let document = locals.get_item("document").unwrap().unwrap(); + let error = document.extract::().err().unwrap(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs new file mode 100644 index 00000000000..66bdfb7583e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -0,0 +1,72 @@ +use litellm_core::error::Error; +use pyo3::prelude::*; + +use crate::errors::{RustUpstreamError, core_error_to_pyerr}; + +pub(super) fn to_pyerr(error: Error) -> PyErr { + let status = error.http_status_code(); + let mapped = match error { + Error::Http { status, body } => RustUpstreamError::new_err((status, body)), + other => core_error_to_pyerr(other), + }; + attach_status(mapped, status) +} + +fn attach_status(error: PyErr, status: Option) -> PyErr { + if let Some(status) = status { + Python::attach(|py| { + let value = error.value(py); + value.setattr("status_code", status).ok(); + value.setattr("message", value.to_string()).ok(); + }); + } + error +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyValueError; + + #[test] + fn preserves_python_validation_and_provider_details() { + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(Error::MissingDocumentUrl); + assert!(mapped.is_instance_of::(py)); + assert_eq!(mapped.value(py).to_string(), "Document URL is required"); + assert_eq!( + mapped + .value(py) + .getattr("status_code") + .unwrap() + .extract::() + .unwrap(), + 500 + ); + let mapped = to_pyerr(Error::Http { + status: 429, + body: r#"{"message":"rate limited"}"#.to_string(), + }); + assert!(mapped.is_instance_of::(py)); + let args: (u16, String) = mapped + .value(py) + .getattr("args") + .and_then(|args| args.extract()) + .expect("OCR failures retain status and unprefixed provider message"); + assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); + + let mapped = to_pyerr(Error::InvalidRequest("invalid format".into())); + assert!(mapped.is_instance_of::(py)); + assert_eq!( + mapped + .value(py) + .getattr("status_code") + .unwrap() + .extract::() + .unwrap(), + 400 + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs new file mode 100644 index 00000000000..12d902a3544 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -0,0 +1,311 @@ +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use litellm_core::auth::ResolvedCredential; +use litellm_core::ocr::hooks::{OcrDuringCallRequest, OcrPostCallRequest, OcrPreCallRequest}; +use litellm_core::ocr::{OcrAdmission, OcrCall, OcrClient, OcrHostOperation, OcrHostResult}; +use litellm_python_interop::{ + from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +}; + +use super::callbacks; +use super::errors::to_pyerr as ocr_error_to_pyerr; +use super::project::{ProjectedOcrFields, admitted_call, project_request}; +use crate::lifecycle::{ + OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call, +}; + +struct PythonOcrHost { + state: PythonCallState, + data: OcrHostData, +} + +enum OcrHostData { + Unprojected { request: Py }, + Projected(Box), + Released, +} + +struct ProjectedOcrHost { + fields: ProjectedOcrFields, + pre_call: Option, + retained_fields: Option>, + body: Option>, + headers: Option>, +} + +impl PythonOcrHost { + fn projected(&self) -> PyResult<&ProjectedOcrHost> { + match &self.data { + OcrHostData::Projected(projected) => Ok(projected), + _ => Err(missing_state()), + } + } + + fn projected_mut(&mut self) -> PyResult<&mut ProjectedOcrHost> { + match &mut self.data { + OcrHostData::Projected(projected) => Ok(projected), + _ => Err(missing_state()), + } + } + + fn pre_call( + &mut self, + py: Python<'_>, + request: OcrPreCallRequest, + ) -> PyResult { + let kwargs = self.state.kwargs.bind(py); + let retained_fields = PyDict::new(py); + for name in request + .optional_params + .as_object() + .ok_or_else(missing_state)? + .keys() + { + if let Some(value) = kwargs.get_item(name)? { + retained_fields.set_item(name, value)?; + } + } + retained_fields.set_item("document", &self.projected()?.fields.document)?; + let projected = self.projected_mut()?; + projected.retained_fields = Some(retained_fields.unbind()); + projected.pre_call = Some((&request).into()); + Ok(request) + } + + fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult { + let provider = self + .projected()? + .fields + .azure_ad_token_provider + .as_ref() + .ok_or_else(missing_state)?; + provider.acquire(py) + } + + fn python_pre_call( + &mut self, + py: Python<'_>, + mut request: OcrDuringCallRequest, + ) -> PyResult { + let projected = self.projected()?; + let pre_call = projected.pre_call.as_ref().ok_or_else(missing_state)?; + self.state.logger()?.update_ocr( + py, + &self.state.kwargs, + pre_call, + &projected.fields.secret_fields, + &request.url, + )?; + if !self.state.logger()?.callbacks_needed(py, "payload")? { + self.state + .logger()? + .object(py) + .call_method0("record_api_call_start_time")?; + return Ok(request); + } + if let Some(body) = request.body.as_object_mut() { + for name in &request.retained_fields { + body.remove(name); + } + } + let body = to_py(py, &request.body)? + .into_bound(py) + .cast_into::()?; + if let Some(retained) = &self.projected()?.retained_fields { + for name in &request.retained_fields { + if let Some(value) = retained.bind(py).get_item(name)? { + body.set_item(name, value)?; + } + } + } + let headers = PyDict::new(py); + for (name, value) in &request.headers { + headers.set_item(name, value)?; + } + let api_key = self.projected()?.fields.api_key.clone_ref(py); + let projected = self.projected_mut()?; + projected.body = Some(body.clone().unbind()); + projected.headers = Some(headers.clone().unbind()); + self.state + .logger()? + .pre_ocr(py, &Some(api_key), &body, &headers, &request.url)?; + let headers = headers + .iter() + .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) + .collect::>>()?; + request.body = from_py(&body)?; + request.headers = headers; + Ok(request) + } + + fn python_post_call( + &mut self, + py: Python<'_>, + request: OcrPostCallRequest, + ) -> PyResult { + let logger = self.state.logger()?; + if logger.callbacks_needed(py, "payload")? { + let projected = self.projected()?; + logger.post_ocr( + py, + &request.original_response, + projected.body.as_ref(), + projected.headers.as_ref(), + )?; + } + Ok(request) + } +} + +impl PythonRoute for PythonOcrHost { + type Call = OcrCall; + + fn state(&self) -> &PythonCallState { + &self.state + } + + fn state_mut(&mut self) -> &mut PythonCallState { + &mut self.state + } + + fn classify(operation: &OcrHostOperation) -> OperationClass { + operation + .phase() + .map_or(OperationClass::Route, OperationClass::Phase) + } + + fn lifecycle_result() -> OcrHostResult { + OcrHostResult::Lifecycle(Ok(())) + } + + fn map_error(error: litellm_core::Error) -> PyErr { + ocr_error_to_pyerr(error) + } + + fn invoke(&mut self, py: Python<'_>, operation: OcrHostOperation) -> PyResult { + Ok(match operation { + OcrHostOperation::ProjectRequest => { + let OcrHostData::Unprojected { request } = &self.data else { + return Err(missing_state()); + }; + let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?; + let has_token_provider = projected.fields.azure_ad_token_provider.is_some(); + let request = projected.request; + self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost { + fields: projected.fields, + pre_call: None, + retained_fields: None, + body: None, + headers: None, + })); + OcrHostResult::Request(Ok((Box::new(request), has_token_provider))) + } + OcrHostOperation::AcquireAzureAdToken => { + OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?)) + } + OcrHostOperation::PreCall(request) => { + OcrHostResult::PreCall(Ok(self.pre_call(py, request)?)) + } + OcrHostOperation::DuringCall(request) => { + OcrHostResult::DuringCall(Ok(self.python_pre_call(py, request)?)) + } + OcrHostOperation::PostCall(request) => { + OcrHostResult::PostCall(Ok(self.python_post_call(py, request)?)) + } + OcrHostOperation::ConstructResponse(response) => { + self.state.end = Some(now(py)?); + self.state.response = Some(callbacks::response(py, response.as_ref())?); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::MapFailure(error) => { + if self.state.error.is_none() { + self.state.retain_error(py, ocr_error_to_pyerr(error)); + } + if self.state.end.is_none() { + self.state.end = Some(now(py)?); + } + let error = self.state.error.as_ref().ok_or_else(missing_state)?; + let (request, provider) = match &self.data { + OcrHostData::Unprojected { request } => (request.bind(py), ""), + OcrHostData::Projected(projected) => ( + projected.fields.boundary_request.bind(py), + projected.fields.provider, + ), + OcrHostData::Released => return Err(missing_state()), + }; + let mapped = callbacks::map_failure(py, error, request, provider)?; + self.state + .retain_error(py, PyErr::from_value(mapped.into_bound(py).into_any())); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => return Err(missing_state()), + }) + } + + fn cleanup(&mut self) { + self.data = OcrHostData::Released; + } + fn traverse(&self, visit: &pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { + match &self.data { + OcrHostData::Unprojected { request } => visit.call(request), + OcrHostData::Projected(projected) => { + visit.call(&projected.fields.boundary_request)?; + visit.call(&projected.fields.document)?; + visit.call(&projected.fields.api_key)?; + if let Some(provider) = &projected.fields.azure_ad_token_provider { + provider.traverse(visit)?; + } + visit.call(&projected.retained_fields)?; + visit.call(&projected.body)?; + visit.call(&projected.headers) + } + OcrHostData::Released => Ok(()), + } + } +} + +pub(super) struct BridgeOcrHooks; + +impl litellm_core::ocr::hooks::OcrHooks for BridgeOcrHooks { + fn intercepts_requests(&self) -> bool { + true + } +} + +#[pyfunction] +fn _ocr_lifecycle( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult> { + let client = OcrClient::shared().map_err(ocr_error_to_pyerr)?; + let call = admitted_call(OcrCall::admit( + client, + OcrAdmission { + asynchronous, + ..OcrAdmission::all() + }, + ))?; + let host = PythonOcrHost { + state: PythonCallState::new( + py, + args.unbind(), + kwargs.copy()?.unbind(), + asynchronous, + if asynchronous { "aocr" } else { "ocr" }, + )?, + data: OcrHostData::Unprojected { + request: request.unbind(), + }, + }; + run_call(py, call, host) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(_ocr_lifecycle, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs new file mode 100644 index 00000000000..10fa40b65ea --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -0,0 +1,19 @@ +mod callbacks; +mod document; +mod errors; +mod lifecycle; +mod project; +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module)?; + document::register(module)?; + lifecycle::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs new file mode 100644 index 00000000000..8b6a1b02e19 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -0,0 +1,579 @@ +use std::sync::Arc; + +use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request}; +use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall}; +use litellm_python_interop::{ + from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::{Map, Value}; + +use super::errors::to_pyerr as ocr_error_to_pyerr; +use super::lifecycle::BridgeOcrHooks; +use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider}; +use crate::errors::RustBridgeDeclined; +use crate::marshal::{project_optional_fields, python_timeout_seconds, request_input_sources}; + +pub(super) struct ProjectedOcrFields { + pub boundary_request: Py, + pub document: Py, + pub api_key: Py, + pub azure_ad_token_provider: Option, + pub provider: &'static str, + pub secret_fields: Vec<&'static str>, +} + +pub(super) struct ProjectedOcrCall { + pub request: LiteLLMOcrRequest, + pub fields: ProjectedOcrFields, +} + +struct OcrArguments<'a, 'py> { + request: &'a Bound<'py, PyAny>, + kwargs: &'a Bound<'py, PyDict>, +} + +impl<'py> OcrArguments<'_, 'py> { + fn lookup(&self, name: &str) -> PyResult> { + match self.kwargs.get_item(name)? { + Some(value) => Ok(value), + None => self.request.getattr(name), + } + } + + fn model(&self) -> PyResult { + self.lookup("model")?.extract() + } + + fn custom_llm_provider(&self) -> PyResult> { + self.lookup("custom_llm_provider")?.extract() + } + + fn document(&self) -> PyResult> { + self.lookup("document") + } + + fn api_key(&self) -> PyResult> { + self.lookup("api_key") + } + + fn api_base(&self) -> PyResult> { + self.lookup("api_base")?.extract() + } + + fn extra_headers(&self) -> PyResult>> { + self.lookup("extra_headers")? + .extract::>>()? + .map(|value| from_py(value.bind(self.request.py()))) + .transpose() + } + + fn timeout_seconds(&self) -> PyResult> { + Ok(self + .lookup("timeout")? + .extract::>>()? + .map(|value| python_timeout_seconds(self.request.py(), value)) + .transpose()? + .flatten()) + } +} + +enum ProjectedDocument { + File { wire: Value, retained: Py }, + Other { wire: Value, retained: Py }, +} + +impl ProjectedDocument { + fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult { + let kind: String = document.get_item("type")?.extract()?; + if kind != "file" { + return Ok(Self::Other { + wire: from_py(document)?, + retained: document.clone().unbind(), + }); + } + let input = document.extract()?; + let encoded = super::document::file_document(py, input)?; + let wire = serde_json::to_value(encoded) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + Ok(Self::File { + retained: to_py(py, &wire)?, + wire, + }) + } + + fn into_parts(self) -> (Value, Py) { + match self { + Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained), + } + } +} + +pub(super) fn project_request( + py: Python<'_>, + request: &Bound<'_, PyAny>, + kwargs: &Bound<'_, PyDict>, +) -> PyResult { + let boundary_request = request.clone().unbind(); + let arguments = OcrArguments { request, kwargs }; + let model = arguments.model()?; + let custom_llm_provider = arguments.custom_llm_provider()?; + let (wire_document, retained_document) = + ProjectedDocument::project(py, &arguments.document()?)?.into_parts(); + let api_key = arguments.api_key()?; + let specs = consumed_optional_params(&model, custom_llm_provider.as_deref()) + .map_err(ocr_error_to_pyerr)?; + let names = specs.iter().map(|spec| spec.name).collect::>(); + let optional_params = project_optional_fields(kwargs, &names)?; + let input_sources = request_input_sources( + kwargs, + names + .iter() + .copied() + .chain(["api_key", "api_base", "extra_headers"]), + )?; + let azure_ad_token_provider = kwargs + .get_item("azure_ad_token_provider")? + .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER)); + let wire = OcrWireRequest { + model, + document: wire_document, + api_key: api_key.extract()?, + api_base: arguments.api_base()?, + custom_llm_provider, + extra_headers: arguments.extra_headers()?, + optional_params, + input_sources, + timeout_seconds: arguments.timeout_seconds()?, + }; + let request = decode_request(wire).map_err(ocr_error_to_pyerr)?; + let provider = request.provider_name(); + Ok(ProjectedOcrCall { + request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None), + fields: ProjectedOcrFields { + boundary_request, + document: retained_document, + api_key: api_key.unbind(), + azure_ad_token_provider, + provider, + secret_fields: specs + .into_iter() + .filter(|spec| spec.secret) + .map(|spec| spec.name) + .collect(), + }, + }) +} + +pub(super) fn admitted_call(outcome: NativeOutcome) -> PyResult { + match outcome { + NativeOutcome::Completed(call) => Ok(call), + NativeOutcome::Declined(reason) => Err(RustBridgeDeclined::new_err(format!( + "native OCR admission declined: {reason:?}" + ))), + } +} + +#[cfg(test)] +mod tests { + use litellm_core::Error; + use litellm_core::ocr::OcrDecline; + use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError}; + + use super::*; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn arguments<'a, 'py>( + request: &'a Bound<'py, PyAny>, + kwargs: &'a Bound<'py, PyDict>, + ) -> OcrArguments<'a, 'py> { + OcrArguments { request, kwargs } + } + + fn project_document( + py: Python<'_>, + document: &Bound<'_, PyAny>, + ) -> PyResult<(Value, Py)> { + ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts) + } + + fn stub_timeout_conversion(py: Python<'_>) { + eval( + py, + c" +import sys +import types +timeouts = types.ModuleType('litellm.rust_bridge.timeouts') +timeouts.timeout_to_seconds = lambda timeout: None if timeout is None else float(timeout) +sys.modules.setdefault('litellm', types.ModuleType('litellm')) +sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) +sys.modules['litellm.rust_bridge.timeouts'] = timeouts +", + ); + } + + #[test] + fn typed_initial_decline_uses_bridge_decline_contract() { + Python::initialize(); + Python::attach(|py| { + let Err(error) = admitted_call(NativeOutcome::Declined(OcrDecline::HostOperations)) + else { + panic!("unsupported host operations should decline admission"); + }; + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + fn post_admission_error_does_not_use_bridge_decline_contract() { + Python::initialize(); + Python::attach(|py| { + let error = ocr_error_to_pyerr(Error::InvalidRequest("callback result".into())); + assert!(error.is_instance_of::(py)); + assert!(!error.is_instance_of::(py)); + }); + } + + #[test] + fn kwargs_override_request_attributes_including_explicit_none() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Request: + def __init__(self): + self.accesses = [] + def __getattribute__(self, name): + if name != 'accesses': + object.__getattribute__(self, 'accesses').append(name) + return object.__getattribute__(self, name) +request = Request() +request.model = 'from-request' +request.custom_llm_provider = 'mistral' +kwargs = {'model': 'from-kwargs', 'custom_llm_provider': None} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let arguments = arguments(&request, &kwargs); + assert_eq!(arguments.model().unwrap(), "from-kwargs"); + assert_eq!(arguments.custom_llm_provider().unwrap(), None); + let accesses: Vec = request.getattr("accesses").unwrap().extract().unwrap(); + assert_eq!(accesses, Vec::::new()); + }); + } + + #[test] + fn missing_kwargs_read_the_request_property_once() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Request: + def __init__(self): + self.reads = 0 + @property + def model(self): + self.reads += 1 + return 'mistral-ocr-latest' +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert_eq!( + arguments(&request, &kwargs).model().unwrap(), + "mistral-ocr-latest" + ); + assert_eq!( + request.getattr("reads").unwrap().extract::().unwrap(), + 1 + ); + }); + } + + #[test] + fn request_property_exceptions_keep_their_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +failure = LookupError('model failed') +class Request: + @property + def model(self): + raise failure +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let error = arguments(&request, &kwargs).model().unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn unused_raising_property_is_never_inspected() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Request: + @property + def unused(self): + raise RuntimeError('unused') + model = 'mistral-ocr-latest' + custom_llm_provider = None +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let arguments = arguments(&request, &kwargs); + assert_eq!(arguments.model().unwrap(), "mistral-ocr-latest"); + assert_eq!(arguments.custom_llm_provider().unwrap(), None); + }); + } + + #[test] + fn document_reader_mutations_are_visible_to_later_field_reads() { + Python::initialize(); + Python::attach(|py| { + stub_timeout_conversion(py); + let locals = eval( + py, + c" +class Request: + api_base = 'original' + timeout = 1 + @property + def document(self): + return document +class Reader: + def read(self): + Request.api_base = 'mutated' + Request.timeout = 9 + return b'abc' +document = {'type': 'file', 'file': Reader()} +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let arguments = arguments(&request, &kwargs); + let document = arguments.document().unwrap(); + project_document(py, &document).unwrap(); + assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated")); + assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0)); + }); + } + + #[test] + fn captured_api_key_keeps_the_original_python_object() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +key = object() +class Request: + api_key = None +request = Request() +kwargs = {'api_key': key} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let captured = arguments(&request, &kwargs).api_key().unwrap(); + assert!( + captured + .unbind() + .bind(py) + .is(locals.get_item("key").unwrap().unwrap()) + ); + }); + } + + #[test] + fn file_documents_are_encoded_and_other_documents_keep_the_python_object() { + Python::initialize(); + Python::attach(|py| { + let file = py + .eval( + c"{'type': 'file', 'file': b'%PDF-1.4', 'mime_type': 'application/pdf'}", + None, + None, + ) + .unwrap(); + assert_eq!( + project_document(py, &file).unwrap().0, + serde_json::json!({ + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + }) + ); + + let original = py + .eval( + c"{'type': 'document_url', 'document_url': 'https://example.com/a.pdf'}", + None, + None, + ) + .unwrap(); + let (wire, retained) = project_document(py, &original).unwrap(); + assert_eq!( + wire, + serde_json::json!({ + "type": "document_url", + "document_url": "https://example.com/a.pdf", + }) + ); + assert!(retained.bind(py).is(&original)); + }); + } + + #[test] + fn unknown_document_types_reach_existing_downstream_validation() { + Python::initialize(); + Python::attach(|py| { + let document = py + .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None) + .unwrap(); + let wire_document = project_document(py, &document).unwrap().0; + assert_eq!( + wire_document, + serde_json::json!({"type": "mystery", "mystery": "x"}) + ); + let error = match decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: wire_document, + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Map::new(), + input_sources: Default::default(), + timeout_seconds: None, + }) { + Ok(_) => panic!("unknown discriminators belong to core validation"), + Err(error) => error, + }; + assert!(error.to_string().contains("document")); + }); + } + + #[test] + fn document_discriminator_errors_keep_their_existing_exceptions() { + Python::initialize(); + Python::attach(|py| { + let missing = py.eval(c"{}", None, None).unwrap(); + assert!( + project_document(py, &missing) + .unwrap_err() + .is_instance_of::(py) + ); + + let non_string = py.eval(c"{'type': 1}", None, None).unwrap(); + assert!( + project_document(py, &non_string) + .unwrap_err() + .is_instance_of::(py) + ); + + let locals = eval( + py, + c" +failure = RuntimeError('type lookup failed') +class Document: + def __getitem__(self, key): + raise failure +document = Document() +", + ); + let error = + project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn document_classification_happens_once() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Document(dict): + def __init__(self): + super().__init__({'file': b'abc'}) + self.reads = [] + def __getitem__(self, key): + self.reads.append(key) + if key == 'type': + return 'file' if self.reads.count('type') == 1 else 'document_url' + return super().__getitem__(key) +document = Document() +", + ); + let document = locals.get_item("document").unwrap().unwrap(); + let (wire, retained) = project_document(py, &document).unwrap(); + assert_eq!(wire["type"], "document_url"); + assert!(!retained.bind(py).is(&document)); + let reads: Vec = document.getattr("reads").unwrap().extract().unwrap(); + assert_eq!(reads, ["type", "mime_type", "file"]); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs similarity index 52% rename from litellm-rust/crates/python-bridge/src/routes/ocr.rs rename to litellm-rust/crates/python-bridge/src/routes/ocr/value.rs index c5def64c2f1..051ac19d4fb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs @@ -1,12 +1,11 @@ use litellm_core::Error; use std::future::Future; -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; -use litellm_core::ocr::wire::{OcrWireRequest, decode_request, is_supported_request}; +use litellm_core::ocr::wire::{OcrWireRequest, decode_request}; use pyo3::prelude::*; use serde_json::Value; -use crate::errors::ocr_error_to_pyerr; +use super::errors::to_pyerr as ocr_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; fn prepare_ocr( @@ -38,37 +37,20 @@ fn prepare_ocr( extra_headers, timeout, } = options; - if is_supported_request(&model, custom_llm_provider.as_deref()) { - let request = decode_request(OcrWireRequest { - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - input_sources, - timeout_seconds: timeout.map(|value| value.as_secs_f64()), - })?; - return litellm_core::ocr::ocr(request) - .await - .map(|response| response.into_json()); - } - run_ocr(OcrRequest { - model: &model, + let request = decode_request(OcrWireRequest { + model, document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), + api_key, + api_base, + custom_llm_provider, extra_headers, optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }) - .await + input_sources, + timeout_seconds: timeout.map(|value| value.as_secs_f64()), + })?; + litellm_core::ocr::ocr(request) + .await + .map(|response| response.into_json()) }) } @@ -96,22 +78,3 @@ bridge_route! { prepare = prepare_ocr, errors = ocr_error_to_pyerr, } - -#[cfg(test)] -mod tests { - use litellm_core::ocr::wire::is_supported_request; - - #[test] - fn native_activation_includes_migrated_providers() { - assert!(is_supported_request("model", Some("mistral"))); - assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); - assert!(is_supported_request( - "documentintelligence/prebuilt-read", - Some("azure_ai") - )); - assert!(is_supported_request("parse-v3", Some("reducto"))); - assert!(is_supported_request("parse-legacy", Some("reducto"))); - assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); - assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); - } -} diff --git a/litellm-rust/crates/python-bridge/tests/lifecycle.py b/litellm-rust/crates/python-bridge/tests/lifecycle.py new file mode 100644 index 00000000000..fd6742102a4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/lifecycle.py @@ -0,0 +1,186 @@ +import asyncio +import gc +import threading +import weakref +from contextvars import ContextVar + + +async def exercise(): + caller = asyncio.current_task() + thread = threading.get_ident() + loop = asyncio.get_running_loop() + marker = ContextVar("driver", default="before") + entered = asyncio.Event() + released = asyncio.Event() + result = object() + + class CustomAwaitable: + def __await__(self): + return operation().__await__() + + async def operation(): + assert asyncio.current_task() is caller + assert threading.get_ident() == thread + assert asyncio.get_running_loop() is loop + marker.set("inside") + entered.set() + await released.wait() + assert asyncio.current_task() is caller + assert marker.get() == "inside" + return result + + async def release(): + await entered.wait() + released.set() + + releaser = asyncio.create_task(release()) + execution = await_execution(CustomAwaitable()) + try: + execution.resume_value(None) + except RuntimeError: + pass + else: + raise AssertionError("resumed an unstarted execution") + wrapped = drive(execution) + try: + wrapped.send(1) + except TypeError: + pass + else: + raise AssertionError("accepted initial value") + assert await wrapped is result + assert marker.get() == "inside" + await releaser + execution.close() + execution.close() + try: + await wrapped + except RuntimeError: + pass + else: + raise AssertionError("accepted coroutine reuse") + + final_awaitable = CustomAwaitable() + assert await drive(calling_execution(lambda: final_awaitable)) is final_awaitable + + cause = KeyError("cause") + failure = ValueError("original") + + async def failing(): + await asyncio.sleep(0) + raise failure from cause + + try: + await drive(await_execution(failing())) + except ValueError as error: + assert error is failure + assert error.__cause__ is cause + names = [] + traceback = error.__traceback__ + while traceback: + names.append(traceback.tb_frame.f_code.co_name) + traceback = traceback.tb_next + assert "failing" in names + else: + raise AssertionError("lost original exception") + + for suppress in (False, True): + pending = asyncio.Event() + cleanup_entered = asyncio.Event() + cleanup_release = asyncio.Event() + cleaned = [] + + async def cancel_operation(): + try: + pending.set() + await asyncio.Event().wait() + except asyncio.CancelledError: + if suppress: + return result + raise + finally: + cleanup_entered.set() + try: + await cleanup_release.wait() + except asyncio.CancelledError: + await cleanup_release.wait() + cleaned.append(asyncio.current_task()) + + task = asyncio.create_task(drive(await_execution(cancel_operation()))) + await pending.wait() + task.cancel() + await cleanup_entered.wait() + assert not task.done() + task.cancel() + await asyncio.sleep(0) + cleanup_release.set() + if suppress: + assert await task is result + else: + try: + await task + except asyncio.CancelledError: + pass + else: + raise AssertionError("lost cancellation") + assert cleaned == [task] + + observed = [] + + def reenter(): + try: + active.start() + except RuntimeError as error: + observed.append(str(error)) + return result + + active = calling_execution(reenter) + assert await drive(active) is result + assert observed == ["execution is already running"] + + class Finalizer: + def __call__(self): + return result + + def __del__(self): + self.owner.close() + observed.append("released") + + def cycle(started): + callback = Finalizer() + execution = calling_execution(callback) + callback.owner = execution + if started: + assert execution.start().value is result + return weakref.ref(callback) + + for started in (False, True): + reference = cycle(started) + gc.collect() + assert reference() is None + assert observed[-2:] == ["released", "released"] + + class Awaitable: + def __await__(self): + try: + yield self + finally: + observed.append("unwound") + + def abandoned(started): + awaitable = Awaitable() + coroutine = drive(await_execution(awaitable)) + awaitable.owner = coroutine + if started: + assert coroutine.send(None) is awaitable + coroutine.close() + return weakref.ref(awaitable) + + for started in (False, True): + reference = abandoned(started) + gc.collect() + assert reference() is None + assert observed[-1] == "unwound" + + +asyncio.run(asyncio.wait_for(exercise(), 10)) diff --git a/litellm-rust/crates/python-interop/AGENTS.md b/litellm-rust/crates/python-interop/AGENTS.md index d1d61e5dfa0..63996d3a92b 100644 --- a/litellm-rust/crates/python-interop/AGENTS.md +++ b/litellm-rust/crates/python-interop/AGENTS.md @@ -1 +1,16 @@ -litellm-python-interop is the domain-neutral PyO3 foundation. Keep generic Python/Serde conversion and interpreter primitives here. Do not add LiteLLM domain crates, route types, API registration, or cdylib build features. +- Target invariants; implementation and runtime validation may lag these rules +- Keep this crate a small, domain-neutral foundation: Python/Serde conversion and interpreter-boundary utilities + - No LiteLLM domain dependencies, route types, callback policy, public API registration or cdylib build features + - Generic code alone does not justify extraction: runtime integration stays in `python-bridge/src/execution.rs`, host adaptation in its `lifecycle.rs` +- Use standard PyO3 ownership and conversion APIs + - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads + - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` + - Preserve `PythonizeError`'s standard conversion into `PyErr`; do not stringify original Python exceptions into new `ValueError`s + - Keep serializer-panic containment in `Pythonized`: async output conversion can run in an unjoined blocking task and otherwise strand delivery +- Use `Python::detach` for Rust-only work; Python operations require attachment + - Keep diagnostic counters in the consumer; wrapper invocations do not measure every interpreter release + - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers; expose retained Python edges to GC without calling Python during traversal +- Keep coroutine driving in the shared Python driver and native adapter + - Driver: `litellm/rust_bridge/lifecycle.py`; handle: `python-bridge/src/lifecycle.rs`; native-backed behavior tests: `python-bridge/tests/lifecycle.py` +- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html) + - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html) diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index 2e562bdae70..79af79e8c61 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -2,4 +2,6 @@ mod gil; mod marshal; pub use gil::{release_count, release_gil}; -pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py}; +pub use marshal::{ + Pythonized, from_py, from_py_preserving_errors, panic_to_pyerr, to_py, to_py_preserving_errors, +}; diff --git a/litellm-rust/crates/python-interop/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs index a16d1e0ae13..ed4cce862c0 100644 --- a/litellm-rust/crates/python-interop/src/marshal.rs +++ b/litellm-rust/crates/python-interop/src/marshal.rs @@ -14,6 +14,13 @@ where pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) } +pub fn from_py_preserving_errors(value: &Bound<'_, PyAny>) -> PyResult +where + T: DeserializeOwned, +{ + pythonize::depythonize(value).map_err(PyErr::from) +} + pub fn to_py(py: Python<'_>, value: &T) -> PyResult> where T: Serialize + ?Sized, @@ -23,6 +30,15 @@ where .map_err(|error| PyValueError::new_err(error.to_string())) } +pub fn to_py_preserving_errors(py: Python<'_>, value: &T) -> PyResult> +where + T: Serialize + ?Sized, +{ + pythonize::pythonize(py, value) + .map(Bound::unbind) + .map_err(PyErr::from) +} + pub struct Pythonized(pub T); impl<'py, T> IntoPyObject<'py> for Pythonized @@ -89,4 +105,49 @@ mod tests { assert_eq!(error.to_string(), "PanicException: serializer panicked"); }); } + + #[test] + fn depythonize_preserves_python_exception_identity_and_traceback() { + Python::initialize(); + Python::attach(|py| { + let locals = pyo3::types::PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +failure = LookupError('conversion failed') +cause = ValueError('cause') +class Broken: + def __index__(self): + raise failure from cause +value = Broken() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let value = locals.get_item("value").unwrap().unwrap(); + let legacy_error = from_py::(&value).unwrap_err(); + assert!(legacy_error.is_instance_of::(py)); + assert!( + !legacy_error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + let error = from_py_preserving_errors::(&value).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + assert!( + error + .cause(py) + .unwrap() + .value(py) + .is(locals.get_item("cause").unwrap().unwrap()) + ); + assert!(error.traceback(py).is_some()); + }); + } } diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 4058a3d72dd..56c9147e066 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -13,6 +13,7 @@ import json import sys import threading import time +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: @@ -34,6 +35,7 @@ class InMemoryCache(BaseCache): default_ttl: int | None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute max_size_per_item: int | None = 1024, # 1MB = 1024KB + clock: Callable[[], float] | None = None, ): """ max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default @@ -49,6 +51,7 @@ class InMemoryCache(BaseCache): self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() + self._clock = clock if clock is not None else lambda: time.time() def check_value_size(self, value: Any): """ @@ -91,7 +94,7 @@ class InMemoryCache(BaseCache): """ Check if a specific key is expired """ - return key in self.ttl_dict and time.time() > self.ttl_dict[key] + return key in self.ttl_dict and self._clock() > self.ttl_dict[key] def _remove_key(self, key: str) -> None: """ @@ -113,7 +116,7 @@ class InMemoryCache(BaseCache): - 3. the size of in-memory cache is bounded """ - current_time: Final = time.time() + current_time: Final = self._clock() # Step 1: Remove expired or outdated items while self.expiration_heap: @@ -147,7 +150,7 @@ class InMemoryCache(BaseCache): Check if ttl is set for a key """ ttl_time: Final = self.ttl_dict.get(key) - if ttl_time is None or float(ttl_time) < time.time(): # if ttl is not set, allow override + if ttl_time is None or float(ttl_time) < self._clock(): # if ttl is not set, allow override return True else: return False @@ -167,10 +170,10 @@ class InMemoryCache(BaseCache): self.cache_dict[key] = value if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: - self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) + self.ttl_dict[key] = self._clock() + float(kwargs["ttl"]) heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) else: - self.ttl_dict[key] = time.time() + self.default_ttl + self.ttl_dict[key] = self._clock() + self.default_ttl heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) async def async_set_cache(self, key, value, **kwargs): diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2c36995c4f8..1a024b6b2b0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1138,6 +1138,46 @@ class RedisCache(BaseCache): ) _record_swallowed_redis_failure(self._circuit_breaker, e) + @_redis_circuit_breaker_guard + async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: + """One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL.""" + if len(cache_list) == 0: + return + commands: Final = tuple( + (self.check_and_fix_namespace(key=cache_key), json.dumps(cache_value), self.get_ttl(ttl=ttl)) + for cache_key, cache_value, ttl in cache_list + ) + start_time: Final = time.time() + try: + async with self.init_async_client().pipeline(transaction=False) as pipe: + for cache_key, json_cache_value, ttl in commands: + pipe.set(name=cache_key, value=json_cache_value, ex=None if ttl is None else timedelta(seconds=ttl)) + await pipe.execute() + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=time.time() - start_time, + call_type=f"async_set_cache_pipeline_with_ttls <- {_get_call_stack_info()}", + start_time=start_time, + end_time=time.time(), + ) + ) + except Exception as e: + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=time.time() - start_time, + error=e, + call_type=f"async_set_cache_pipeline_with_ttls <- {_get_call_stack_info()}", + start_time=start_time, + end_time=time.time(), + ) + ) + verbose_logger.error( + "LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS %s", str(e) + ) + _record_swallowed_redis_failure(self._circuit_breaker, e) + async def _set_cache_sadd_helper( self, redis_client: async_redis_client, @@ -1233,6 +1273,24 @@ class RedisCache(BaseCache): if len(self.redis_batch_writing_buffer) >= self.redis_flush_size: await self.flush_cache_buffer() # logging done in here + @staticmethod + async def _incrbyfloat_with_ttl( + _redis_client: "Redis", key: str, value: float, ttl: int | None, refresh_ttl: bool + ) -> float: + """INCRBYFLOAT plus its TTL command in one round trip; a third only when an unexpiring key needs an EXPIRE.""" + if ttl is None: + return await _redis_client.incrbyfloat(name=key, amount=value) + async with _redis_client.pipeline(transaction=False) as pipe: + pipe.incrbyfloat(name=key, amount=value) + if refresh_ttl: + pipe.expire(key, ttl) + else: + pipe.ttl(key) + result, ttl_or_expire = await pipe.execute() + if not refresh_ttl and ttl_or_expire == -1: + await _redis_client.expire(key, ttl) + return float(result) + @_redis_circuit_breaker_guard async def async_increment( self, @@ -1249,14 +1307,9 @@ class RedisCache(BaseCache): _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: - result: Final = await _redis_client.incrbyfloat(name=key, amount=value) - if _used_ttl is not None: - if refresh_ttl: - await _redis_client.expire(key, _used_ttl) - else: - current_ttl: Final = await _redis_client.ttl(key) - if current_ttl == -1: - await _redis_client.expire(key, _used_ttl) + result: Final = await self._incrbyfloat_with_ttl( + _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl + ) ## LOGGING ## end_time = time.time() diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 6cecc1e6157..ee01a53ecb3 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -4,6 +4,8 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 +import hashlib +import json import os from collections.abc import Awaitable, Callable, Generator from contextlib import AbstractAsyncContextManager @@ -343,6 +345,22 @@ class MCPClient: if auth_value: self.update_auth_value(auth_value) + async def discovery_auth_fingerprint(self) -> str: + request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers()) + if self._resolved_auth is None: + return self._hash_discovery_auth(request) + flow: Final = self._resolved_auth.async_auth_flow(request) + try: + authenticated: Final = await flow.__anext__() + return self._hash_discovery_auth(authenticated) + finally: + await flow.aclose() + + @staticmethod + def _hash_discovery_auth(request: httpx.Request) -> str: + material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items())))) + return hashlib.sha256(material.encode()).hexdigest() + def _create_transport_context( self, ) -> tuple[_TransportContext, httpx.AsyncClient | None]: @@ -781,7 +799,7 @@ class MCPClient: # Return a default error result instead of raising return self.error_tool_result(e) - async def list_prompts(self) -> list[Prompt]: + async def list_prompts(self, *, raise_on_error: bool = False) -> list[Prompt]: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") @@ -811,6 +829,8 @@ class MCPClient: verbose_logger.warning("MCP client list_prompts was cancelled") raise except Exception as e: + if raise_on_error: + raise error_type: Final = type(e).__name__ verbose_logger.error( "MCP client list_prompts failed - Error Type: %s, Error: %s, Server: %s, Transport: %s", @@ -869,7 +889,7 @@ class MCPClient: ) raise - async def list_resources(self) -> list[Resource]: + async def list_resources(self, *, raise_on_error: bool = False) -> list[Resource]: """List available resources from the server.""" verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") @@ -899,6 +919,8 @@ class MCPClient: verbose_logger.warning("MCP client list_resources was cancelled") raise except Exception as e: + if raise_on_error: + raise error_type: Final = type(e).__name__ verbose_logger.error( "MCP client list_resources failed - Error Type: %s, Error: %s, Server: %s, Transport: %s", @@ -916,7 +938,7 @@ class MCPClient: # Return empty list instead of raising to allow graceful degradation return [] - async def list_resource_templates(self) -> list[ResourceTemplate]: + async def list_resource_templates(self, *, raise_on_error: bool = False) -> list[ResourceTemplate]: """List available resource templates from the server.""" verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") @@ -949,6 +971,8 @@ class MCPClient: verbose_logger.warning("MCP client list_resource_templates was cancelled") raise except Exception as e: + if raise_on_error: + raise error_type: Final = type(e).__name__ verbose_logger.error( "MCP client list_resource_templates failed - Error Type: %s, Error: %s, Server: %s, Transport: %s", diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index dc41c7dadc8..caac8e888fd 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -673,7 +673,7 @@ class SlackAlerting(CustomBatchLogger): Create a standard message for a budget alert """ _all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True) - _all_fields_as_dict.pop("token") + _all_fields_as_dict.pop("token", None) msg = "" for k, v in _all_fields_as_dict.items(): if isinstance(v, Litellm_EntityType): diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bb54767edef..8a976a966a6 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1217,6 +1217,7 @@ class CustomGuardrail(CustomLogger): _, metadata_bucket = get_or_create_metadata_bucket(request_data) _append_guardrail_info(metadata_bucket) + _sync_guardrail_info_to_logging_obj(request_data, request_data.get("litellm_logging_obj")) _guardrail_self_recorded.set(True) diff --git a/litellm/integrations/otel/langfuse_logger.py b/litellm/integrations/otel/langfuse_logger.py index ed47533e700..f8fd417392f 100644 --- a/litellm/integrations/otel/langfuse_logger.py +++ b/litellm/integrations/otel/langfuse_logger.py @@ -3,7 +3,12 @@ from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger from litellm.integrations.otel.logger import OpenTelemetryV2 -from litellm.integrations.otel.mappers.langfuse import LANGFUSE_OBSERVATION_INPUT, LANGFUSE_OBSERVATION_OUTPUT +from litellm.integrations.otel.mappers.langfuse import ( + LANGFUSE_OBSERVATION_INPUT, + LANGFUSE_OBSERVATION_OUTPUT, + LANGFUSE_TRACE_NAME, +) +from litellm.integrations.otel.model.metadata import caller_trace_name from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output from litellm.integrations.otel.plumbing.context import request_root_span @@ -13,6 +18,18 @@ if TYPE_CHECKING: class LangfuseOpenTelemetryV2(OpenTelemetryV2): + """Names the trace from the request. Langfuse reads ``langfuse.trace.name`` off the root observation, + and the proxy's root span is still recording when the LLM call starts.""" + + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + root: Final = request_root_span() + name: Final = caller_trace_name(kwargs) + if root is not None and root.is_recording() and name is not None: + root.set_attribute(LANGFUSE_TRACE_NAME, name) + super().log_pre_api_call(model, messages, kwargs) + + +class LangfuseContentOpenTelemetryV2(LangfuseOpenTelemetryV2): """Stamps the request's input and output on the root observation while it is still recording. Langfuse shows a trace's input and output from its root observation. The proxy's root span ends diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 630aa313dc9..9ac748b231c 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -554,6 +554,7 @@ class OpenTelemetryV2(CustomLogger): capture_content=self.config.capture_span_content, time_to_first_chunk_seconds=call.time_to_first_chunk_seconds, request_route=request_root_http_route(), + trace_name=call.trace_name, ) end_time_ns: Final = to_ns(end_time) if carrier is not None and carrier.span is not None: @@ -984,8 +985,8 @@ def build_otel_v2_logger( def _logger_class(config: OpenTelemetryV2Config) -> type[OpenTelemetryV2]: - if "langfuse" not in config.mapper_names or not config.capture_span_content: + if "langfuse" not in config.mapper_names: return OpenTelemetryV2 - from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2 + from litellm.integrations.otel.langfuse_logger import LangfuseContentOpenTelemetryV2, LangfuseOpenTelemetryV2 - return LangfuseOpenTelemetryV2 + return LangfuseContentOpenTelemetryV2 if config.capture_span_content else LangfuseOpenTelemetryV2 diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index 01063d85355..98ff0f155a1 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -28,6 +28,7 @@ from litellm.integrations.otel.model.payloads import ( LANGFUSE_OBSERVATION_INPUT: Final = "langfuse.observation.input" LANGFUSE_OBSERVATION_OUTPUT: Final = "langfuse.observation.output" +LANGFUSE_TRACE_NAME: Final = "langfuse.trace.name" class LangfuseMapper: @@ -36,6 +37,7 @@ class LangfuseMapper: "langfuse.observation.model.name": lambda d: d.request_model or None, "langfuse.observation.metadata.provider": lambda d: d.provider or None, "langfuse.observation.id": lambda d: d.identity.call_id or None, + LANGFUSE_TRACE_NAME: lambda d: d.trace_name or None, "langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None, "langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None, } diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index ee116aca46b..cc81b689708 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -48,6 +48,8 @@ from litellm.integrations.otel.model.utils import as_str, to_seconds if TYPE_CHECKING: from litellm.types.utils import StandardLoggingPayload +LANGFUSE_TRACE_NAME_HEADER: Final = "langfuse_trace_name" + @dataclass(frozen=True) class RequestIdentity: @@ -215,6 +217,7 @@ class LLMCallEvent: # needs to be reasonable for a span that never gets closed (a leak). provisional_span_name: str time_to_first_chunk_seconds: float | None + trace_name: str | None @classmethod def from_dict(cls, kwargs: Mapping[str, Any]) -> LLMCallEvent: @@ -231,9 +234,30 @@ class LLMCallEvent: upstream_started=kwargs.get("api_call_start_time") is not None, provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), + trace_name=caller_trace_name(kwargs), ) +def caller_trace_name(kwargs: Mapping[str, object]) -> str | None: + request: Final = _as_str_mapping(kwargs.get("litellm_params")) + if request is None: + return None + proxy_request: Final = _as_str_mapping(request.get("proxy_server_request")) + headers: Final = _as_str_mapping(proxy_request.get("headers")) if proxy_request is not None else None + from_header: Final = as_str(headers.get(LANGFUSE_TRACE_NAME_HEADER)) if headers is not None else None + if from_header: + return from_header + return next( + ( + name + for key in ("metadata", "litellm_metadata") + if (metadata := _as_str_mapping(request.get(key))) is not None + and (name := as_str(metadata.get("trace_name"))) + ), + None, + ) + + def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: """Seconds from the upstream request being issued (``api_call_start_time``) to the first streamed chunk (``completion_start_time``); ``None`` for diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d0959a6c2e9..c11c4a7a27d 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -387,6 +387,7 @@ class LLMCallSpanData: output_type: GenAIOutputType | None = None call_type: str | None = None request_route: str | None = None + trace_name: str | None = None @classmethod def from_standard_logging_payload( @@ -395,6 +396,7 @@ class LLMCallSpanData: capture_content: bool = False, time_to_first_chunk_seconds: float | None = None, request_route: str | None = None, + trace_name: str | None = None, ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) # The single parse of the request's metadata — the request-vs-provider @@ -436,6 +438,7 @@ class LLMCallSpanData: output_type=resolve_output_type(call_type), call_type=call_type or None, request_route=request_route or context.identity.request_route, + trace_name=trace_name, ) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 540ce6738fc..2528f07f92c 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -52,6 +52,12 @@ from litellm.types.integrations.prometheus import ( _sanitize_prometheus_label_value, validate_prometheus_deployment_and_latency_caller_identity, ) +from litellm.types.proxy.carried_budget_state import ( + KeyBudgetSnapshot, + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -1941,6 +1947,8 @@ class PrometheusLogger(CustomLogger): _user_spend: Final = _metadata.get("user_api_key_user_spend", None) _user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None) + _user_email: Final = _metadata.get("user_api_key_user_email", None) + _org_alias: Final = _metadata.get("user_api_key_org_alias", None) # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -1957,6 +1965,7 @@ class PrometheusLogger(CustomLogger): response_cost=response_cost, key_max_budget=_api_key_max_budget, key_spend=_api_key_spend, + carried=KeyBudgetSnapshot.from_metadata(_metadata), ), self._set_team_budget_metrics_after_api_request( user_api_team=user_api_team, @@ -1964,16 +1973,21 @@ class PrometheusLogger(CustomLogger): team_spend=_team_spend, team_max_budget=_team_max_budget, response_cost=response_cost, + carried=TeamBudgetSnapshot.from_metadata(_metadata), ), self._set_user_budget_metrics_after_api_request( user_id=user_id, user_spend=_user_spend, user_max_budget=_user_max_budget, response_cost=response_cost, + carried=UserBudgetSnapshot.from_metadata(_metadata), + user_email=_user_email if isinstance(_user_email, str) else None, ), self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, response_cost=response_cost, + carried=OrgBudgetSnapshot.from_metadata(_metadata), + org_alias=_org_alias if isinstance(_org_alias, str) else None, ), return_exceptions=True, ) @@ -3821,6 +3835,7 @@ class PrometheusLogger(CustomLogger): team_spend: float | None, team_max_budget: float | None, response_cost: float, + carried: TeamBudgetSnapshot | None = None, ): """ Set team budget metrics after an LLM API request @@ -3839,6 +3854,7 @@ class PrometheusLogger(CustomLogger): spend=team_spend, max_budget=team_max_budget, response_cost=response_cost, + carried=carried, ) self._set_team_budget_metrics(team_object) @@ -3850,18 +3866,26 @@ class PrometheusLogger(CustomLogger): spend: float | None, max_budget: float | None, response_cost: float, + carried: TeamBudgetSnapshot | None = None, ) -> LiteLLM_TeamTable: """ Assemble a LiteLLM_TeamTable object - for fields not available in metadata, we fetch from db - Fields not available in metadata: - - `budget_reset_at` + ``budget_reset_at`` comes from the auth-carried snapshot when the request has one, + otherwise from the team lookup """ from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _total_team_spend: Final = (spend or 0) + response_cost + if carried is not None: + return LiteLLM_TeamTable( + team_id=team_id, + team_alias=team_alias, + spend=_total_team_spend, + max_budget=max_budget if max_budget is not None else carried.max_budget, + budget_reset_at=carried.budget_reset_at, + ) team_object: Final = LiteLLM_TeamTable( team_id=team_id, team_alias=team_alias, @@ -3946,11 +3970,13 @@ class PrometheusLogger(CustomLogger): self, org_id: str | None, response_cost: float, + carried: OrgBudgetSnapshot | None = None, + org_alias: str | None = None, ): """ Set org budget metrics after an LLM API request - - Fetches org info via cache (get_org_object) + - Uses the auth-carried org budget when the request has one, else fetches via get_org_object - Sets org budget metrics """ if isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric): @@ -3959,6 +3985,16 @@ class PrometheusLogger(CustomLogger): if not org_id: return + if carried is not None: + self._set_org_budget_metrics( + org_id=org_id, + org_alias=org_alias or "", + spend=carried.spend + response_cost, + max_budget=carried.max_budget, + budget_reset_at=None, + ) + return + from litellm.proxy.auth.auth_checks import get_org_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -3979,7 +4015,6 @@ class PrometheusLogger(CustomLogger): if org_info is None: return - org_alias: Final = org_info.organization_alias or "" _total_org_spend: Final = (org_info.spend or 0.0) + response_cost budget_table: Final = org_info.litellm_budget_table max_budget: Final = budget_table.max_budget if budget_table else None @@ -3987,7 +4022,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, - org_alias=org_alias, + org_alias=org_info.organization_alias or "", spend=_total_org_spend, max_budget=max_budget, budget_reset_at=budget_reset_at, @@ -4084,6 +4119,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, key_max_budget: float | None, key_spend: float | None, + carried: KeyBudgetSnapshot | None = None, ): if isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric): return @@ -4095,6 +4131,7 @@ class PrometheusLogger(CustomLogger): key_max_budget=key_max_budget, key_spend=key_spend, response_cost=response_cost, + carried=carried, ) self._set_key_budget_metrics(user_api_key_dict) @@ -4105,6 +4142,7 @@ class PrometheusLogger(CustomLogger): key_max_budget: float | None, key_spend: float | None, response_cost: float, + carried: KeyBudgetSnapshot | None = None, ) -> UserAPIKeyAuth: """ Assemble a UserAPIKeyAuth object @@ -4113,6 +4151,14 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _total_key_spend: Final = (key_spend or 0) + response_cost + if carried is not None: + return UserAPIKeyAuth( + token=user_api_key, + key_alias=user_api_key_alias, + max_budget=key_max_budget, + spend=_total_key_spend, + budget_reset_at=carried.budget_reset_at, + ) user_api_key_dict: Final = UserAPIKeyAuth( token=user_api_key, key_alias=user_api_key_alias, @@ -4140,6 +4186,8 @@ class PrometheusLogger(CustomLogger): user_spend: float | None, user_max_budget: float | None, response_cost: float, + carried: UserBudgetSnapshot | None = None, + user_email: str | None = None, ): """ Set user budget metrics after an LLM API request @@ -4157,6 +4205,8 @@ class PrometheusLogger(CustomLogger): spend=user_spend, max_budget=user_max_budget, response_cost=response_cost, + carried=carried, + user_email=user_email, ) self._set_user_budget_metrics(user_object) @@ -4167,18 +4217,28 @@ class PrometheusLogger(CustomLogger): spend: float | None, max_budget: float | None, response_cost: float, + carried: UserBudgetSnapshot | None = None, + user_email: str | None = None, ) -> LiteLLM_UserTable: """ Assemble a LiteLLM_UserTable object - for fields not available in metadata, we fetch from db - Fields not available in metadata: - - `budget_reset_at` + ``budget_reset_at`` and ``user_alias`` come from the auth-carried snapshot when the + request has one, otherwise from the user lookup """ from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _total_user_spend: Final = (spend or 0) + response_cost + if carried is not None: + return LiteLLM_UserTable( + user_id=user_id, + spend=_total_user_spend, + max_budget=max_budget if max_budget is not None else carried.max_budget, + budget_reset_at=carried.budget_reset_at, + user_alias=carried.user_alias, + user_email=user_email, + ) user_object: Final = LiteLLM_UserTable( user_id=user_id, spend=_total_user_spend, diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 54b116639e7..cdc108a6b4e 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.websearch_interception.tools import is_web_search_tool_responses from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata from litellm.litellm_core_utils.llm_judge import ( @@ -335,6 +336,16 @@ def _forwards_nothing(value: object) -> bool: return value is None or (isinstance(value, list) and len(value) == 0) +def _request_has_hosted_web_search(request: Mapping[str, object]) -> bool: + if request.get("web_search_options") is not None: + return True + tools: Final = request.get("tools") + return isinstance(tools, Sequence) and any( + isinstance(tool, Mapping) and tool.get("type") != "function" and is_web_search_tool_responses(tool) + for tool in tools + ) + + def _judgeable_sample( ops: _SurfaceOps, kwargs: Mapping[str, object], @@ -343,9 +354,14 @@ def _judgeable_sample( ) -> tuple[tuple[Mapping[str, object], ...], Mapping[str, object], str] | None: """The normalized chat conversation, the forwardable generation params, and the judgeable final text; None when this request's shapes cannot be sampled (no text and no - tool call to serialize, or a shape the owner transformations reject).""" + tool call to serialize, hosted web search the shadow cannot replay comparably, + or a shape the owner transformations reject).""" + if _request_has_hosted_web_search(_proxy_wire_body(kwargs) if ops.wire_params else model_parameters): + return None try: request: Final = ops.chat_request(kwargs, model_parameters) + if _request_has_hosted_web_search(request): + return None items: Final = _MESSAGE_ITEMS_ADAPTER.validate_python(request.get("messages")) messages: Final = _CHAT_MESSAGES_ADAPTER.validate_python( tuple(m.model_dump(exclude_none=True) if isinstance(m, BaseModel) else m for m in items) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 498d662a906..2d3a99abe81 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -23,7 +23,6 @@ import litellm from litellm import ( _custom_logger_compatible_callbacks_literal, json_logs, - log_raw_request_response, turn_off_message_logging, ) from litellm._logging import ( @@ -563,6 +562,7 @@ class Logging(LiteLLMLoggingBaseClass): self.streaming_chunks: list[Any] = [] # for generating complete stream response self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response + self._native_callback_fast_path: bool = False # Initialize dynamic callbacks self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks @@ -1236,6 +1236,11 @@ class Logging(LiteLLMLoggingBaseClass): additional_args.get("api_base", "") ) + def record_api_call_start_time(self) -> None: + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + if self.model_call_details.get("first_api_call_start_time") is None: + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] + def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API try: @@ -1253,7 +1258,7 @@ class Logging(LiteLLMLoggingBaseClass): additional_args=additional_args, ) # log raw request to provider (like LangFuse) -- if opted in. - if self.log_raw_request_response is True or log_raw_request_response is True: + if self.log_raw_request_response is True or litellm.log_raw_request_response is True: _litellm_params: Final = self.model_call_details.get("litellm_params", {}) _metadata: Final = _litellm_params.get("metadata", {}) or {} try: @@ -1300,15 +1305,7 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - # Set-once first provider-handoff instant. api_call_start_time - # is overwritten on every retry, so it can't measure one-time - # preprocessing; pinning the first attempt excludes retry loops - # + backoff. Logging object only — must NOT go into - # litellm_params["metadata"] (caller request metadata, typed - # Dict[str, str], echoed downstream; a datetime breaks it). - if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] + self.record_api_call_start_time() # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1442,16 +1439,21 @@ class Logging(LiteLLMLoggingBaseClass): """ return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) + def record_post_call( + self, original_response: object, input: object, api_key: object, additional_args: dict[str, object] + ) -> None: + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" + def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" + self.record_post_call(original_response, input, api_key, additional_args) attr: Literal["warning", "debug"] if self.litellm_request_debug: @@ -2116,6 +2118,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, + build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2140,6 +2143,9 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) + if not build_logging_payload: + return + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2201,6 +2207,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, + build_logging_payload: bool = True, ): try: if start_time is None: @@ -2238,6 +2245,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, + build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3261,7 +3269,9 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): + def _failure_handler_helper_fn( + self, exception, traceback_exception, start_time=None, end_time=None, build_logging_payload: bool = True + ): if start_time is None: start_time = self.start_time if end_time is None: @@ -3296,6 +3306,9 @@ class Logging(LiteLLMLoggingBaseClass): metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) + if not build_logging_payload: + return start_time, end_time + ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 75046f2cf87..4923bdda305 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -90,12 +90,12 @@ class _ResponseDoneBody(TypedDict, total=False): output: ReadOnly[Sequence[Mapping[str, object]]] -class _ScopedWebSocket(Protocol): +class ScopedWebSocket(Protocol): @property def scope(self) -> _ASGIScope: ... -class _ClientWebSocket(_ScopedWebSocket, Protocol): +class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... async def close(self, code: int = 1000, reason: str | None = None) -> None: ... @@ -1149,7 +1149,7 @@ class RealTimeStreaming: ) @staticmethod - def _detect_beta_header(websocket: _ScopedWebSocket) -> bool: + def _detect_beta_header(websocket: ScopedWebSocket) -> bool: """Return True if the client sent 'OpenAI-Beta: realtime=v1'. Checks the raw ASGI scope headers so it works for both FastAPI WebSocket @@ -1584,6 +1584,6 @@ class RealTimeStreaming: verbose_logger.debug("Could not relay the upstream close to the client: %s", e) -def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: +def client_sent_openai_beta_realtime_header(websocket: ScopedWebSocket) -> bool: """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``.""" return RealTimeStreaming._detect_beta_header(websocket) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 82189461403..5d3ae444b42 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -44,6 +44,7 @@ from litellm.types.llms.openai import ( from litellm.types.responses.main import ( OutputCodeInterpreterCall, build_code_interpreter_log_outputs, + build_web_search_call, ) from litellm.types.utils import ( Delta, @@ -649,6 +650,7 @@ class ModelResponseIterator: # Accumulate web_search_tool_result blocks for multi-turn reconstruction # See: https://github.com/BerriAI/litellm/issues/17737 self.web_search_results: list[dict[str, object]] = [] + self._web_search_calls: dict[str, object] = {} # mutable-ok: provider call state by id # Accumulate compaction blocks for multi-turn reconstruction self.compaction_blocks: list[dict[str, object]] = [] @@ -822,6 +824,19 @@ class ModelResponseIterator: return content_block_start + def _web_search_call_snapshot(self) -> dict[str, object]: + return dict(self._web_search_calls) # mutable-ok: stream payload snapshot + + def _complete_web_search_call(self, result: dict[str, object]) -> None: + tool_use_id: Final = result.get("tool_use_id") + if not isinstance(tool_use_id, str) or tool_use_id not in self._web_search_calls: + return + self._web_search_calls[tool_use_id] = build_web_search_call( + tool_id=tool_use_id, + tool_input=self._server_tool_inputs.get(tool_use_id, {}), # mutable-ok: empty provider input + result=result, + ) + def _build_code_interpreter_results(self) -> list: """Convert accumulated tool_results to OutputCodeInterpreterCall objects. @@ -923,6 +938,14 @@ class ModelResponseIterator: self._current_server_tool_id = content_block_start["content_block"]["id"] tool_input: Final = content_block_start["content_block"].get("input", {}) self._server_tool_inputs[self._current_server_tool_id] = tool_input + if _stream_tool_name == "web_search": + self._web_search_calls[self._current_server_tool_id] = build_web_search_call( + self._current_server_tool_id, + tool_input, + {"content": []}, # mutable-ok: no provider result yet + status="in_progress", + ) + provider_specific_fields["web_search_calls"] = self._web_search_call_snapshot() # Include caller information if present (for programmatic tool calling) if "caller" in content_block_start["content_block"]: caller_data: Final = content_block_start["content_block"]["caller"] @@ -957,7 +980,9 @@ class ModelResponseIterator: # The full content comes in content_block_start, not in deltas # See: https://github.com/BerriAI/litellm/issues/17737 self.web_search_results.append(content_block_start["content_block"]) + self._complete_web_search_call(content_block_start["content_block"]) provider_specific_fields["web_search_results"] = self.web_search_results + provider_specific_fields["web_search_calls"] = self._web_search_call_snapshot() elif content_type == "web_fetch_tool_result": # Capture web_fetch_tool_result for multi-turn reconstruction # The full content comes in content_block_start, not in deltas diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 5463f1862ad..0f99441a115 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -70,6 +70,7 @@ from litellm.types.llms.openai import ( from litellm.types.responses.main import ( OutputCodeInterpreterCall, build_code_interpreter_log_outputs, + build_web_search_call, ) from litellm.types.utils import ( CacheCreationTokenDetails, @@ -2464,6 +2465,35 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) return code_interpreter_results + def _build_web_search_calls( + self, + web_search_results: Sequence[object], + completion_response: Mapping[str, object], + ) -> list[object]: + content: Final = completion_response.get("content") + blocks: Final = content if isinstance(content, Sequence) else () + inputs: Final = { # mutable-ok: indexes provider server inputs + call_id: tool_input + for block in blocks + if isinstance(block, Mapping) + and block.get("type") == "server_tool_use" + and block.get("name") == "web_search" + and isinstance((call_id := block.get("id")), str) + and isinstance((tool_input := block.get("input")), Mapping) + } + return [ # mutable-ok: provider-neutral response items + build_web_search_call( + tool_id=tool_use_id, + tool_input=inputs.get(tool_use_id, {}), # mutable-ok: empty provider input + result=result, + ) + for result in web_search_results + if isinstance(result, dict) + and result.get("type") == "web_search_tool_result" + and isinstance((tool_use_id := result.get("tool_use_id")), str) + and tool_use_id in inputs + ] + def _build_provider_specific_fields( self, completion_response: dict, @@ -2485,6 +2515,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if web_search_results is not None: provider_specific_fields["web_search_results"] = web_search_results + provider_specific_fields["web_search_calls"] = self._build_web_search_calls( + web_search_results, + completion_response, + ) if tool_results is not None: provider_specific_fields["tool_results"] = tool_results diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 9b3a57cc422..8ff9f2e0679 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -400,7 +400,7 @@ class LiteLLMAnthropicMessagesAdapter: Anthropic web search tools have: - type starting with "web_search" (e.g., "web_search_20260209") - - name = "web_search" + - legacy name = "web_search" without a client input_schema Args: tool: Tool definition dict @@ -410,7 +410,9 @@ class LiteLLMAnthropicMessagesAdapter: """ tool_type: Final = tool.get("type", "") tool_name: Final = tool.get("name", "") - return (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search" + return (isinstance(tool_type, str) and tool_type.startswith("web_search")) or ( + tool_name == "web_search" and "input_schema" not in tool + ) def translate_anthropic_messages_to_openai( self, diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index e9913f0108d..146915dd6fd 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -13,7 +13,11 @@ from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from ....litellm_core_utils.realtime_streaming import RealTimeStreaming +from ....litellm_core_utils.realtime_streaming import ( + RealTimeStreaming, + ScopedWebSocket, + client_sent_openai_beta_realtime_header, +) from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion @@ -31,6 +35,19 @@ async def forward_messages(client_ws: Any, backend_ws: Any): pass +def azure_realtime_protocol_for_client( + configured_protocol: object, + *, + query_params: RealtimeQueryParams | None, + websocket: ScopedWebSocket, +) -> str: + if isinstance(configured_protocol, str) and configured_protocol: + return configured_protocol + if (query_params or {}).get("intent") == "transcription": + return "GA" + return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" + + class _ProxyClientWebSocket(Protocol): """Client-facing websocket handle: this path only closes it after a failed handshake.""" diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index f5126f81006..3a2af8a5aba 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -92,6 +92,18 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def get_api_key_env_var(self) -> str | None: return AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR + def resolve_connection_params( + self, + *, + api_key: str | None, + api_base: str | None, + dynamic_api_key: str | None, + dynamic_api_base: str | None, + ) -> tuple[str | None, str | None]: + explicit_api_key: Final = None if api_key is None else dynamic_api_key or api_key + explicit_api_base: Final = None if api_base is None else dynamic_api_base or api_base + return explicit_api_key, explicit_api_base + def get_supported_ocr_params(self, model: str) -> list: """ Get supported OCR parameters for Azure Document Intelligence. @@ -618,7 +630,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): except SSRFError as ssrf_err: raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") - poll_headers = {"Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "")} + poll_headers: Final = { + header: raw_response.request.headers[header] + for header in ("Ocp-Apim-Subscription-Key", "Authorization") + if header in raw_response.request.headers + } return operation_url, poll_headers @staticmethod diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 8111f9a194a..bd67dbf1a2a 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -144,9 +144,15 @@ class BaseOCRConfig: """ return None - def supports_rust_bridge(self) -> bool: - """Whether the Rust OCR bridge may serve this config when it is enabled for the provider.""" - return True + def resolve_connection_params( + self, + *, + api_key: str | None, + api_base: str | None, + dynamic_api_key: str | None, + dynamic_api_base: str | None, + ) -> tuple[str | None, str | None]: + return dynamic_api_key or api_key, dynamic_api_base or api_base def get_health_check_document(self) -> DocumentType: return { # mutable-ok: litellm.aocr rejects any document that is not a dict diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index bbbda4d14b6..53a3e634adf 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -56,6 +56,7 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset( ) _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) +_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"}) _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" @@ -187,6 +188,43 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return {key: value for key, value in params.items() if key != "service_tier"} + def _handle_unsupported_reasoning_summary( + self, params: dict[str, object], model: str, drop_params: bool + ) -> dict[str, object]: + reasoning: Final = params.get("reasoning") + if not self.use_openai_path or not isinstance(reasoning, dict): + return params + summary: Final = reasoning.get("summary") + if summary is None or ( + isinstance(summary, str) and summary in _BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES + ): + return params + if not drop_params: + raise litellm.utils.UnsupportedParamsError( + status_code=400, + message=( + f"bedrock_mantle does not support reasoning.summary={summary!r} for {model!r}; the Bedrock Mantle " + "OpenAI Responses path only accepts 'auto'. Set `drop_params: true` (litellm_settings or this " + 'deployment\'s litellm_params) to have LiteLLM drop it, or set `model_reasoning_summary = "auto"` ' + "in the client (Codex CLI: ~/.codex/config.toml)." + ), + ) + verbose_logger.warning( + "Bedrock Mantle Responses API: dropping unsupported reasoning.summary %r (supported: %s).", + summary, + sorted(_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES), + ) + stripped: Final = { # mutable-ok: map_openai_params contract returns a plain dict + key: value for key, value in reasoning.items() if key != "summary" + } + return ( + {**params, "reasoning": stripped} # mutable-ok: map_openai_params contract returns a plain dict + if stripped + else { # mutable-ok: map_openai_params contract returns a plain dict + key: value for key, value in params.items() if key != "reasoning" + } + ) + def transform_responses_api_request( self, model: str, @@ -343,12 +381,16 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI model: str, drop_params: bool, ) -> dict: - params: Final = self._handle_unsupported_service_tier( - super().map_openai_params( - response_api_optional_params=response_api_optional_params, - model=model, + params: Final = self._handle_unsupported_reasoning_summary( + self._handle_unsupported_service_tier( + super().map_openai_params( + response_api_optional_params=response_api_optional_params, + model=model, + drop_params=drop_params, + ), drop_params=drop_params, ), + model=model, drop_params=drop_params, ) diff --git a/litellm/llms/cohere/ocr/transformation.py b/litellm/llms/cohere/ocr/transformation.py index dd15d5360a6..b55ff4a3cbf 100644 --- a/litellm/llms/cohere/ocr/transformation.py +++ b/litellm/llms/cohere/ocr/transformation.py @@ -144,9 +144,6 @@ class CohereParseConfig(BaseOCRConfig): def get_api_key_env_var(self) -> str | None: return COHERE_API_KEY_ENV_VAR - def supports_rust_bridge(self) -> bool: - return False - def get_health_check_document(self) -> DocumentType: return { # mutable-ok: litellm.aocr rejects any document that is not a dict "type": "image_url", diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 26ad9a02a79..b4e1856f499 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -511,7 +511,6 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): m = cast(dict, message) m.pop("provider_specific_fields", None) m.pop("thinking_blocks", None) - m.pop("reasoning_content", None) return messages diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 33e0a0a923f..2fe11d9f7bd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -497,9 +497,14 @@ class OpenAIResponsesHandler(BaseTranslation): 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: + rewritten_texts: Final = guardrailed_inputs.get("texts") or () + if len(rewritten_texts) != len(extracted.task_mappings): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") await self._apply_guardrail_responses_to_input( messages=input_data, - responses=guardrailed_inputs.get("texts") or (), + responses=rewritten_texts, task_mappings=extracted.task_mappings, ) verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) @@ -635,10 +640,12 @@ class OpenAIResponsesHandler(BaseTranslation): """ Apply guardrail responses back to input messages. + ``responses`` pairs positionally with ``task_mappings``; the caller rejects + the request when the two disagree, so this never has to guess an alignment. + Override this method to customize how responses are applied. """ - for task_idx, guardrail_response in enumerate(responses): - mapping = task_mappings[task_idx] + for guardrail_response, mapping in zip(responses, task_mappings): msg_idx = cast(int, mapping[0]) content_idx_optional = cast(int | None, mapping[1]) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index aadd9bd3028..7fa09951eae 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57695,6 +57695,23 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "input_cost_per_token": 2.2e-07, @@ -57745,6 +57762,23 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "input_cost_per_token": 2.2e-07, diff --git a/litellm/ocr/input.py b/litellm/ocr/input.py new file mode 100644 index 00000000000..bcb448371c4 --- /dev/null +++ b/litellm/ocr/input.py @@ -0,0 +1,112 @@ +from collections.abc import Mapping +from os import PathLike +from typing import Final, Literal, Protocol, cast # noqa: TID251 # native callables are validated when loaded + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.configuration import rust_ocr_enabled + + +class FileReader(Protocol): + def read(self) -> bytes | str: ... + + +class FileDocument(TypedDict): + type: ReadOnly[Literal["file"]] + file: ReadOnly[bytes | PathLike[str] | FileReader] + mime_type: ReadOnly[NotRequired[str]] + + +class NativeFileDocument(Protocol): + def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ... + + +class NativeUploadDocument(Protocol): + def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ... + + +class NativeMimeType(Protocol): + def __call__(self, file_name: str) -> str: ... + + +_FILE_DOCUMENT: Final = NativeBinding( + "_ocr_file_document", + validate=lambda value: ( + cast( # cast-ok: native export owns the callable signature + NativeFileDocument, value + ) + if callable(value) + else None + ), +) +_UPLOAD_DOCUMENT: Final = NativeBinding( + "_ocr_upload_document", + validate=lambda value: ( + cast( # cast-ok: native export owns the callable signature + NativeUploadDocument, value + ) + if callable(value) + else None + ), +) +_MAX_FILE_BYTES: Final = NativeBinding( + "_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None +) +_MIME_TYPE: Final = NativeBinding( + "_ocr_mime_type", + validate=lambda value: ( + cast( # cast-ok: native export owns the callable signature + NativeMimeType, value + ) + if callable(value) + else None + ), +) +_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024 + + +def get_mime_type(file_path: str) -> str: + native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None + if native is None: + from litellm.ocr import legacy + + return legacy.get_mime_type(file_path) + return native(file_path) + + +def get_max_file_bytes() -> int: + limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None + if limit is None: + return _PYTHON_MAX_FILE_BYTES + return limit + + +def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]: + native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None + if native is None: + from litellm.ocr import legacy + + return legacy.convert_file_document_to_url_document(document) + return native(document) + + +def convert_upload_to_url_document( + file_content: bytes, filename: str | None, content_type: str | None +) -> dict[str, str]: + native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None + if native is None: + from litellm.ocr import legacy + + if len(file_content) > _PYTHON_MAX_FILE_BYTES: + raise ValueError("OCR file exceeds the size limit") + content_mime: Final = content_type.split(";")[0].strip() if content_type else None + mime_type: Final = ( + legacy.get_mime_type(filename) + if filename and (not content_mime or content_mime == "application/octet-stream") + else content_mime or "application/octet-stream" + ) + return legacy.convert_file_document_to_url_document( + {"type": "file", "file": file_content, "mime_type": mime_type} + ) + return native(file_content, filename, content_type) diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py new file mode 100644 index 00000000000..ddf6016dce3 --- /dev/null +++ b/litellm/ocr/legacy.py @@ -0,0 +1,411 @@ +""" +Main OCR function for LiteLLM. +""" + +import asyncio +import base64 +import mimetypes +import os +import re +from collections.abc import Coroutine, Mapping +from dataclasses import dataclass +from io import IOBase +from types import MappingProxyType +from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.constants import request_timeout +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.ocr.transformation import ( + OCR_REQUEST_FORMAT_PARAM, + BaseOCRConfig, + OCRResponse, + parse_ocr_request_format, +) +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.ocr.input import FileReader +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager, client + +base_llm_http_handler: Final = BaseLLMHTTPHandler() + + +@dataclass(frozen=True, slots=True) +class _PreparedOCRRequest: + model: str + document: Mapping[str, object] + api_key: str | None + api_base: str | None + custom_llm_provider: str + extra_headers: dict[str, object] | None + provider_config: BaseOCRConfig + optional_params: dict[str, object] + litellm_params: dict[str, object] + effective_timeout: float | httpx.Timeout + litellm_logging_obj: LiteLLMLoggingObj + + +def _prepare_ocr_request( + model: str, + document: Mapping[str, object], + api_key: str | None, + api_base: str | None, + timeout: float | httpx.Timeout | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + kwargs: dict[str, object], +) -> _PreparedOCRRequest: + litellm_logging_obj: Final = cast( # cast-ok: @client supplies the logging object; preserve legacy failure behavior + LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj") + ) + litellm_call_id: Final = cast( # cast-ok: @client supplies the call id without coercion + str | None, kwargs.get("litellm_call_id", None) + ) + + if not isinstance(document, dict): + raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}") + + doc_type = document.get("type") + + if doc_type == "file": + document = convert_file_document_to_url_document(document) + doc_type = document.get("type") + + if doc_type not in ["document_url", "image_url"]: + raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") + + ( + model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + ) + + ocr_provider_config: Final = ProviderConfigManager.get_provider_ocr_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) + + if ocr_provider_config is None: + raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}") + + resolved_api_key, resolved_api_base = ocr_provider_config.resolve_connection_params( + api_key=api_key, + api_base=api_base, + dynamic_api_key=dynamic_api_key, + dynamic_api_base=dynamic_api_base, + ) + + verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider) + + litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) + + supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model) + requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM) + if requested_format is not None: + try: + parsed_format: Final = parse_ocr_request_format(requested_format) + except ValueError as e: + raise litellm.exceptions.UnsupportedParamsError( + message=f"{e}", model=model, llm_provider=custom_llm_provider + ) from e + if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native": + raise litellm.exceptions.UnsupportedParamsError( + message=( + f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, " + f"model: {model}" + ), + model=model, + llm_provider=custom_llm_provider, + ) + + non_default_params: Final = {} + for param in supported_params: + if param in kwargs: + non_default_params[param] = kwargs.pop(param) + + optional_params: Final = ocr_provider_config.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + ) + + verbose_logger.debug("OCR optional_params after mapping: %s", optional_params) + + effective_timeout: Final = timeout or request_timeout + + litellm_logging_obj.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params={ + "litellm_call_id": litellm_call_id, + "api_base": resolved_api_base, + }, + custom_llm_provider=custom_llm_provider, + ) + + return _PreparedOCRRequest( + model=model, + document=document, + api_key=resolved_api_key, + api_base=resolved_api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + provider_config=ocr_provider_config, + optional_params=cast( + dict[str, object], optional_params + ), # cast-ok: provider configs return heterogeneous OCR options + litellm_params=dict(litellm_params), + effective_timeout=effective_timeout, + litellm_logging_obj=litellm_logging_obj, + ) + + +def _error_provider(model: str, custom_llm_provider: str | None) -> str | None: + if custom_llm_provider is not None: + return custom_llm_provider + prefix: Final = model.partition("/")[0] + if prefix in {"mistral", "azure_ai", "vertex_ai"}: + return prefix + return "mistral" if model.startswith("mistral-ocr") else None + + +@client +async def aocr( + model: str, + document: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + **kwargs: object, # kwargs-ok: public OCR accepts provider-specific options +) -> OCRResponse: + completion_kwargs: Final[dict[str, object]] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } + try: + prepared: Final = _prepare_ocr_request( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) + + response = base_llm_http_handler.ocr( + model=prepared.model, + document=cast( # cast-ok: preserve legacy document fields for provider validation + dict[str, str], prepared.document + ), + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=True, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + if asyncio.iscoroutine(response): + response = await response + + if response is None: + raise ValueError(f"Got an unexpected None response from the OCR API: {response}") + + return response + except Exception as e: + error_provider: Final = _error_provider(model, custom_llm_provider) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model + raise litellm.exception_type( + model=error_model, + custom_llm_provider=error_provider, + original_exception=e, + completion_kwargs=completion_kwargs, + extra_kwargs=kwargs, + ) + + +_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$") + +_MIME_TYPE_MAP: Final = MappingProxyType( + { + ".pdf": "application/pdf", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".bmp": "image/bmp", + } +) + + +def get_mime_type(file_path: str) -> str: + ext: Final = os.path.splitext(file_path)[1].lower() + mime: Final = _MIME_TYPE_MAP.get(ext) + if mime: + return mime + guessed, _ = mimetypes.guess_type(file_path) + return guessed or "application/octet-stream" + + +def _read_file(file_input: object) -> tuple[bytes, str, str | None]: + if isinstance(file_input, str): + raise ValueError( + "OCR file input does not accept bare str values. Pass bytes, " + "a pathlib.Path, or a file-like object. To OCR a local file " + "from a path, call open(path, 'rb') yourself." + ) + if isinstance(file_input, os.PathLike): + file_path: Final = str(cast(object, file_input)) # cast-ok: preserve staging's str(PathLike) conversion + if not os.path.isfile(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + mime_type: Final = get_mime_type(file_path) + with open(file_path, "rb") as stream: + return stream.read(), mime_type, os.path.basename(file_path) + if isinstance(file_input, bytes): + return file_input, "application/octet-stream", None + if isinstance(file_input, IOBase) or hasattr(file_input, "read"): + file_name: Final = cast( # cast-ok: retain legacy validation and errors for file-like metadata + str | None, getattr(file_input, "name", None) + ) + inferred_mime: Final = get_mime_type(file_name) if file_name else "application/octet-stream" + reader: Final = cast(FileReader, file_input) # cast-ok: legacy accepts duck-typed file readers + content: Final = reader.read() + return content.encode("utf-8") if isinstance(content, str) else content, inferred_mime, file_name + raise ValueError( + f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object." + ) + + +def convert_file_document_to_url_document(document: Mapping[str, object]) -> dict[str, str]: + file_input: Final = document.get("file") + if file_input is None: + raise ValueError( + "document with type='file' must include a 'file' field containing " + "a pathlib.Path, file-like object, or bytes" + ) + file_bytes, inferred_mime, file_name = _read_file(file_input) + if not file_bytes: + raise ValueError("File is empty or could not be read") + mime_type: Final = cast( # cast-ok: keep staging's MIME validation errors + str, document.get("mime_type", inferred_mime) + ) + if not _MIME_PATTERN.match(mime_type): + raise ValueError(f"Invalid MIME type: {mime_type}") + + base64_data: Final = base64.b64encode(file_bytes).decode("utf-8") + data_uri: Final = f"data:{mime_type};base64,{base64_data}" + + if mime_type.startswith("image/"): + verbose_logger.debug( + "OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)", + mime_type, + len(file_bytes), + file_name, + ) + return {"type": "image_url", "image_url": data_uri} + + verbose_logger.debug( + "OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)", + mime_type, + len(file_bytes), + file_name, + ) + return {"type": "document_url", "document_url": data_uri} + + +@client +def ocr( + model: str, + document: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + **kwargs: object, # kwargs-ok: public OCR accepts provider-specific options +) -> OCRResponse | Coroutine[object, object, OCRResponse]: + completion_kwargs: Final[dict[str, object]] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } + try: + _is_async: Final = kwargs.pop("aocr", False) is True + completion_kwargs["aocr"] = _is_async + prepared: Final = _prepare_ocr_request( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + kwargs=kwargs, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout=timeout, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) + + response: Final = base_llm_http_handler.ocr( + model=prepared.model, + document=cast( # cast-ok: preserve legacy document fields for provider validation + dict[str, str], prepared.document + ), + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=_is_async, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + return response + except Exception as e: + error_provider: Final = _error_provider(model, custom_llm_provider) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model + raise litellm.exception_type( + model=error_model, + custom_llm_provider=error_provider, + original_exception=e, + completion_kwargs=completion_kwargs, + extra_kwargs=kwargs, + ) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 56bfd98895d..382c5d6aae4 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -1,460 +1,20 @@ -""" -Main OCR function for LiteLLM. -""" - -import asyncio -import base64 -import mimetypes -import os -import re -from collections.abc import Callable, Coroutine, Mapping -from dataclasses import dataclass -from io import IOBase -from types import MappingProxyType -from typing import Any, Final, cast +from collections.abc import Awaitable, Callable, Coroutine, Mapping +from typing import Final, cast # noqa: TID251 # native binding selects a sync result or an async awaitable import httpx -import litellm -from litellm._logging import verbose_logger -from litellm.constants import request_timeout -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.azure_ai.ocr.common_utils import ( - is_azure_cohere_parse_model, - is_azure_document_intelligence_model, -) -from litellm.llms.base_llm.ocr.transformation import ( - OCR_REQUEST_FORMAT_PARAM, - BaseOCRConfig, - OCRResponse, - parse_ocr_request_format, -) -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import legacy +from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type from litellm.rust_bridge.bindings import native_exception_types -from litellm.rust_bridge.configuration import rust_enabled -from litellm.types.router import GenericLiteLLMParams -from litellm.utils import ProviderConfigManager, client +from litellm.rust_bridge.configuration import rust_ocr_enabled +from litellm.rust_bridge.ocr import LiteLLMOcrRequest +from litellm.rust_bridge.ocr_lifecycle import select -####### ENVIRONMENT VARIABLES ################### -base_llm_http_handler = BaseLLMHTTPHandler() -################################################# +__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr") -@dataclass -class _PreparedOCRRequest: - model: str - document: dict[str, Any] - api_key: str | None - api_base: str | None - custom_llm_provider: str - extra_headers: dict[str, object] | None - provider_config: BaseOCRConfig - optional_params: dict[str, object] - litellm_params: dict[str, object] - effective_timeout: float | httpx.Timeout - litellm_logging_obj: LiteLLMLoggingObj - caller_supplied_api_key: bool = True - caller_supplied_api_base: bool = True - - -_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) -_RUST_OCR_CONFIG_FIELDS: Final = frozenset( - { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_scope", - "azure_authority_host", - "azure_credential", - "azure_federated_token_file", - "vertex_credentials", - "vertex_ai_credentials", - "vertex_project", - "vertex_ai_project", - "vertex_location", - "vertex_ai_location", - } -) -_RUST_OCR_SECRET_FIELDS: Final = frozenset( - {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} -) - - -def _prepare_ocr_request( - model: str, - document: Mapping[str, object], - api_key: str | None, - api_base: str | None, - timeout: float | httpx.Timeout | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - kwargs: dict[str, object], -) -> _PreparedOCRRequest: - litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) - litellm_call_id: Final = cast(str | None, kwargs.get("litellm_call_id", None)) - - if not isinstance(document, dict): - raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}") - - doc_type = document.get("type") - - if doc_type == "file": - document = convert_file_document_to_url_document(document) - doc_type = document.get("type") - - if doc_type not in ["document_url", "image_url"]: - raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") - - caller_supplied_api_key: Final = api_key is not None - caller_supplied_api_base: Final = api_base is not None - - ( - model, - custom_llm_provider, - dynamic_api_key, - dynamic_api_base, - ) = litellm.get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - api_key=api_key, - ) - - suppress_dynamic_api_base: Final = ( - not caller_supplied_api_base - and custom_llm_provider == "azure_ai" - and is_azure_document_intelligence_model(model) - ) - if dynamic_api_key: - api_key = dynamic_api_key - if dynamic_api_base and not suppress_dynamic_api_base: - api_base = dynamic_api_base - - ocr_provider_config: Final = ProviderConfigManager.get_provider_ocr_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) - - if ocr_provider_config is None: - raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}") - - verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider) - - litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) - - supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model) - requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM) - if requested_format is not None: - try: - parsed_format: Final = parse_ocr_request_format(requested_format) - except ValueError as e: - raise litellm.exceptions.UnsupportedParamsError( - message=f"{e}", model=model, llm_provider=custom_llm_provider - ) from e - if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native": - raise litellm.exceptions.UnsupportedParamsError( - message=( - f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, " - f"model: {model}" - ), - model=model, - llm_provider=custom_llm_provider, - ) - - non_default_params: Final = {} - for param in supported_params: - if param in kwargs: - non_default_params[param] = kwargs.pop(param) - - optional_params: Final = ocr_provider_config.map_ocr_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - ) - - verbose_logger.debug("OCR optional_params after mapping: %s", optional_params) - - effective_timeout: Final = timeout or request_timeout - - litellm_logging_obj.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params={ - "litellm_call_id": litellm_call_id, - "api_base": api_base, - }, - custom_llm_provider=custom_llm_provider, - ) - - return _PreparedOCRRequest( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - provider_config=ocr_provider_config, - optional_params=cast(dict[str, object], optional_params), - litellm_params=dict(litellm_params), - effective_timeout=effective_timeout, - litellm_logging_obj=litellm_logging_obj, - caller_supplied_api_key=caller_supplied_api_key, - caller_supplied_api_base=caller_supplied_api_base, - ) - - -def _rust_ocr_provider(request: rust_ocr_bridge.LiteLLMOcrRequest) -> str | None: - if request.custom_llm_provider is not None: - return request.custom_llm_provider - prefix: Final = request.model.partition("/")[0] - if prefix in _RUST_OCR_PROVIDERS: - return prefix - if request.model.startswith("mistral-ocr"): - return "mistral" - return None - - -def _rust_ocr_supported(request: rust_ocr_bridge.LiteLLMOcrRequest) -> bool: - provider: Final = _rust_ocr_provider(request) - if provider not in _RUST_OCR_PROVIDERS or request.kwargs.get(OCR_REQUEST_FORMAT_PARAM) == "native": - return False - if provider == "azure_ai": - return ( - not is_azure_cohere_parse_model(request.model) - and not callable(request.kwargs.get("azure_ad_token_provider")) - and request.kwargs.get("azure_username") is None - and request.kwargs.get("azure_password") is None - ) - return True - - -def _rust_bridge_optional_params( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], -) -> Mapping[str, object]: - optional_params: Final = MappingProxyType( - { - name: value - for name, value in request.kwargs.items() - if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) - and name not in {"litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request"} - } - ) - provider: Final = _rust_ocr_provider(request) - if provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: - return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) - if provider != "vertex_ai": - return optional_params - project: Final = ( - request.kwargs.get("vertex_project") - or request.kwargs.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - location: Final = ( - request.kwargs.get("vertex_location") - or request.kwargs.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - credentials: Final = ( - request.kwargs.get("vertex_credentials") - or request.kwargs.get("vertex_ai_credentials") - or resolve_secret("VERTEXAI_CREDENTIALS") - ) - vertex_params: Final = MappingProxyType( - { - name: value - for name, value in ( - ("vertex_project", project), - ("vertex_location", location), - ("vertex_credentials", credentials), - ) - if value is not None - } - ) - return MappingProxyType({**optional_params, **vertex_params}) - - -def _rust_bridge_input_sources( - request: rust_ocr_bridge.LiteLLMOcrRequest, - optional_params: Mapping[str, object], -) -> Mapping[str, str]: - proxy_request: Final = request.kwargs.get("proxy_server_request") - if not isinstance(proxy_request, Mapping): - return MappingProxyType({}) - proxy_request_mapping: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types - Mapping[object, object], proxy_request - ) - body_value: Final = proxy_request_mapping.get("body") - if not isinstance(body_value, Mapping): - return MappingProxyType({}) - body: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types - Mapping[object, object], body_value - ) - credential_fields_value: Final = proxy_request_mapping.get("credential_fields", ()) - credential_fields: Final = ( - frozenset(name for name in credential_fields_value if isinstance(name, str)) - if isinstance(credential_fields_value, (list, tuple, set, frozenset)) - else frozenset() - ) - names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) - request_sources: Final = MappingProxyType( - {name: "request" for name in names if name in body or name in credential_fields} - ) - if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: - return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) - return request_sources - - -def _marshal_rust_ocr_request( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], -) -> rust_ocr_bridge.LiteLLMOcrRequest: - if not isinstance(request.document, dict): - raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") - document: Final = ( - convert_file_document_to_url_document(request.document) - if request.document.get("type") == "file" - else request.document - ) - provider: Final = _rust_ocr_provider(request) - api_key: Final = request.api_key or resolve_secret("MISTRAL_API_KEY") if provider == "mistral" else request.api_key - optional_params: Final = _rust_bridge_optional_params(request, resolve_secret) - input_sources: Final = _rust_bridge_input_sources(request, optional_params) - logged_optional_params: Final = MappingProxyType( - {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} - ) - logged_kwargs: Final = MappingProxyType( - { - name: "****" if name in _RUST_OCR_SECRET_FIELDS else value - for name, value in request.kwargs.items() - if name != "proxy_server_request" - } - ) - logging_obj: Final = cast( # cast-ok: bridge kwargs carry the prepared logging object - LiteLLMLoggingObj, request.kwargs["litellm_logging_obj"] - ) - logging_obj.update_from_kwargs( - kwargs=dict(logged_kwargs), # mutable-ok: logging API requires an owned dict - model=request.model, - optional_params=dict(logged_optional_params), # mutable-ok: logging API requires an owned dict - litellm_params={ - "litellm_call_id": request.kwargs.get("litellm_call_id"), - "api_base": request.api_base, - }, # mutable-ok: legacy logging requires a concrete params dict - custom_llm_provider=provider, - ) - logging_obj.pre_call( - input="OCR document processing", - api_key=api_key, - additional_args={ # mutable-ok: pre_call mutates the additional_args dict - "complete_input_dict": { - "model": request.model, - "document": document, - **logged_optional_params, - }, # mutable-ok: callbacks consume a JSON-serializable request dict - "api_base": request.api_base or "", - "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict - }, - ) - return rust_ocr_bridge.LiteLLMOcrRequest( - model=request.model, - document=document, - api_key=api_key, - api_base=request.api_base, - timeout=request.timeout if request.timeout is not None else request_timeout, - custom_llm_provider=request.custom_llm_provider, - extra_headers=request.extra_headers, - kwargs=optional_params, - input_sources=input_sources, - ) - - -def _map_rust_ocr_error( - error: Exception, - request: rust_ocr_bridge.LiteLLMOcrRequest, - exception_types: tuple[type[BaseException], type[BaseException]] | None, -) -> Exception: - if exception_types is None or not isinstance(error, exception_types[1]): - return error - provider: Final = _rust_ocr_provider(request) - if provider is None: - return error - provider_config: Final = ProviderConfigManager.get_provider_ocr_config( - model=request.model.removeprefix(f"{provider}/"), provider=litellm.LlmProviders(provider) - ) - if provider_config is None: - return error - error_args: Final = cast( # cast-ok: Python exceptions expose positional args as a tuple - tuple[object, ...], error.args - ) - status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 - message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) - error_factory: Final = cast( # cast-ok: provider configs expose heterogeneous exception factories - Callable[..., Exception], provider_config.get_error_class - ) - return error_factory( - error_message=message, status_code=status or 500, headers={} - ) # mutable-ok: provider error factories require a concrete headers dict - - -def _run_rust_ocr( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_ocr() is None: - return None - marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) - input_sources: Final = marshalled.input_sources - try: - response: Final = rust_ocr_bridge.ocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_rust_ocr_error(error, request, native_exception_types()) from error - return OCRResponse.model_validate(response) if response is not None else None - - -async def _run_rust_aocr( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_aocr() is None: - return None - marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) - input_sources: Final = marshalled.input_sources - try: - response: Final = await rust_ocr_bridge.aocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_rust_ocr_error(error, request, native_exception_types()) from error - return OCRResponse.model_validate(response) if response is not None else None - - -@client -async def aocr( +def _bind_request( model: str, document: Mapping[str, object], api_key: str | None = None, @@ -462,77 +22,9 @@ async def aocr( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, - **kwargs: object, -) -> OCRResponse: - """ - Async OCR function. - - Args: - model: Model name (e.g., "mistral/mistral-ocr-latest") - document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs, - {"type": "image_url", "image_url": "https://..."} for images, or - {"type": "file", "file": } for local files - api_key: Optional API key - api_base: Optional API base URL - timeout: Optional timeout - custom_llm_provider: Optional custom LLM provider - extra_headers: Optional extra headers - **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - - Returns: - OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - - Example: - ```python - import litellm - - # OCR with PDF - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": "https://arxiv.org/pdf/2201.04234" - }, - include_image_base64=True - ) - - # OCR with image - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "image_url", - "image_url": "https://example.com/image.png" - } - ) - - # OCR with base64 encoded PDF - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": f"data:application/pdf;base64,{base64_pdf}" - } - ) - - # OCR with local file - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": "/path/to/document.pdf"} - ) - ``` - """ - completion_kwargs: Final[dict[str, object]] = { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "timeout": timeout, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "kwargs": kwargs, - } - request: Final = rust_ocr_bridge.LiteLLMOcrRequest( + **kwargs: object, # kwargs-ok: public OCR accepts provider-specific options +) -> LiteLLMOcrRequest: + return LiteLLMOcrRequest( model=model, document=document, api_key=api_key, @@ -542,340 +34,50 @@ async def aocr( extra_headers=extra_headers, kwargs=kwargs, ) + + +def _public_request(name: str, args: tuple[object, ...], kwargs: dict[str, object]) -> LiteLLMOcrRequest: try: - if rust_enabled() and _rust_ocr_supported(request): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = await _run_rust_aocr( - request=request, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - - prepared: Final = _prepare_ocr_request( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, - ) - model = prepared.model - custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - - response = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=True, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) - - if asyncio.iscoroutine(response): - response = await response - - if response is None: - raise ValueError(f"Got an unexpected None response from the OCR API: {response}") - - return response - except Exception as e: - error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) - error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model - raise litellm.exception_type( - model=error_model, - custom_llm_provider=error_provider, - original_exception=e, - completion_kwargs=completion_kwargs, - extra_kwargs=kwargs, - ) + return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation + except TypeError as error: + raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None -################################################# -# Public utilities — used by the SDK and the proxy -################################################# - -_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$") - -_MIME_TYPE_MAP: Final = { - ".pdf": "application/pdf", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".webp": "image/webp", - ".tiff": "image/tiff", - ".tif": "image/tiff", - ".bmp": "image/bmp", -} - - -def get_mime_type(file_path: str) -> str: - """ - Determine MIME type from file path extension. - - Falls back to mimetypes.guess_type, then to 'application/octet-stream'. - """ - ext: Final = os.path.splitext(file_path)[1].lower() - mime: Final = _MIME_TYPE_MAP.get(ext) - if mime: - return mime - guessed, _ = mimetypes.guess_type(file_path) - return guessed or "application/octet-stream" - - -def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]: - """ - Convert a file-type document dict to a document_url-type document dict - with an inline base64 data URI. - - Accepts document dicts like: - {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path - {"type": "file", "file": } # file-like object (BinaryIO) - {"type": "file", "file": b"raw bytes"} # raw bytes - - Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or - ``open(path, "rb")`` instead. See the str check below for the rationale. - - Returns: - {"type": "document_url", "document_url": "data:;base64,"} - or {"type": "image_url", "image_url": "data:;base64,"} - """ - file_input: Final = document.get("file") - if file_input is None: - raise ValueError( - "document with type='file' must include a 'file' field containing " - "a pathlib.Path, file-like object, or bytes" - ) - - file_bytes: bytes - mime_type: str = "application/octet-stream" - file_name: str | None = None - - if isinstance(file_input, str): - # Bare strings are rejected here. The OCR ``document`` accepts a - # ``{"type": "file", "file": }`` shape, and when this helper - # runs in a proxy request handler ```` is attacker-controlled. - # Opening it as a path is an arbitrary local file read on the proxy - # host, which is then base64-encoded and forwarded to the OCR - # provider — an exfiltration primitive. - raise ValueError( - "OCR file input does not accept bare str values. Pass bytes, " - "a pathlib.Path, or a file-like object. To OCR a local file " - "from a path, call open(path, 'rb') yourself." - ) - if isinstance(file_input, os.PathLike): - # os.PathLike (pathlib.Path and custom __fspath__ classes) is a - # Python-level type that HTTP form values can't fabricate. - file_path: Final = str(file_input) - if not os.path.isfile(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - mime_type = get_mime_type(file_path) - file_name = os.path.basename(file_path) - with open(file_path, "rb") as f: - file_bytes = f.read() - elif isinstance(file_input, bytes): - file_bytes = file_input - elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): - if hasattr(file_input, "name"): - file_name = getattr(file_input, "name", None) - if file_name: - mime_type = get_mime_type(file_name) - file_bytes = file_input.read() - if isinstance(file_bytes, str): - file_bytes = file_bytes.encode("utf-8") - else: - raise ValueError( - f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object." - ) - - if not file_bytes: - raise ValueError("File is empty or could not be read") - - if "mime_type" in document: - mime_type = document["mime_type"] - - if not _MIME_PATTERN.match(mime_type): - raise ValueError(f"Invalid MIME type: {mime_type}") - - base64_data: Final = base64.b64encode(file_bytes).decode("utf-8") - data_uri: Final = f"data:{mime_type};base64,{base64_data}" - - if mime_type.startswith("image/"): - verbose_logger.debug( - "OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)", - mime_type, - len(file_bytes), - file_name, - ) - return {"type": "image_url", "image_url": data_uri} - - verbose_logger.debug( - "OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)", - mime_type, - len(file_bytes), - file_name, - ) - return {"type": "document_url", "document_url": data_uri} - - -@client def ocr( - model: str, - document: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - timeout: float | httpx.Timeout | None = None, - custom_llm_provider: str | None = None, - extra_headers: dict[str, object] | None = None, - **kwargs: object, + *args: object, + **kwargs: object, # kwargs-ok: preserve the public OCR call shape ) -> OCRResponse | Coroutine[object, object, OCRResponse]: - """ - Synchronous OCR function. - - Args: - model: Model name (e.g., "mistral/mistral-ocr-latest") - document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs, - {"type": "image_url", "image_url": "https://..."} for images, or - {"type": "file", "file": } for local files - api_key: Optional API key - api_base: Optional API base URL - timeout: Optional timeout - custom_llm_provider: Optional custom LLM provider - extra_headers: Optional extra headers - **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - - Returns: - OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - - Example: - ```python - import litellm - - # OCR with PDF - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": "https://arxiv.org/pdf/2201.04234" - }, - include_image_base64=True - ) - - # OCR with image - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "image_url", - "image_url": "https://example.com/image.png" - } - ) - - # OCR with base64 encoded PDF - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": f"data:application/pdf;base64,{base64_pdf}" - } - ) - - # OCR with local file - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": "/path/to/document.pdf"} - ) - - # Access pages - for page in response.pages: - print(f"Page {page.index}: {page.markdown}") - ``` - """ - completion_kwargs: Final[dict[str, object]] = { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "timeout": timeout, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "kwargs": kwargs, - } - request: Final = rust_ocr_bridge.LiteLLMOcrRequest( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, - ) - try: - _is_async: Final = kwargs.pop("aocr", False) is True - completion_kwargs["aocr"] = _is_async - if rust_enabled() and _rust_ocr_supported(request): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = _run_rust_ocr( - request=request, - resolve_api_key=get_secret_str, + request: Final = _public_request("ocr", args, kwargs) + native: Final = select(request) if rust_ocr_enabled() else None + if native is not None: + try: + return cast( # cast-ok: False selects the synchronous result + OCRResponse, native(request, args, kwargs, False) ) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response + except _decline_types(): + pass + fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator + Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]], legacy.ocr + ) + return fallback(*args, **kwargs) - prepared: Final = _prepare_ocr_request( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - kwargs=kwargs, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout=timeout, - ) - model = prepared.model - custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - response: Final = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=_is_async, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) +async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape + request: Final = _public_request("aocr", args, kwargs) + native: Final = select(request) if rust_ocr_enabled() else None + if native is not None: + try: + return await cast( # cast-ok: True selects the asynchronous result + Awaitable[OCRResponse], native(request, args, kwargs, True) + ) + except _decline_types(): + pass + fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator + Callable[..., Awaitable[OCRResponse]], legacy.aocr + ) + return await fallback(*args, **kwargs) - return response - except Exception as e: - error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) - error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model - raise litellm.exception_type( - model=error_model, - custom_llm_provider=error_provider, - original_exception=e, - completion_kwargs=completion_kwargs, - extra_kwargs=kwargs, - ) + +def _decline_types() -> tuple[type[BaseException], ...]: + exception_types: Final = native_exception_types() + return (exception_types[0],) if exception_types is not None else () diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7e6de474b0b..cd6739777dd 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -2162,6 +2162,10 @@ class MCPRequestHandler: # No team restrictions → use key restrictions allowed_tools = cast(list[str], key_tools) + allowed_tools = _as_list( + await MCPRequestHandler._apply_end_user_tool_ceiling(allowed_tools, server_id, user_api_key_auth) + ) + allowed_tools = _as_list( await MCPRequestHandler._apply_user_tool_ceiling( allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source @@ -3027,6 +3031,38 @@ class MCPRequestHandler: return list(user_tools) return list(set(allowed_tools) & set(user_tools)) + @staticmethod + async def _apply_end_user_tool_ceiling( + allowed_tools: Sequence[str] | None, + server_id: str, + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> Sequence[str] | None: + """Narrow a key/team tool allowlist by the end user's (customer's) tool entitlement.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_auth is None or not user_api_key_auth.end_user_id or prisma_client is None: + return allowed_tools + + object_permissions: Final = await MCPRequestHandler._get_end_user_object_permission( + user_api_key_auth, prisma_client + ) + if object_permissions is None: + return allowed_tools + + end_user_direct_tools: Final = global_mcp_server_manager.expand_tool_permissions( + object_permissions.mcp_tool_permissions + ).get(server_id) + end_user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(object_permissions, server_id) + end_user_tools: Final = MCPRequestHandler._union_tool_grants(end_user_direct_tools, end_user_toolset_tools) + if end_user_tools is None: + return allowed_tools + if allowed_tools is None: + return list(end_user_tools) + return list(set(allowed_tools) & set(end_user_tools)) + # Sentinel stored in cache when an agent has no object_permission, so we # don't re-query the DB on every MCP request for that agent. _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 28f3ec6521a..0ab76588b1f 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -7,7 +7,6 @@ just a form that asks the user for their API key — not a full identity-provide Endpoints implemented here: GET /.well-known/oauth-authorization-server — OAuth authorization server metadata - GET /.well-known/oauth-protected-resource — OAuth protected resource metadata GET /v1/mcp/oauth/authorize — Shows HTML form to collect the API key POST /v1/mcp/oauth/authorize — Stores temp auth code and redirects POST /v1/mcp/oauth/token — Exchanges code for a bearer JWT token @@ -19,7 +18,7 @@ import html as _html_module import time import uuid from typing import Final, cast -from urllib.parse import urlencode +from urllib.parse import urlencode, urlparse import jwt from fastapi import APIRouter, Depends, Form, HTTPException, Request @@ -27,14 +26,15 @@ from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import store_user_credential -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, -) from litellm.proxy._experimental.mcp_server.oauth_utils import ( + BYOK_RESOURCE_METADATA_PATH, TOKEN_NO_CACHE_HEADERS, + get_request_base_url, validate_loopback_redirect_uri, + well_known_root_suffix, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.middleware.per_request_root_path_middleware import get_server_root_paths # --------------------------------------------------------------------------- # In-memory store for pending authorization codes. @@ -596,13 +596,10 @@ def _build_authorize_html( # --------------------------------------------------------------------------- -@router.get("/.well-known/oauth-authorization-server", include_in_schema=False) -async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: - """RFC 8414 Authorization Server Metadata for the BYOK OAuth flow.""" - base_url: Final = get_request_base_url(request) +def _byok_authorization_server_response(base_url: str, issuer: str) -> JSONResponse: return JSONResponse( { - "issuer": base_url, + "issuer": issuer, "authorization_endpoint": f"{base_url}/v1/mcp/oauth/authorize", "token_endpoint": f"{base_url}/v1/mcp/oauth/token", "response_types_supported": ["code"], @@ -612,14 +609,36 @@ async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: ) -@router.get("/.well-known/oauth-protected-resource", include_in_schema=False) -async def oauth_protected_resource_metadata(request: Request) -> JSONResponse: - """RFC 9728 Protected Resource Metadata pointing back at this server.""" +@router.get("/.well-known/oauth-authorization-server", include_in_schema=False) +async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: base_url: Final = get_request_base_url(request) + return _byok_authorization_server_response(base_url, base_url) + + +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/v1/mcp/oauth", include_in_schema=False) +async def byok_authorization_server_metadata(request: Request) -> JSONResponse: + base_url: Final = get_request_base_url(request) + return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth") + + +@router.get("/.well-known/oauth-authorization-server/{root_path:path}/v1/mcp/oauth", include_in_schema=False) +async def byok_prefixed_authorization_server_metadata(request: Request, root_path: str) -> JSONResponse: + prefix: Final = f"/{root_path}" + if prefix not in get_server_root_paths(): + raise HTTPException(status_code=404, detail="Unknown proxy root path") + parsed: Final = urlparse(get_request_base_url(request)) + base_url: Final = f"{parsed.scheme}://{parsed.netloc}{prefix}" + return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth") + + +@router.get(BYOK_RESOURCE_METADATA_PATH, include_in_schema=False) +async def byok_protected_resource_metadata(request: Request) -> JSONResponse: + base_url: Final = get_request_base_url(request) + parsed: Final = urlparse(base_url) return JSONResponse( { - "resource": base_url, - "authorization_servers": [base_url], + "resource": f"{parsed.scheme}://{parsed.netloc}", + "authorization_servers": (f"{base_url}/v1/mcp/oauth",), } ) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 42fbe82531c..bafe33d0a6b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1842,6 +1842,30 @@ async def register_client_with_server( return JSONResponse(token_response) +@router.get("/authorize/mcp-session") +async def authorize_mcp_session( + request: Request, + redirect_uri: str, + client_id: str, + state: str = "", + code_challenge: str | None = None, + code_challenge_method: str | None = None, + response_type: str | None = None, + resource: str | None = None, +) -> Response: + return aggregate_authorize( + request=request, + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + response_type=response_type, + session_user_id=_session_cookie_user_id(request), + resource=resource, + ) + + @router.get("/{mcp_server_name}/authorize") @router.get("/authorize") async def authorize( @@ -2393,8 +2417,7 @@ async def _build_oauth_protected_resource_response( per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint supports and is admitted with a gateway session bearer. The per-server relay authorize/token endpoints stay registered for the keyed interactive flow (which - is challenged with an explicit ``authorization_uri``), and the root-resolved - (unnamed) legacy shape keeps the relay authorization server. + is challenged with an explicit ``authorization_uri``). Args: request: FastAPI Request object @@ -2405,15 +2428,11 @@ async def _build_oauth_protected_resource_response( Returns: OAuth protected resource metadata dict """ + if mcp_server_name is None: + return oauth_protected_resource_root(request) + 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 - - # When no server name provided, try to resolve the single OAuth2 server - if mcp_server_name is None: - resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) - if resolved: - mcp_server_name = resolved.server_name or resolved.name mcp_server: MCPServer | None = None if mcp_server_name: @@ -2478,7 +2497,7 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response - if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server: + if mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": resource_url, @@ -2542,6 +2561,17 @@ def _jwt_auth_issuers() -> list: return issuers +@router.get("/.well-known/oauth-protected-resource") +def oauth_protected_resource_root(request: Request) -> dict[str, str | tuple[str, ...]]: + request_base_url: Final = get_request_base_url(request) + parsed: Final = urlparse(request_base_url) + return { + "resource": f"{parsed.scheme}://{parsed.netloc}", + "authorization_servers": (f"{request_base_url}/mcp",), + "scopes_supported": (), + } + + def _build_aggregate_protected_resource_response(request: Request) -> dict: """RFC 9728 metadata for the aggregate /mcp resource: the gateway itself is the authorization server. No per-server names or scopes leak here; access @@ -2568,14 +2598,14 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: The issuer is ``{base}/mcp`` and must stay equal to the value the aggregate protected-resource document advertises: spec clients verify the issuer in the metadata matches the one that derived the well-known URL. - Advertises the root /authorize, /token, and /register endpoints and + Advertises the MCP session authorize endpoint, root /token and /register endpoints, and ``token_endpoint_auth_methods_supported: ["none", ...]`` because DCR clients (Claude Desktop, MCP Inspector) register as public clients; PKCE S256 is mandatory in the gateway's authorize flow.""" request_base_url: Final = get_request_base_url(request) return { "issuer": f"{request_base_url}/mcp", - "authorization_endpoint": f"{request_base_url}/authorize", + "authorization_endpoint": f"{request_base_url}/authorize/mcp-session", "token_endpoint": f"{request_base_url}/token", "introspection_endpoint": f"{request_base_url}/introspect", "registration_endpoint": f"{request_base_url}/register", @@ -2645,7 +2675,6 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments @router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp") -@router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp(request: Request, mcp_server_name: str | None = None): """ OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. @@ -2666,6 +2695,8 @@ async def oauth_protected_resource_mcp(request: Request, mcp_server_name: str | def _build_oauth_authorization_server_response( request: Request, mcp_server_name: str | None, + *, + issuer_path: str | None = None, ) -> dict: """Build OAuth authorization server metadata response (gateway-as-AS shape). @@ -2694,7 +2725,13 @@ def _build_oauth_authorization_server_response( _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth authorization server") - issuer: Final = f"{request_base_url}/{mcp_server_name}" if explicitly_named else request_base_url + issuer: Final = ( + f"{request_base_url}/{issuer_path}" + if issuer_path is not None + else f"{request_base_url}/{mcp_server_name}" + if explicitly_named + else request_base_url + ) return { "issuer": issuer, @@ -2724,6 +2761,7 @@ async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_n return _build_oauth_authorization_server_response( request=request, mcp_server_name=mcp_server_name, + issuer_path=f"mcp/{mcp_server_name}", ) @@ -2802,7 +2840,7 @@ async def jwks_json(request: Request): # Additional legacy pattern support -@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp") +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}/mcp") async def oauth_authorization_server_legacy(request: Request, mcp_server_name: str): """ OAuth authorization server discovery for legacy /{server_name}/mcp pattern. @@ -2810,6 +2848,7 @@ async def oauth_authorization_server_legacy(request: Request, mcp_server_name: s return _build_oauth_authorization_server_response( request=request, mcp_server_name=mcp_server_name, + issuer_path=f"{mcp_server_name}/mcp", ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cc7ea0c1ea5..3291d5effc2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -10,6 +10,7 @@ import asyncio import datetime import hashlib import json +import math import os import re import time @@ -26,8 +27,9 @@ from collections.abc import ( from contextlib import asynccontextmanager from dataclasses import dataclass, replace from functools import lru_cache +from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeAlias, TypedDict, TypeVar, cast from urllib.parse import ParseResult, urlparse import anyio @@ -44,11 +46,12 @@ from mcp.types import ( ResourceTemplate, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl, BaseModel +from pydantic import AnyUrl, BaseModel, TypeAdapter from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( MCP_CLIENT_TIMEOUT, MCP_HEALTH_CHECK_TIMEOUT, @@ -91,6 +94,7 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, canonicalize_url_identity, + get_byok_www_authenticate, ) from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, @@ -193,7 +197,6 @@ if TYPE_CHECKING: from mcp.shared.context import RequestContext from mcp.types import CreateMessageRequestParams - from litellm.caching.caching import InMemoryCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.mcp_server.mcp_toolset import MCPToolset @@ -1232,7 +1235,7 @@ async def _resolve_byok_mcp_auth_header( "Complete the OAuth authorization flow to provide your API key." ), }, - headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) return byok_cred @@ -1677,6 +1680,105 @@ def _record_mcp_guardrail_evaluations( verbose_logger.warning("Failed to record MCP guardrail evaluation for logging: %s", e) +_DiscoveryItem = TypeVar("_DiscoveryItem", bound=BaseModel) +_DiscoveryKey: TypeAlias = tuple[str, str | None] +_DISCOVERY_CACHE_LIMIT: Final = 1024 + + +class _DiscoveryCache(Generic[_DiscoveryItem]): + def __init__( + self, ttl: float, clock: Callable[[], float], adapter: TypeAdapter[tuple[_DiscoveryItem, ...]] + ) -> None: + self._ttl = ttl + self._adapter = adapter + self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, max_size_per_item=64, clock=clock) + self._pending: dict[ + _DiscoveryKey, asyncio.Task[list[_DiscoveryItem]] + ] = {} # mutable-ok: constant-time fetch registration + self._waiters: dict[asyncio.Task[list[_DiscoveryItem]], int] = {} # mutable-ok: constant-time waiter accounting + + def invalidate(self, server_id: str) -> None: + prefix: Final = f"[{json.dumps(server_id)}," + keys: Final = cast( # cast-ok: private cache contains only JSON string keys + "tuple[str, ...]", tuple(self._entries.cache_dict) + ) + for entry_key in keys: + if entry_key.startswith(prefix): + self._entries.delete_cache(entry_key) + for key in tuple(self._pending): + if key[0] == server_id: + self._pending.pop(key) + + @staticmethod + def _observe_completion(task: asyncio.Task[list[_DiscoveryItem]]) -> None: + if not task.cancelled(): + task.exception() + + async def get( + self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[list[_DiscoveryItem]]] + ) -> tuple[_DiscoveryItem, ...]: + if self._ttl <= 0: + return tuple(await fetch()) + entry: Final[object] = self._entries.get_cache(json.dumps(key)) + if entry is not None: + return self._adapter.validate_python(entry) + pending: Final = self._pending.get(key) + if pending is not None: + return await self._await_fetch(key, pending) + if len(self._pending) >= _DISCOVERY_CACHE_LIMIT: + return tuple(await fetch()) + task: Final = asyncio.create_task(self._fetch(key, fetch)) + self._pending[key] = task + task.add_done_callback(self._observe_completion) + return await self._await_fetch(key, task) + + async def _await_fetch( + self, key: _DiscoveryKey, task: asyncio.Task[list[_DiscoveryItem]] + ) -> tuple[_DiscoveryItem, ...]: + self._waiters[task] = self._waiters.get(task, 0) + 1 + try: + return tuple(item.model_copy(deep=True) for item in await asyncio.shield(task)) + finally: + remaining: Final = self._waiters[task] - 1 + if remaining: + self._waiters[task] = remaining + else: + self._waiters.pop(task) + if self._pending.get(key) is task: + self._pending.pop(key) + if not task.done(): + task.cancel() + + async def _fetch( + self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[list[_DiscoveryItem]]] + ) -> list[_DiscoveryItem]: + try: + items: Final = await fetch() + if self._pending.get(key) is asyncio.current_task(): + self._entries.set_cache( + json.dumps(key), + self._adapter.dump_json(tuple(items)), + ttl=self._ttl, + ) + return items + finally: + if self._pending.get(key) is asyncio.current_task(): + self._pending.pop(key) + + +def _mcp_discovery_cache_ttl() -> float: + raw: Final = os.environ.get("LITELLM_MCP_DISCOVERY_CACHE_TTL", "60") + try: + ttl: Final = float(raw) + except ValueError: + verbose_logger.warning("Invalid LITELLM_MCP_DISCOVERY_CACHE_TTL; using 60 seconds") + return 60.0 + if not math.isfinite(ttl) or ttl < 0: + verbose_logger.warning("Invalid LITELLM_MCP_DISCOVERY_CACHE_TTL; using 60 seconds") + return 60.0 + return ttl + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -1793,6 +1895,7 @@ class MCPServerManager: cred_provider: UpstreamCredentialProvider | None = None, per_user_oauth_token_store: InvalidatableOAuthTokenStore | None = None, per_user_token_cache: MCPPerUserTokenCache | None = None, + discovery_clock: Callable[[], float] = time.monotonic, ): self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( self.get_mcp_server_by_id @@ -1802,6 +1905,16 @@ class MCPServerManager: oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), ) + discovery_ttl: Final = _mcp_discovery_cache_ttl() + self._prompt_discovery_cache = _DiscoveryCache[Prompt]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[Prompt, ...]) + ) + self._resource_discovery_cache = _DiscoveryCache[Resource]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[Resource, ...]) + ) + self._template_discovery_cache = _DiscoveryCache[ResourceTemplate]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[ResourceTemplate, ...]) + ) self.registry: dict[str, MCPServer] = {} self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} @@ -2529,6 +2642,7 @@ class MCPServerManager: self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) + self._invalidate_discovery_lists(server_id) self.config_mcp_servers[server_id] = new_server self._set_oauth_discovery_deferred( server_id, @@ -2730,6 +2844,7 @@ class MCPServerManager: global_mcp_tool_registry, ) + self._invalidate_discovery_lists(server.server_id) prefix_root: Final = normalize_server_name(get_server_prefix(server)) if server.spec_path and prefix_root: openapi_key_prefix: Final = prefix_root + MCP_TOOL_PREFIX_SEPARATOR @@ -3106,6 +3221,7 @@ class MCPServerManager: # env_vars_are_encrypted=False. new_server: Final = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False) self._assign_unique_short_prefix(new_server) + self._invalidate_discovery_lists(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) @@ -3142,6 +3258,7 @@ class MCPServerManager: previous_server=self.registry[mcp_server.server_id], ) self._assign_unique_short_prefix(new_server) + self._invalidate_discovery_lists(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) @@ -4375,6 +4492,41 @@ class MCPServerManager: ) raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) + def _invalidate_discovery_lists(self, server_id: str) -> None: + self._prompt_discovery_cache.invalidate(server_id) + self._resource_discovery_cache.invalidate(server_id) + self._template_discovery_cache.invalidate(server_id) + + def _discovery_key( + self, + server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | dict[str, str] | None, + extra_headers: dict[str, str] | None, + stdio_env: dict[str, str] | None, + subject_token: str | None, + credential_fingerprint: str | None = None, + ) -> _DiscoveryKey: + per_user: Final = ( + server.requires_per_user_auth + or self._references_per_user_env_var(server) + or server.delegate_auth_to_upstream + or server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) + ) + if not (per_user or mcp_auth_header or extra_headers or stdio_env or subject_token): + return server.server_id, None + identity: Final = ( + (user_api_key_auth.user_id, user_api_key_auth.api_key) + if per_user and user_api_key_auth is not None + else None + ) + material: Final = json.dumps( + (identity, mcp_auth_header, extra_headers, stdio_env, subject_token, credential_fingerprint), + sort_keys=True, + separators=(",", ":"), + ) + return server.server_id, hashlib.sha256(material.encode()).hexdigest() + async def get_prompts_from_server( self, server: MCPServer, @@ -4384,47 +4536,38 @@ class MCPServerManager: add_prefix: bool = True, raw_headers: dict[str, str] | None = None, ) -> list[Prompt]: - """ - Helper method to get prompts from a single MCP server with prefixed names. - - Args: - server (MCPServer): The server to query prompts from - mcp_auth_header: Optional auth header for MCP server - - Returns: - List[Prompt]: List of prompts available on the server with prefixed names - """ - - verbose_logger.debug("Connecting to url: %s", server.url) - verbose_logger.info("get_prompts_from_server for %s...", server.name) - - client = None - try: - if server.static_headers: - if extra_headers is None: - extra_headers = {} - extra_headers.update(server.static_headers) - + headers: Final = ( + dict( + chain( + extra_headers.items() if extra_headers else (), + server.static_headers.items() if server.static_headers else (), + ) + ) + or None + ) stdio_env: Final = self._build_stdio_env(server, raw_headers) subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) - - client = await self._create_mcp_client( + client: Final = await self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, - extra_headers=extra_headers, + extra_headers=headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + ) + credential_fingerprint: Final = await client.discovery_auth_fingerprint() + key: Final = self._discovery_key( + server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token, credential_fingerprint ) - prompts: Final = await client.list_prompts() + async def fetch() -> list[Prompt]: + return await client.list_prompts(raise_on_error=True) - prefixed_or_original_prompts: Final = self._create_prefixed_prompts(prompts, server, add_prefix=add_prefix) - - return prefixed_or_original_prompts - - except Exception as e: - verbose_logger.warning("Failed to get prompts from server %s: %s", server.name, e) + items: Final = await self._prompt_discovery_cache.get(key, fetch) + return self._create_prefixed_prompts(items, server, add_prefix=add_prefix) + except Exception as error: + verbose_logger.warning("Failed to get prompts from server %s: %s", server.name, error) return [] async def get_resources_from_server( @@ -4436,38 +4579,38 @@ class MCPServerManager: add_prefix: bool = True, raw_headers: dict[str, str] | None = None, ) -> list[Resource]: - """Fetch available resources from a single MCP server.""" - - verbose_logger.debug("Connecting to url: %s", server.url) - verbose_logger.info("get_resources_from_server for %s...", server.name) - - client = None - try: - if server.static_headers: - if extra_headers is None: - extra_headers = {} - extra_headers.update(server.static_headers) - + headers: Final = ( + dict( + chain( + extra_headers.items() if extra_headers else (), + server.static_headers.items() if server.static_headers else (), + ) + ) + or None + ) stdio_env: Final = self._build_stdio_env(server, raw_headers) subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) - - client = await self._create_mcp_client( + client: Final = await self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, - extra_headers=extra_headers, + extra_headers=headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + ) + credential_fingerprint: Final = await client.discovery_auth_fingerprint() + key: Final = self._discovery_key( + server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token, credential_fingerprint ) - resources: Final = await client.list_resources() + async def fetch() -> list[Resource]: + return await client.list_resources(raise_on_error=True) - prefixed_resources: Final = self._create_prefixed_resources(resources, server, add_prefix=add_prefix) - - return prefixed_resources - - except Exception as e: - verbose_logger.warning("Failed to get resources from server %s: %s", server.name, e) + items: Final = await self._resource_discovery_cache.get(key, fetch) + return self._create_prefixed_resources(items, server, add_prefix=add_prefix) + except Exception as error: + verbose_logger.warning("Failed to get resources from server %s: %s", server.name, error) return [] async def get_resource_templates_from_server( @@ -4479,40 +4622,38 @@ class MCPServerManager: add_prefix: bool = True, raw_headers: dict[str, str] | None = None, ) -> list[ResourceTemplate]: - """Fetch available resource templates from a single MCP server.""" - - verbose_logger.debug("Connecting to url: %s", server.url) - verbose_logger.info("get_resource_templates_from_server for %s...", server.name) - - client = None - try: - if server.static_headers: - if extra_headers is None: - extra_headers = {} - extra_headers.update(server.static_headers) - + headers: Final = ( + dict( + chain( + extra_headers.items() if extra_headers else (), + server.static_headers.items() if server.static_headers else (), + ) + ) + or None + ) stdio_env: Final = self._build_stdio_env(server, raw_headers) subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) - - client = await self._create_mcp_client( + client: Final = await self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, - extra_headers=extra_headers, + extra_headers=headers, stdio_env=stdio_env, subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + ) + credential_fingerprint: Final = await client.discovery_auth_fingerprint() + key: Final = self._discovery_key( + server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token, credential_fingerprint ) - resource_templates: Final = await client.list_resource_templates() + async def fetch() -> list[ResourceTemplate]: + return await client.list_resource_templates(raise_on_error=True) - prefixed_templates: Final = self._create_prefixed_resource_templates( - resource_templates, server, add_prefix=add_prefix - ) - - return prefixed_templates - - except Exception as e: - verbose_logger.warning("Failed to get resource templates from server %s: %s", server.name, e) + items: Final = await self._template_discovery_cache.get(key, fetch) + return self._create_prefixed_resource_templates(items, server, add_prefix=add_prefix) + except Exception as error: + verbose_logger.warning("Failed to get resource_templates from server %s: %s", server.name, error) return [] async def read_resource_from_server( @@ -5220,7 +5361,7 @@ class MCPServerManager: return prefixed_tools def _create_prefixed_prompts( - self, prompts: list[Prompt], server: MCPServer, add_prefix: bool = True + self, prompts: Sequence[Prompt], server: MCPServer, add_prefix: bool = True ) -> list[Prompt]: """ Create prefixed prompts and update prompt mapping. @@ -5247,7 +5388,7 @@ class MCPServerManager: return prefixed_prompts def _create_prefixed_resources( - self, resources: list[Resource], server: MCPServer, add_prefix: bool = True + self, resources: Sequence[Resource], server: MCPServer, add_prefix: bool = True ) -> list[Resource]: """Prefix resource names and track origin server for read requests.""" @@ -5264,7 +5405,7 @@ class MCPServerManager: def _create_prefixed_resource_templates( self, - resource_templates: list[ResourceTemplate], + resource_templates: Sequence[ResourceTemplate], server: MCPServer, add_prefix: bool = True, ) -> list[ResourceTemplate]: @@ -6001,6 +6142,7 @@ class MCPServerManager: failure is logged, never raised, because the DB write already succeeded and the TTL remains the backstop. """ + self._invalidate_discovery_lists(server_id) try: await self._per_user_oauth_token_store.invalidate(user_id, server_id) except Exception as exc: # noqa: BLE001 - cache drop is best-effort; TTL is the backstop @@ -6466,6 +6608,9 @@ class MCPServerManager: for registry_key in dropped_registry_keys: self._invalidate_oauth_discovery_state(previous_registry[registry_key].server_id) + for server_id in previous_registry.keys() | registered_registry.keys(): + if previous_registry.get(server_id) != registered_registry.get(server_id): + self._invalidate_discovery_lists(server_id) self.registry = registered_registry # A discovery task may have published into ``previous_registry`` while # this replacement was being staged. Reconcile every published entry diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 1ca2ffc703d..39865a35ec6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -16,6 +16,7 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( normalize_token_endpoint_auth_method, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -126,6 +127,14 @@ def _resolve_proxy_base_url_env() -> str | None: return None +BYOK_RESOURCE_METADATA_PATH: Final = "/v1/mcp/oauth/protected-resource" + + +def get_byok_www_authenticate() -> str: + base_url: Final = _resolve_proxy_base_url_env() or get_request_root_path().rstrip("/") + return f'Bearer resource_metadata="{base_url}{BYOK_RESOURCE_METADATA_PATH}"' + + def get_request_base_url(request: Request) -> str: """ Get the base URL for the request, considering X-Forwarded-* headers. diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 43a7f894db8..fc87db69e16 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import ( ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, + get_byok_www_authenticate, get_passthrough_www_authenticate, get_route_relative_request_path, well_known_root_suffix, @@ -2852,7 +2853,7 @@ if MCP_AVAILABLE: "server_name": mcp_server.server_name or mcp_server.name, "message": "User identity is required for BYOK servers", }, - headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) # Check shared credential cache before hitting the DB. @@ -2873,9 +2874,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) return @@ -2914,7 +2913,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) async def execute_mcp_tool( @@ -3068,9 +3067,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) mcp_auth_header = byok_cred elif mcp_server.is_byok: diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 53af85baac6..cb7a18cd107 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -22415,47 +22415,34 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource_get", - "parameters": [ - { - "in": "query", - "name": "mcp_server_name", - "required": false, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Mcp Server Name" - } - } - ], + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get", "responses": { "200": { - "content": { - "application/json": { - "schema": {} - } - }, - "description": "Successful Response" - }, - "422": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/HTTPValidationError" + "additionalProperties": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ] + }, + "title": "Response Oauth Protected Resource Root Well Known Oauth Protected Resource Get", + "type": "object" } } }, - "description": "Validation Error" + "description": "Successful Response" } }, - "summary": "Oauth Protected Resource Mcp", + "summary": "Oauth Protected Resource Root", "tags": [ "mcp_byok_oauth" ] @@ -24636,47 +24623,34 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource_get_2", - "parameters": [ - { - "in": "query", - "name": "mcp_server_name", - "required": false, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Mcp Server Name" - } - } - ], + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2", "responses": { "200": { - "content": { - "application/json": { - "schema": {} - } - }, - "description": "Successful Response" - }, - "422": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/HTTPValidationError" + "additionalProperties": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ] + }, + "title": "Response Oauth Protected Resource Root Well Known Oauth Protected Resource Get", + "type": "object" } } }, - "description": "Validation Error" + "description": "Successful Response" } }, - "summary": "Oauth Protected Resource Mcp", + "summary": "Oauth Protected Resource Root", "tags": [ "mcp_discoverable" ] @@ -25052,6 +25026,129 @@ ] } }, + "/authorize/mcp-session": { + "get": { + "operationId": "authorize_mcp_session_authorize_mcp_session_get", + "parameters": [ + { + "in": "query", + "name": "redirect_uri", + "required": true, + "schema": { + "title": "Redirect Uri", + "type": "string" + } + }, + { + "in": "query", + "name": "client_id", + "required": true, + "schema": { + "title": "Client Id", + "type": "string" + } + }, + { + "in": "query", + "name": "state", + "required": false, + "schema": { + "default": "", + "title": "State", + "type": "string" + } + }, + { + "in": "query", + "name": "code_challenge", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Code Challenge" + } + }, + { + "in": "query", + "name": "code_challenge_method", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Code Challenge Method" + } + }, + { + "in": "query", + "name": "response_type", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Response Type" + } + }, + { + "in": "query", + "name": "resource", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Resource" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Authorize Mcp Session", + "tags": [ + "mcp_discoverable" + ] + } + }, "/callback": { "get": { "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7b0e92d81ae..ae6c042ab3a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -40,6 +40,11 @@ from litellm.types.mcp import ( MCPTransportType, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.proxy.carried_budget_state import ( + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem @@ -1302,6 +1307,11 @@ class UpdateKeyRequest(KeyRequestBase): rotation_interval: str | None = None organization_id: str | None = None + project_id: str | None = Field( + default=None, + description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.", + ) + @model_validator(mode="before") @classmethod def drop_blank_team_id(cls, values: object) -> object: @@ -3106,6 +3116,9 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) + user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) + org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used @@ -4388,7 +4401,12 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): agent_ids: tuple[str, ...] = () +class TeamInfoMember(Member): + user_alias: str | None = None + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): + members_with_roles: tuple[TeamInfoMember, ...] = () team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1efc9611fe6..9c175242a9a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -102,6 +102,7 @@ from litellm.proxy.guardrails.tool_name_extraction import ( ) from litellm.proxy.route_llm_request import route_request from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository @@ -1022,6 +1023,19 @@ async def common_checks( fallback_spend=user_object.spend or 0.0, max_budget=user_budget, ) + call_info: Final = CallInfo( + spend=user_spend, + max_budget=user_budget, + user_id=user_object.user_id, + user_email=user_object.user_email, + event_group=Litellm_EntityType.USER, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="user_budget", + user_info=call_info, + ) + ) if math.isfinite(user_budget) and user_spend >= user_budget: raise litellm.BudgetExceededError( current_cost=user_spend, @@ -5622,6 +5636,8 @@ async def _organization_max_budget_check( if org_table is None: return + carry_organization_budget_state(valid_token=valid_token, org_table=org_table) + # Get max_budget from organization's budget table org_max_budget: float | None = None if org_table.litellm_budget_table is not None: diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py new file mode 100644 index 00000000000..52e26e885c9 --- /dev/null +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -0,0 +1,307 @@ +"""Warm the user, team, membership, org and project cache entries auth reads: one MGET, one DB query, one +pipeline write instead of one Redis GET (and one DB query when cold) per object. The per-object getters stay +the readers and the fallback, so enforcement never depends on this running.""" + +from __future__ import annotations + +import time +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, Protocol, TypeAlias + +from pydantic import BaseModel, TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_cache import RedisCache +from litellm.constants import DEFAULT_IN_MEMORY_TTL +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.team import LiteLLM_TeamTableCachedObj +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import LiteLLM_ProjectTableCachedObj, UserAPIKeyAuth +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + get_management_object_ttl, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, +) +from litellm.proxy.utils import PrismaClient + +_RowKind: TypeAlias = Literal["user_row", "team_row", "membership_row", "organization_row", "project_row"] + +_TEAM_MEMBERSHIP_AUTH_TTL: Final = 5 +_RowValues: Final = TypeAdapter(dict[str, object]) +_NO_ROWS: Final[Mapping[str, object]] = MappingProxyType({}) +_TEAM_BOUND_ROWS: Final = frozenset({"team_row", "membership_row"}) +_REFRESH_STAMPED_ROWS: Final = frozenset({"team_row", "project_row"}) + + +def _lists_as_json(alias: str, columns: Sequence[str]) -> str: + """Prisma reads a NULL scalar list as ``[]``; ``to_jsonb`` reads it as ``null``, which the models reject.""" + return ", ".join(f"'{column}', COALESCE(to_jsonb({alias}.{column}), '[]'::jsonb)" for column in columns) + + +_USER_LISTS: Final = _lists_as_json("u", ("teams", "models", "allowed_cache_controls", "policies")) +_TEAM_LISTS: Final = _lists_as_json( + "t", + ( + "admins", + "members", + "models", + "team_member_permissions", + "access_group_ids", + "policies", + "default_team_member_models", + ), +) +_ORG_LISTS: Final = _lists_as_json("o", ("models",)) +_PROJECT_LISTS: Final = _lists_as_json("p", ("models",)) +_PERMISSION_LISTS: Final = _lists_as_json( + "op", + ( + "mcp_servers", + "mcp_access_groups", + "mcp_toolsets", + "blocked_tools", + "vector_stores", + "agents", + "agent_access_groups", + "models", + "search_tools", + "skills", + ), +) +_BUDGET_LISTS: Final = _lists_as_json("b", ("allowed_models",)) + + +def _budget_json(owner_alias: str) -> str: + return ( + f"(SELECT to_jsonb(b) || jsonb_build_object({_BUDGET_LISTS}) " + f'FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = {owner_alias}.budget_id)' + ) + + +def _permission_json(owner_alias: str) -> str: + return ( + f"(SELECT to_jsonb(op) || jsonb_build_object({_PERMISSION_LISTS}) " + f'FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = {owner_alias}.object_permission_id)' + ) + + +_SQL: Final = f""" +SELECT + ( + SELECT to_jsonb(u) || jsonb_build_object( + {_USER_LISTS}, + 'organization_memberships', + COALESCE(( + SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id + ), '[]'::jsonb) + ) + FROM "LiteLLM_UserTable" u WHERE u.user_id = $1 + ) AS user_row, + ( + SELECT to_jsonb(t) || jsonb_build_object( + {_TEAM_LISTS}, + 'litellm_model_table', ( + SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) + FROM "LiteLLM_ModelTable" m WHERE m.id = t.model_id + ), + 'object_permission', {_permission_json("t")} + ) + FROM "LiteLLM_TeamTable" t WHERE t.team_id = $2 + ) AS team_row, + ( + SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', {_budget_json("tm")}) + FROM "LiteLLM_TeamMembership" tm WHERE tm.user_id = $3 AND tm.team_id = $2 + ) AS membership_row, + ( + SELECT to_jsonb(o) || jsonb_build_object( + {_ORG_LISTS}, + 'litellm_budget_table', {_budget_json("o")}, + 'object_permission', {_permission_json("o")} + ) + FROM "LiteLLM_OrganizationTable" o WHERE o.organization_id = $4 + ) AS organization_row, + ( + SELECT to_jsonb(p) || jsonb_build_object( + {_PROJECT_LISTS}, + 'litellm_budget_table', {_budget_json("p")}, + 'object_permission', {_permission_json("p")} + ) + FROM "LiteLLM_ProjectTable" p WHERE p.project_id = $5 + ) AS project_row +""" + + +@dataclass(frozen=True, slots=True) +class AuthObjectRefs: + """Ids of the objects a request's auth checks will read. ``None`` means not referenced.""" + + user_id: str | None = None + team_id: str | None = None + membership_user_id: str | None = None + organization_id: str | None = None + project_id: str | None = None + + @classmethod + def from_token(cls, token: UserAPIKeyAuth) -> AuthObjectRefs: + has_membership: Final = token.team_id is not None and token.user_id is not None + return cls( + user_id=token.user_id, + team_id=token.team_id, + membership_user_id=token.user_id if has_membership else None, + organization_id=token.org_id, + project_id=token.project_id, + ) + + +class _InMemoryCache(Protocol): + def get_cache(self, key: str) -> object: ... + def set_cache(self, key: str, value: object, *, ttl: float | None = ...) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _CacheEntry: + cache_key: str + row: _RowKind + model_type: type[BaseModel] + ttl: float | None + + +def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_CacheEntry]: + if refs.user_id is not None: + yield _CacheEntry(refs.user_id, "user_row", LiteLLM_UserTable, management_ttl) + if refs.team_id is not None: + yield _CacheEntry(f"team_id:{refs.team_id}", "team_row", LiteLLM_TeamTableCachedObj, management_ttl) + if refs.team_id is not None and refs.membership_user_id is not None: + yield _CacheEntry( + team_membership_auth_cache_key(team_id=refs.team_id, user_id=refs.membership_user_id), + "membership_row", + LiteLLM_TeamMembership, + _TEAM_MEMBERSHIP_AUTH_TTL, + ) + yield _CacheEntry( + team_membership_reservation_cache_key(user_id=refs.membership_user_id, team_id=refs.team_id), + "membership_row", + LiteLLM_TeamMembership, + None, + ) + if refs.organization_id is not None: + yield _CacheEntry( + f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL + ) + yield _CacheEntry( + f"org_id:{refs.organization_id}:with_budget", + "organization_row", + LiteLLM_OrganizationTable, + DEFAULT_IN_MEMORY_TTL, + ) + if refs.project_id is not None: + yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl) + + +def _entries(refs: AuthObjectRefs, cache: UserApiKeyCache) -> tuple[_CacheEntry, ...]: + return tuple(_iter_entries(refs, get_management_object_ttl(cache))) + + +def _missing_in_memory(entries: Sequence[_CacheEntry], memory: _InMemoryCache) -> tuple[_CacheEntry, ...]: + return tuple(entry for entry in entries if memory.get_cache(key=entry.cache_key) is None) + + +def _set_in_memory(memory: _InMemoryCache, cache_key: str, value: object, ttl: float | None) -> None: + if ttl is None: + memory.set_cache(key=cache_key, value=value) + else: + memory.set_cache(key=cache_key, value=value, ttl=ttl) + + +async def _fill_from_redis(entries: Sequence[_CacheEntry], redis_cache: RedisCache, memory: _InMemoryCache) -> None: + if not entries: + return + found: Final = _RowValues.validate_python( + await redis_cache.async_batch_get_cache(key_list=sorted(entry.cache_key for entry in entries)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped cache API + ) + for entry, value in ((entry, found.get(entry.cache_key)) for entry in entries): + if value is not None: + _set_in_memory(memory, entry.cache_key, value, entry.ttl) + + +def _validate_row( + row_value: object, model_type: type[BaseModel], row: _RowKind, refreshed_at: float +) -> BaseModel | None: + if row_value is None: + return None + try: + columns: Final = _RowValues.validate_python(row_value) + if row in _REFRESH_STAMPED_ROWS: + stamped: Final = {**columns, "last_refreshed_at": refreshed_at} # mutable-ok: validators write into it + return model_type.model_validate(stamped) + return model_type.model_validate(columns) + except ValidationError as e: + verbose_proxy_logger.warning("auth prefetch: %s did not validate as %s: %s", row, model_type.__name__, e) + return None + + +async def _fetch_rows( + refs: AuthObjectRefs, kinds: frozenset[_RowKind], prisma_client: PrismaClient +) -> Mapping[str, object]: + row: Final[object] = await prisma_client.db.query_first( # pyright: ignore[reportAny] # prisma types query_first as Any + _SQL, + refs.user_id if "user_row" in kinds else None, + refs.team_id if kinds & _TEAM_BOUND_ROWS else None, + refs.membership_user_id if "membership_row" in kinds else None, + refs.organization_id if "organization_row" in kinds else None, + refs.project_id if "project_row" in kinds else None, + ) + return _RowValues.validate_python(row) if row is not None else _NO_ROWS + + +async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: UserApiKeyCache) -> None: + payloads: Final = tuple( + (entry.cache_key, CacheCodec.serialize(value, model_type=entry.model_type), entry.ttl) + for entry, value in entries + ) + memory: Final[_InMemoryCache] = cache.in_memory_cache + for cache_key, payload, ttl in payloads: + _set_in_memory(memory, cache_key, payload, cache.default_in_memory_ttl if ttl is None else ttl) + if cache.redis_cache is not None: + await cache.redis_cache.async_set_cache_pipeline_with_ttls(payloads) + + +async def _fill_from_db( + refs: AuthObjectRefs, entries: Sequence[_CacheEntry], cache: UserApiKeyCache, prisma_client: PrismaClient +) -> None: + if not entries: + return + model_for: Final[Mapping[_RowKind, type[BaseModel]]] = MappingProxyType( + {entry.row: entry.model_type for entry in entries} + ) + rows: Final = await _fetch_rows(refs, frozenset(model_for), prisma_client) + refreshed_at: Final = time.time() + objects: Final[Mapping[_RowKind, BaseModel | None]] = MappingProxyType( + {row: _validate_row(rows.get(row), model_type, row, refreshed_at) for row, model_type in model_for.items()} + ) + writes: Final = tuple((entry, value) for entry in entries if (value := objects[entry.row]) is not None) + if writes: + await _write_back(writes, cache) + + +async def prefetch_auth_objects( + refs: AuthObjectRefs, + user_api_key_cache: UserApiKeyCache, + prisma_client: PrismaClient | None, +) -> None: + """Best effort: any failure leaves the per-object getters to fetch as before.""" + try: + memory: Final[_InMemoryCache] = user_api_key_cache.in_memory_cache + missing: Final = _missing_in_memory(_entries(refs, user_api_key_cache), memory) + if user_api_key_cache.redis_cache is not None: + await _fill_from_redis(missing, user_api_key_cache.redis_cache, memory) + if prisma_client is None: + return + await _fill_from_db(refs, _missing_in_memory(missing, memory), user_api_key_cache, prisma_client) + except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own + verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 20ab9904f46..9828311112e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -24,6 +24,7 @@ from starlette.exceptions import WebSocketException import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.caching.redis_cache import RedisCache from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, INVALID_VIRTUAL_KEY_ERROR_MARKER, @@ -63,6 +64,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, get_end_user_id_from_request_body, @@ -101,6 +103,12 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state +from litellm.proxy.spend_tracking.spend_counter_batch import ( + bind_admission_counter_keys, + release_spend_counter_batch, + spend_counter_batch_scope, +) from litellm.proxy.utils import ( PrismaClient, ProxyLogging, @@ -1970,6 +1978,9 @@ async def _user_api_key_auth_builder( llm_model_list=llm_model_list, llm_router=llm_router, ) + await _prefetch_referenced_auth_objects( + valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client + ) # Check 2. If user_id for this token is in budget - done in common_checks() if valid_token.user_id is not None: @@ -2621,6 +2632,11 @@ async def _run_centralized_common_checks( None if isinstance(end_user_result, BaseException) else end_user_result ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result + carry_team_and_user_budget_state( + valid_token=user_api_key_auth_obj, + team_object=team_object, + user_object=user_object, + ) if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id @@ -2672,21 +2688,25 @@ async def _run_centralized_common_checks( user_api_key_dict=user_api_key_auth_obj, ) - _ = await common_checks( - request=request, - request_body=request_data, - team_object=team_object, - user_object=user_object, - end_user_object=end_user_object, - general_settings=general_settings, - global_proxy_spend=global_proxy_spend, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=user_api_key_auth_obj, - skip_budget_checks=skip_budget_checks, - project_object=project_object, - ) + bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id) + try: + _ = await common_checks( + request=request, + request_body=request_data, + team_object=team_object, + user_object=user_object, + end_user_object=end_user_object, + general_settings=general_settings, + global_proxy_spend=global_proxy_spend, + route=route, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=user_api_key_auth_obj, + skip_budget_checks=skip_budget_checks, + project_object=project_object, + ) + finally: + release_spend_counter_batch() await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, @@ -2864,6 +2884,28 @@ async def _authorize_authenticated_request( return None +def _spend_counter_redis_cache() -> RedisCache | None: + from litellm.proxy.proxy_server import spend_counter_cache + + return spend_counter_cache.redis_cache + + +async def _prefetch_referenced_auth_objects( + valid_token: UserAPIKeyAuth, + end_user_id: str | None, + user_api_key_cache: UserApiKeyCache, + prisma_client: PrismaClient | None, +) -> None: + """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold). + Runs after the key's model access check so a denied request costs no more than it did before.""" + bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None) + await prefetch_auth_objects( + refs=AuthObjectRefs.from_token(valid_token), + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + + def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None: """Anchor the OTLP destinations this key or team overrides its traces to. @@ -2928,7 +2970,7 @@ async def user_api_key_auth( # Run the whole auth phase inside a live ``auth`` span so the DB lookups it # triggers (key/user/team object reads) nest under it instead of flattening # onto the server span. No-op when OTel V2 isn't active. - with phase_span(f"auth {route}"): + with phase_span(f"auth {route}"), spend_counter_batch_scope(_spend_counter_redis_cache()): try: user_api_key_auth_obj: Final = await _user_api_key_auth_builder( request=request, diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 2071576a943..3b0ff9d7add 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -536,7 +536,30 @@ It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABL The key in the file is the login's own, so it expires with it (24h by default): run `lite login --config-claude` again after that, which rewrites the key in place. Earlier versions wrote an `apiKeyHelper` that ran `lite auth print-token` instead, so a later login refreshed Claude Code by itself; that meant Claude Code spawning a full `lite` start, keychain check included, on every credential refresh, so the helper is no longer written and a stale one is stripped by the next `--config-claude` or `configure claude`. Like `lite up`, the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first -#### Configuring Claude Code Once, With a Virtual Key +#### Configuring Claude Code or Codex Once, With a Virtual Key + +Run the setup wizard with your gateway URL and a long-lived virtual key: + +```bash +lite configure --api-key sk-... --gateway-url https://your-proxy.example.com +``` + +Select Claude Code, Codex, or both, then choose a gateway model for each selected agent. The wizard validates the key and reads the models your key can access before changing settings. Start either configured agent normally with `claude` or `codex`; the gateway connection persists across terminals without a wrapper or exported API key + +`--gateway-url` also accepts a deployment path prefix and a trailing `/v1`. `--base-url` is an alias. If omitted, setup uses `lite --base-url`, `LITELLM_PROXY_URL`, or the saved CLI URL; the wizard asks for a URL when none was provided + +For a scripted setup, name the agent and model: + +```bash +lite configure --gateway-url https://your-proxy.example.com codex --api-key sk-... --model my-coding-model +lite unconfigure codex +``` + +Codex setup requires an installed stable Codex version of [0.129.0 or newer](https://github.com/openai/codex/releases/tag/rust-v0.129.0), which prevents repository settings from redirecting requests carrying your saved key. Setup checks `codex --version` before fetching models or changing either selected agent's settings. Undo remains available without Codex installed + +Codex setup updates `~/.codex/config.toml` (or `$CODEX_HOME/config.toml`) with the selected model and a LiteLLM Responses provider. The gateway key lives in that provider's static Authorization header, in a file written atomically with owner-only permissions. Other providers, hooks, MCP servers and comments are preserved. A default profile selection is removed so it cannot override the gateway settings; its contents are preserved, and undo restores the selection. Explicit Codex flags and supported project settings still follow Codex's normal precedence + +The Codex undo receipt is kept in a private `.litellm` directory beside the resolved config file. `lite unconfigure codex` restores only values still holding what configure wrote, preserving later edits. The provider URL and credential are restored together. Symlinks are followed and their targets become owner-only; keep these credential-bearing files out of version control `lite configure claude` wires Claude Code up persistently with a long-lived virtual key, a pinned model and an undo, and `lite unconfigure claude` puts things back: @@ -548,7 +571,7 @@ claude The key comes from `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) and is written into `env.ANTHROPIC_AUTH_TOKEN`; without one the command refuses, since a `lite login` credential expires within a day and keeping it fresh would mean Claude Code running `lite` through `apiKeyHelper` on every credential refresh. The command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key and as `env.ANTHROPIC_MODEL`, both of which have to be on `/v1/models` for the key. The second one matters for `claude -c` and `claude --resume`: a resumed session otherwise re-sends the model its transcript recorded, which behind an auto-router with `return_raw_model_name: true` is the tier model that answered, and a key scoped to the router alias gets a 403 for it; `ANTHROPIC_MODEL` outranks the transcript on resume. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control -Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt +Plain `lite configure`, with no agent named, asks which agents to wire and which gateway model each starts on, picked from `/v1/models` with a type-to-filter prompt. All choices and selected config files are checked before the first settings write. If a later filesystem write fails, the output identifies each agent already configured and its undo command What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any request diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 0354aefa95c..f424e07968e 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -125,6 +125,18 @@ def build_agent_env( return env +def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]: + return MappingProxyType( + { + "name": "LiteLLM proxy", + "base_url": base_url.rstrip("/") + "/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": False, + } + ) + + def _codex_proxy_args(base_url: str) -> list[str]: """Codex `-c` overrides that point it at the proxy. @@ -134,21 +146,19 @@ def _codex_proxy_args(base_url: str) -> list[str]: because the proxy does not speak the Responses WebSocket protocol. The key is read from OPENAI_API_KEY, which build_agent_env already exports. """ - root: Final = base_url.rstrip("/") + "/v1" provider: Final = f"model_providers.{CODEX_PROXY_PROVIDER}" return [ "-c", f'model_provider="{CODEX_PROXY_PROVIDER}"', - "-c", - f'{provider}.name="LiteLLM proxy"', - "-c", - f'{provider}.base_url="{root}"', + *( + argument + for key, value in codex_proxy_provider(base_url).items() + for argument in ("-c", f"{provider}.{key}={json.dumps(value)}") + ), "-c", f'{provider}.env_key="{OPENAI_API_KEY_ENV}"', "-c", - f'{provider}.wire_api="responses"', - "-c", - f"{provider}.supports_websockets=false", + f"{provider}.http_headers={{}}", ] diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index e6231f3cac9..1473e40070f 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -458,11 +458,17 @@ def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None: return ConfigureReceipt.model_validate_json(state_path.read_bytes()) except (OSError, ValidationError) as e: raise ClaudeSettingsError( - f"{state_path} is not a readable `lite configure claude` receipt ({e}). " + f"{state_path} is not a readable `lite configure claude` receipt. " "Remove it and edit Claude Code's settings by hand if they still point at the proxy." ) from e +def preflight_claude_settings(settings_path: Path) -> None: + refuse_while_owned(settings_path, settings_file_owners(settings_path)) + _env_object(load_json_or_empty(settings_path), settings_path) + read_configure_receipt(configure_state_path(settings_path)) + + def configure_claude_settings( base_url: str, credential: StaticToken, diff --git a/litellm/proxy/client/cli/commands/codex_settings.py b/litellm/proxy/client/cli/commands/codex_settings.py new file mode 100644 index 00000000000..686eaa47ff0 --- /dev/null +++ b/litellm/proxy/client/cli/commands/codex_settings.py @@ -0,0 +1,307 @@ +import hashlib +import json +import re +import subprocess +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import reduce +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import tomlkit +from pydantic import BaseModel, ConfigDict, ValidationError +from tomlkit.container import OutOfOrderTableProxy +from tomlkit.exceptions import TOMLKitError +from tomlkit.items import InlineTable, Table +from tomlkit.toml_document import TOMLDocument + +from litellm.litellm_core_utils.private_json import ( + commit_staged_json, + discard_staged_json, + ensure_private_dir, + stage_private_bytes, + stage_private_json, +) + +from .agents import CODEX_PROXY_PROVIDER, codex_proxy_provider + +_PROVIDER_PATH: Final = f"model_providers.{CODEX_PROXY_PROVIDER}" +_OWNED_PATHS: Final = ("model_provider", "model", "profile", _PROVIDER_PATH) +_Table: TypeAlias = TOMLDocument | Table | InlineTable | OutOfOrderTableProxy +_EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +_MIN_CODEX_VERSION: Final = (0, 129, 0) + + +class CodexSettingsError(Exception): + pass + + +class _Receipt(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + version: Literal[1] = 1 + settings_path: str + file_existed: bool + providers_existed: bool + previous: Mapping[str, str | None] + written: Mapping[str, str] + + +@dataclass(frozen=True, slots=True) +class CodexUnconfigureOutcome: + restored: tuple[str, ...] + kept: tuple[str, ...] + file_removed: bool + + +def codex_configure_state_path(settings_path: Path) -> Path: + target: Final = settings_path.resolve() + digest: Final = hashlib.sha256(str(target).encode()).hexdigest() + return target.parent / ".litellm" / f"codex_configure_{digest}.json" + + +def _read(settings_path: Path) -> TOMLDocument: + try: + document: Final = tomlkit.parse(settings_path.read_bytes()) if settings_path.exists() else tomlkit.document() + except (OSError, UnicodeError, TOMLKitError) as error: + raise CodexSettingsError( + f"Could not read Codex settings at {settings_path}; no settings were changed" + ) from error + providers: Final = _mapping(document).get("model_providers") + parent: Final = _table(providers) + if providers is not None and parent is None: + raise CodexSettingsError("Codex model_providers must be a TOML table; no settings were changed") + entries: Final = _mapping(parent) if parent is not None else _EMPTY + configured: Final = entries.get(CODEX_PROXY_PROVIDER) + if configured is not None and _table(configured) is None: + raise CodexSettingsError("Codex model_providers.litellm must be a TOML table; no settings were changed") + return document + + +def _mapping(value: Mapping[str, object]) -> Mapping[str, object]: + return value + + +def _table(value: object) -> _Table | None: + return value if isinstance(value, (TOMLDocument, Table, InlineTable, OutOfOrderTableProxy)) else None + + +def _snapshot(document: TOMLDocument, path: str) -> str | None: + section, _, key = path.rpartition(".") + parent: Final = _table(_mapping(document).get(section)) if section else document + if parent is None or key not in parent: + return None + values: Final = _mapping(parent) + return tomlkit.dumps(MappingProxyType({"value": values[key]})) + + +def _fingerprint(value: str | None) -> str: + normalized: Final = "missing" if value is None else json.dumps(tomlkit.parse(value), sort_keys=True, default=str) + return hashlib.sha256(normalized.encode()).hexdigest() + + +def _with(document: TOMLDocument, path: str, snapshot: str | None) -> TOMLDocument: + section, _, key = path.rpartition(".") + if section and section not in document and snapshot is not None: + contents: Final = tomlkit.parse(tomlkit.dumps(MappingProxyType({key: tomlkit.parse(snapshot).item("value")}))) + return tomlkit.parse(document.as_string() + "\n" + tomlkit.dumps(MappingProxyType({section: contents}))) + # mutable-ok: TOMLKit editing requires private node mutation to preserve comments and order + updated: Final = tomlkit.parse(document.as_string()) + parent: Final = _table(_mapping(updated).get(section)) if section else updated + if parent is None: + return updated + if snapshot is None: + if key in parent: + del parent[key] + else: + parent[key] = tomlkit.parse(snapshot).item("value") + return updated + + +def _receipt(settings_path: Path) -> _Receipt | None: + path: Final = codex_configure_state_path(settings_path) + if not path.exists(): + return None + try: + receipt: Final = _Receipt.model_validate_json(path.read_bytes()) + if receipt.settings_path != str(settings_path.resolve()) or frozenset(receipt.previous) != frozenset( + receipt.written + ): + raise ValueError("invalid receipt scope") + if not frozenset(receipt.written) <= frozenset(_OWNED_PATHS): + raise ValueError("invalid receipt ownership") + for snapshot in receipt.previous.values(): + if snapshot is not None and tuple(tomlkit.parse(snapshot)) != ("value",): + raise ValueError("invalid receipt snapshot") + except (OSError, UnicodeError, TOMLKitError, ValidationError, ValueError) as error: + raise CodexSettingsError( + f"Could not read the Codex configure receipt at {path}; no settings were changed" + ) from error + return receipt + + +def _codex_version() -> str | None: + try: + result: Final = subprocess.run(("codex", "--version"), capture_output=True, text=True, timeout=5, check=False) + except (OSError, subprocess.SubprocessError, UnicodeError): + return None + return result.stdout if result.returncode == 0 else None + + +def require_safe_codex(*, version: Callable[[], str | None] = _codex_version) -> None: + output: Final = version() + matched: Final = re.fullmatch(r"codex-cli (\d+)\.(\d+)\.(\d+)", output.strip()) if output is not None else None + if matched is not None and tuple(int(part) for part in matched.groups()) >= _MIN_CODEX_VERSION: + return + raise CodexSettingsError( + "Codex 0.129.0 or newer (stable) must be installed before saving a gateway key. " + "Older versions allow repository settings to redirect authenticated requests. " + "Install or update Codex, check `codex --version`, then retry." + ) + + +def preflight_codex_settings(settings_path: Path) -> None: + require_safe_codex() + _read(settings_path) + _receipt(settings_path) + + +def _ours(document: TOMLDocument, path: str, receipt: _Receipt) -> bool: + return receipt.written.get(path) == _fingerprint(_snapshot(document, path)) + + +def _stage_settings(path: Path, document: TOMLDocument) -> str: + try: + return stage_private_bytes(str(path), document.as_string().encode()) + except OSError as error: + raise CodexSettingsError(f"Could not stage Codex settings at {path}; no settings were changed") from error + + +def _commit(path: Path, staged: str | None, commit: Callable[[str, str], None]) -> None: + if staged is None: + path.unlink(missing_ok=True) + else: + commit(staged, str(path)) + + +def configure_codex_settings( + base_url: str, + api_key: str, + model: str, + settings_path: Path, + *, + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + require_safe_codex() + current: Final = _read(settings_path) + earlier: Final = _receipt(settings_path) + headers: Final = tomlkit.parse(tomlkit.dumps(MappingProxyType({"Authorization": f"Bearer {api_key}"}))) + provider_table: Final = tomlkit.parse( + tomlkit.dumps(MappingProxyType({**codex_proxy_provider(base_url), "http_headers": headers})) + ) + provider: Final = tomlkit.dumps(MappingProxyType({"value": provider_table})) + selections: Final = tomlkit.parse( + tomlkit.dumps(MappingProxyType({"model_provider": CODEX_PROXY_PROVIDER, "model": model})) + ) + merged: Final = _with( + _with( + _with(_with(current, "profile", None), "model", _snapshot(selections, "model")), + "model_provider", + _snapshot(selections, "model_provider"), + ), + _PROVIDER_PATH, + provider, + ) + owned: Final = tuple( + path + for path in _OWNED_PATHS + if _fingerprint(_snapshot(current, path)) != _fingerprint(_snapshot(merged, path)) + or (earlier is not None and _ours(current, path, earlier)) + ) + receipt: Final = _Receipt( + settings_path=str(settings_path.resolve()), + file_existed=settings_path.exists() if earlier is None else earlier.file_existed, + providers_existed="model_providers" in current if earlier is None else earlier.providers_existed, + previous=MappingProxyType( + { + path: earlier.previous[path] + if earlier is not None and _ours(current, path, earlier) + else _snapshot(current, path) + for path in owned + } + ), + written=MappingProxyType({path: _fingerprint(_snapshot(merged, path)) for path in owned}), + ) + target: Final = settings_path.resolve() + state_path: Final = codex_configure_state_path(settings_path) + try: + ensure_private_dir(state_path.parent) + staged_receipt: Final = stage_private_json(str(state_path), receipt.model_dump(mode="json")) + except OSError as error: + raise CodexSettingsError(f"Could not stage the Codex configure receipt at {state_path}") from error + try: + staged_settings: Final = _stage_settings(target, merged) + except CodexSettingsError: + discard_staged_json(staged_receipt) + raise + try: + commit(staged_receipt, str(state_path)) + except OSError as error: + discard_staged_json(staged_receipt) + discard_staged_json(staged_settings) + raise CodexSettingsError( + f"Could not write the Codex configure receipt at {state_path}; no settings were changed" + ) from error + try: + commit(staged_settings, str(target)) + except OSError as error: + discard_staged_json(staged_settings) + try: + _commit( + state_path, + None if earlier is None else stage_private_json(str(state_path), earlier.model_dump(mode="json")), + commit_staged_json, + ) + except OSError as rollback_error: + raise CodexSettingsError( + f"Codex settings were not written and its receipt at {state_path} could not be restored" + ) from rollback_error + raise CodexSettingsError( + f"Could not write Codex settings at {settings_path}; the earlier receipt was restored" + ) from error + + +def unconfigure_codex_settings( + settings_path: Path, *, commit: Callable[[str, str], None] = commit_staged_json +) -> CodexUnconfigureOutcome: + current: Final = _read(settings_path) + receipt: Final = _receipt(settings_path) + if receipt is None: + raise CodexSettingsError("Codex is not configured by `lite configure codex`; nothing to undo") + ours: Final = tuple(path for path in receipt.written if settings_path.exists() and _ours(current, path, receipt)) + restored_owned: Final = reduce(lambda document, path: _with(document, path, receipt.previous[path]), ours, current) + providers: Final = _table(_mapping(restored_owned).get("model_providers")) + restored: Final = ( + _with(restored_owned, "model_providers", None) + if providers is not None and not providers and not receipt.providers_existed + else restored_owned + ) + target: Final = settings_path.resolve() + file_removed: Final = not restored.as_string().strip() and not (receipt.file_existed and target.exists()) + staged: Final = None if file_removed else _stage_settings(target, restored) + state_path: Final = codex_configure_state_path(settings_path) + try: + _commit(target, staged, commit) + state_path.unlink() + except OSError as error: + if staged is not None: + discard_staged_json(staged) + raise CodexSettingsError( + "Could not finish undoing Codex configuration; the receipt was kept for retry" + ) from error + return CodexUnconfigureOutcome( + restored=tuple(path for path in ours if _snapshot(current, path) != _snapshot(restored, path)), + kept=tuple(path for path in receipt.written if path not in ours and _snapshot(current, path) is not None), + file_removed=file_removed, + ) diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 2715a0a9a38..de22251ea8c 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -62,12 +62,16 @@ def hidden_command_names() -> frozenset[str]: return parse_hidden_commands(get_config_value(HIDDEN_COMMANDS_KEY)) -def _normalize_base_url(value: str) -> str: +def normalize_base_url(value: str) -> str: + if any(ord(char) <= 32 or ord(char) == 127 for char in value): + raise click.UsageError("base_url must not contain whitespace or control characters") parsed: Final = urlparse(value) if parsed.scheme not in ("http", "https") or not parsed.netloc: raise click.UsageError("base_url must be a full http:// or https:// URL including a host") if "?" in value or "#" in value: raise click.UsageError("base_url must not include a query string or fragment") + if parsed.username is not None or parsed.password is not None: + raise click.UsageError("base_url must not contain credentials; pass --api-key separately") return value.rstrip("/") @@ -86,7 +90,7 @@ def _normalize_hidden_commands(value: str) -> str: _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType( { - "base_url": _normalize_base_url, + "base_url": normalize_base_url, HIDDEN_COMMANDS_KEY: _normalize_hidden_commands, } ) diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py index 4acf94e16f9..7988f8aef3c 100644 --- a/litellm/proxy/client/cli/commands/configure.py +++ b/litellm/proxy/client/cli/commands/configure.py @@ -1,4 +1,4 @@ -"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" +"""Persistent Claude Code and Codex gateway configuration.""" import os import sys @@ -11,6 +11,7 @@ from typing import Final import click from InquirerPy import inquirer from InquirerPy.base.control import Choice +from pydantic import BaseModel from litellm.proxy.common_utils.model_listing_utils import ( CLAUDE_CODE_CLIENT, @@ -18,6 +19,7 @@ from litellm.proxy.common_utils.model_listing_utils import ( GATEWAY_CLIENT_HEADER, ) +from .agents import codex_config_path from .auth import CliContextObj from .claude_settings import ( STARTING_MODEL_ROLE, @@ -30,15 +32,23 @@ from .claude_settings import ( claude_settings_path, configure_claude_settings, configure_state_path, - refuse_while_owned, + preflight_claude_settings, settings_file_owners, unconfigure_claude_settings, ) +from .codex_settings import ( + CodexSettingsError, + configure_codex_settings, + preflight_codex_settings, + unconfigure_codex_settings, +) +from .config import normalize_base_url from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing _LISTED_MODELS_SHOWN: Final = 20 _CLAUDE_TARGET: Final = "claude" -_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),) +_CODEX_TARGET: Final = "codex" +_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"), (_CODEX_TARGET, "Codex (CLI)")) _KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default" _CLAUDE_CODE_VIEW: Final = MappingProxyType( {"anthropic-version": "2023-06-01", GATEWAY_CLIENT_HEADER: CLAUDE_CODE_CLIENT} @@ -60,10 +70,12 @@ def resolve_credential(ctx: click.Context, api_key: str | None) -> StaticToken: explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) if not explicit: raise ClaudeSettingsError( - "`lite configure claude` needs a long-lived virtual key: pass --api-key, `lite --api-key`, or set " + "`lite configure` needs a long-lived virtual key: pass --api-key, `lite --api-key`, or set " "LITELLM_PROXY_API_KEY. Your `lite login` credential expires within a day, so it is not written " - "into Claude Code's settings." + "into agent settings." ) + if not explicit.strip() or any(ord(char) <= 32 or ord(char) == 127 for char in explicit): + raise ClaudeSettingsError("The virtual key must not be blank or contain whitespace or control characters.") return StaticToken(explicit) @@ -76,33 +88,45 @@ class _Listing: return tuple(model.id for model in self.models) -def _start(ctx: click.Context, api_key: str | None) -> tuple[StaticToken, _Listing]: - """Every configure path begins the same way: the local ownership check first, so a `lite up` - session is refused before any request, then the credential, then the listing.""" - settings_path: Final = claude_settings_path(os.environ) +def _preflight(target: str) -> None: + try: + if target == _CLAUDE_TARGET: + preflight_claude_settings(claude_settings_path(os.environ)) + else: + preflight_codex_settings(codex_config_path(os.environ)) + except (ClaudeSettingsError, CodexSettingsError) as e: + raise click.ClickException(str(e)) from e + + +def _start(ctx: click.Context, api_key: str | None, target: str = _CLAUDE_TARGET) -> tuple[StaticToken, _Listing]: + _preflight(target) try: - refuse_while_owned(settings_path, settings_file_owners(settings_path)) credential: Final = resolve_credential(ctx, api_key) except ClaudeSettingsError as e: raise click.ClickException(str(e)) - return credential, _listed_models(ctx.obj["base_url"], credential.token) + return credential, _listed_models(ctx.obj["base_url"], credential.token, target) -def _listing_error(base_url: str, error: PiSyncError) -> str: +def _listing_error(base_url: str, error: PiSyncError, target: str) -> str: """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question.""" if error.kind is ListingFailure.REJECTED: return f"LiteLLM rejected your key (HTTP {error.status}). Pass a valid --api-key." if error.kind is ListingFailure.UNREACHABLE: - return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" + return ( + f"Could not connect. Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" + ) if error.kind is ListingFailure.EMPTY: - return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model." - return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." + name: Final = "Claude Code" if target == _CLAUDE_TARGET else "Codex" + return f"{error.message} {name} would have nothing to run; give the key access to at least one model." + return f"The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." -def _listed_models(base_url: str, key: str) -> _Listing: - listed: Final = fetch_model_listing(base_url, key, headers=_CLAUDE_CODE_VIEW) +def _listed_models(base_url: str, key: str, target: str = _CLAUDE_TARGET) -> _Listing: + listed: Final = fetch_model_listing( + base_url, key, headers=_CLAUDE_CODE_VIEW if target == _CLAUDE_TARGET else MappingProxyType({}) + ) if isinstance(listed, PiSyncError): - raise click.ClickException(_listing_error(base_url, listed)) + raise click.ClickException(_listing_error(base_url, listed, target)) return _Listing(listed) @@ -115,17 +139,19 @@ def _model_choice(model: str | None) -> ModelChoice: return StartOn(model) if model is not None else UnpinModel() +def _validated_model(model: str | None, listing: _Listing, base_url: str) -> str | None: + starting: Final = _starting_model(model, listing) if model is not None else None + if model is not None and starting is None: + shown: Final = ", ".join(listing.ids[:_LISTED_MODELS_SHOWN]) + raise click.ClickException(f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}.") + return starting + + def _apply_claude(ctx: click.Context, credential: StaticToken, listing: _Listing, model: str | None) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] listed: Final = listing.ids - starting: Final = _starting_model(model, listing) if model is not None else None - if model is not None and starting is None: - shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) - more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" - raise click.ClickException( - f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}." - ) + starting: Final = _validated_model(model, listing, base_url) settings_path: Final = claude_settings_path(os.environ) try: configure_claude_settings( @@ -178,40 +204,131 @@ def _pick_model(listed: Sequence[str]) -> str | None: picked: Final = inquirer.fuzzy( message="Model Claude Code starts on (type to filter; /model switches any time):", choices=[_KEEP_DEFAULT_MODEL, *listed], + default=listed[0] if listed else _KEEP_DEFAULT_MODEL, ).execute() return None if picked == _KEEP_DEFAULT_MODEL else str(picked) +def _pick_codex_model(listed: Sequence[str]) -> str: + choices: Final = list(listed) # mutable-ok: InquirerPy's choices parameter requires a list + return str(inquirer.fuzzy(message="Model Codex starts on (type to filter):", choices=choices).execute()) + + +def _apply_codex(ctx: click.Context, credential: StaticToken, listing: _Listing, model: str) -> None: + base_url: Final[str] = ctx.obj["base_url"] + _validated_model(model, listing, base_url) + settings_path: Final = codex_config_path(os.environ) + try: + configure_codex_settings(base_url, credential.token, model, settings_path) + except CodexSettingsError as e: + raise click.ClickException(str(e)) from e + click.echo(f"Configured Codex: {settings_path} now routes through {base_url}.") + click.echo(f"Starting model: {model}. Credential: your virtual key, stored in the private provider settings.") + click.echo("Start `codex` from any terminal. Undo with `lite unconfigure codex`.") + if settings_path.is_symlink(): + click.echo(f"Note: your key now lives in {settings_path.resolve()}; keep it out of version control.", err=True) + + +@dataclass(frozen=True, slots=True) +class _Setup: + target: str + listing: _Listing + model: str | None + + +def _choose_setup( + ctx: click.Context, + target: str, + credential: StaticToken, + pick_model: Callable[[Sequence[str]], str | None], + pick_codex_model: Callable[[Sequence[str]], str], +) -> _Setup: + base_url: Final[str] = ctx.obj["base_url"] + listing: Final = _listed_models(base_url, credential.token, target) + model: Final = ( + pick_model(tuple(item.source_model or item.id for item in listing.models)) + if target == _CLAUDE_TARGET + else pick_codex_model(listing.ids) + ) + _validated_model(model, listing, base_url) + return _Setup(target, listing, model) + + def interactive_configure( ctx: click.Context, pick_targets: Callable[[], tuple[str, ...]] = _pick_targets, pick_model: Callable[[Sequence[str]], str | None] = _pick_model, + pick_codex_model: Callable[[Sequence[str]], str] = _pick_codex_model, ) -> None: """`lite configure` with no agent named: ask which agents to wire and which model to pin.""" targets: Final = pick_targets() - if _CLAUDE_TARGET not in targets: + if not targets: return - credential, listing = _start(ctx, None) - _apply_claude( - ctx, credential, listing, pick_model(tuple(model.source_model or model.id for model in listing.models)) + for target in targets: + _preflight(target) + try: + credential: Final = resolve_credential(ctx, None) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) from e + setups: Final = tuple(_choose_setup(ctx, target, credential, pick_model, pick_codex_model) for target in targets) + for setup in setups: + if setup.target == _CLAUDE_TARGET: + _apply_claude(ctx, credential, setup.listing, setup.model) + elif setup.model is not None: + _apply_codex(ctx, credential, setup.listing, setup.model) + + +class _ConnectionOptions(BaseModel): + api_key: str | None = None + gateway_url: str | None = None + + +def _connection_context(ctx: click.Context, api_key: str | None, gateway_url: str | None) -> click.Context: + ctx_obj: Final[CliContextObj] = ctx.obj + group: Final = ( + _ConnectionOptions.model_validate(ctx.parent.params) + if ctx.parent is not None and ctx.parent.command.name == "configure" + else _ConnectionOptions() ) + key: Final = api_key if api_key is not None else group.api_key + url: Final = gateway_url if gateway_url is not None else group.gateway_url + normalized: Final = normalize_base_url(url if url is not None else ctx_obj["base_url"]) + connection: Final[CliContextObj] = { + **ctx_obj, + "base_url": normalized.removesuffix("/v1"), + "base_url_explicit": url is not None or ctx_obj.get("base_url_explicit", False), + "api_key": key if key is not None else ctx_obj.get("api_key"), + "api_key_from_token_file": False if key is not None else ctx_obj.get("api_key_from_token_file", False), + } + return click.Context(ctx.command, parent=ctx.parent, obj=connection) @click.group(name="configure", invoke_without_command=True) +@click.option("--api-key", default=None, help="Long-lived LiteLLM virtual key to store in the selected agents.") +@click.option( + "--gateway-url", "--base-url", default=None, help="Gateway URL; defaults to `lite --base-url` / LITELLM_PROXY_URL." +) @click.pass_context -def configure_group(ctx: click.Context) -> None: +def configure_group(ctx: click.Context, api_key: str | None, gateway_url: str | None) -> None: """Persistently route a coding agent through your LiteLLM proxy. With no agent named, asks which agents to wire and which proxy model to pin. """ if ctx.invoked_subcommand is not None: return + connection: Final = _connection_context(ctx, api_key, gateway_url) if not sys.stdin.isatty(): raise click.ClickException( "`lite configure` asks questions, so it needs a terminal. Non-interactively, run " - "`lite configure claude --api-key --model `." + "`lite configure claude --api-key --model ` or " + "`lite configure codex --api-key --model `." ) - interactive_configure(ctx) + prompted: Final = ( + connection + if connection.obj.get("base_url_explicit") + else _connection_context(connection, None, click.prompt("Gateway URL", default=connection.obj["base_url"])) + ) + interactive_configure(prompted) @click.group(name="unconfigure") @@ -228,8 +345,9 @@ def unconfigure_group() -> None: "LITELLM_PROXY_API_KEY value; required, since a `lite login` credential expires within a day.", ) @click.option("--model", default=None, help=_MODEL_OPTION_HELP) +@click.option("--gateway-url", "--base-url", default=None, help="Gateway URL, including any deployment path prefix.") @click.pass_context -def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None: +def configure_claude(ctx: click.Context, api_key: str | None, model: str | None, gateway_url: str | None) -> None: """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`. Patches ~/.claude/settings.json in place: the proxy URL, your virtual key as a static token, @@ -238,8 +356,39 @@ def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. Assumes the proxy is already running. """ - credential, listing = _start(ctx, api_key) - _apply_claude(ctx, credential, listing, model) + connection: Final = _connection_context(ctx, api_key, gateway_url) + credential, listing = _start(connection, api_key) + _apply_claude(connection, credential, listing, model) + + +@configure_group.command(name="codex") +@click.option("--api-key", default=None, help="Long-lived LiteLLM virtual key to store in Codex's user config.") +@click.option("--gateway-url", "--base-url", default=None, help="Gateway URL, including any deployment path prefix.") +@click.option("--model", required=True, help="Gateway model Codex starts on, as listed by /v1/models for your key.") +@click.pass_context +def configure_codex(ctx: click.Context, api_key: str | None, gateway_url: str | None, model: str) -> None: + """Route plain `codex` through the gateway until `lite unconfigure codex`.""" + connection: Final = _connection_context(ctx, api_key, gateway_url) + credential, listing = _start(connection, api_key, _CODEX_TARGET) + _apply_codex(connection, credential, listing, model) + + +@unconfigure_group.command(name="codex") +def unconfigure_codex() -> None: + """Restore only Codex settings still holding what configure wrote.""" + settings_path: Final = codex_config_path(os.environ) + try: + outcome: Final = unconfigure_codex_settings(settings_path) + except CodexSettingsError as e: + raise click.ClickException(str(e)) from e + if outcome.file_removed: + click.echo(f"Removed {settings_path}; it held only settings created by `lite configure codex`.") + elif outcome.restored: + click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.") + else: + click.echo(f"Nothing in {settings_path} was still ours to restore.") + if outcome.kept: + click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.") @unconfigure_group.command(name="claude") diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index b0e81a222c0..05fb877d0f1 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -95,7 +95,7 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s # If no API key provided via flag or environment variable, try to load from saved token. # Pass base_url so we only use the stored key when it was issued for this server. - api_key_from_token_file: Final = api_key is None + api_key_from_token_file: Final = api_key is None and ctx.invoked_subcommand not in ("configure", "unconfigure") resolved_api_key: Final = ( get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) if api_key_from_token_file diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9f58aaf24f1..289fe086379 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3181,6 +3181,11 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can exercise the production gating logic directly rather than reimplementing the finally block. """ + if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"): + pending: Final = getattr(logging_obj, "_native_pending_logging", None) + if pending is not None: + logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once + pending.release(not exception_raised) _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is None: return diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 1d4024bb84e..dc329e55e31 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -3,6 +3,7 @@ import asyncio import gc import json import os +import socket import sys import tracemalloc from collections import Counter @@ -237,6 +238,80 @@ def _process_memory_usage(process: _ProcessHandle) -> _ProcessMemoryUsage: ) +PROC_STATM_PATH: Final = "/proc/self/statm" +PROC_MEMINFO_PATH: Final = "/proc/meminfo" +PSUTIL_MISSING_ERROR: Final = "Install psutil for memory monitoring: pip install psutil" + + +class _ProcMemoryInfo(NamedTuple): + rss: int + vms: int + + +class _ProcFilesystemProcess: + """Memory of the running process read from the Linux proc filesystem, for images without psutil.""" + + def __init__( + self, + statm_path: str = PROC_STATM_PATH, + meminfo_path: str = PROC_MEMINFO_PATH, + page_size: int | None = None, + ) -> None: + self._statm_path: Final = statm_path + self._meminfo_path: Final = meminfo_path + self._page_size: Final = os.sysconf("SC_PAGE_SIZE") if page_size is None else page_size + + def memory_info(self) -> _ProcMemoryInfo: + with open(self._statm_path, encoding="ascii") as statm: + size_pages, resident_pages = statm.read().split()[:2] + return _ProcMemoryInfo(rss=int(resident_pages) * self._page_size, vms=int(size_pages) * self._page_size) + + def memory_percent(self) -> float: + with open(self._meminfo_path, encoding="ascii") as meminfo: + total_kilobytes: Final = next(int(line.split()[1]) for line in meminfo if line.startswith("MemTotal:")) + return self.memory_info().rss / (total_kilobytes * 1024) * 100 + + +def _process_handle() -> _ProcessHandle | None: + try: + import psutil + except ImportError: + return _ProcFilesystemProcess() if os.path.exists(PROC_STATM_PATH) else None + return psutil.Process() + + +def _health_status(memory_percent: float) -> str: + if memory_percent > 80: + return "critical" + if memory_percent > 60: + return "warning" + return "healthy" + + +class _SummaryProcessMemory(TypedDict, total=False): + summary: ReadOnly[str] + ram_usage_mb: ReadOnly[float] + system_memory_percent: ReadOnly[float] + error: ReadOnly[str] + + +def _summary_process_memory(process: _ProcessHandle | None) -> tuple[_SummaryProcessMemory, str]: + if process is None: + missing: Final[_SummaryProcessMemory] = {"error": PSUTIL_MISSING_ERROR} + return missing, "healthy" + try: + usage: Final = _process_memory_usage(process) + except Exception as e: + unreadable: Final[_SummaryProcessMemory] = {"error": str(e)} + return unreadable, "healthy" + memory: Final[_SummaryProcessMemory] = { + "summary": f"{usage.resident_megabytes:.1f} MB ({usage.percent:.1f}% of system memory)", + "ram_usage_mb": round(usage.resident_megabytes, 2), + "system_memory_percent": round(usage.percent, 2), + } + return memory, _health_status(usage.percent) + + @router.get("/debug/memory/summary", include_in_schema=False) async def get_memory_summary( _: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -246,6 +321,7 @@ async def get_memory_summary( Returns: - worker_pid: Process ID + - hostname: Host (the pod on Kubernetes) the worker runs on - status: Overall health based on memory usage - memory: Process memory usage and RAM info - caches: Cache item counts and descriptions @@ -263,35 +339,7 @@ async def get_memory_summary( user_api_key_cache, ) - # Get process memory info - process_memory = {} - health_status = "healthy" - - try: - import psutil - - usage: Final = _process_memory_usage(psutil.Process()) - memory_mb: Final = usage.resident_megabytes - memory_percent: Final = usage.percent - - process_memory = { - "summary": f"{memory_mb:.1f} MB ({memory_percent:.1f}% of system memory)", - "ram_usage_mb": round(memory_mb, 2), - "system_memory_percent": round(memory_percent, 2), - } - - # Check memory health status - if memory_percent > 80: - health_status = "critical" - elif memory_percent > 60: - health_status = "warning" - else: - health_status = "healthy" - - except ImportError: - process_memory["error"] = "Install psutil for memory monitoring: pip install psutil" - except Exception as e: - process_memory["error"] = str(e) + process_memory, health_status = _summary_process_memory(_process_handle()) # Get cache information caches: Final[dict[str, object]] = {} @@ -347,6 +395,7 @@ async def get_memory_summary( return { "worker_pid": os.getpid(), + "hostname": socket.gethostname(), "status": health_status, "memory": process_memory, "caches": { diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index e93bc96da69..54021e68980 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -74,6 +74,8 @@ DisablePreparedStatementsFlag = Annotated[ bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=DISABLE_PREPARED_STATEMENTS_ENV_VAR)) ] MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR: Final = "DATABASE_MAX_IDLE_CONNECTION_LIFETIME" +DATABASE_SSLMODE_ENV_VAR: Final = "DATABASE_SSLMODE" +DATABASE_SSLROOTCERT_ENV_VAR: Final = "DATABASE_SSLROOTCERT" # schema.prisma pins `provider = "postgresql"`, so these are the only schemes # Prisma can actually connect with. @@ -135,6 +137,7 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"}) +PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset({"sslmode", "sslcert", "sslaccept"}) PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----" PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103) TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0 @@ -263,6 +266,19 @@ def connection_params_from_url(url: str) -> Mapping[str, str | int | float]: ) +def token_refresh_params_from_url(url: str) -> Mapping[str, str | int | float]: + """Return the params a re-minted token URL carries over from the URL it replaces. + + The pool and timeout params plus Prisma's TLS params (already translated from + libpq spelling), so a refreshed URL keeps verifying the server the way the + first one did. + """ + kept: Final = CONNECTION_PARAM_KEYS | PRISMA_TLS_PARAM_KEYS + return MappingProxyType( + {key: value for key, value in urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query) if key in kept} + ) + + def unsupported_db_scheme(database_url: str) -> str | None: """Return the connection URL scheme when it is not PostgreSQL, else None. @@ -312,6 +328,9 @@ class DatabaseURLSettings(BaseSettings): default=None, validation_alias=MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR ) + database_sslmode: str | None = Field(default=None, validation_alias=DATABASE_SSLMODE_ENV_VAR) + database_sslrootcert: str | None = Field(default=None, validation_alias=DATABASE_SSLROOTCERT_ENV_VAR) + # Writer database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") direct_url: str | None = Field(default=None, validation_alias="DIRECT_URL") @@ -353,6 +372,25 @@ class DatabaseURLSettings(BaseSettings): azure_postgresql_auth=self.azure_postgresql_auth, ) + def tls_params(self) -> Mapping[str, str]: + """``sslmode`` / ``sslrootcert`` query params for every URL assembled from the discrete vars. + + A root cert on its own means ``verify-full``: under libpq's default + ``prefer`` the CA would never be consulted, and PgBouncer would dial + Postgres unverified with the bundle loaded. + """ + sslmode: Final = self.database_sslmode or ("verify-full" if self.database_sslrootcert else None) + return MappingProxyType( + { + key: value + for key, value in ( + ("sslmode", sslmode), + ("sslrootcert", self.database_sslrootcert), + ) + if value + } + ) + def build_writer_url(self) -> str | None: """Return the writer URL to set, or ``None`` to leave it as-is. @@ -362,6 +400,12 @@ class DatabaseURLSettings(BaseSettings): A ``DATABASE_URL`` the supervisor pointed at the in-container PgBouncer is kept even under token auth: the pooler renews the token upstream. """ + assembled: Final = self._assemble_writer_url() + if assembled is None: + return None + return add_missing_query_params(assembled, self.tls_params()) + + def _assemble_writer_url(self) -> str | None: auth: Final = self.token_auth() if auth is not None and database_url_is_pooled(): return None @@ -411,6 +455,12 @@ class DatabaseURLSettings(BaseSettings): pre-existing ``DATABASE_URL_READ_REPLICA``. Reader fields fall back to the writer's values. """ + assembled: Final = self._assemble_reader_url() + if assembled is None: + return None + return add_missing_query_params(assembled, self.tls_params()) + + def _assemble_reader_url(self) -> str | None: if not self.database_host_read_replica: return None # reader is opt-in if self.database_url_read_replica: diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 21b73f27a80..acd01b0e99e 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -16,7 +16,7 @@ from datetime import datetime, timedelta from typing import Any, Final, Protocol from litellm._logging import verbose_proxy_logger -from litellm.proxy.db.db_url_settings import add_missing_query_params, connection_params_from_url +from litellm.proxy.db.db_url_settings import add_missing_query_params, token_refresh_params_from_url from litellm.proxy.db.token_auth import ( DEFAULT_POSTGRES_PORT, DatabaseTokenAuth, @@ -441,7 +441,7 @@ class PrismaWrapper: endpoint: Final = self._iam_endpoint if self._iam_endpoint is not None else self._endpoint_from_env() db_url: Final = add_missing_query_params( endpoint.build_url(mint_database_token(auth, endpoint)), - connection_params_from_url(os.environ.get(self._db_url_env_var, "")), + token_refresh_params_from_url(os.environ.get(self._db_url_env_var, "")), ) os.environ[self._db_url_env_var] = db_url return db_url diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index e35b1c8c82b..89a07234c6c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -24,6 +24,7 @@ from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import db_lookup_gate +from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, @@ -104,6 +105,15 @@ class SpendCounterReseed: SpendCounterReseed._locks.popitem(last=False) return lock + @staticmethod + async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: + """Apply local deltas after an in-flight reseed establishes the spend balance.""" + lock: Final = await SpendCounterReseed._get_lock(counter_key) + async with lock: + return await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, local_only=True, refresh_ttl=True + ) + @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: """ @@ -185,6 +195,11 @@ class SpendCounterReseed: return True return False + @staticmethod + async def _read_active_batch(counter_key: str) -> tuple[float | None, bool] | None: + """The request's MGET answers for this counter; a Redis miss there is authoritative.""" + return await read_batched_spend_counter(counter_key) + @staticmethod async def coalesced( prisma_client: Optional["PrismaClient"], @@ -202,10 +217,13 @@ class SpendCounterReseed: """ lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: + batched: Final = await SpendCounterReseed._read_active_batch(counter_key) + if batched is not None and batched[0] is not None: + return batched[0] # Re-check after acquiring the lock. Skip in-memory on a clean # Redis miss - in-memory is per-pod-stale. - redis_clean_miss = False - if spend_counter_cache.redis_cache is not None: + redis_clean_miss = batched is not None + if spend_counter_cache.redis_cache is not None and not redis_clean_miss: try: val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: @@ -244,8 +262,12 @@ class SpendCounterReseed: key=counter_key, value=current_value, ) + record_spend_counter_value(counter_key, current_value) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -392,8 +414,11 @@ class SpendCounterReseed: ) -> float | None: lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: - redis_clean_miss = False - if spend_counter_cache.redis_cache is not None: + batched: Final = await SpendCounterReseed._read_active_batch(counter_key) + if batched is not None and batched[0] is not None: + return batched[0] + redis_clean_miss = batched is not None + if spend_counter_cache.redis_cache is not None and not redis_clean_miss: try: val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: @@ -437,12 +462,18 @@ class SpendCounterReseed: key=counter_key, value=current_value, ) + record_spend_counter_value(counter_key, float(current_value)) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=window_spend) + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = ( + max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend + ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", counter_key, ) raise - return window_spend + return current_value diff --git a/litellm/proxy/guardrails/guardrail_hooks/conduct/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/conduct/__init__.py new file mode 100644 index 00000000000..9eac143be88 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/conduct/__init__.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .conduct import ConductGuardrail + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import Guardrail, LitellmParams + +DEFAULT_TIMEOUT_SECONDS: Final = 8.0 +_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def initialize_guardrail( + litellm_params: LitellmParams, + guardrail: Guardrail, + guardrail_cls: type[CustomGuardrail] = ConductGuardrail, +) -> CustomGuardrail: + import litellm + + extras: Final = litellm_params.model_extra or _NO_EXTRAS + _callback: Final = guardrail_cls( + api_url=litellm_params.api_base, + agent_token=litellm_params.api_key, + workspace_id=extras.get("workspace_id"), + tool_name=extras.get("tool_name", "llm_call"), + unreachable_fallback=litellm_params.unreachable_fallback, + timeout=DEFAULT_TIMEOUT_SECONDS if litellm_params.timeout is None else litellm_params.timeout, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + supported_event_hooks=guardrail_cls.get_supported_event_hooks(), + ) + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback + + +guardrail_initializer_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.CONDUCT.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.CONDUCT.value: ConductGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/conduct/conduct.py b/litellm/proxy/guardrails/guardrail_hooks/conduct/conduct.py new file mode 100644 index 00000000000..c87f8c016b1 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/conduct/conduct.py @@ -0,0 +1,158 @@ +"""Conduct Guard as a LiteLLM guardrail, backed by the ``conduct-litellm-guard`` PyPI package. + +Install: ``pip install "conduct-litellm-guard>=0.2.5"`` +Source: https://github.com/sseshachala/conductai/tree/main/packages/conduct-litellm-guard +""" + +from __future__ import annotations + +import inspect +from collections.abc import Awaitable, Callable, Mapping +from functools import partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, Protocol + +from pydantic import BaseModel, ConfigDict + +from litellm.integrations.custom_guardrail import CustomGuardrail, log_guardrail_information +from litellm.types.llms.openai import ChatCompletionUserMessage +from litellm.types.proxy.guardrails.guardrail_hooks.conduct import ConductGuardrailConfigModel + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus + +MISSING_PACKAGE_MESSAGE: Final = ( + "conduct-litellm-guard>=0.2.5 is required for the Conduct guardrail. " + 'Install it with: pip install "conduct-litellm-guard>=0.2.5"' +) + +BLOCKING_VERDICTS: Final = frozenset({"block", "approval"}) +FLAGGED_VERDICTS: Final = frozenset({"warning", "advisory"}) + + +class ConductDecision(Protocol): + @property + def verdict(self) -> str: ... + + @property + def rule_id(self) -> str | None: ... + + +class ConductCheck(Protocol): + def __call__(self, *, data: Mapping[str, object], call_type: str) -> Awaitable[ConductDecision]: ... + + +def request_payload( + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], +) -> Mapping[str, object] | None: + if input_type != "request": + return None + messages: Final = inputs.get("structured_messages") or tuple( + ChatCompletionUserMessage(role="user", content=text) for text in inputs.get("texts") or () + ) + return MappingProxyType({**request_data, "prompt": None, "messages": messages}) + + +def decision_status(decision: ConductDecision) -> GuardrailStatus: + return "guardrail_flagged" if decision.verdict in FLAGGED_VERDICTS else "success" + + +class ConductVerdict(BaseModel): + model_config = ConfigDict(frozen=True) + + verdict: str + rule_id: str | None = None + + +def record_decision( + guardrail: CustomGuardrail, + request_data: dict[str, object], # mutable-ok: the logging helper writes metadata into it + decision: ConductDecision, +) -> None: + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=ConductVerdict(verdict=decision.verdict, rule_id=decision.rule_id).model_dump(), + request_data=request_data, + guardrail_status=decision_status(decision), + ) + + +async def apply_conduct_guardrail( + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], + check: ConductCheck, + blocked: Callable[[ConductDecision], Exception], + record: Callable[[ConductDecision], None], +) -> GenericGuardrailAPIInputs: + payload: Final = request_payload(inputs, request_data, input_type) + if payload is None: + return inputs + decision: Final = await check(data=payload, call_type=input_type) + if decision.verdict in BLOCKING_VERDICTS: + raise blocked(decision) + record(decision) + return inputs + + +def binds_unreachable_fallback(guardrail_cls: type[object]) -> bool: + return "unreachable_fallback" in inspect.signature(guardrail_cls.__init__).parameters + + +try: + from conduct_litellm_guard.guardrail import ConductGuard, ConductGuardBlocked + + if not binds_unreachable_fallback(ConductGuard): + raise ImportError(MISSING_PACKAGE_MESSAGE) +except ImportError as import_error: + _import_error: Final = import_error + + class ConductGuardrail(CustomGuardrail): + def __init__(self, **kwargs: object) -> None: # kwargs-ok: mirrors the plugin constructor, only raises + raise ImportError(MISSING_PACKAGE_MESSAGE) from _import_error + + @staticmethod + def get_config_model() -> type[ConductGuardrailConfigModel]: + return ConductGuardrailConfigModel + +else: + + class ConductGuardrail(ConductGuard): # pyright: ignore[reportUntypedBaseClass] # optional dep, absent at type-check + @staticmethod + def get_config_model() -> type[ConductGuardrailConfigModel]: + return ConductGuardrailConfigModel + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail contract + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return await apply_conduct_guardrail( + inputs, + request_data, + input_type, + self.check, + ConductGuardBlocked, + partial(record_decision, self, request_data), + ) + + +__all__ = ( + "BLOCKING_VERDICTS", + "FLAGGED_VERDICTS", + "MISSING_PACKAGE_MESSAGE", + "ConductCheck", + "ConductDecision", + "ConductGuardrail", + "ConductVerdict", + "apply_conduct_guardrail", + "binds_unreachable_fallback", + "decision_status", + "record_decision", + "request_payload", +) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index f7f500b1adc..8fed1f906e5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -18,6 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_tool_message_for_guardrail, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) @@ -261,6 +262,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): fail_on_error: bool | None = True, streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, + async_handler: AsyncHTTPHandler | None = None, **kwargs, ) -> None: """ @@ -273,9 +275,13 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of every streaming_sampling_rate chunks. Defaults to False. streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5. + async_handler (AsyncHTTPHandler | None): HTTP client to call AI Guard with. Defaults to the shared + guardrail-callback client. **kwargs: Additional arguments passed to the CustomGuardrail base class. """ - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.fail_on_error = True if fail_on_error is None else fail_on_error self._set_streaming_params( CrowdStrikeAIDRGuardrailConfigModelOptionalParams( 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 64a47f4f4ff..ee5cd7c4cb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -315,7 +315,6 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is None: call_type = _infer_call_type(call_type=None, completion_response=response) - # Fallback: resolve call_type from logging_obj for pass-through endpoints if call_type is None: litellm_logging_obj: Final = data.get("litellm_logging_obj") logging_call_type: Final = ( @@ -324,6 +323,8 @@ class UnifiedLLMGuardrails(CustomLogger): if logging_call_type in ( CallTypes.pass_through.value, CallTypes.allm_passthrough_route.value, + CallTypes.ocr.value, + CallTypes.aocr.value, ): call_type = logging_call_type diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 219f6f270ed..b1e4f6fd9c3 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -32,32 +32,36 @@ from litellm.router_utils.auto_router_model_naming import ( strategy_router_dependencies, ) -ILLEGAL_DISPLAY_PARAMS: Final = [ - "messages", - "api_key", - "prompt", - "input", - "client_secret", - "azure_ad_token", - "azure_username", - "azure_password", - "vertex_credentials", - "vertex_ai_credentials", - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - "aws_web_identity_token", - "extra_headers", - "headers", - "exception", # internal; not JSON-serializable, never for display - "litellm_metadata", # internal tracking metadata with auth objects; not for display -] # Provider routing fields. Allowed for proxy admins so they can see which # region/version a deployment is checking; gated at the endpoint layer for # non-admin callers (see _strip_admin_only_fields_from_health_result). -ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = ("api_base", "api_version") +ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = ("api_base", "api_version", "aws_bedrock_runtime_endpoint") -MINIMAL_DISPLAY_PARAMS: Final = ["model", "mode_error"] +MINIMAL_DISPLAY_PARAMS: Final = frozenset({"model", "mode_error"}) + +HEALTH_DISPLAY_PARAMS: Final = ( + MINIMAL_DISPLAY_PARAMS + | frozenset(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS) + | frozenset( + { + "custom_llm_provider", + "mode", + "base_model", + "aws_region_name", + "region_name", + "watsonx_region_name", + "vertex_project", + "vertex_location", + "tpm", + "rpm", + "error", + "raw_request_typed_dict", + "x-ratelimit-remaining-requests", + "x-ratelimit-remaining-tokens", + "x-ms-region", + } + ) +) # Modes whose health-check probe is a chat-style completion call and # therefore accept `max_tokens`. Other modes (embedding, image_generation, @@ -143,14 +147,10 @@ def _get_random_llm_message(): def _clean_endpoint_data(endpoint_data: dict, details: bool | None = True): """ - Clean the endpoint data for display to users. + Keep only the explicitly approved, JSON-safe diagnostic fields for display to users. """ - endpoint_data.pop("litellm_logging_obj", None) - return ( - {k: v for k, v in endpoint_data.items() if k not in ILLEGAL_DISPLAY_PARAMS} - if details is not False - else {k: v for k, v in endpoint_data.items() if k in MINIMAL_DISPLAY_PARAMS} - ) + displayed: Final = HEALTH_DISPLAY_PARAMS if details is not False else MINIMAL_DISPLAY_PARAMS + return {k: v for k, v in endpoint_data.items() if k in displayed} def health_check_filter_kwargs_from_general_settings( @@ -258,8 +258,52 @@ def _deployment_model(deployment: Mapping[str, object]) -> str | None: return params.get("model") if isinstance(params, Mapping) else None +def _owner_team_id(deployment: Mapping[str, object]) -> str | None: + info: Final = deployment.get("model_info") + owner: Final = info.get("team_id") if isinstance(info, Mapping) else None + return owner if isinstance(owner, str) else None + + +def _team_public_model_name(deployment: Mapping[str, object]) -> str | None: + info: Final = deployment.get("model_info") + name: Final = info.get("team_public_model_name") if isinstance(info, Mapping) else None + return name if isinstance(name, str) else None + + +def _deployments_routed_by_name( + model_list: Sequence[Mapping[str, object]], model_name: str, team_id: str | None +) -> tuple[Mapping[str, object], ...]: + """The deployments a request for ``model_name`` from this caller routes to. + + A team's own copies published under that name win, then deployments carrying it as + ``model_name``. A caller with no team reaches a public name only when nothing carries + it as ``model_name``, and only an admin still has another team's deployment in a + scoped ``model_list`` by then. + """ + own_copies: Final = tuple( + x + for x in model_list + if team_id is not None and _owner_team_id(x) == team_id and _team_public_model_name(x) == model_name + ) + if own_copies: + return own_copies + by_name: Final = tuple(x for x in model_list if x.get("model_name") == model_name) + if by_name or team_id is not None: + return by_name + return tuple(x for x in model_list if _team_public_model_name(x) == model_name) + + +def deployments_targeted_by_name( + model_list: Sequence[Mapping[str, object]], model: str, team_id: str | None +) -> tuple[Mapping[str, object], ...]: + """``model`` targets deployments the way a request for it routes, else by ``litellm_params.model``.""" + return _deployments_routed_by_name(model_list, model, team_id) or tuple( + x for x in model_list if _deployment_model(x) == model + ) + + def _narrow_to_target( - model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None + model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None, team_id: str | None ) -> tuple[Mapping[str, object], ...]: """Narrow to the requested deployment. An id matching nothing keeps the whole list.""" if model_id is not None: @@ -267,8 +311,7 @@ def _narrow_to_target( return by_id or tuple(model_list) if model is None: return tuple(model_list) - by_param: Final = tuple(x for x in model_list if _deployment_model(x) == model) - return by_param or tuple(x for x in model_list if x.get("model_name") == model) + return deployments_targeted_by_name(model_list, model, team_id) def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool: @@ -813,13 +856,18 @@ async def perform_health_check( instrumentation_context: dict | None = None, health_check_skip_disabled_background_models: bool = False, router: "Router | None" = None, + team_id: str | None = None, ): """ Perform a health check on the system. When model_id is provided, only the deployment with that id is checked (so models that share the same name but have different ids are checked separately). - When model (name) is provided, all deployments matching that name are checked. + When model (name) is provided, the deployments a request for that name from the + caller (``team_id``) would route to are checked: the caller's team copies published + under that name, else the deployments named that way, else a public name that only + another team's deployment carries, else the deployments whose ``litellm_params.model`` + is that string. When ``health_check_skip_disabled_background_models`` is True (via ``general_settings.health_check_skip_disabled_background_models``), deployments @@ -850,7 +898,7 @@ async def perform_health_check( cycle_start_time: Final = time.monotonic() requested_model_count: Final = len(model_list) skip_disabled: Final = health_check_skip_disabled_background_models - narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id), skip_disabled) + narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id, team_id), skip_disabled) if not narrowed: if instrumentation_enabled: logger.debug( diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index bf527e1e868..b9964c0e342 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -36,9 +36,13 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.proxy.auth.auth_checks import ( + _resolve_key_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request +) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check ) +from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.health_check_latest import LatestHealthCheckRow @@ -47,6 +51,7 @@ from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, _update_litellm_params_for_health_check, + deployments_targeted_by_name, health_check_filter_kwargs_from_general_settings, perform_health_check, run_with_timeout, @@ -58,6 +63,7 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( get_in_flight_requests, ) from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager +from litellm.router import Router from litellm.router_utils.clientside_credential_handler import ( _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path clientside_credential_keys, @@ -917,7 +923,7 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: def _strip_admin_only_fields_from_health_result(result: dict) -> dict: """ Return a copy of the /health response with provider routing fields - (``api_base``, ``api_version``) removed from each healthy/unhealthy + (``ADMIN_ONLY_HEALTH_DISPLAY_PARAMS``) removed from each healthy/unhealthy endpoint entry. Used to hide those fields from non-admin callers while still showing them which deployments they own and whether each one is healthy. Proxy admins receive the unmodified result. @@ -931,41 +937,68 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: return out -def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: str | None) -> set | None: +def _health_accessible_model_names( + user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None +) -> frozenset[str] | None: + """Model names the caller may health-check, or None when the key is unrestricted.""" + granted_models: Final = _resolve_key_models_for_auth_check(user_api_key_dict) + if not granted_models or SpecialModelNames.all_proxy_models.value in granted_models: + return None + if llm_router is None: + return frozenset(granted_models) + return frozenset( + get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=llm_router.get_model_names(team_id=user_api_key_dict.team_id), + model_access_groups=llm_router.get_model_access_groups(), + ) + ) + + +def _caller_may_probe_deployment( + deployment: Mapping[str, object], + allowed_models: frozenset[str] | None, + llm_router: Router | None, + team_id: str | None, + caller_is_admin: bool, +) -> bool: + """Same deployment visibility rule as routing: another team's deployment is never in scope, team-less callers included.""" + if not caller_is_admin and not Router._deployment_usable_by_team(deployment, team_id): + return False + if allowed_models is None: + return True + if llm_router is None: + return deployment.get("model_name") in allowed_models + model: Final = dict(deployment) + return any( + llm_router.should_include_deployment(model_name=name, model=model, team_id=team_id) for name in allowed_models + ) + + +def _resolve_targeted_model_ids( + model_list: list, model: str | None, model_id: str | None, team_id: str | None +) -> set | None: """ Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of - deployment IDs the response should be scoped to. + deployment IDs the response should be scoped to, mirroring the live-path + narrowing in ``perform_health_check()``: ``model_id`` wins when given and + matches ``model_info.id`` only; ``model`` targets the deployments a request + for that name from the caller would route to, else those whose + ``litellm_params.model`` provider string is that value (``deployments_targeted_by_name``). - Mirrors the live-path semantics in ``perform_health_check()``: ``model`` - matches either the deployment's ``model_name`` alias or its - ``litellm_params.model`` provider string. ``model_id`` matches - ``model_info.id``. - - Both query params are validated against the supplied ``model_list``. - Callers pass an already-scoped list (filtered to the caller's allowed - models for non-admins, full list for admins), so a ``model_id`` that - isn't present resolves to an empty set rather than a single-element - set — preventing a non-admin from reading another deployment's cached - health entry by guessing its ID. - - Returns ``None`` when no targeting is requested — callers should treat - that as "no filter." + Callers pass an already-scoped list, so a ``model_id`` outside the + caller's scope resolves to an empty set and never to the unvalidated id. + Returns ``None`` when no targeting is requested. """ - if not model and not model_id: + if model_id: + return {i for m in model_list if (i := (m.get("model_info") or {}).get("id")) == model_id} + if not model: return None - target_ids: Final[set] = set() - for m in model_list: - deployment_id = (m.get("model_info") or {}).get("id") - if not deployment_id: - continue - if model_id and deployment_id == model_id: - target_ids.add(deployment_id) - continue - if model: - litellm_model = (m.get("litellm_params") or {}).get("model") - if m.get("model_name") == model or litellm_model == model: - target_ids.add(deployment_id) - return target_ids + return { + i + for m in deployments_targeted_by_name(model_list, model, team_id) + if (i := (m.get("model_info") or {}).get("id")) + } def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: set) -> dict: @@ -1046,8 +1079,12 @@ def _health_endpoint_resolve_target_model_name( model_id: str | None, llm_router, ) -> str | None: - """Map ``model_id`` (without ``model``) to ``model_name`` for live health checks.""" - if not model_id or model: + """Map ``model_id`` to its deployment's ``model_name`` for live health checks. + + ``model_id`` wins over ``model``, so an id no deployment carries is a 404 even + when it is paired with a known name. + """ + if not model_id: return model if llm_router is None: raise HTTPException( @@ -1133,7 +1170,9 @@ async def health_endpoint( response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE if is_admin: return result - response.headers["Litellm-Health-Field-Notice"] = "api_base and api_version are admin-only on this endpoint" + response.headers["Litellm-Health-Field-Notice"] = ( + f"{', '.join(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS)} are admin-only on this endpoint" + ) return _strip_admin_only_fields_from_health_result(result) try: @@ -1157,32 +1196,24 @@ async def health_endpoint( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) - _llm_model_list = copy.deepcopy(llm_model_list) - ### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ### - # Live path: scope by model_name (every deployment has one). - # Cache path: scope by model_id (the cache is keyed on model_id). - # Consequence: a deployment whose model_name the caller can access - # but which lacks model_info.id will appear in the live /health - # response but NOT in the background-cache /health response. This is - # surfaced via the "warnings" field below so operators can fix the - # missing model_info.id rather than guess at the discrepancy. - # Keys granted SpecialModelNames.all_proxy_models carry the literal - # "all-proxy-models" entry, which matches no real model_name; treat - # them as unrestricted instead of filtering the list down to nothing. - # Keys granted SpecialModelNames.all_team_models inherit the parent - # team's allowlist (same semantics as get_key_models in - # model_checks.py). Without a team_id the sentinel cannot resolve and - # stays in the list, matching nothing; denied rather than - # unrestricted, mirroring _resolve_key_models_for_auth_check. - accessible_models = list(user_api_key_dict.models) - if SpecialModelNames.all_team_models.value in accessible_models and user_api_key_dict.team_id is not None: - accessible_models = list(user_api_key_dict.team_models) - restrict_to_allowed_models: Final = ( - len(accessible_models) > 0 and SpecialModelNames.all_proxy_models.value not in accessible_models - ) - if restrict_to_allowed_models: - allowed_models: Final = set(accessible_models) - _llm_model_list = [m for m in _llm_model_list if m.get("model_name") in allowed_models] + allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) + restrict_to_allowed_models: Final = not is_admin or allowed_models is not None + _llm_model_list: Final = [ + m + for m in copy.deepcopy(llm_model_list) + if not restrict_to_allowed_models + or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin) + ] + targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id, user_api_key_dict.team_id) + if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": f"key not allowed to health-check model_id {model_id}" + if model_id + else f"key not allowed to health-check model {model}" + }, + ) if use_background_health_checks: # The cached background result covers every model. When the # caller targets a specific model/model_id we have to narrow the @@ -1190,7 +1221,6 @@ async def health_endpoint( # healthy_count, otherwise an unhealthy "foo" combined with any # other healthy model would still report healthy_count > 0 and # the targeted-503 path would never fire. - targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id) if restrict_to_allowed_models: allowed_model_ids: Final = { (m.get("model_info") or {}).get("id") @@ -1202,7 +1232,7 @@ async def health_endpoint( # intersection of "targeted" and "allowed." filter_ids: Final = targeted_ids if targeted_ids is not None else allowed_model_ids filtered: Final = _filter_health_check_results_by_model_ids(health_check_results, filter_ids) - if targeted_ids is None and not allowed_model_ids: + if targeted_ids is None and _llm_model_list and not allowed_model_ids: # Caller has accessible model_names but none of the # matching deployments expose a model_info.id, so the # cache filter (which keys on model_id) drops every @@ -1241,6 +1271,7 @@ async def health_endpoint( model_id=model_id, max_concurrency=health_check_concurrency, router=llm_router, + team_id=user_api_key_dict.team_id, **_hc_filter, ) return _post_process(router_result) diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index cdaa6d5a81c..5cfef11df8d 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -3,6 +3,8 @@ import json from datetime import datetime, timezone from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -108,30 +110,32 @@ class KeyManagementEventHooks: from litellm.proxy.proxy_server import litellm_proxy_admin_name if is_audit_logging_enabled(): - _updated_values: Final = json.dumps(data.json(exclude_none=True), default=str) - - _before_value = existing_key_row.json(exclude_none=True) - _before_value = json.dumps(_before_value, default=str) - - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=_hash_token_if_needed(data.key), - action="updated", - updated_values=_updated_values, - before_value=_before_value, - ) - ) + updated_fields: Final = { + **data.model_dump(exclude_none=True), + **({"project_id": data.project_id} if "project_id" in data.model_fields_set else {}), + } + audit_log: Final = LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=_hash_token_if_needed(data.key), + action="updated", + updated_values=json.dumps(updated_fields, default=str), + before_value=json.dumps(existing_key_row.json(exclude_none=True), default=str), ) + masked_values: Final = TypeAdapter(dict[str, object]).validate_json(str(audit_log.updated_values)) + request_data: Final = ( + audit_log.model_copy(update={"updated_values": json.dumps({**masked_values, "project_id": None})}) + if "project_id" in data.model_fields_set and data.project_id is None + else audit_log + ) + asyncio.create_task(create_audit_log_for_update(request_data=request_data)) @staticmethod async def async_key_rotated_hook( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1e946cc2e23..64da2ad00f6 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -180,7 +180,7 @@ class _ProxyDBLogger(CustomLogger): # here because the input above is constructed non-None. _error_information = cast( StandardLoggingPayloadErrorInformation, - _sanitize_error_information_for_spend_logs(_error_information), + _sanitize_error_information_for_spend_logs(_error_information, original_exception=original_exception), ) _metadata["error_information"] = _error_information diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8f7b515c22a..59971e54e46 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -64,6 +64,7 @@ from litellm.proxy.common_utils.callback_utils import ( strip_callback_config, ) from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers +from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY # Cache special headers as a frozenset for O(1) lookup performance @@ -2300,6 +2301,7 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget user_model_budget: Final = user_api_key_dict.user_model_max_budget data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param + data[_metadata_variable_name].update(carried_budget_metadata(user_api_key_dict)) data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d066f1e9138..10c11119006 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -22,6 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -30,6 +31,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.auth.password_policy import validate_password_policy from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, user_object_permission_id_cache_key, @@ -86,6 +88,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIM_ENTITLEMENTS_METADATA_KEY, SCIM_ROLES_METADATA_KEY, ) +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from prisma import models as prisma_models @@ -96,6 +99,8 @@ if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() +_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) +_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 def _user_table( @@ -1252,6 +1257,13 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda if k == "max_budget": if "max_budget" in fields_set: non_default_values[k] = v + elif k == "model_max_budget": + if k in fields_set: + try: + _USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v) + except ValidationError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + non_default_values[k] = {} if v is None else v elif ( v is not None and v @@ -1421,7 +1433,7 @@ async def _update_single_user_helper( Returns the updated user data or raises an exception on failure. """ - from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client, user_api_key_cache if prisma_client is None: raise Exception("Not connected to DB!") @@ -1464,7 +1476,7 @@ async def _update_single_user_helper( # because `_update_internal_user_params` drops empty values, and `object_permission: {}` is # precisely the clear-my-own-ceiling case this must refuse. _sent_fields: Final = user_request.fields_set() if hasattr(user_request, "fields_set") else set() - _protected_fields: Final = ("max_budget", "soft_budget", "spend", "object_permission") + _protected_fields: Final = ("max_budget", "model_max_budget", "soft_budget", "spend", "object_permission") for _field in _protected_fields: if _field in non_default_values or _field in _sent_fields: raise HTTPException( @@ -1548,6 +1560,12 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) + if "model_max_budget" in non_default_values: + await evict_and_broadcast( + cache_keys=(non_default_values["user_id"],), + user_api_key_cache=user_api_key_cache, + ) + if "object_permission_id" in non_default_values: await _invalidate_cached_user_entitlement( user_id=non_default_values.get("user_id"), @@ -1802,7 +1820,7 @@ async def bulk_user_update( }' ``` """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException( @@ -1867,9 +1885,22 @@ async def bulk_user_update( # Perform bulk database update await UserRepository(prisma_client).table.update_many( where={}, - data=non_default_values, # Update all users + data=( + {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} + if "model_max_budget" in non_default_values + else non_default_values + ), ) + if "model_max_budget" in non_default_values: + for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE): + await asyncio.gather( + *( + evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache) + for user in all_users_in_db[start : start + _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE] + ) + ) + # Create individual success results for user in all_users_in_db: results.append( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index db4467dda5b..1fe174763c4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2718,6 +2718,12 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, ) + if data.project_id is not None and data.project_id != existing_key_row.project_id: + raise HTTPException( + status_code=400, detail="Project reassignment is not supported. Use null to detach the key." + ) + is_project_change: Final = "project_id" in data.model_fields_set and data.project_id != existing_key_row.project_id + common_key_access_checks( user_api_key_dict=user_api_key_dict, data=data, @@ -2810,7 +2816,9 @@ async def _validate_update_key_data( # non-budget change means the caller was authorized — skip the redundant # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key: Final = getattr(existing_key_row, "team_id", None) is not None - can_skip_admin_check: Final = (caller_is_creator or _key_is_team_key) and not _is_budget_change + can_skip_admin_check: Final = (caller_is_creator or _key_is_team_key) and not ( + _is_budget_change or is_project_change + ) if (not _is_proxy_admin) and not can_skip_admin_check: hashed_key: Final = existing_key_row.token await _check_key_admin_access( @@ -2853,7 +2861,9 @@ async def _validate_update_key_data( ) # Validate key against project limits if project_id is being set - _project_id_to_check: Final = getattr(data, "project_id", None) or getattr(existing_key_row, "project_id", None) + _project_id_to_check: Final = ( + data.project_id if "project_id" in data.model_fields_set else existing_key_row.project_id + ) if _project_id_to_check is not None and (data.models is not None or data.max_budget is not None): await _check_project_key_limits( project_id=_project_id_to_check, @@ -2962,6 +2972,7 @@ async def update_key_fn( - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. + - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. - organization_id: Optional[str] - The organization id of the key. - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - models: Optional[list] - Model_name's a user is allowed to call diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 94d2b773e14..8484279c69a 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -763,6 +763,15 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: merged_litellm_params.pop(field, None) merged_model_info.pop(field, None) + elif ( + field + in ( + "auto_router_routing_compression", + "auto_router_model_compression", + ) + and getattr(updated_patch.litellm_params, field) is None + ): + merged_litellm_params.pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c050368b3fe..e9e37540dd8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy._types import ( SpecialProxyStrings, TeamAccessGroupModelGrant, TeamAddMemberResponse, + TeamInfoMember, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, TeamListResponseObject, @@ -4292,37 +4293,35 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _hydrate_member_emails( +async def _hydrate_member_user_details( prisma_client: PrismaClient, members: Sequence[Member], -) -> tuple[Member, ...]: - """Fill in ``user_email`` for roster entries that were stored without one. - - ``members_with_roles`` is a denormalized snapshot written at add-time, so an entry - stored with ``user_email=None`` keeps that null even once the user row has an email. - Look the missing ones up in ``LiteLLM_UserTable`` (one indexed query) and fill them - in. A stored email is never overwritten - the snapshot stays the source of truth - wherever it has a value. - """ - missing_user_ids: Final = frozenset(m.user_id for m in members if not m.user_email and m.user_id is not None) - if not missing_user_ids: - return tuple(members) - - user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await _user_db(prisma_client).find_many( - where={ # mutable-ok: Prisma query filters are dict-shaped - "user_id": { # mutable-ok: Prisma query filters are dict-shaped - "in": sorted(missing_user_ids) +) -> tuple[TeamInfoMember, ...]: + """Attach ``user_alias`` and fill in a missing ``user_email`` from ``LiteLLM_UserTable`` in one query.""" + user_ids: Final = frozenset(m.user_id for m in members if m.user_id is not None) + user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = ( + await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(user_ids) + } } - } + ) + if user_ids + else () ) - email_by_user_id: Final = MappingProxyType({u.user_id: u.user_email for u in user_rows if u.user_email}) + user_by_id: Final = MappingProxyType({u.user_id: u for u in user_rows}) - return tuple( - m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload - if not m.user_email and m.user_id is not None and m.user_id in email_by_user_id - else m - for m in members - ) + def hydrate(m: Member) -> TeamInfoMember: + user_row: Final = user_by_id.get(m.user_id) if m.user_id is not None else None + return TeamInfoMember( + role=m.role, + user_id=m.user_id, + user_email=m.user_email or (user_row.user_email if user_row is not None else None), + user_alias=user_row.user_alias if user_row is not None else None, + ) + + return tuple(hydrate(m) for m in members) async def _resolve_team_access_group_resources( @@ -4462,17 +4461,12 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) - # Fill in emails the add-time roster snapshot never captured - hydrated_members: Final = await _hydrate_member_emails( + hydrated_members: Final = await _hydrate_member_user_details( prisma_client=prisma_client, members=resolved_team_info.members_with_roles, ) hydrated_team_info: Final = resolved_team_info.model_copy( - update={ # mutable-ok: pydantic update payload - # list(), not the tuple: model_copy skips validation, so the field has - # to be handed the list[Member] the response model declares. - "members_with_roles": list(hydrated_members) # mutable-ok: declared list[Member] - } + update={"members_with_roles": hydrated_members} # mutable-ok: pydantic update payload ) response_object: Final = TeamInfoResponseObject( diff --git a/litellm/proxy/mcp_registry.json b/litellm/proxy/mcp_registry.json index f37fc39813e..b117f35600d 100644 --- a/litellm/proxy/mcp_registry.json +++ b/litellm/proxy/mcp_registry.json @@ -66,7 +66,7 @@ "name": "slack", "title": "Slack", "description": "Channel management, messaging, and Slack workspace integration", - "icon_url": "https://cdn.simpleicons.org/slack", + "icon_url": "/ui/assets/logos/slack.svg", "category": "Communication", "registry_url": null, "transport": "stdio", @@ -249,7 +249,7 @@ "name": "exa", "title": "Exa", "description": "Fast, intelligent web search and web crawling", - "icon_url": "https://cdn.simpleicons.org/exa", + "icon_url": "/ui/assets/logos/exa_ai.png", "category": "Search", "registry_url": "https://registry.modelcontextprotocol.io/servers/ai.exa%2Fexa", "transport": "http", @@ -262,7 +262,7 @@ "name": "tavily", "title": "Tavily", "description": "AI-optimized search engine for research and retrieval", - "icon_url": "https://cdn.simpleicons.org/tavily", + "icon_url": "/ui/assets/logos/tavily.png", "category": "Search", "registry_url": null, "transport": "stdio", @@ -288,7 +288,7 @@ "name": "playwright", "title": "Playwright", "description": "Browser automation and testing with Playwright", - "icon_url": "https://cdn.simpleicons.org/playwright", + "icon_url": "https://raw.githubusercontent.com/microsoft/playwright/2f6148bcd1a96ec687d55ce08645fc6315b1514e/packages/recorder/public/playwright-logo.svg", "category": "Web & Browser", "registry_url": null, "transport": "stdio", @@ -300,7 +300,7 @@ "name": "browserbase", "title": "Browserbase", "description": "Cloud browser automation and session management", - "icon_url": "https://cdn.simpleicons.org/browserbase", + "icon_url": "https://www.browserbase.com/favicon.svg", "category": "Web & Browser", "registry_url": null, "transport": "stdio", @@ -315,7 +315,7 @@ "name": "aws", "title": "AWS", "description": "Interact with Amazon Web Services resources and APIs", - "icon_url": "https://cdn.simpleicons.org/amazonaws", + "icon_url": "/ui/assets/logos/aws.svg", "category": "Cloud", "registry_url": null, "transport": "stdio", @@ -392,7 +392,7 @@ "name": "twilio", "title": "Twilio", "description": "Send SMS, make calls, and manage communication via Twilio", - "icon_url": "https://cdn.simpleicons.org/twilio", + "icon_url": "/ui/assets/logos/twilio.svg", "category": "Communication", "registry_url": null, "transport": "stdio", diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index ebf4d988fdd..53ebbe91b54 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -15,7 +15,7 @@ from litellm.llms.base_llm.ocr.transformation import ( OCRResponse, parse_ocr_request_format, ) -from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type +from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes 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 @@ -28,24 +28,7 @@ def _build_document_from_upload( filename: str | None, content_type: str | None, ) -> dict[str, str]: - """ - Convert uploaded file bytes into a Mistral-format document dict with base64 data URI. - - Delegates to convert_file_document_to_url_document after resolving MIME type - from the upload's content_type header or filename. - """ - mime_type = content_type.split(";")[0].strip() if content_type else None - if not mime_type or mime_type == "application/octet-stream": - if filename: - mime_type = get_mime_type(filename) - - return convert_file_document_to_url_document( - { - "type": "file", - "file": file_content, - "mime_type": mime_type or "application/octet-stream", - } - ) + return convert_upload_to_url_document(file_content, filename, content_type) def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]: @@ -120,7 +103,7 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: # Seek to start in case the file was already partially read by middleware await uploaded_file.seek(0) - file_content: Final = await uploaded_file.read() + file_content: Final = await uploaded_file.read(get_max_file_bytes() + 1) if not file_content: raise ValueError("Uploaded file is empty") diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ed193c7f434..ad45781d5d2 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -58,6 +58,15 @@ class UndeliverableStreamRewrite(Exception): self.guardrail_name: Final = guardrail_name +class UnappliableRequestRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, " + "so the request was rejected rather than sent unrewritten" + ) + self.guardrail_name: Final = guardrail_name + + def _tool_call_shape(tool_call: object) -> tuple[object, object]: plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call function: Final = plain.get("function") if isinstance(plain, Mapping) else None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c8fd7edfed6..606a590c24b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,7 +26,6 @@ from collections.abc import ( MutableMapping, Sequence, ) -from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -659,6 +658,15 @@ from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.spend_counter_batch import ( + PendingSpendIncrement, + active_spend_counter_batch, + forget_spend_counter, + post_call_counter_keys, + read_batched_spend_counter, + record_spend_counter_value, + spend_counter_batch_scope, +) from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) @@ -2627,6 +2635,7 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None if needs_update: spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) if spend_counter_cache.redis_cache is not None: + forget_spend_counter(counter_key) try: await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend) except Exception: @@ -2724,6 +2733,12 @@ async def read_spend_counter_cache_value(counter_key: str) -> tuple[float | None """Return (value, authoritative) for the live counter, None when absent. A clean Redis miss is final: the per-pod in-memory copy outlives the Redis TTL and only holds this pod's writes, so it is consulted only when Redis is unreachable.""" + batch: Final = active_spend_counter_batch() + if batch is not None: + batched: Final = await batch.read(counter_key) + if batched is not None: + return batched + if spend_counter_cache.redis_cache is not None: try: redis_val: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) @@ -2761,12 +2776,6 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) return fallback_spend, False -@dataclass(frozen=True, slots=True) -class _PendingSpendIncrement: - counter_key: str - increment: float - - async def increment_spend_counters( token: str | None, team_id: str | None, @@ -2789,6 +2798,45 @@ async def increment_spend_counters( Awaited (not create_task) in the cost callback, so the counter is updated before the next request's auth check runs. """ + with spend_counter_batch_scope( + spend_counter_cache.redis_cache, + counter_keys=post_call_counter_keys( + token=token, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + tags=tags, + model_access_groups=model_access_groups, + ), + ): + await _increment_spend_counters_batched( + token=token, + team_id=team_id, + user_id=user_id, + response_cost=response_cost, + org_id=org_id, + budget_reservation=budget_reservation, + end_user_id=end_user_id, + tags=tags, + request_started_at=request_started_at, + model_access_groups=model_access_groups, + ) + + +async def _increment_spend_counters_batched( + token: str | None, + team_id: str | None, + user_id: str | None, + response_cost: float | None, + org_id: str | None, + budget_reservation: dict | None, + end_user_id: str | None, + tags: list[str] | None, + request_started_at: datetime | None, + model_access_groups: Sequence[str] | None, +): + """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET.""" reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update( budget_reservation=budget_reservation, response_cost=response_cost, @@ -2801,7 +2849,7 @@ async def increment_spend_counters( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + async def _key_scope(key_token: str) -> tuple[PendingSpendIncrement | BaseException, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -2812,7 +2860,7 @@ async def increment_spend_counters( hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token ) key_counter_key: Final = f"spend:key:{hashed_token}" - key_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + key_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () if key_counter_key in reserved_counter_keys else ( @@ -2824,7 +2872,7 @@ async def increment_spend_counters( ) ) - async def _key_window_increment(window: object) -> _PendingSpendIncrement | None: + async def _key_window_increment(window: object) -> PendingSpendIncrement | None: duration = ( window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) ) @@ -2871,9 +2919,9 @@ async def increment_spend_counters( ) return key_pending + tuple(item for item in window_pending if item is not None) - async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + async def _team_scope(scope_team_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: team_counter_key: Final = f"spend:team:{scope_team_id}" - team_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + team_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () if team_counter_key in reserved_counter_keys else ( @@ -2885,7 +2933,7 @@ async def increment_spend_counters( ) ) - async def _team_window_increment(window: object) -> _PendingSpendIncrement | None: + async def _team_window_increment(window: object) -> PendingSpendIncrement | None: duration = ( window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) ) @@ -2934,7 +2982,7 @@ async def increment_spend_counters( async def _team_member_scope( scope_user_id: str, scope_team_id: str - ) -> tuple[_PendingSpendIncrement | BaseException, ...]: + ) -> tuple[PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" if team_member_counter_key in reserved_counter_keys: return () @@ -2946,7 +2994,7 @@ async def increment_spend_counters( ), ) - async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" if user_counter_key in reserved_counter_keys: return () @@ -3056,7 +3104,7 @@ async def _prepare_end_user_and_tag_spend_increments( tags: list[str] | None, response_cost: float, reserved_counter_keys: set[str], -) -> tuple[_PendingSpendIncrement | BaseException, ...]: +) -> tuple[PendingSpendIncrement | BaseException, ...]: unique_tags: Final = ( tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else () ) @@ -3093,7 +3141,7 @@ async def _prepare_model_access_group_spend_increments( model_access_groups: Sequence[object], response_cost: float, reserved_counter_keys: set[str], -) -> tuple[_PendingSpendIncrement | BaseException, ...]: +) -> tuple[PendingSpendIncrement | BaseException, ...]: """Charge the model access groups that authorized this request. Without this the counter auth reads is written only by the reservation path, so @@ -3126,7 +3174,7 @@ async def _prepare_org_spend_increment( org_id: str | None, response_cost: float, reserved_counter_keys: set[str], -) -> tuple[_PendingSpendIncrement, ...]: +) -> tuple[PendingSpendIncrement, ...]: if org_id is None: return () @@ -3144,7 +3192,7 @@ async def _prepare_unreserved_spend_counter_increment( source_cache_key: str | list[str], increment: float, reserved_counter_keys: set[str], -) -> _PendingSpendIncrement | None: +) -> PendingSpendIncrement | None: if counter_key in reserved_counter_keys: return None @@ -3159,7 +3207,7 @@ async def _prepare_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, -) -> _PendingSpendIncrement: +) -> PendingSpendIncrement: """ Initialize counter from the authoritative DB spend value if not yet set, then return the pending increment for the caller to apply in one @@ -3181,7 +3229,7 @@ async def _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key=source_cache_key, ) - return _PendingSpendIncrement(counter_key=counter_key, increment=increment) + return PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _enqueue_window_spend_row_update( @@ -3240,7 +3288,7 @@ async def _prepare_window_spend_counter_increment( window_duration: str | None, window_start: datetime | None, increment: float, -) -> _PendingSpendIncrement | None: +) -> PendingSpendIncrement | None: if window_start is None: verbose_proxy_logger.warning( "Skipping spend counter increment for invalid budget window %s", @@ -3257,7 +3305,7 @@ async def _prepare_window_spend_counter_increment( ) if initialized is False: return None - return _PendingSpendIncrement(counter_key=counter_key, increment=increment) + return PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _ensure_spend_counter_initialized( @@ -3324,6 +3372,14 @@ async def _ensure_window_spend_counter_initialized( async def _is_spend_counter_cache_warm(counter_key: str) -> bool: + batched: Final = await read_batched_spend_counter(counter_key) + if batched is not None: + batched_value, _ = batched + if batched_value is None: + return False + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value) + return True + if spend_counter_cache.redis_cache is not None: try: current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache( @@ -3368,16 +3424,16 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float): key=counter_key, value=current_value, ) + record_spend_counter_value(counter_key, float(current_value)) return current_value - return await spend_counter_cache.async_increment_cache( - key=counter_key, - value=increment, - refresh_ttl=True, + return await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment ) async def _invalidate_spend_counter(counter_key: str): + forget_spend_counter(counter_key) spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) if spend_counter_cache.redis_cache is not None: try: @@ -3390,16 +3446,23 @@ async def _invalidate_spend_counter(counter_key: str): ) -async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None: +async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: + try: + await increment_spend_counters_pipeline(pending=pending) + except RedisCircuitBreakerOpenError: + return + + +async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> None: + """One INCRBYFLOAT+EXPIRE pipeline for every pending counter; on failure every counter is invalidated + before the error propagates, so no caller can read a half-applied batch.""" if not pending: return redis_cache: Final = spend_counter_cache.redis_cache if redis_cache is None: for item in pending: - await spend_counter_cache.async_increment_cache( - key=item.counter_key, - value=item.increment, - refresh_ttl=True, + await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment ) return ttl: Final = redis_cache.get_ttl() @@ -3409,13 +3472,12 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme ] try: results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) - except Exception as e: + except Exception: await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) - if isinstance(e, RedisCircuitBreakerOpenError): - return raise for item, current_value in zip(pending, results or ()): spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) + record_spend_counter_value(item.counter_key, float(current_value)) async def update_cache( diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 3d0bd5e61c9..20b4708c193 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -145,12 +145,14 @@ ROUTE_ENDPOINT_MAPPING: Final = { } +_AVAILABLE_MODELS_HINT: Final = "Call `/v1/models` to view available models for your key." + + class ProxyModelNotFoundError(HTTPException): def __init__(self, route: str, model_name: str, retryable_with_model_read_through: bool = True): self.retryable_with_model_read_through: Final = retryable_with_model_read_through - detail: Final = { - "error": f"{route}: Invalid model name passed in model={model_name}. Call `/v1/models` to view available models for your key." - } + self.spend_log_error_message: Final = f"{route}: Invalid model name passed in. {_AVAILABLE_MODELS_HINT}" + detail: Final = {"error": f"{route}: Invalid model name passed in model={model_name}. {_AVAILABLE_MODELS_HINT}"} super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index eedec0619db..6074a50a69b 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router from litellm.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer @@ -257,46 +258,47 @@ async def reserve_budget_for_request( applied_entries: Final[list[dict[str, float | str]]] = [] try: - for counter in counters: - entry = _counter_to_reservation_entry( - counter=counter, - reserved_cost=reservation_cost, - ) - applied_entries.append(entry) - try: - reserved_value = await _reserve_counter( + with _counters_batch_scope(frozenset(counter.counter_key for counter in counters)): + for counter in counters: + entry = _counter_to_reservation_entry( counter=counter, - reservation_cost=reservation_cost, + reserved_cost=reservation_cost, ) - except _CounterReservationUnavailable as exc: - if exc.touched_counter and not exc.counter_invalidated: - await _release_applied_entries_best_effort( - entries=[entry], - default_reserved_cost=reservation_cost, + applied_entries.append(entry) + try: + reserved_value = await _reserve_counter( + counter=counter, + reservation_cost=reservation_cost, ) - applied_entries.remove(entry) - if fail_closed_budget_enforcement: - _raise_reservation_unavailable(counter_key=counter.counter_key) - continue + except _CounterReservationUnavailable as exc: + if exc.touched_counter and not exc.counter_invalidated: + await _release_applied_entries_best_effort( + entries=[entry], + default_reserved_cost=reservation_cost, + ) + applied_entries.remove(entry) + if fail_closed_budget_enforcement: + _raise_reservation_unavailable(counter_key=counter.counter_key) + continue - if reserved_value is not None: - current_spend = reserved_value - else: - cached_spend = current_spend_by_counter_key.get(counter.counter_key) - if cached_spend is None: - cached_spend = await _get_current_counter_value(counter=counter) - current_spend = cached_spend + reservation_cost - if current_spend > counter.max_budget: - reservation_cost = await _apply_over_budget_reservation_policy( - counter=counter, - valid_token=valid_token, - entry=entry, - applied_entries=applied_entries, - reservation_cost=reservation_cost, - current_spend=current_spend, - fail_closed_budget_enforcement=fail_closed_budget_enforcement, - ) - continue + if reserved_value is not None: + current_spend = reserved_value + else: + cached_spend = current_spend_by_counter_key.get(counter.counter_key) + if cached_spend is None: + cached_spend = await _get_current_counter_value(counter=counter) + current_spend = cached_spend + reservation_cost + if current_spend > counter.max_budget: + reservation_cost = await _apply_over_budget_reservation_policy( + counter=counter, + valid_token=valid_token, + entry=entry, + applied_entries=applied_entries, + reservation_cost=reservation_cost, + current_spend=current_spend, + fail_closed_budget_enforcement=fail_closed_budget_enforcement, + ) + continue except Exception: await _release_applied_entries_best_effort( entries=applied_entries, @@ -878,67 +880,92 @@ async def _get_current_counter_value(counter: _BudgetCounter) -> float: ) +def _counters_batch_scope(counter_keys: frozenset[str]) -> spend_counter_batch_scope: + """Each counter is read once, then written, so one MGET up front serves every read in the loop.""" + from litellm.proxy.proxy_server import spend_counter_cache + + return spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys) + + +@dataclass(frozen=True, slots=True) +class _EntryAdjustment: + entry: dict[str, float | str] + counter_key: str + target_adjustment: float + adjustment: float + + +def _entry_adjustment( + entry: dict[str, float | str], actual_cost: float, default_reserved_cost: float +) -> _EntryAdjustment | None: + counter_key: Final = entry.get("counter_key") + if counter_key is None: + return None + target_adjustment: Final = actual_cost - _get_entry_reserved_cost( + entry=entry, default_reserved_cost=default_reserved_cost + ) + adjustment: Final = target_adjustment - float(entry.get("applied_adjustment") or 0.0) + if adjustment == 0: + return None + return _EntryAdjustment( + entry=entry, counter_key=str(counter_key), target_adjustment=target_adjustment, adjustment=adjustment + ) + + async def _set_reserved_entries_actual_cost( entries: list[dict], actual_cost: float, default_reserved_cost: float, reseed_on_inconsistent: bool = True, ) -> None: - for entry in entries: - await _set_reserved_entry_actual_cost( - entry=entry, - actual_cost=actual_cost, - default_reserved_cost=default_reserved_cost, - reseed_on_inconsistent=reseed_on_inconsistent, + """Every reserved counter is read from one MGET and the consistent adjustments go out in one pipeline. + A counter that was flushed or reseeded since reservation is settled on its own after the pipeline.""" + from litellm.proxy.proxy_server import increment_spend_counters_pipeline + + with _counters_batch_scope(frozenset(str(entry["counter_key"]) for entry in entries if "counter_key" in entry)): + adjustments: Final = tuple( + adjustment + for entry in entries + if (adjustment := _entry_adjustment(entry, actual_cost, default_reserved_cost)) is not None ) - - -async def _set_reserved_entry_actual_cost( - entry: dict, - actual_cost: float, - default_reserved_cost: float, - reseed_on_inconsistent: bool = True, -) -> None: - from litellm.proxy.proxy_server import ( - _increment_spend_counter_cache, - reseed_spend_counter_from_db, - ) - - counter_key: Final = entry.get("counter_key") - if counter_key is None: - return - reserved_cost: Final = _get_entry_reserved_cost( - entry=entry, - default_reserved_cost=default_reserved_cost, - ) - target_adjustment: Final = actual_cost - reserved_cost - applied_adjustment: Final = float(entry.get("applied_adjustment") or 0.0) - adjustment: Final = target_adjustment - applied_adjustment - if adjustment == 0: - return - if await _counter_can_apply_adjustment( - counter_key=counter_key, - adjustment=adjustment, - ): - await _increment_spend_counter_cache( - counter_key=counter_key, - increment=adjustment, + consistent: Final = tuple( + await asyncio.gather( + *( + _counter_can_apply_adjustment(counter_key=item.counter_key, adjustment=item.adjustment) + for item in adjustments + ) + ) ) - elif reseed_on_inconsistent: - # Post-call reconcile / release: the counter was flushed, expired or reseeded - # between reservation and reconcile, so the optimistic delta no longer applies. - # Reseed from the DB floor (which cannot include this request's cost yet) and - # add the settled cost, since increment_spend_counters skips reserved keys. - reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key) - if reseeded and actual_cost > 0: - await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost) - else: - # Pre-call admission resize: the in-flight reservation cost is not yet - # persisted, so the DB floor would discard it. Keep the original - # fail-closed behavior (raise -> reserve_budget_for_request releases and - # denies) rather than admitting against an inconsistent counter. - raise RuntimeError(f"Cannot resize budget reservation against inconsistent counter {counter_key}") - entry["applied_adjustment"] = target_adjustment + inconsistent: Final = tuple(item for item, ok in zip(adjustments, consistent) if not ok) + if inconsistent and not reseed_on_inconsistent: + # Pre-call admission resize: the in-flight reservation cost is not yet + # persisted, so the DB floor would discard it. Keep the original + # fail-closed behavior (raise -> reserve_budget_for_request releases and + # denies) rather than admitting against an inconsistent counter. + raise RuntimeError( + f"Cannot resize budget reservation against inconsistent counter {inconsistent[0].counter_key}" + ) + applicable: Final = tuple(item for item, ok in zip(adjustments, consistent) if ok) + await increment_spend_counters_pipeline( + pending=tuple( + PendingSpendIncrement(counter_key=item.counter_key, increment=item.adjustment) for item in applicable + ) + ) + for item in inconsistent: + await _reseed_reserved_entry(item=item, actual_cost=actual_cost) + for item in adjustments: + item.entry["applied_adjustment"] = item.target_adjustment + + +async def _reseed_reserved_entry(item: _EntryAdjustment, actual_cost: float) -> None: + """Post-call reconcile / release of a counter that was flushed, expired or reseeded between reservation and + reconcile: the optimistic delta no longer applies, so reseed from the DB floor (which cannot include this + request's cost yet) and add the settled cost, since increment_spend_counters skips reserved keys.""" + from litellm.proxy.proxy_server import _increment_spend_counter_cache, reseed_spend_counter_from_db + + reseeded: Final = await reseed_spend_counter_from_db(counter_key=item.counter_key) + if reseeded and actual_cost > 0: + await _increment_spend_counter_cache(counter_key=item.counter_key, increment=actual_cost) async def _counter_can_apply_adjustment( @@ -963,8 +990,8 @@ async def _release_applied_entries_best_effort( ) -> None: for entry in entries: try: - await _set_reserved_entry_actual_cost( - entry=entry, + await _set_reserved_entries_actual_cost( + entries=[entry], # mutable-ok: the reconcile takes the reservation's list of entries actual_cost=0.0, default_reserved_cost=default_reserved_cost, ) diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py new file mode 100644 index 00000000000..efd3a78d211 --- /dev/null +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -0,0 +1,60 @@ +"""Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import ( + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) + + +def carry_team_and_user_budget_state( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, +) -> None: + if team_object is not None: + valid_token.team_budget_snapshot = TeamBudgetSnapshot( # rebind-ok: the request credential is pinned in place + budget_reset_at=team_object.budget_reset_at, + max_budget=team_object.max_budget, + ) + if user_object is not None: + valid_token.user_budget_snapshot = UserBudgetSnapshot( # rebind-ok: same object the caller keeps using + budget_reset_at=user_object.budget_reset_at, + max_budget=user_object.max_budget, + user_alias=user_object.user_alias, + ) + + +def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: + budget_table: Final = org_table.litellm_budget_table + valid_token.organization_alias = ( + org_table.organization_alias + ) # rebind-ok: the request credential is pinned in place + valid_token.org_budget_snapshot = OrgBudgetSnapshot( # rebind-ok: same object the caller keeps using + spend=org_table.spend, + max_budget=budget_table.max_budget if budget_table is not None else None, + ) + + +def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object]: + snapshots: Final = ( + valid_token.team_budget_snapshot, + valid_token.user_budget_snapshot, + valid_token.org_budget_snapshot, + ) + return MappingProxyType( + { + key: value + for snapshot in snapshots + if snapshot is not None + for key, value in snapshot.metadata_entries().items() + } + ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py new file mode 100644 index 00000000000..7106d88c655 --- /dev/null +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -0,0 +1,217 @@ +"""One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each.""" + +import asyncio +from collections.abc import Iterator, Mapping, Sequence +from contextvars import ContextVar, Token +from dataclasses import dataclass +from types import MappingProxyType, TracebackType +from typing import Final + +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key + +_CounterValues: Final = TypeAdapter(dict[str, float | None]) +_NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) + + +@dataclass(frozen=True, slots=True) +class PendingSpendIncrement: + counter_key: str + increment: float + + +class SpendCounterBatch: + """Bound counters are read with one MGET on first use; counters bound later join the next MGET. + ``async_batch_get_cache`` maps a clean miss to ``None`` and drops keys only when Redis failed, so an absent + key means "read it yourself" and a present ``None`` is an authoritative miss.""" + + __slots__ = ("_fetched", "_keys", "_loaded", "_lock", "_open", "_redis_cache") + + def __init__(self, redis_cache: RedisCache) -> None: + self._redis_cache: Final = redis_cache + self._lock: Final = asyncio.Lock() + self._open = True + self._keys: frozenset[str] = frozenset() + self._fetched: frozenset[str] = frozenset() + self._loaded: Mapping[str, float | None] = _NO_VALUES + + @property + def counter_keys(self) -> frozenset[str]: + return self._keys + + @property + def is_open(self) -> bool: + return self._open + + def bind(self, counter_keys: frozenset[str]) -> None: + if self._open: + self._keys = self._keys | counter_keys + + def close(self) -> None: + """Later reads go to Redis directly; call before any read-then-write on the counters.""" + self._open = False + + async def read(self, counter_key: str) -> tuple[float | None, bool] | None: + """(value, authoritative) for a bound counter, None when the caller must read Redis itself.""" + if not self._open or counter_key not in self._keys: + return None + loaded: Final = await self._load() + if counter_key not in loaded: + return None + return loaded[counter_key], True + + def record(self, counter_key: str, value: float) -> None: + """A write returned the counter's new value; later reads in this scope see it instead of the MGET value.""" + if not self._open: + return + key: Final = frozenset((counter_key,)) + self._keys = self._keys | key + self._fetched = self._fetched | key + self._loaded = MappingProxyType({**self._loaded, counter_key: value}) + + def forget(self, counter_key: str) -> None: + """A write left the counter's value unknown; later reads in this scope go to Redis.""" + key: Final = frozenset((counter_key,)) + self._keys = self._keys - key + self._fetched = self._fetched - key + self._loaded = MappingProxyType({k: v for k, v in self._loaded.items() if k != counter_key}) + + async def _load(self) -> Mapping[str, float | None]: + async with self._lock: + pending: Final = self._keys - self._fetched + if pending: + self._fetched = self._fetched | pending + fetched: Final = await self._fetch(pending) + self._loaded = MappingProxyType({**fetched, **self._loaded}) + return self._loaded + + async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]: + try: + return _CounterValues.validate_python( + await self._redis_cache.async_batch_get_cache(key_list=sorted(keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped cache API + ) + except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback + verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e) + return _NO_VALUES + + +_active_batch: Final[ContextVar[SpendCounterBatch | None]] = ContextVar("spend_counter_batch", default=None) + + +def active_spend_counter_batch() -> SpendCounterBatch | None: + return _active_batch.get() + + +class spend_counter_batch_scope: + """Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it. + Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET.""" + + __slots__ = ("_counter_keys", "_redis_cache", "_token") + + def __init__(self, redis_cache: RedisCache | None, counter_keys: frozenset[str] = frozenset()) -> None: + self._redis_cache: Final = redis_cache + self._counter_keys: Final = counter_keys + self._token: Token[SpendCounterBatch | None] | None = None + + def __enter__(self) -> None: + if self._redis_cache is None: + return + outer: Final = _active_batch.get() + if outer is not None and outer.is_open: + outer.bind(self._counter_keys) + return + batch: Final = SpendCounterBatch(self._redis_cache) + batch.bind(self._counter_keys) + self._token = _active_batch.set(batch) + + def __exit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + if self._token is not None: + _active_batch.reset(self._token) + + +def release_spend_counter_batch() -> None: + batch: Final = _active_batch.get() + if batch is not None: + batch.close() + + +def _iter_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> Iterator[str]: + if token.token is not None: + yield f"spend:key:{token.token}" + if token.team_id is not None: + yield f"spend:team:{token.team_id}" + if token.user_id is not None: + yield f"spend:team_member:{token.user_id}:{token.team_id}" + if token.user_id is not None: + yield f"spend:user:{token.user_id}" + if end_user_id is not None: + yield f"spend:end_user:{end_user_id}" + if token.org_id is not None: + yield f"spend:org:{token.org_id}" + + +def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]: + return frozenset(_iter_admission_counter_keys(token, end_user_id)) + + +def post_call_counter_keys( + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + tags: Sequence[object] | None, + model_access_groups: Sequence[object] | None, +) -> frozenset[str]: + """Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read.""" + entity_keys: Final = admission_counter_keys( + UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id + ) + tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str)) + group_keys: Final = frozenset( + model_access_group_spend_counter_key(group) + for group in model_access_groups or () + if group and isinstance(group, str) + ) + return entity_keys | tag_keys | group_keys + + +def bind_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> None: + """Idempotent: call again after the token gains ids (end user, team org) so those counters join the MGET.""" + bind_spend_counter_keys(admission_counter_keys(token, end_user_id)) + + +def bind_spend_counter_keys(counter_keys: frozenset[str]) -> None: + batch: Final = _active_batch.get() + if batch is None: + return + batch.bind(counter_keys) + + +def record_spend_counter_value(counter_key: str, value: float) -> None: + batch: Final = _active_batch.get() + if batch is None: + return + batch.record(counter_key, value) + + +def forget_spend_counter(counter_key: str) -> None: + batch: Final = _active_batch.get() + if batch is None: + return + batch.forget(counter_key) + + +async def read_batched_spend_counter(counter_key: str) -> tuple[float | None, bool] | None: + """Bind-on-read for counters only known at read time (budget windows); the first reader pays the MGET.""" + batch: Final = _active_batch.get() + if batch is None: + return None + batch.bind(frozenset((counter_key,))) + return await batch.read(counter_key) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 01398d38687..3fd20bb7e81 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1167,6 +1167,7 @@ def _redact_prompt_fields_in_guardrail_entry( def _sanitize_error_information_for_spend_logs( error_information: StandardLoggingPayloadErrorInformation | None, + original_exception: BaseException | None = None, ) -> StandardLoggingPayloadErrorInformation | None: """ Sanitize ``error_information`` before it lands in ``LiteLLM_SpendLogs.metadata``. @@ -1188,7 +1189,12 @@ def _sanitize_error_information_for_spend_logs( if error_information is None: return None - sanitized = cast(dict, {**error_information}) + persisted: Final = ( + {**error_information, "error_message": original_exception.spend_log_error_message} + if isinstance(original_exception, ProxyModelNotFoundError) + else error_information + ) + sanitized = cast(dict, {**persisted}) if not should_store_prompts_and_responses_in_spend_logs(): for field in ("error_message", "traceback"): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8e6142090a4..89625021e37 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -121,6 +121,11 @@ from litellm.proxy.db.create_views import ( should_create_missing_views, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_url_settings import ( + DatabaseURLSettings, + add_missing_query_params, + token_refresh_params_from_url, +) from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, @@ -4054,7 +4059,10 @@ class PrismaClient: # loop and times out after 30s. if token_auth is not None and reader_iam_endpoint is not None: reader_token: Final = mint_database_token(token_auth, reader_iam_endpoint) - read_replica_url = reader_iam_endpoint.build_url(reader_token) + read_replica_url = add_missing_query_params( + reader_iam_endpoint.build_url(reader_token), + token_refresh_params_from_url(read_replica_url), + ) os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url reader_kwargs: Final[dict[str, Any]] = {"datasource": {"url": read_replica_url}} if http_client is not None: @@ -7807,7 +7815,7 @@ def construct_database_url_from_env_vars() -> str | None: if database_schema: database_url += f"?schema={database_schema}" - return database_url + return add_missing_query_params(database_url, DatabaseURLSettings.from_env().tls_params()) return None diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index b824a5928c6..44c47af57f4 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -33,7 +33,7 @@ from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..llms.azure.common_utils import get_azure_ad_token -from ..llms.azure.realtime.handler import AzureOpenAIRealtime +from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..llms.openai.realtime.handler import OpenAIRealtime @@ -413,14 +413,14 @@ async def _arealtime( api_version = api_version or litellm_params.api_version or "2024-10-01-preview" - realtime_protocol = ( + configured_realtime_protocol: Final = ( kwargs.get("realtime_protocol") or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - if realtime_protocol is None and (query_params or {}).get("intent") == "transcription": - realtime_protocol = "GA" - realtime_protocol = realtime_protocol or "beta" + realtime_protocol: Final = azure_realtime_protocol_for_client( + configured_realtime_protocol, query_params=query_params, websocket=websocket + ) resolved_azure_ad_token: Final = ( None if api_key else get_azure_ad_token(GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token)) ) @@ -586,9 +586,7 @@ def _azure_realtime_health_protocol( configured: Final = configured_raw if isinstance(configured_raw, str) else None if configured is not None: return configured, query_params - if query_params is not None: - return "GA", query_params - return "beta", None + return "GA", query_params def _realtime_health_check_auth_headers( @@ -621,8 +619,8 @@ async def _realtime_health_check( api_key: str - api key custom_llm_provider: str - custom llm provider realtime_protocol: Optional[str] - protocol version ("GA"/"v1" for GA path, "beta" for beta path); - None resolves it for Azure from model_params/env, with transcription-only models probing GA - plus intent=transcription the way real calls do + None resolves it for Azure from model_params/env and otherwise probes GA, the upstream a client + without the OpenAI-Beta header is bridged to, with transcription-only models adding intent=transcription Returns: bool - True if connection is successful, False otherwise diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 660dd8f0c92..126b976e2c5 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -40,6 +40,9 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, + WebSearchCallCompletedEvent, + WebSearchCallInProgressEvent, + WebSearchCallSearchingEvent, ) from litellm.types.utils import Delta as ChatCompletionDelta from litellm.types.utils import ( @@ -135,6 +138,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( self.responses_api_request.get("tools") ) + self._web_search_calls: dict[str, object] = {} # mutable-ok: latest call by provider id + self._queued_web_search_call_ids: set[str] = set() # mutable-ok: emitted call ids def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing: Final = self._tool_output_index_by_call_id.get(call_id) @@ -172,6 +177,43 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return delta.content or delta.function_call or delta.tool_calls or chunk.choices[0].finish_reason is not None + def _reserve_web_search_indexes(self, provider_fields: object) -> None: + if not isinstance(provider_fields, dict): + return + calls: Final = provider_fields.get("web_search_calls") + items: Final = calls.values() if isinstance(calls, dict) else calls if isinstance(calls, list) else () + for item in items: + try: + call_id = item.id.removeprefix("ws_") + status = item.status + except AttributeError: + call_id = str(item.get("id", "")).removeprefix("ws_") if isinstance(item, dict) else "" + status = item.get("status") if isinstance(item, dict) else None + if call_id: + output_index = self._get_or_assign_tool_output_index(call_id) + self._web_search_calls[call_id] = item + if status == "in_progress": + self._pending_tool_events = [ # mutable-ok: replaces speculative function events + event + for event in self._pending_tool_events + if getattr(event, "output_index", None) != output_index + ] + + def _tool_call_id(self, tool_call: object) -> str: + index: Final = self._normalize_tool_call_index(tool_call) + call_id_raw: Final = tool_call.get("id") if isinstance(tool_call, dict) else getattr(tool_call, "id", None) + if call_id_raw: + call_id: Final = str(call_id_raw) + if index is not None: + existing: Final = self._tool_call_id_by_index.get(index) + if existing is not None and existing != call_id: + self._ambiguous_tool_call_indexes.add(index) + self._tool_call_id_by_index[index] = call_id + return call_id + if index is None or index in self._ambiguous_tool_call_indexes: + return "" + return self._tool_call_id_by_index.get(index, "") + def _queue_tool_call_delta_events(self, tool_calls: object) -> None: """ Convert chat-completions streaming `tool_calls` deltas into Responses API streaming events. @@ -187,28 +229,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return for tc in tool_calls: - tc_index = self._normalize_tool_call_index(tc) - call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) - call_id = "" - - if call_id_raw: - call_id = str(call_id_raw) - if tc_index is not None: - existing_call_id = self._tool_call_id_by_index.get(tc_index) - if existing_call_id is not None and existing_call_id != call_id: - # Reusing the same index for multiple call_ids is ambiguous for id-less deltas. - # Guard against silent misrouting by disabling index fallback for this index. - self._ambiguous_tool_call_indexes.add(tc_index) - self._tool_call_id_by_index[tc_index] = call_id - elif tc_index is not None: - if tc_index in self._ambiguous_tool_call_indexes: - continue - mapped_call_id = self._tool_call_id_by_index.get(tc_index) - if mapped_call_id: - call_id = mapped_call_id - + call_id = self._tool_call_id(tc) if not call_id: continue + if call_id in self._web_search_calls: + continue fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) fn_name = "" @@ -220,7 +245,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_name = str(getattr(fn, "name", "") or "") fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) - output_index = self._get_or_assign_tool_output_index(call_id) if call_id not in self._tool_args_by_call_id: @@ -292,9 +316,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_name = str(getattr(fn, "name", "") or "") fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) + web_search_call = self._web_search_calls.get(call_id) + if web_search_call is not None: + if call_id not in self._queued_web_search_call_ids: + self._queue_web_search_events(call_id, web_search_call) + self._queued_web_search_call_ids.add(call_id) + continue # Track if this is a new tool call that wasn't streamed - is_new_tool_call = call_id not in self._tool_args_by_call_id + is_new_tool_call = call_id not in self._tool_item_id_by_call_id # If we never sent output_item.added for this call_id, emit it now. if is_new_tool_call: @@ -359,6 +389,49 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) self._pending_tool_events.append(item_done_event) + def _queue_web_search_events(self, call_id: str, web_search_call: object) -> None: + from openai.types.responses import ResponseFunctionWebSearch + + item: Final = ( + web_search_call + if isinstance(web_search_call, ResponseFunctionWebSearch) + else ResponseFunctionWebSearch.model_validate(web_search_call) + ) + output_index: Final = self._get_or_assign_tool_output_index(call_id) + self._sequence_number += 1 + added: Final = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=output_index, + item=BaseLiteLLMOpenAIResponseObject( + **{ # mutable-ok: BaseLiteLLM object accepts dynamic item fields + "id": item.id, + "type": item.type, + "status": "in_progress", + "action": None, + } + ), + ) + added.__dict__["sequence_number"] = self._sequence_number + self._pending_tool_events.append(added) + for event_type, event_class in ( + (ResponsesAPIStreamEvents.WEB_SEARCH_CALL_IN_PROGRESS, WebSearchCallInProgressEvent), + (ResponsesAPIStreamEvents.WEB_SEARCH_CALL_SEARCHING, WebSearchCallSearchingEvent), + (ResponsesAPIStreamEvents.WEB_SEARCH_CALL_COMPLETED, WebSearchCallCompletedEvent), + ): + self._sequence_number += 1 + event = event_class(type=event_type, output_index=output_index, item_id=item.id) + event.__dict__["sequence_number"] = self._sequence_number + self._pending_tool_events.append(event) + self._sequence_number += 1 + self._pending_tool_events.append( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + sequence_number=self._sequence_number, + item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), + ) + ) + def _adopt_response_id_from_chunk(self, chunk: ModelResponseStream) -> None: if self._cached_response_id is not None: return @@ -915,8 +988,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: chunk = cast(ModelResponseStream, chunk) - self._ensure_output_item_for_chunk(chunk) - # Accumulate provider_specific_fields from chunk and delta for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -927,6 +998,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ): if src and isinstance(src, dict): self._merge_provider_specific_fields(src) + self._reserve_web_search_indexes(src) + self._ensure_output_item_for_chunk(chunk) # Proceed to transformation self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(chunk) @@ -1021,8 +1094,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopIteration else: chunk = self.litellm_custom_stream_wrapper.__next__() - self._ensure_output_item_for_chunk(chunk) - # Accumulate provider_specific_fields from chunk and delta for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1033,6 +1104,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ): if src and isinstance(src, dict): self._merge_provider_specific_fields(src) + self._reserve_web_search_indexes(src) + self._ensure_output_item_for_chunk(chunk) # Always snapshot before returning any pending events so that # finish_reason (e.g. content_filter) is captured even when # _ensure_output_item_for_chunk queues events on the same chunk. @@ -1168,7 +1241,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "message", self._cached_item_id, ) - return _output_items_with_id(message_aligned, "reasoning", self._cached_reasoning_item_id) + reasoning_aligned: Final = _output_items_with_id( + message_aligned, + "reasoning", + self._cached_reasoning_item_id, + ) + return reasoning_aligned def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None: if litellm_model_response: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..cc594f167c7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -25,7 +25,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( from openai.types.chat.chat_completion_named_tool_choice_param import ( Function as NamedToolChoiceFunction, ) -from openai.types.responses import ResponseFunctionToolCall +from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam @@ -45,6 +45,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ) from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionAssistantMessage, ChatCompletionImageObject, ChatCompletionImageUrlObject, ChatCompletionRedactedThinkingBlock, @@ -635,6 +636,7 @@ class LiteLLMCompletionResponsesConfig: merged_assistant = LiteLLMCompletionResponsesConfig._merged_trailing_assistant_message( messages=messages, chat_completion_messages=chat_completion_messages, + hosted_search=_input.get("type") == "web_search_call", ) if merged_assistant is not None: messages[-1] = merged_assistant @@ -807,29 +809,44 @@ class LiteLLMCompletionResponsesConfig: chat_completion_messages: Sequence[ AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage ], - ) -> ChatCompletionResponseMessage | None: - """Fold an assistant content message into a directly preceding assistant - tool_calls message. Providers like DeepSeek and Anthropic require tool - results immediately after the tool_calls message, so an assistant message - between them is rejected.""" + hosted_search: bool = False, + ) -> ChatCompletionAssistantMessage | None: + """Keep replayed search context on the assistant turn so client tool results + still immediately follow the assistant that requested them.""" if not messages or len(chat_completion_messages) != 1: return None - last_message = messages[-1] - new_message = chat_completion_messages[0] - if not isinstance(last_message, dict): + if not isinstance(messages[-1], dict): return None + last_message: Final = _STR_KEY_DICT_ADAPTER.validate_python(messages[-1]) + new_message: Final = _STR_KEY_DICT_ADAPTER.validate_python(chat_completion_messages[0]) if last_message.get("role") != "assistant" or new_message.get("role") != "assistant": return None - if not last_message.get("tool_calls") or last_message.get("content") or new_message.get("tool_calls"): + if not (last_message.get("tool_calls") or hosted_search) or new_message.get("tool_calls"): return None - new_content = new_message.get("content") + new_content: Final = new_message.get("content") if new_content is None: return None + previous_content: Final = last_message.get("content") + content: Final = ( + new_content + if not previous_content + else [ # mutable-ok: outbound chat content uses JSON arrays + block + for value in (previous_content, new_content) + for block in ( + (ChatCompletionTextObject(type="text", text=value),) + if isinstance(value, str) + else _OBJECT_LIST_ADAPTER.validate_python(value) + ) + ] + ) merged: Final = { # mutable-ok: json.dumps rejects MappingProxyType in outbound chat messages **last_message, - "content": new_content, + "content": content, } - return cast(ChatCompletionResponseMessage, merged) # cast-ok: TypedDict spread widens to dict[str, object] + return cast( # cast-ok: preserves the assistant fields and content blocks + ChatCompletionAssistantMessage, merged + ) @staticmethod def _deduplicate_tool_call_output_messages( @@ -1252,6 +1269,14 @@ class LiteLLMCompletionResponsesConfig: - ResponseReasoningItemParam - ItemReference """ + if input_item.get("type") == "web_search_call": + search: Final = ResponseFunctionWebSearch.model_validate(input_item) + return [ # mutable-ok: input conversion returns chat message lists + GenericChatCompletionMessage( + role="assistant", + content="Hosted web search: " + search.model_dump_json(exclude_none=True), + ) + ] if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(input_item): # handle executed tool call results return ( @@ -1438,7 +1463,6 @@ class LiteLLMCompletionResponsesConfig: return input_item.get("type") in [ "function_call_output", "custom_tool_call_output", - "web_search_call", "computer_call_output", "tool_result", # Anthropic/MCP format ] @@ -2041,7 +2065,7 @@ class LiteLLMCompletionResponsesConfig: def transform_chat_completion_tools_to_responses_tools( chat_completion_response: ModelResponse, responses_api_request: ResponsesAPIOptionalRequestParams | None = None, - ) -> list[ResponseFunctionToolCall | CustomToolCallOutputItem]: + ) -> list[ResponseFunctionToolCall | ResponseFunctionWebSearch | CustomToolCallOutputItem]: """ Transform a Chat Completion tools into a Responses API tools. @@ -2064,7 +2088,12 @@ class LiteLLMCompletionResponsesConfig: custom_tool_names: Final = extract_custom_tool_names(request_tools) namespace_tool_names: Final = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(request_tools) - responses_tools: Final[list[ResponseFunctionToolCall | CustomToolCallOutputItem]] = [] + web_search_calls: Final = LiteLLMCompletionResponsesConfig._web_search_calls_by_call_id( + chat_completion_response + ) + responses_tools: Final[ + list[ResponseFunctionToolCall | ResponseFunctionWebSearch | CustomToolCallOutputItem] + ] = [] # mutable-ok: preserves provider tool-call order for tool in all_chat_completion_tools: if tool.type == "function": function_definition = tool.function @@ -2072,8 +2101,10 @@ class LiteLLMCompletionResponsesConfig: tool_id = tool.id or "" tool_arguments = serialize_tool_call_arguments(function_definition.get("arguments")) - # Check if this is a custom tool - if is_custom_tool_call(tool_name, custom_tool_names): + web_search_call = web_search_calls.get(tool_id) + if web_search_call is not None: + responses_tools.append(web_search_call) + elif is_custom_tool_call(tool_name, custom_tool_names): # Build custom_tool_call output item input_str = unwrap_custom_tool_arguments(tool_arguments) custom_item = CustomToolCallOutputItem( @@ -2128,6 +2159,35 @@ class LiteLLMCompletionResponsesConfig: responses_tools.append(output_tool_call) return responses_tools + @staticmethod + def _web_search_calls_by_call_id( + chat_completion_response: ModelResponse, + ) -> Mapping[str, ResponseFunctionWebSearch]: + calls: Final[dict[str, ResponseFunctionWebSearch]] = {} # mutable-ok: indexes provider-built calls + for choice in chat_completion_response.choices: + provider_fields = getattr(choice.message, "provider_specific_fields", None) + if not isinstance(provider_fields, Mapping): + continue + web_search_calls = provider_fields.get("web_search_calls") + items = ( + web_search_calls.values() + if isinstance(web_search_calls, Mapping) + else web_search_calls + if isinstance(web_search_calls, Sequence) + else () + ) + for item in items: + try: + call = ( + item + if isinstance(item, ResponseFunctionWebSearch) + else ResponseFunctionWebSearch.model_validate(item) + ) + except (TypeError, ValueError): + continue + calls[call.id.removeprefix("ws_")] = call + return MappingProxyType(calls) + @staticmethod def _map_chat_completion_finish_reason_to_responses_status( finish_reason: str | None, @@ -2326,6 +2386,7 @@ class LiteLLMCompletionResponsesConfig: | OutputFunctionToolCall | OutputImageGenerationCall | ResponseFunctionToolCall + | ResponseFunctionWebSearch | CustomToolCallOutputItem ]: responses_output: list[ @@ -2334,6 +2395,7 @@ class LiteLLMCompletionResponsesConfig: | OutputFunctionToolCall | OutputImageGenerationCall | ResponseFunctionToolCall + | ResponseFunctionWebSearch | CustomToolCallOutputItem ] = [] diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 5582027bb5d..ff2e389a6bb 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -42,6 +42,17 @@ def rust_enabled() -> bool: ) +def rust_ocr_enabled() -> bool: + environment: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) + if environment is False: + return False + return resolve_rust_enabled( + process_override=_CONFIGURATION.override, + environment_override=environment, + release_default=True, + ) + + def reset_rust_configuration() -> None: _CONFIGURATION.override = None diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py new file mode 100644 index 00000000000..f5e0c1b0fc6 --- /dev/null +++ b/litellm/rust_bridge/lifecycle.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import datetime +import os +import uuid +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + + +@dataclass(frozen=True, slots=True) +class Await: + awaitable: Awaitable[object] + + +@dataclass(frozen=True, slots=True) +class Complete: + value: object + + +class Execution(Protocol): + def start(self) -> Await | Complete: ... + + def resume_value(self, value: object) -> Await | Complete: ... + + def resume_error(self, error: BaseException) -> Await | Complete: ... + + def close(self) -> None: ... + + +async def drive(execution: Execution) -> object: + try: + step = execution.start() # rebind-ok: the execution protocol advances after each selected await + while isinstance(step, Await): + try: + value = await step.awaitable # rebind-ok: each selected await produces the next protocol input + except GeneratorExit: + raise + except BaseException as error: + step = execution.resume_error(error) # rebind-ok: advance the execution protocol + else: + step = execution.resume_value(value) # rebind-ok: advance the execution protocol + return step.value + finally: + execution.close() + + +class MetadataUpdater(Protocol): + def __call__( + self, + result: object, + logging_obj: Logging, + model: str | None, + kwargs: dict[str, object], + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: ... + + +@dataclass(frozen=True, slots=True) +class CallSetup: + logger: Logging + kwargs: dict[str, object] + + +def setup( + call_type: str, + args: tuple[object, ...], + kwargs: Mapping[str, object], + start_time: datetime.datetime, + asynchronous: bool, +) -> CallSetup: + from litellm import utils + from litellm.litellm_core_utils.litellm_logging import Logging + + arguments: Final = { # mutable-ok: function_setup consumes an owned kwargs dict + "litellm_call_id": str(uuid.uuid4()), + **kwargs, + } + supplied: Final = arguments.get("litellm_logging_obj") + if isinstance(supplied, Logging): + supplied._native_callback_fast_path = False # pyright: ignore[reportPrivateUsage] # supplied loggers retain all dispatch contracts + return CallSetup(supplied, arguments) + logger, prepared = utils.function_setup( + call_type, utils.Rules(), start_time, *args, is_async_call=asynchronous, **arguments + ) + if type(logger) is Logging and call_type in ("ocr", "aocr"): + logger._native_callback_fast_path = True # pyright: ignore[reportPrivateUsage] # only bridge-created OCR loggers opt into callback elision + return CallSetup(logger, prepared) + + +def check_limits(kwargs: Mapping[str, object]) -> None: + import litellm + + current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor + if litellm.max_budget and current_cost > litellm.max_budget: + raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) + metadata: Final = kwargs.get("metadata") + if isinstance(metadata, Mapping): + typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata + Mapping[str, object], metadata + ) + previous: Final = typed_metadata.get("previous_models") + if ( + isinstance(previous, list) + and litellm.num_retries_per_request is not None + and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history + >= litellm.num_retries_per_request + ): + raise RuntimeError("Max retries per request hit!") + + +def finalize( + response: object, + logger: Logging, + kwargs: dict[str, object], + start_time: datetime.datetime, + end_time: datetime.datetime, +) -> None: + from litellm.litellm_core_utils.llm_response_utils import response_metadata + + model: Final = kwargs.get("model") + update: Final = cast( # cast-ok: legacy metadata function accepts concrete kwargs + MetadataUpdater, response_metadata.update_response_metadata + ) + update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) + + +def deployment_callbacks_needed() -> bool: + import litellm + from litellm.integrations.custom_logger import CustomLogger + + return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) + + +def callbacks_needed(logger: Logging, phase: str) -> bool: + import litellm + from litellm._logging import ( + _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging + ) + + if ( + _is_debugging_on() + or getattr(logger, "litellm_request_debug", False) + or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") + ): + return True + input_needed: Final = bool( + litellm.input_callback + or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + or logger.dynamic_input_callbacks + or callable(getattr(logger, "logger_fn", None)) + or logger.log_raw_request_response + or litellm.log_raw_request_response + ) + match phase: + case "input": + return input_needed + case "sync_success": + return bool(litellm.success_callback or logger.dynamic_success_callbacks) + case "sync_success_async": + return bool( + (litellm.success_callback or logger.dynamic_success_callbacks) + and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks + ) + case "async_success": + return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + case "sync_failure": + return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) + case "async_failure": + return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + case "payload": + return bool( + input_needed + or litellm.success_callback + or litellm.failure_callback + or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + or logger.dynamic_success_callbacks + or logger.dynamic_async_success_callbacks + or logger.dynamic_failure_callbacks + or logger.dynamic_async_failure_callbacks + ) + case _: + return True + + +def success_bookkeeping( + logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> None: + phase: Final = "async_success" if asynchronous else "sync_success" + if logger.should_run_logging(phase): + logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload + result=response, start_time=start, end_time=end, build_logging_payload=False + ) + logger.has_run_logging(phase) + + +def failure_bookkeeping( + logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> None: + phase: Final = "async_failure" if asynchronous else "sync_failure" + if logger.should_run_logging(phase): + logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload + error, "", start, end, build_logging_payload=False + ) + logger.has_run_logging(phase) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 89eab71ccba..de8a93dd8b1 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,44 +2,16 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Mapping from dataclasses import dataclass from types import MappingProxyType from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx -import litellm -from litellm.constants import request_timeout -from litellm.llms.azure_ai.ocr.common_utils import is_azure_cohere_parse_model from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse -from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds -from litellm.types.router import GenericLiteLLMParams -from litellm.utils import ProviderConfigManager - -_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) -_RUST_OCR_CONFIG_FIELDS: Final = frozenset( - { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_scope", - "azure_authority_host", - "azure_credential", - "azure_federated_token_file", - "vertex_credentials", - "vertex_ai_credentials", - "vertex_project", - "vertex_ai_project", - "vertex_location", - "vertex_ai_location", - } -) -_RUST_OCR_SECRET_FIELDS: Final = frozenset( - {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} -) @dataclass(frozen=True, slots=True) @@ -87,26 +59,6 @@ class RustAocr(Protocol): raise NotImplementedError -class _OCRLogging(Protocol): - def update_from_kwargs( - self, - *, - kwargs: dict[str, object], - model: str, - optional_params: dict[str, object], - litellm_params: dict[str, object], - custom_llm_provider: str | None, - ) -> None: ... - - def pre_call( - self, - *, - input: str, - api_key: str | None, - additional_args: dict[str, object], - ) -> None: ... - - def _as_ocr(value: object) -> RustOcr | None: return cast(RustOcr, value) if callable(value) else None @@ -127,204 +79,6 @@ def load_rust_aocr() -> RustAocr | None: return _AOCR.load() -def provider(request: LiteLLMOcrRequest) -> str | None: - if request.custom_llm_provider is not None: - return request.custom_llm_provider - prefix: Final = request.model.partition("/")[0] - if prefix in _RUST_OCR_PROVIDERS: - return prefix - if request.model.startswith("mistral-ocr"): - return "mistral" - return None - - -def supported(request: LiteLLMOcrRequest) -> bool: - request_provider: Final = provider(request) - if request_provider not in _RUST_OCR_PROVIDERS: - return False - if request_provider == "azure_ai": - return ( - not is_azure_cohere_parse_model(request.model) - and not callable(request.kwargs.get("azure_ad_token_provider")) - and request.kwargs.get("azure_username") is None - and request.kwargs.get("azure_password") is None - ) - return True - - -def _optional_params(request: LiteLLMOcrRequest, resolve_secret: Callable[[str], str | None]) -> Mapping[str, object]: - optional_params: Final = MappingProxyType( - { - name: value - for name, value in request.kwargs.items() - if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) - and name not in ("litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request") - } - ) - request_provider: Final = provider(request) - if request_provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: - return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) - if request_provider != "vertex_ai": - return optional_params - project: Final = ( - request.kwargs.get("vertex_project") - or request.kwargs.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - location: Final = ( - request.kwargs.get("vertex_location") - or request.kwargs.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - credentials: Final = ( - request.kwargs.get("vertex_credentials") - or request.kwargs.get("vertex_ai_credentials") - or resolve_secret("VERTEXAI_CREDENTIALS") - ) - vertex_params: Final = MappingProxyType( - { - name: value - for name, value in ( - ("vertex_project", project), - ("vertex_location", location), - ("vertex_credentials", credentials), - ) - if value is not None - } - ) - return MappingProxyType({**optional_params, **vertex_params}) - - -def _input_sources(request: LiteLLMOcrRequest, optional_params: Mapping[str, object]) -> Mapping[str, str]: - proxy_request_value: Final = request.kwargs.get("proxy_server_request") - if not isinstance(proxy_request_value, Mapping): - return MappingProxyType({}) - proxy_request: Final = cast( # cast-ok: runtime Mapping check narrows metadata with unknown key and value types - Mapping[object, object], proxy_request_value - ) - credential_fields_value: Final = proxy_request.get("credential_fields", ()) - credential_fields: Final = ( - frozenset(name for name in credential_fields_value if isinstance(name, str)) - if isinstance(credential_fields_value, (list, tuple, set, frozenset)) - else frozenset() - ) - request_fields_value: Final = proxy_request.get("body_fields") - request_fields: Sequence[object] - if isinstance(request_fields_value, Sequence) and not isinstance(request_fields_value, (str, bytes)): - request_fields = cast( # cast-ok: runtime Sequence check excludes scalar strings and bytes - Sequence[object], request_fields_value - ) - else: - body_value: Final = proxy_request.get("body") - request_fields = ( - tuple(cast(Mapping[object, object], body_value)) # cast-ok: runtime Mapping check establishes iterable keys - if isinstance(body_value, Mapping) - else () - ) - names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) - request_sources: Final = MappingProxyType( - {name: "request" for name in names if name in request_fields or name in credential_fields} - ) - if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: - return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) - return request_sources - - -def _marshal( - request: LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], - convert_file_document: Callable[[dict[str, object]], dict[str, str]], -) -> LiteLLMOcrRequest: - if not isinstance(request.document, dict): - raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") - document: Final = ( - convert_file_document(request.document) if request.document.get("type") == "file" else request.document - ) - request_provider: Final = provider(request) - api_key: Final = ( - request.api_key or resolve_secret("MISTRAL_API_KEY") if request_provider == "mistral" else request.api_key - ) - optional_params: Final = _optional_params(request, resolve_secret) - input_sources: Final = _input_sources(request, optional_params) - logged_optional_params: Final = MappingProxyType( - {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} - ) - logged_kwargs: Final = MappingProxyType( - { - name: "****" if name in _RUST_OCR_SECRET_FIELDS else value - for name, value in request.kwargs.items() - if name != "proxy_server_request" - } - ) - logging_obj: Final = cast( # cast-ok: client decorator injects the logging object through untyped kwargs - _OCRLogging, request.kwargs["litellm_logging_obj"] - ) - logging_obj.update_from_kwargs( - kwargs=dict(logged_kwargs), # mutable-ok: legacy logging mutates its kwargs copy - model=request.model, - optional_params=dict(logged_optional_params), # mutable-ok: legacy logging requires concrete dict params - litellm_params={ # mutable-ok: legacy logging requires a concrete params dict - "litellm_call_id": request.kwargs.get("litellm_call_id"), - "api_base": request.api_base, - }, - custom_llm_provider=request_provider, - ) - logging_obj.pre_call( - input="OCR document processing", - api_key=api_key, - additional_args={ # mutable-ok: pre_call mutates the additional_args dict - "complete_input_dict": { # mutable-ok: callbacks consume a JSON-serializable request dict - "model": request.model, - "document": document, - **logged_optional_params, - }, - "api_base": request.api_base or "", - "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict - }, - ) - return LiteLLMOcrRequest( - model=request.model, - document=document, - api_key=api_key, - api_base=request.api_base, - timeout=request.timeout if request.timeout is not None else request_timeout, - custom_llm_provider=request.custom_llm_provider, - extra_headers=request.extra_headers, - kwargs=optional_params, - input_sources=input_sources, - ) - - -def _map_error(error: Exception, request: LiteLLMOcrRequest) -> Exception: - exception_types: Final = native_exception_types() - if exception_types is None or not isinstance(error, exception_types[1]): - return error - request_provider: Final = provider(request) - if request_provider is None: - return error - provider_config: Final = ProviderConfigManager.get_provider_ocr_config( - model=request.model.removeprefix(f"{request_provider}/"), provider=litellm.LlmProviders(request_provider) - ) - if provider_config is None: - return error - error_args: Final = cast( # cast-ok: BaseException.args exposes Any while native errors carry scalar args - tuple[object, ...], error.args - ) - status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 - message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) - error_factory: Final = cast( # cast-ok: legacy provider error factories have untyped callable parameters - Callable[..., Exception], provider_config.get_error_class - ) - return error_factory( - error_message=message, - status_code=status or 500, - headers={}, # mutable-ok: provider error factories require a concrete headers dict - ) - - def _response(response: Mapping[str, object]) -> OCRResponse: provider_native_response: Final = response.get(PROVIDER_NATIVE_RESPONSE_KEY) normalized: Final = OCRResponse.model_validate( @@ -335,56 +89,6 @@ def _response(response: Mapping[str, object]) -> OCRResponse: return normalized -def run( - request: LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], - convert_file_document: Callable[[dict[str, object]], dict[str, str]], -) -> OCRResponse | None: - if load_rust_ocr() is None: - return None - marshalled: Final = _marshal(request, resolve_secret, convert_file_document) - try: - response: Final = ocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=marshalled.input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_error(error, request) from error - return _response(response) if response is not None else None - - -async def arun( - request: LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], - convert_file_document: Callable[[dict[str, object]], dict[str, str]], -) -> OCRResponse | None: - if load_rust_aocr() is None: - return None - marshalled: Final = _marshal(request, resolve_secret, convert_file_document) - try: - response: Final = await aocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=marshalled.input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_error(error, request) from error - return _response(response) if response is not None else None - - def ocr( *, model: str, diff --git a/litellm/rust_bridge/ocr_lifecycle.py b/litellm/rust_bridge/ocr_lifecycle.py new file mode 100644 index 00000000000..5ca584e1c11 --- /dev/null +++ b/litellm/rust_bridge/ocr_lifecycle.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Mapping, Sequence +from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.ocr import LiteLLMOcrRequest + + +class NativeOcrLifecycle(Protocol): + def __call__( + self, + request: LiteLLMOcrRequest, + args: Sequence[object], + kwargs: Mapping[str, object], + asynchronous: bool, + ) -> OCRResponse | Awaitable[OCRResponse]: ... + + +class ExceptionMapper(Protocol): + def __call__( + self, + *, + model: str, + custom_llm_provider: str | None, + original_exception: Exception, + completion_kwargs: dict[str, object], + extra_kwargs: dict[str, object], + ) -> Exception: ... + + +def _binding(value: object) -> NativeOcrLifecycle | None: + if not callable(value): + return None + return cast("NativeOcrLifecycle", value) # cast-ok: callable validated at the native binding boundary + + +NATIVE_OCR_LIFECYCLE: Final = NativeBinding("_ocr_lifecycle", validate=_binding) + + +def select(request: LiteLLMOcrRequest) -> NativeOcrLifecycle | None: + if request.kwargs.get("aocr"): + return None + return NATIVE_OCR_LIFECYCLE.load() + + +def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]: + return request.kwargs + + +def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception: + mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper + ExceptionMapper, litellm.exception_type + ) + try: + return mapper( + model=request.model.removeprefix(f"{request_provider}/"), + custom_llm_provider=request_provider, + original_exception=error, + completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs + extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs + ) + except Exception as public_error: + public_error.__context__ = error + return public_error diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 02dee40f2a3..69cb88bfa2f 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -137,6 +137,7 @@ class SupportedGuardrailIntegrations(Enum): COMPRESR = "compresr" STRAIKER = "straiker" ALICE = "alice" + CONDUCT = "conduct" class Role(Enum): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b7c4371f32f..6612227f532 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -61,6 +61,7 @@ from openai.types.responses.response_create_params import ( ToolParam, ) from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall +from openai.types.responses.response_function_web_search import ResponseFunctionWebSearch from pydantic import ( BaseModel, ConfigDict, @@ -1358,6 +1359,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | OutputFunctionToolCall | OutputImageGenerationCall | ResponseFunctionToolCall + | ResponseFunctionWebSearch | CustomToolCallOutputItem ] ) diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py new file mode 100644 index 00000000000..0e64b91645d --- /dev/null +++ b/litellm/types/proxy/carried_budget_state.py @@ -0,0 +1,64 @@ +"""Budget fields auth already resolved, carried on the request so success logging does no object lookups. + +Auth pins one snapshot per entity on ``UserAPIKeyAuth`` (request-scoped, never cached), the pre-call +setup writes them into the request metadata under the aliased ``user_api_key_*`` names, and a logger +reads them back with ``from_metadata``. ``None`` means this request never carried that entity +(unauthenticated route, custom auth, budget check skipped) and the logger keeps its own lookup. +""" + +from collections.abc import Mapping +from datetime import datetime +from types import MappingProxyType + +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from typing_extensions import Self + + +class _BudgetSnapshot(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True, extra="ignore") + + def metadata_entries(self) -> Mapping[str, object]: + return MappingProxyType(self.model_dump(by_alias=True, mode="json")) + + @classmethod + def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None: + try: + return cls.model_validate(metadata) + except ValidationError: + return None + + +class KeyBudgetSnapshot(_BudgetSnapshot): + """Read-only view of the ``user_api_key_budget_reset_at`` entry ``add_user_api_key_auth_to_request_metadata`` writes.""" + + budget_reset_at: datetime | None = Field( + validation_alias="user_api_key_budget_reset_at", serialization_alias="user_api_key_budget_reset_at" + ) + + +class TeamBudgetSnapshot(_BudgetSnapshot): + budget_reset_at: datetime | None = Field( + validation_alias="user_api_key_team_budget_reset_at", serialization_alias="user_api_key_team_budget_reset_at" + ) + max_budget: float | None = Field( + validation_alias="user_api_key_team_table_max_budget", serialization_alias="user_api_key_team_table_max_budget" + ) + + +class UserBudgetSnapshot(_BudgetSnapshot): + budget_reset_at: datetime | None = Field( + validation_alias="user_api_key_user_budget_reset_at", serialization_alias="user_api_key_user_budget_reset_at" + ) + max_budget: float | None = Field( + validation_alias="user_api_key_user_table_max_budget", serialization_alias="user_api_key_user_table_max_budget" + ) + user_alias: str | None = Field( + validation_alias="user_api_key_user_alias", serialization_alias="user_api_key_user_alias" + ) + + +class OrgBudgetSnapshot(_BudgetSnapshot): + spend: float = Field(validation_alias="user_api_key_org_spend", serialization_alias="user_api_key_org_spend") + max_budget: float | None = Field( + validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget" + ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/conduct.py b/litellm/types/proxy/guardrails/guardrail_hooks/conduct.py new file mode 100644 index 00000000000..fbff4363351 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/conduct.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class ConductGuardrailConfigModelOptionalParams(BaseModel): + workspace_id: str | None = Field( + default=None, + description="Conduct workspace id, sent as the X-Workspace-Id header. Env: CONDUCT_WORKSPACE_ID.", + ) + tool_name: str | None = Field( + default="llm_call", + description="Conduct tool name the prompt is evaluated under. Match the tool your rules target.", + ) + timeout: float | None = Field( + default=8.0, + gt=0.0, + description="Timeout in seconds for the Conduct check.", + ) + unreachable_fallback: Literal["fail_open", "fail_closed"] | None = Field( + default="fail_closed", + description="Behavior when Conduct is unreachable, times out, or rejects the token.", + ) + + +class ConductGuardrailConfigModel(GuardrailConfigModel[ConductGuardrailConfigModelOptionalParams]): + api_key: str = Field( + min_length=1, + description="Conduct agent token. Env: CONDUCT_AGENT_TOKEN.", + ) + api_base: str | None = Field( + default="https://api.conductai.ai", + description="Conduct API base URL. The MCP endpoint is derived as /mcp.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Conduct Guard" diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 00635a8e1ef..26cc5c4c6cc 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import Final, Literal, Optional, Union from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall +from openai.types.responses.response_function_web_search import ActionSearchSource, ResponseFunctionWebSearch from pydantic import PrivateAttr from typing_extensions import Any, TypedDict @@ -39,6 +41,36 @@ class OutputFunctionToolCall(BaseLiteLLMOpenAIResponseObject): phase: Phase = None +def build_web_search_call( + tool_id: str, + tool_input: object, + result: object, + status: Literal["in_progress", "searching", "completed", "failed"] | None = None, +) -> ResponseFunctionWebSearch: + query: Final = tool_input.get("query", "") if isinstance(tool_input, Mapping) else "" + content: Final = result.get("content") if isinstance(result, Mapping) else None + result_items: Final = content if isinstance(content, Sequence) and not isinstance(content, (str, bytes)) else () + sources: Final = [ # mutable-ok: official SDK expects a source list + ActionSearchSource(type="url", url=url) + for item in result_items + if isinstance(item, Mapping) + and item.get("type") == "web_search_result" + and isinstance((url := item.get("url")), str) + ] + failed: Final = isinstance(content, Mapping) and content.get("type") == "web_search_tool_result_error" + return ResponseFunctionWebSearch( + id=f"ws_{tool_id}", + type="web_search_call", + status=status or ("failed" if failed else "completed"), + action={ # mutable-ok: official SDK expects an action mapping + "type": "search", + "query": query if isinstance(query, str) else "", + "queries": [query] if isinstance(query, str) and query else [], # mutable-ok: SDK list field + "sources": sources, + }, + ) + + class OutputImageGenerationCall(BaseLiteLLMOpenAIResponseObject): """An image generation call output""" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index aadd9bd3028..7fa09951eae 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57695,6 +57695,23 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "input_cost_per_token": 2.2e-07, @@ -57745,6 +57762,23 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://fireworks.ai/models/deepseek-ai/deepseek-v4p1-flash", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "input_cost_per_token": 2.2e-07, diff --git a/pyproject.toml b/pyproject.toml index 29609ce5ca1..69ecf2a42a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,13 +72,14 @@ proxy = [ "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", + "tomlkit>=0.13.3,<1.0", "polars>=1.38.1,<2.0", "soundfile>=0.12.1,<1.0", "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "expression>=5.6.0,<6.0", ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy -# imports are all guarded, so it runs on the base SDK plus just these five, and +# imports are all guarded, so it runs on the base SDK plus these packages, and # none of the server runtime in `proxy` is pulled in. On Linux, # keyring reaches the Secret Service through secretstorage, which brings # cryptography with it. @@ -88,6 +89,7 @@ cli = [ "requests>=2.32.0,<3.0", "InquirerPy>=0.3.4,<1.0", "keyring>=25.6.0,<26.0", + "tomlkit>=0.13.3,<1.0", ] extra_proxy = [ "prisma>=0.11.0,<1.0", diff --git a/scripts/benchmark_ocr_callbacks.py b/scripts/benchmark_ocr_callbacks.py new file mode 100644 index 00000000000..5db182c3d0e --- /dev/null +++ b/scripts/benchmark_ocr_callbacks.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +"""Measure serial sync/async OCR latency through a loopback HTTP provider + +Run each callback mode in a fresh process against an installed release wheel: +python -I scripts/benchmark_ocr_callbacks.py --callbacks none --label before \ + --expected-transport rust --iterations 200 --warmup 20 --output before-none.json +Repeat with --callbacks noop and with the candidate wheel in a separate venv +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import importlib.metadata +import json +import statistics +import sys +import threading +import time +from collections.abc import Sequence +from dataclasses import asdict, dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final, cast + +SIZES: Final = ( + 1024, + 4 * 1024, + 16 * 1024, + 64 * 1024, + 256 * 1024, + 1024 * 1024, +) +MODEL: Final = "mistral/mistral-ocr-latest" +EXPECTED_MARKDOWN: Final = "mock remote OCR response" +RESPONSE: Final = json.dumps( + { + "pages": [{"index": 0, "markdown": EXPECTED_MARKDOWN, "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, + }, + separators=(",", ":"), +).encode() + + +class Server(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), Handler) + self.user_agents: set[str] = set() + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + server: Final = cast(Server, self.server) + server.user_agents.add(self.headers.get("User-Agent", "")) + length: Final = int(self.headers["Content-Length"]) + body: Final = self.rfile.read(length) + request: Final = json.loads(body) + if self.path != "/v1/ocr" or request.get("model") != "mistral-ocr-latest": + self.send_error(400) + return + document: Final = request.get("document", {}) + if not isinstance(document, dict) or not str(document.get("document_url", "")).startswith( + "data:application/pdf;base64," + ): + self.send_error(400) + return + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(RESPONSE))) + self.end_headers() + self.wfile.write(RESPONSE) + + def log_message(self, format: str, *args: object) -> None: + return + + +@dataclass(frozen=True, slots=True) +class Result: + label: str + mode: str + size: int + iterations: int + median_ms: float + mean_ms: float + p95_ms: float + requests_per_second: float + + +def document(size: int) -> dict[str, str]: + payload: Final = b"%PDF-1.4\n" + b"x" * max(0, size - 9) + encoded: Final = base64.b64encode(payload[:size]).decode("ascii") + return {"type": "document_url", "document_url": f"data:application/pdf;base64,{encoded}"} + + +def percentile(values: Sequence[float], quantile: float) -> float: + ordered: Final = sorted(values) + index: Final = min(len(ordered) - 1, round((len(ordered) - 1) * quantile)) + return ordered[index] + + +def verify(response: object) -> None: + pages: Final = getattr(response, "pages", ()) + if len(pages) != 1 or getattr(pages[0], "markdown", None) != EXPECTED_MARKDOWN: + raise RuntimeError(f"unexpected OCR response: {response!r}") + + +def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Result: + median: Final = statistics.median(samples) + return Result( + label=label, + mode=mode, + size=size, + iterations=len(samples), + median_ms=median * 1000, + mean_ms=statistics.fmean(samples) * 1000, + p95_ms=percentile(samples, 0.95) * 1000, + requests_per_second=1 / median, + ) + + +def sync_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]: + samples: list[float] = [] + for _ in range(count): + started: Final = time.perf_counter() + response: Final = litellm.ocr( + model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30 + ) + samples.append(time.perf_counter() - started) + verify(response) + return tuple(samples) + + +async def async_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]: + samples: list[float] = [] + for _ in range(count): + started: Final = time.perf_counter() + response: Final = await litellm.aocr( + model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30 + ) + samples.append(time.perf_counter() - started) + verify(response) + return tuple(samples) + + +async def main() -> int: + parser: Final = argparse.ArgumentParser(description="E2E OCR benchmark against a local remote-style HTTP server") + parser.add_argument("--callbacks", choices=("none", "noop"), required=True) + parser.add_argument("--label", required=True) + parser.add_argument("--expected-transport", choices=("python", "rust"), required=True) + parser.add_argument("--iterations", type=int, default=30) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--sizes", type=int, nargs="+", default=SIZES) + parser.add_argument("--output", type=Path, required=True) + args: Final = parser.parse_args() + + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class NoopCallback(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.pre_calls = 0 + self.sync_successes = 0 + self.async_successes = 0 + + def log_pre_api_call(self, model, messages, kwargs): + self.pre_calls += 1 + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.sync_successes += 1 + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.async_successes += 1 + + registry_names: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ) + if any(getattr(litellm, name) for name in registry_names): + raise RuntimeError("benchmark requires initially empty callback registrations") + callback: Final = NoopCallback() + if args.callbacks == "noop": + litellm.callbacks.append(callback) + + rust_toggle: Final = getattr(litellm, "rust", None) + if callable(rust_toggle): + rust_toggle(False) + package: Final = Path(litellm.__file__).resolve() + version: Final = importlib.metadata.version("litellm") + native_path: str | None = None + native_sha256: str | None = None + try: + from litellm.rust_bridge import _native + + native: Final = Path(_native.__file__).resolve() + native_path = str(native) + native_sha256 = hashlib.file_digest(native.open("rb"), "sha256").hexdigest() + except ImportError: + pass + + server: Final = Server() + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + results: list[Result] = [] + try: + for size in args.sizes: + request_document: Final = document(size) + sync_samples(litellm, url, request_document, args.warmup) + sync_result: Final = summarize( + args.label, "sync", size, sync_samples(litellm, url, request_document, args.iterations) + ) + results.append(sync_result) + await async_samples(litellm, url, request_document, args.warmup) + async_result: Final = summarize( + args.label, + "async", + size, + await async_samples(litellm, url, request_document, args.iterations), + ) + results.append(async_result) + sys.stdout.write(json.dumps(asdict(sync_result)) + "\n") + sys.stdout.write(json.dumps(asdict(async_result)) + "\n") + sys.stdout.flush() + finally: + server.shutdown() + server.server_close() + thread.join() + + from litellm.litellm_core_utils.litellm_logging import executor + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.to_thread(executor.shutdown, wait=True) + per_mode: Final = len(args.sizes) * (args.iterations + args.warmup) + if args.callbacks == "noop": + if (callback.pre_calls, callback.sync_successes, callback.async_successes) != ( + 2 * per_mode, + per_mode, + per_mode, + ): + raise RuntimeError(f"callback delivery mismatch: {vars(callback)}") + elif any(getattr(litellm, name) for name in registry_names): + raise RuntimeError("callback registrations appeared in the no-callback case") + await GLOBAL_LOGGING_WORKER.stop() + + user_agents: Final = tuple(sorted(server.user_agents)) + python_transport: Final = any( + value.startswith("python-httpx") or value.startswith("litellm/") for value in user_agents + ) + if (args.expected_transport == "python") != python_transport: + raise RuntimeError(f"unexpected transport for {args.label}: user_agents={user_agents}") + artifact: Final = { + "label": args.label, + "callbacks": args.callbacks, + "python": sys.executable, + "callback_counts": { + "pre": callback.pre_calls, + "sync_success": callback.sync_successes, + "async_success": callback.async_successes, + }, + "version": version, + "package": str(package), + "native": native_path, + "native_sha256": native_sha256, + "user_agents": user_agents, + "results": tuple(asdict(result) for result in results), + } + args.output.write_text(json.dumps(artifact, indent=2) + "\n") + sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package", "user_agents")}) + "\n") + sys.stdout.write(f"results={args.output}\n") + sys.stdout.flush() + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index e5e0a164a83..8c0ef5a8b15 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -37,6 +37,7 @@ longer signal it. ### Fixed +- **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state - **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected - **key**: Read now picks up `model_rpm_limit`, `model_tpm_limit`, `guardrails`, `tags`, `enforced_params`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type` and `prompts` from `info.metadata`, where the proxy actually stores them; previously they stayed empty in state, so a matching config showed a permanent phantom diff on them and out-of-band changes to them were never detected diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 018d01f75a8..39546d588df 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -3,6 +3,7 @@ package litellm import ( "context" "encoding/json" + "errors" "fmt" "log" @@ -321,19 +322,42 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ metadata, err := plannedKeyMetadata(c, d) if err != nil { - d.Partial(true) - return diag.FromErr(fmt.Errorf("error updating key: %s", err)) + return failedKeyUpdate(ctx, d, m, err) } key.Metadata = metadata if _, err := c.UpdateKey(key); err != nil { - d.Partial(true) - return diag.FromErr(fmt.Errorf("error updating key: %s", err)) + return failedKeyUpdate(ctx, d, m, err) } return resourceKeyRead(ctx, d, m) } +// Deleting a team cascade-deletes its keys, so an apply that moves a key onto a +// replacement team can find the key already gone, and recreating it is the only +// way forward. Confirming it is really gone keeps an unrelated 404 (a rejected +// project_id, say) a hard failure rather than silently orphaning a live key. +func failedKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{}, err error) diag.Diagnostics { + c := m.(*Client) + if d.HasChange("team_id") && keyIsGone(c, d.Id(), err) { + log.Printf("[WARN] Key %q no longer exists, most likely cascade-deleted with its previous team; recreating it under the new team_id", d.Id()) + return resourceKeyCreate(ctx, d, m) + } + d.Partial(true) + return diag.FromErr(fmt.Errorf("error updating key: %s", err)) +} + +func keyIsGone(c *Client, keyID string, err error) bool { + if errors.Is(err, errKeyGone) { + return true + } + if !isNotFound(err) { + return false + } + key, getErr := c.GetKey(keyID) + return getErr == nil && key == nil +} + func changedMap(d *schema.ResourceData, name string) map[string]interface{} { if !d.HasChange(name) { return nil @@ -341,6 +365,8 @@ func changedMap(d *schema.ResourceData, name string) map[string]interface{} { return d.Get(name).(map[string]interface{}) } +var errKeyGone = errors.New("no longer exists") + func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) { if !d.HasChange("metadata") { return nil, nil @@ -350,7 +376,7 @@ func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface return nil, err } if current == nil { - return nil, fmt.Errorf("key %s no longer exists", d.Id()) + return nil, fmt.Errorf("key %s %w", d.Id(), errKeyGone) } oldDeclared, newDeclared := d.GetChange("metadata") return mergeKeyMetadata(current.Metadata, oldDeclared.(map[string]interface{}), newDeclared.(map[string]interface{})), nil diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 66291eadcc5..fe708edd3d3 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -7,8 +7,10 @@ import ( "net/http" "net/http/httptest" "reflect" + "sync/atomic" "testing" + "github.com/hashicorp/terraform-plugin-sdk/v2/diag" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) @@ -686,3 +688,191 @@ func TestKeyUpdateOmitsUnchangedDuration(t *testing.T) { t.Errorf("update payload unexpectedly contains duration = %v", v) } } + +// newKeyUpdateResourceData builds a *schema.ResourceData reflecting a real +// state -> config diff for team_id (unlike schema.TestResourceDataRaw, which +// has no notion of prior state), so d.HasChange("team_id") behaves the way it +// does during a real Update call. +func newKeyUpdateResourceData(t *testing.T, id, oldTeamID, newTeamID string) *schema.ResourceData { + t.Helper() + state := &terraform.InstanceState{ID: id, Attributes: map[string]string{"team_id": oldTeamID}} + diff := &terraform.InstanceDiff{Attributes: map[string]*terraform.ResourceAttrDiff{ + "team_id": {Old: oldTeamID, New: newTeamID}, + }} + d, err := schema.InternalMap(resourceKey().Schema).Data(state, diff) + if err != nil { + t.Fatalf("building ResourceData returned error: %v", err) + } + return d +} + +// keyRecoveryProxy fakes the two responses the cascade-delete recovery path +// turns on: what POST /key/update returns, and whether GET /key/info still +// finds the key afterwards. +type keyRecoveryProxy struct { + updateStatus int + updateBody string + staleKeyGone bool + updateCalls int32 + generateCalls int32 +} + +const keyNotFoundBody = `{"error":{"message":"Key not found.","type":"not_found_error","param":"key","code":"404"}}` + +func (p *keyRecoveryProxy) handler() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/key/update": + atomic.AddInt32(&p.updateCalls, 1) + w.WriteHeader(p.updateStatus) + io.WriteString(w, p.updateBody) + case "/key/generate": + atomic.AddInt32(&p.generateCalls, 1) + io.WriteString(w, `{"key": "sk-new", "token_id": "new-token"}`) + case "/key/info": + requested := r.URL.Query().Get("key") + if p.staleKeyGone && requested != "new-token" { + w.WriteHeader(http.StatusNotFound) + io.WriteString(w, keyNotFoundBody) + return + } + json.NewEncoder(w).Encode(map[string]interface{}{ + "key": requested, + "info": map[string]interface{}{"team_id": "team-b"}, + }) + default: + http.NotFound(w, r) + } + } +} + +func runKeyUpdate(t *testing.T, p *keyRecoveryProxy, d *schema.ResourceData) diag.Diagnostics { + t.Helper() + srv := httptest.NewServer(p.handler()) + defer srv.Close() + return resourceKeyUpdate(context.Background(), d, NewClient(srv.URL, "test-key", true)) +} + +// Reassigning a key between two teams that both still exist is a plain +// in-place /key/update and must not be turned into a destroy/recreate. +func TestResourceKeyUpdateTeamReassignmentStaysInPlace(t *testing.T) { + proxy := &keyRecoveryProxy{updateStatus: http.StatusOK, updateBody: `{"key": "hash-1"}`} + d := newKeyUpdateResourceData(t, "hash-1", "team-a", "team-b") + + if diags := runKeyUpdate(t, proxy, d); diags.HasError() { + t.Fatalf("update returned error: %v", diags) + } + if got := atomic.LoadInt32(&proxy.generateCalls); got != 0 { + t.Errorf("a benign team reassignment must not recreate the key, got %d /key/generate calls", got) + } + if d.Id() != "hash-1" { + t.Errorf("Id = %q, want hash-1 unchanged", d.Id()) + } +} + +// The reported bug: the key was cascade-deleted along with its old team, so +// /key/update 404s and the apply must recover by recreating it. +func TestResourceKeyUpdateRecreatesCascadeDeletedKey(t *testing.T) { + proxy := &keyRecoveryProxy{updateStatus: http.StatusNotFound, updateBody: keyNotFoundBody, staleKeyGone: true} + d := newKeyUpdateResourceData(t, "stale-token", "team-a", "team-b") + + if diags := runKeyUpdate(t, proxy, d); diags.HasError() { + t.Fatalf("a cascade-deleted key must be recreated, not error: %v", diags) + } + if got := atomic.LoadInt32(&proxy.updateCalls); got != 1 { + t.Errorf("expected 1 /key/update attempt before recovering, got %d", got) + } + if got := atomic.LoadInt32(&proxy.generateCalls); got != 1 { + t.Errorf("expected exactly 1 /key/generate recreate, got %d", got) + } + if d.Id() != "new-token" { + t.Errorf("Id = %q, want the recreated key's new-token", d.Id()) + } +} + +// /key/update 404s for reasons other than a missing key, a rejected +// project_id among them. Recovering on the status code alone would orphan a +// key that is still live on the proxy, so the key's absence must be confirmed. +func TestResourceKeyUpdateNotFoundWithLiveKeyFailsLoudly(t *testing.T) { + proxy := &keyRecoveryProxy{ + updateStatus: http.StatusNotFound, + updateBody: `{"error":{"message":"Project not found, project_id=proj-1"}}`, + } + d := newKeyUpdateResourceData(t, "hash-1", "team-a", "team-b") + + if diags := runKeyUpdate(t, proxy, d); !diags.HasError() { + t.Fatal("a 404 on a key that still exists must stay an error") + } + if got := atomic.LoadInt32(&proxy.generateCalls); got != 0 { + t.Errorf("expected no recreate while the key is still live, got %d /key/generate calls", got) + } + if d.Id() != "hash-1" { + t.Errorf("Id = %q, want hash-1 untouched on a hard failure", d.Id()) + } +} + +// A key gone for some reason unrelated to a team move still fails loudly. +func TestResourceKeyUpdateNotFoundWithoutTeamChangeFailsLoudly(t *testing.T) { + proxy := &keyRecoveryProxy{updateStatus: http.StatusNotFound, updateBody: keyNotFoundBody, staleKeyGone: true} + d := newKeyUpdateResourceData(t, "gone-token", "team-a", "team-a") + + if diags := runKeyUpdate(t, proxy, d); !diags.HasError() { + t.Fatal("expected an error when team_id did not change") + } + if got := atomic.LoadInt32(&proxy.generateCalls); got != 0 { + t.Errorf("expected no recreate when team_id is unchanged, got %d /key/generate calls", got) + } + if d.Id() != "gone-token" { + t.Errorf("Id = %q, want gone-token untouched on a hard failure", d.Id()) + } +} + +// A transient failure must never be mistaken for a cascade-deleted key. +func TestResourceKeyUpdateServerErrorDoesNotRecreate(t *testing.T) { + proxy := &keyRecoveryProxy{ + updateStatus: http.StatusInternalServerError, + updateBody: `{"error":{"message":"Internal Server Error"}}`, + staleKeyGone: true, + } + d := newKeyUpdateResourceData(t, "hash-1", "team-a", "team-b") + + if diags := runKeyUpdate(t, proxy, d); !diags.HasError() { + t.Fatal("expected a 500 to surface as an error") + } + if got := atomic.LoadInt32(&proxy.generateCalls); got != 0 { + t.Errorf("expected no recreate for a transient error, got %d /key/generate calls", got) + } +} + +// The metadata pre-read fails before /key/update is ever reached when the key +// is gone, so that path needs the same recovery. +func TestResourceKeyUpdateRecreatesCascadeDeletedKeyWithMetadataChange(t *testing.T) { + proxy := &keyRecoveryProxy{updateStatus: http.StatusOK, updateBody: `{"key": "hash-1"}`, staleKeyGone: true} + state := &terraform.InstanceState{ID: "stale-token", Attributes: map[string]string{ + "team_id": "team-a", + "metadata.%": "1", + "metadata.tier": "gold", + }} + diff := &terraform.InstanceDiff{Attributes: map[string]*terraform.ResourceAttrDiff{ + "team_id": {Old: "team-a", New: "team-b"}, + "metadata.tier": {Old: "gold", New: "silver"}, + }} + d, err := schema.InternalMap(resourceKey().Schema).Data(state, diff) + if err != nil { + t.Fatalf("building ResourceData returned error: %v", err) + } + + if diags := runKeyUpdate(t, proxy, d); diags.HasError() { + t.Fatalf("a cascade-deleted key must be recreated, not error: %v", diags) + } + if got := atomic.LoadInt32(&proxy.updateCalls); got != 0 { + t.Errorf("expected the metadata pre-read to short-circuit /key/update, got %d calls", got) + } + if got := atomic.LoadInt32(&proxy.generateCalls); got != 1 { + t.Errorf("expected exactly 1 /key/generate recreate, got %d", got) + } + if d.Id() != "new-token" { + t.Errorf("Id = %q, want the recreated key's new-token", d.Id()) + } +} diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index d2b842364a6..588402e3996 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -144,6 +144,8 @@ def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them( "tests/e2e/pytest.ini", "tests/e2e/gateway/stage_mirror_ci_config.yml", ".github/e2e-stack/up.sh", + ".github/e2e-stack/start-idp.sh", + "tests/e2e/idp_realm.json", ".github/workflows/test-e2e-changed.yml", ), ) diff --git a/tests/code_coverage_tests/test_e2e_idp_stack.py b/tests/code_coverage_tests/test_e2e_idp_stack.py new file mode 100644 index 00000000000..93596af915b --- /dev/null +++ b/tests/code_coverage_tests/test_e2e_idp_stack.py @@ -0,0 +1,115 @@ +import json +import os +import subprocess +import sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from threading import Thread + +import pytest + +ROOT = Path(__file__).resolve().parents[2] +START_IDP = ROOT / ".github/e2e-stack/start-idp.sh" + + +def run_start(tmp_path: Path, *, platform: str = "Linux", failure: str = "", port: int = 8181, real_curl: bool = False): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + calls = tmp_path / "docker.jsonl" + programs = { + "docker": """import json, os, sys +with open(os.environ['DOCKER_LOG'], 'a') as out: + out.write(json.dumps(sys.argv[1:]) + '\\n') +if os.environ['FAILURE'] == 'schema' and 'psql' in sys.argv: + sys.exit(17) +if os.environ['FAILURE'] == 'launch' and '--name' in sys.argv: + sys.exit(18) +""", + "curl": "import os, sys; sys.exit(1 if os.environ['FAILURE'] == 'readiness' else 0)\n", + "uname": "import os; print(os.environ['PLATFORM'])\n", + } + if real_curl: + del programs["curl"] + for name, source in programs.items(): + program = bin_dir / name + program.write_text(f"#!{sys.executable}\n{source}") + program.chmod(0o755) + result = subprocess.run( + ["bash", str(START_IDP)], + env={ + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "DOCKER_LOG": str(calls), + "PLATFORM": platform, + "FAILURE": failure, + "DATABASE_HOST": "127.0.0.1", + "DATABASE_PORT": "5544", + "DATABASE_USER": "fixture_user", + "DATABASE_PASSWORD": "fixture_password", + "DATABASE_NAME": "fixture_db", + "E2E_KEYCLOAK_PORT": str(port), + "E2E_KEYCLOAK_STARTUP_TIMEOUT": "0", + }, + capture_output=True, + text=True, + timeout=10, + ) + return result, [json.loads(line) for line in calls.read_text().splitlines()] + + +@pytest.mark.parametrize("platform", ("Linux", "Darwin")) +def test_idp_uses_existing_database_and_imports_runner_realm(tmp_path: Path, platform: str) -> None: + result, calls = run_start(tmp_path, platform=platform) + + assert result.returncode == 0, result.stderr + schema, _, launch = calls + host = "127.0.0.1" if platform == "Linux" else "host.docker.internal" + assert schema[schema.index("-h") + 1] == host + assert schema[schema.index("-p") + 1] == "5544" + assert "ON_ERROR_STOP=1" in schema + assert "CREATE SCHEMA IF NOT EXISTS keycloak" in schema + assert f"KC_DB_URL_HOST={host}" in launch + assert "KC_DB_URL_PORT=5544" in launch + assert "KC_DB_SCHEMA=keycloak" in launch + assert "KC_DB_POOL_MAX_SIZE=10" in launch + assert f"{ROOT}/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" in launch + assert "KC_HTTP_PORT=8181" in launch + if platform == "Linux": + assert launch[launch.index("--network") + 1] == "host" + else: + assert launch[launch.index("-p") + 1] == "127.0.0.1:8181:8181" + assert "Keycloak realm is up" in result.stdout + + +@pytest.mark.parametrize(("failure", "code"), (("schema", 17), ("launch", 18), ("readiness", 1))) +def test_idp_failure_stops_stack_startup(tmp_path: Path, failure: str, code: int) -> None: + result, calls = run_start(tmp_path, failure=failure) + + assert result.returncode == code + assert "Keycloak realm is up" not in result.stdout + if failure == "schema": + assert len(calls) == 1, "do not replace an IdP when its database is unavailable" + + +def test_readiness_requires_the_imported_realm_on_the_configured_port(tmp_path: Path) -> None: + expected_path = "/realms/litellm-e2e/.well-known/openid-configuration" + observed_paths: list[str] = [] + + class Discovery(BaseHTTPRequestHandler): + def do_GET(self) -> None: + observed_paths.append(self.path) + self.send_response(200 if self.path == expected_path else 404) + self.end_headers() + + with ThreadingHTTPServer(("127.0.0.1", 0), Discovery) as server: + worker = Thread(target=server.serve_forever, daemon=True) + worker.start() + try: + result, _ = run_start(tmp_path, port=server.server_port, real_curl=True) + finally: + server.shutdown() + worker.join(timeout=5) + + assert result.returncode == 0, result.stderr + assert observed_paths == [expected_path] + assert "Keycloak realm is up" in result.stdout diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 8be2e7b4ce2..a58c13d6a1c 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,9 +17,9 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory regression test (`test_reliability_memory_e2e.py`: a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic -- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite +- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` @@ -163,7 +163,7 @@ reliability... behavior : fallback | retry | cooldown | timeout | routing | cache | circuit_breaker | perf variant : 5xx | context_window | content_policy | 429 | timeout simple_shuffle | usage_based | latency_based | cost_based | least_busy - latency | throughput | session_anomaly (perf only; SLO/threshold assertion, not binary) + latency | throughput | session_anomaly | memory (perf only; SLO/threshold assertion, not binary) assertion : routes_to_fallback | succeeds_within_retries | picks_under_tpm | returns_cached | trips_then_recovers | under_slo e.g. reliability.fallback.context_window.routes_to_fallback exercised_on=[chat_completions] diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index f1f7b17d86b..44564a51e26 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -27,14 +27,50 @@ The suites run against a live proxy, so bring one up first by running the litell 2. Bring up a Postgres and a Redis for the proxy to use. The repo-root `docker-compose.yml` already defines a Postgres on `5432`; a `docker run -p 6379:6379 redis:7` covers Redis. Point `DATABASE_URL` / `REDIS_HOST` / `REDIS_PORT` at whatever you run. Tests that read Redis directly default to the deployed shape (TLS + cluster mode) whenever `REDIS_HOST` is set, so for a local standalone Redis also set `REDIS_CLUSTER=false` and `REDIS_SSL=false` (plus `REDIS_PASSWORD` when your Redis requires auth) -3. Start the litellm proxy locally against your config and confirm it is live: +3. Start the identity provider the JWT API tests authenticate against, then the litellm proxy against your config, and confirm both are live. It is a real Keycloak, running the realm in `tests/e2e/idp_realm.json`, and the proxy trusts it because `JWT_PUBLIC_KEY_URL` points at that realm's JWKS. The proxy caches the JWKS for `public_key_ttl` (600s) and does not refetch on an unknown `kid`, so keep its data volume across restarts; restart the proxy if you deliberately replace that volume: ```bash - set -a && source .env && set +a + docker run -d --name litellm-e2e-idp -p 8480:8080 \ + -e KC_BOOTSTRAP_ADMIN_USERNAME=admin -e KC_BOOTSTRAP_ADMIN_PASSWORD=admin \ + -v "$PWD/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \ + -v litellm-e2e-idp-data:/opt/keycloak/data \ + quay.io/keycloak/keycloak:26.7.3 start-dev --import-realm + curl -fs --retry 30 --retry-delay 2 --retry-all-errors http://127.0.0.1:8480/realms/litellm-e2e/.well-known/openid-configuration + export JWT_ISSUER=http://127.0.0.1:8480/realms/litellm-e2e + export JWT_AUDIENCE=litellm-e2e + export JWT_PUBLIC_KEY_URL="$JWT_ISSUER/protocol/openid-connect/certs" litellm --config .yml --port 4000 curl -fs http://localhost:4000/health/liveliness ``` + The tests reach Keycloak at `E2E_KEYCLOAK_URL` (default `http://127.0.0.1:8480`) and provision their identities through its admin API, so they also need `E2E_KEYCLOAK_ADMIN_USER` and `E2E_KEYCLOAK_ADMIN_PASSWORD` (`admin` / `admin` for the throwaway container above; the deployed stacks take theirs from a secret). JWT auth is an enterprise feature, so the proxy needs `LITELLM_LICENSE` in its environment, and its config needs the JWT block below. `enable_jwt_auth` only routes bearer tokens with three dot-separated segments into the JWT path, so `sk-` virtual keys and the master key keep working for every other suite. `proxy_batch_write_at` is lowered so the JWT spend-attribution test sees its row well inside the poll deadline: + + ```yaml + general_settings: + proxy_batch_write_at: 5 + enable_jwt_auth: true + litellm_jwtauth: + user_id_jwt_field: sub + user_email_jwt_field: email + team_ids_jwt_field: groups + user_id_upsert: true + ``` + + Set `JWT_ISSUER` to the exact realm URL used by the test runner and `JWT_AUDIENCE=litellm-e2e`. The realm explicitly maps this audience, `sub`, `email`, and `groups`; the proxy fetches real signing keys from its JWKS endpoint. The rejection tests obtain signed tokens with a different audience or issuer and verify the corresponding rejection reason. The issuer test uses a different HTTP Host when requesting a token from the isolated, dynamically named test IdP. + + Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack. + + Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write. + + Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`: + + ```bash + E2E_KEYCLOAK_ADMIN_USER=admin E2E_KEYCLOAK_ADMIN_PASSWORD=admin \ + uv run pytest tests/e2e/other/test_jwt_auth_e2e.py tests/e2e/management/test_jwt_management_e2e.py --reruns 0 -v + ``` + + Buildkite runs this suite against a Keycloak deployed beside the ephemeral stack by project-releaser. It fetches the realm from the test-runner revision even when it reuses a gateway image from another commit. The GitHub Actions changed-test stack starts the same digest-pinned Keycloak through `.github/e2e-stack/start-idp.sh`, imports the checked-out realm, and exports the IdP URL and credentials in `stack.env`. Both runners configure issuer/audience validation and store the realm, keys and users in a separate schema in the stack's PostgreSQL, so replacing Keycloak preserves token validity. Both wait for realm discovery before running tests. Losing the whole ephemeral database invalidates the stack. Keycloak skips imports into an existing realm, so changes to the realm export require a fresh stack (or deliberately replacing the local data volume). A stack without it fails the JWT tests rather than skipping them + 4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): ```bash @@ -65,7 +101,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start. The Redis chaos test under `load/` needs a proxy it can pause the Redis of on the same host (`gateway/redis_chaos_ci_config.yml`), which `.github/workflows/test-e2e-redis-chaos.yml` boots, and which the Buildkite `e2e-redis-chaos` step in project-releaser runs co-located with Postgres and Valkey in one pod; it is deselected unless `E2E_REDIS_CHAOS` is set +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite and both JWT suites as canaries, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Keycloak, Jaeger, and TLS cluster-mode Valkey. Realm-only edits also trigger these canaries. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start. The Redis chaos test under `load/` needs a proxy it can pause the Redis of on the same host (`gateway/redis_chaos_ci_config.yml`), which `.github/workflows/test-e2e-redis-chaos.yml` boots, and which the Buildkite `e2e-redis-chaos` step in project-releaser runs co-located with Postgres and Valkey in one pod; it is deselected unless `E2E_REDIS_CHAOS` is set Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 700dc822d59..36569896125 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -17,14 +17,18 @@ import functools import os from collections.abc import Generator, Iterator from datetime import datetime, timezone +from typing import Final import pytest import requests -from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL +from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL, unique_marker from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup +from e2e_http import unwrap from fixture_mode import fixture_mode_collection_error, fixture_report_lines +from idp import Identity, Keycloak, keycloak_from_env from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager +from models import TeamNewBody, UserNewBody, UserNewResponse from provider_edge import replay_leftover_error from proxy_client import ProxyClient, build_proxy_client @@ -32,6 +36,33 @@ _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() +@pytest.fixture(scope="session") +def idp() -> Keycloak: + return keycloak_from_env() + + +@pytest.fixture +def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) -> Identity: + marker: Final = unique_marker() + identity: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer) + resources.defer(lambda: proxy.delete_user(identity.user_id)) + # Seed the canonical user before any JWT call populates the auth cache. + # Group claims grant team access; management membership is added by the test. + unwrap( + proxy.transport.post( + "/user/new", + headers=proxy.transport.master, + json=UserNewBody( + user_id=identity.user_id, user_email=f"{identity.username}@example.com", user_role="internal_user" + ), + response_type=UserNewResponse, + ) + ) + team_id: Final = proxy.create_team(TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=identity.group)) + resources.defer(lambda: proxy.delete_team(team_id)) + return identity + + def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line( "markers", diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index c8d7037d2fd..d1227fe7c0c 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,6 +76,10 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} + +- {id: mgmt.key.jwt.lifecycle, module: mgmt, tier: P0, surface: api, assertions: [lifecycle], source: "management_endpoints/key_management_endpoints.py", rationale: "An IdP-issued admin JWT creates, reads, updates, clears and deletes a key; omitted fields survive updates"} +- {id: mgmt.key.jwt.member_denied, module: mgmt, tier: P0, surface: api, assertions: [member_denied], source: "auth/handle_jwt.py", rationale: "A valid member JWT cannot update an admin-managed key and denial leaves it unchanged"} +- {id: mgmt.key.jwt.other_team_denied, module: mgmt, tier: P0, surface: api, assertions: [other_team_denied], source: "auth/handle_jwt.py", rationale: "A valid JWT for another existing team cannot read the key"} - {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} - {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} - {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 814ebae2e0b..9292e5f07db 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -9,9 +9,12 @@ - {id: other.auth.llm_chat.not_bearer_scheme_denied, module: other, tier: P0, area: auth, assertions: [not_bearer_scheme_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "NotBearer scheme on chat is 401/403"} - {id: other.auth.realtime.missing_header_denied, module: other, tier: P1, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §9.19 / LIT-4778", rationale: "Realtime client-secret and calls routes reject requests without Authorization"} - {id: other.config.responses.metadata_redis_ttl_bounded, module: other, tier: P0, area: config, assertions: [ttl_bounded], source: "responses + redis cache", rationale: "Responses store+metadata must not leave TTL-unbounded Redis entries (LIT-1201)"} -- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"} -- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"} -- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"} +- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:1217-1256 auth_jwt / user_api_key_auth.py:1365-1377", rationale: "An access token issued by the configured IdP whose groups claim names an existing team is accepted on /chat/completions"} +- {id: other.auth.jwt.spend_attributed_to_claims, module: other, tier: P0, area: auth, assertions: [spend_attributed_to_claims], source: "handle_jwt.py:2224 auth_builder / user_api_key_auth.py:1438-1474", rationale: "The spend log row for a JWT-authenticated call carries the team_id from the groups claim and the user_id from sub, which for a real IdP is an opaque uuid, not a virtual key's identity. Single-group claim only: the proxy picks the team from a set, so attribution over several groups is unordered"} +- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:1244-1250", rationale: "A token the IdP issued with a one-second lifespan is rejected 401 (Token Expired) once it lapses, even though its signature still verifies; leeway is 0"} +- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"} +- {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"} +- {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} @@ -48,3 +51,6 @@ - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} - {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} + +- {id: other.auth.jwt.wrong_issuer_denied, module: other, tier: P0, area: auth, assertions: [wrong_issuer_denied], source: "auth/handle_jwt.py", rationale: "A signed token with the correct audience and an unexpected issuer is rejected"} +- {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index cfbedf3d236..e95d27bab84 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -35,4 +35,5 @@ - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} +- {id: reliability.perf.memory.under_slo, module: reliability, tier: P1, behavior: perf, variant: memory, assertions: [under_slo], exercised_on: [chat_completions], source: grammar, rationale: "Proxy RSS and the stored request snapshot stay within fixed budgets across a few hundred failing requests with retries and fallbacks, the v1.100.0 retry-breadcrumb leak shape (MAT-335)"} - {id: reliability.perf.session_anomaly.under_slo, module: reliability, tier: P1, behavior: perf, variant: session_anomaly, assertions: [under_slo], exercised_on: [messages], source: grammar, rationale: "Weekly Claude Code-shaped multi-turn session load against real providers; ceilings on error rate, warm-turn cache read/write, p95 turn time, and gateway-recorded spend (LIT-4562)"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 7344a2b6cec..e15a0cd0f8f 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -160,6 +160,14 @@ ANOMALY_MAX_KEY_SPEND_USD = float( ANOMALY_SPEND_SETTLE_SECONDS = float( os.environ.get("E2E_ANOMALY_SPEND_SETTLE_SECONDS", "75") ) +MEMORY_REQUESTS_PER_PHASE = int(os.environ.get("E2E_MEMORY_REQUESTS_PER_PHASE", "300")) +MEMORY_RETRIES_PER_REQUEST = int(os.environ.get("E2E_MEMORY_RETRIES_PER_REQUEST", "2")) +MEMORY_TRANSCRIPT_TURNS = int(os.environ.get("E2E_MEMORY_TRANSCRIPT_TURNS", "40")) +MEMORY_CONCURRENCY = int(os.environ.get("E2E_MEMORY_CONCURRENCY", "4")) +MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES", "15")) +MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1")) +MEMORY_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_RSS_BUDGET_MB", "48")) +MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64")) def ws_base_url() -> str: diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 415c72bbb3c..ce069720c6e 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -130,6 +130,20 @@ class ProbeResult(BaseModel): return 200 <= self.status_code < 500 and self.status_code != 404 +class ExternalWrite(BaseModel): + """Outcome of a write to a non-proxy API (an identity provider's admin API) + that answers with a status and, on create, a Location header naming the new + resource rather than a JSON body.""" + + status_code: int + location: str = "" + body: str = "" + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 300 + + class StreamingResponse(BaseModel): """Raw outcome for calls whose body is provider-native or streamed: status, the x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging @@ -257,23 +271,23 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None: f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}" ) - def wire_body(json: BaseModel) -> dict[str, object]: if isinstance(json, PartialBody): return json.model_dump(by_alias=True, exclude_unset=True) return json.model_dump(by_alias=True, exclude_none=True) -def _headers(headers: BaseModel) -> dict[str, str]: - dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) +def _flat(model: BaseModel) -> dict[str, str]: + dumped: dict[str, object] = model.model_dump(by_alias=True, exclude_none=True) return {key: str(value) for key, value in dumped.items()} +def _headers(headers: BaseModel) -> dict[str, str]: + return _flat(headers) + + def _params(params: BaseModel | None) -> dict[str, str]: - if params is None: - return {} - dumped: dict[str, object] = params.model_dump(by_alias=True, exclude_none=True) - return {key: str(value) for key, value in dumped.items()} + return _flat(params) if params is not None else {} TRANSIENT_STATUSES: frozenset[int] = frozenset({529}) @@ -416,6 +430,62 @@ def get_external[R: BaseModel]( return classify(resp, response_type) +def post_form_external[R: BaseModel]( + url: str, + *, + form: BaseModel, + response_type: type[R], + headers: BaseModel | None = None, + timeout: float = 30.0, +) -> Result[R]: + """POST an absolute URL outside the proxy as `application/x-www-form-urlencoded`, + the encoding OAuth 2 token endpoints take. Like get_external: no proxy base url, + no proxy auth, and the same tagged-union classification as every other call.""" + try: + resp = requests.post( + url, + data=_flat(form), + headers=_headers(headers) if headers is not None else None, + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return classify(resp, response_type) + + +def post_json_external( + url: str, + *, + headers: BaseModel, + json: BaseModel, + timeout: float = 30.0, +) -> ExternalWrite: + """POST an absolute URL outside the proxy under its own bearer, for an API that + answers a create with a status and a Location header rather than a JSON body.""" + try: + resp = requests.post( + url, + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return ExternalWrite(status_code=-1, body=str(exc)) + return ExternalWrite( + status_code=resp.status_code, + location=resp.headers.get("Location", ""), + body=resp.text, + ) + + +def delete_external(url: str, *, headers: BaseModel, timeout: float = 30.0) -> ExternalWrite: + try: + resp = requests.delete(url, headers=_headers(headers), timeout=timeout) + except requests.RequestException as exc: + return ExternalWrite(status_code=-1, body=str(exc)) + return ExternalWrite(status_code=resp.status_code, body=resp.text) + + def delete[R: BaseModel]( url: URL, *, diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 229e8514dee..8c8e64443cb 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -1,4 +1,11 @@ general_settings: + proxy_batch_write_at: 5 + enable_jwt_auth: true + litellm_jwtauth: + user_id_jwt_field: sub + user_email_jwt_field: email + team_ids_jwt_field: groups + user_id_upsert: true proxy_config_reload_interval_seconds: 7 store_prompts_in_spend_logs: true database_connection_pool_limit: 10 diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py new file mode 100644 index 00000000000..6d2fc84eb27 --- /dev/null +++ b/tests/e2e/idp.py @@ -0,0 +1,228 @@ +"""Provision isolated identities and obtain signed tokens from the test Keycloak realm.""" + +from __future__ import annotations + +import os +import secrets +import warnings +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Final, Literal + +import pytest +from e2e_http import ( + AuthHeaders, + ExternalWrite, + NetworkError, + Result, + Success, + delete_external, + post_form_external, + post_json_external, +) +from pydantic import BaseModel, Field + +KEYCLOAK_URL_ENV: Final = "E2E_KEYCLOAK_URL" +KEYCLOAK_REALM_ENV: Final = "E2E_KEYCLOAK_REALM" +KEYCLOAK_ADMIN_USER_ENV: Final = "E2E_KEYCLOAK_ADMIN_USER" +KEYCLOAK_ADMIN_PASSWORD_ENV: Final = "E2E_KEYCLOAK_ADMIN_PASSWORD" + +DEFAULT_KEYCLOAK_URL: Final = "http://127.0.0.1:8480" +DEFAULT_REALM: Final = "litellm-e2e" +TESTS_CLIENT_ID: Final = "litellm-e2e-tests" +SHORT_LIVED_CLIENT_ID: Final = "litellm-e2e-shortlived" +ADMIN_CLIENT_ID: Final = "litellm-e2e-admin" +WRONG_AUDIENCE_CLIENT_ID: Final = "litellm-e2e-other-app" + +_START_HINT: Final = ( + "Start it with the `docker run ... quay.io/keycloak/keycloak` command in tests/e2e/CONTRIBUTING.md, " + f"and point {KEYCLOAK_URL_ENV} / {KEYCLOAK_ADMIN_USER_ENV} / {KEYCLOAK_ADMIN_PASSWORD_ENV} at it" +) + + +class TokenGrantForm(BaseModel): + """The direct-access (password) grant an OAuth 2 token endpoint takes, form encoded.""" + + grant_type: Literal["password"] = "password" + client_id: str + username: str + password: str + + +class TokenResponse(BaseModel): + access_token: str = Field(repr=False) + + +class TokenRequestHeaders(BaseModel): + host: str | None = None + + +class GroupCreateBody(BaseModel): + name: str + + +class PasswordCredential(BaseModel): + type: Literal["password"] = "password" + value: str + temporary: bool = False + + +class UserCreateBody(BaseModel): + """Keycloak's admin representation of a new user. `firstName` / `lastName` and + an empty `requiredActions` matter: a realm's default VERIFY_PROFILE action + otherwise leaves the account "not fully set up" and every grant fails.""" + + username: str + email: str + email_verified: bool = Field(default=True, alias="emailVerified") + first_name: str = Field(default="E2E", alias="firstName") + last_name: str = Field(default="Tester", alias="lastName") + enabled: bool = True + groups: tuple[str, ...] + credentials: tuple[PasswordCredential, ...] + required_actions: tuple[str, ...] = Field(default=(), alias="requiredActions") + + +def created_id(write: ExternalWrite, context: str) -> str: + """The new resource's id, which Keycloak returns only as the last segment of + the Location header on a 201.""" + if write.status_code != 201: + pytest.fail(f"Keycloak refused to create {context}: HTTP {write.status_code} {write.body[:300]}") + if not write.location or write.location.endswith("/"): + pytest.fail(f"Keycloak created {context} without a resource id in its Location header") + return write.location.rsplit("/", 1)[-1] + + +@dataclass(frozen=True, slots=True) +class Identity: + """One provisioned IdP user: the `sub` the proxy will see, the credential the + test signs in with, and the group whose name the litellm team carries.""" + + user_id: str + username: str + password: str = field(repr=False) + group: str + group_id: str + + +@dataclass(frozen=True, slots=True) +class Keycloak: + base_url: str + realm: str + admin_username: str + admin_password: str = field(repr=False) + + @property + def issuer(self) -> str: + return f"{self.base_url}/realms/{self.realm}" + + @property + def jwks_url(self) -> str: + return f"{self.issuer}/protocol/openid-connect/certs" + + def token_url(self, realm: str) -> str: + return f"{self.base_url}/realms/{realm}/protocol/openid-connect/token" + + def _admin_url(self, path: str) -> str: + return f"{self.base_url}/admin/realms/{self.realm}{path}" + + def _admin_headers(self) -> AuthHeaders: + """A fresh admin token per call: the master realm's tokens are short lived, + and a cached one would expire in the middle of a slow test.""" + form: Final = TokenGrantForm(client_id="admin-cli", username=self.admin_username, password=self.admin_password) + result: Final = post_form_external(self.token_url("master"), form=form, response_type=TokenResponse) + return AuthHeaders(authorization=f"Bearer {self._token(result, 'the Keycloak admin credential')}") + + def _token(self, result: Result[TokenResponse], context: str) -> str: + match result: + case Success(data=granted): + return granted.access_token + case NetworkError(message=message): + return pytest.fail(f"No live Keycloak at {self.base_url} for {context}: {message}. {_START_HINT}") + case _: + return pytest.fail(f"Keycloak refused {context}: {result}") + + def create_group(self, name: str) -> str: + return created_id( + post_json_external( + self._admin_url("/groups"), headers=self._admin_headers(), json=GroupCreateBody(name=name) + ), + f"group {name}", + ) + + def create_user(self, *, username: str, email: str, password: str, group: str) -> str: + return created_id( + post_json_external( + self._admin_url("/users"), + headers=self._admin_headers(), + json=UserCreateBody( + username=username, + email=email, + groups=(group,), + credentials=(PasswordCredential(value=password),), + ), + ), + f"user {username}", + ) + + def delete_user(self, user_id: str) -> None: + self._delete(f"/users/{user_id}") + + def delete_group(self, group_id: str) -> None: + self._delete(f"/groups/{group_id}") + + def _delete(self, path: str) -> None: + try: + headers: Final = self._admin_headers() + except pytest.fail.Exception as exc: + warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2) + return + result: Final = delete_external(self._admin_url(path), headers=headers) + if result.status_code not in (204, 404): + warnings.warn( + f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}", + RuntimeWarning, + stacklevel=2, + ) + + def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity: + """Create `group` and a user in it, credentialed with a password generated + for this test alone, and hand back the identity a token can be minted for.""" + group_id: Final = self.create_group(group) + defer(lambda: self.delete_group(group_id)) + username: Final = f"e2e-jwt-user-{marker}" + password: Final = secrets.token_urlsafe(24) + user_id: Final = self.create_user( + username=username, email=f"{username}@example.com", password=password, group=group + ) + defer(lambda: self.delete_user(user_id)) + return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id) + + def access_token( + self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None + ) -> str: + """Sign `identity` in through the direct-access grant and hand back the + access token Keycloak signed, exactly as it came off the wire.""" + result: Final = post_form_external( + self.token_url(self.realm), + form=TokenGrantForm(client_id=client_id, username=identity.username, password=identity.password), + response_type=TokenResponse, + headers=TokenRequestHeaders(host=issuer_host), + ) + return self._token(result, f"a token for {identity.username}") + + +def keycloak_from_env() -> Keycloak: + admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip() + admin_password: Final = os.environ.get(KEYCLOAK_ADMIN_PASSWORD_ENV, "").strip() + if not admin_username or not admin_password: + pytest.fail( + f"The JWT suite needs {KEYCLOAK_ADMIN_USER_ENV} and {KEYCLOAK_ADMIN_PASSWORD_ENV} to provision " + f"identities in its Keycloak realm, and neither may be empty. {_START_HINT}" + ) + return Keycloak( + base_url=os.environ.get(KEYCLOAK_URL_ENV, DEFAULT_KEYCLOAK_URL).rstrip("/"), + realm=os.environ.get(KEYCLOAK_REALM_ENV, "").strip() or DEFAULT_REALM, + admin_username=admin_username, + admin_password=admin_password, + ) diff --git a/tests/e2e/idp_realm.json b/tests/e2e/idp_realm.json new file mode 100644 index 00000000000..3b747e7a5dd --- /dev/null +++ b/tests/e2e/idp_realm.json @@ -0,0 +1,210 @@ +{ + "realm": "litellm-e2e", + "enabled": true, + "sslRequired": "none", + "registrationAllowed": false, + "accessTokenLifespan": 300, + "clients": [ + { + "clientId": "litellm-e2e-tests", + "enabled": true, + "publicClient": true, + "standardFlowEnabled": false, + "directAccessGrantsEnabled": true, + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": false, + "config": { + "claim.name": "groups", + "full.path": "false", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + }, + { + "name": "litellm-audience", + "protocol": "openid-connect", + "protocolMapper": "oidc-audience-mapper", + "consentRequired": false, + "config": { + "included.custom.audience": "litellm-e2e", + "access.token.claim": "true", + "id.token.claim": "false" + } + } + ], + "defaultClientScopes": [ + "email", + "basic" + ] + }, + { + "clientId": "litellm-e2e-shortlived", + "enabled": true, + "publicClient": true, + "standardFlowEnabled": false, + "directAccessGrantsEnabled": true, + "attributes": { + "access.token.lifespan": "1" + }, + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": false, + "config": { + "claim.name": "groups", + "full.path": "false", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + }, + { + "name": "litellm-audience", + "protocol": "openid-connect", + "protocolMapper": "oidc-audience-mapper", + "consentRequired": false, + "config": { + "included.custom.audience": "litellm-e2e", + "access.token.claim": "true", + "id.token.claim": "false" + } + } + ], + "defaultClientScopes": [ + "email", + "basic" + ] + }, + { + "clientId": "litellm-e2e-admin", + "enabled": true, + "publicClient": true, + "standardFlowEnabled": false, + "directAccessGrantsEnabled": true, + "defaultClientScopes": [ + "email", + "litellm_proxy_admin", + "basic" + ], + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": false, + "config": { + "claim.name": "groups", + "full.path": "false", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + }, + { + "name": "litellm-audience", + "protocol": "openid-connect", + "protocolMapper": "oidc-audience-mapper", + "consentRequired": false, + "config": { + "included.custom.audience": "litellm-e2e", + "access.token.claim": "true", + "id.token.claim": "false" + } + } + ] + }, + { + "clientId": "litellm-e2e-other-app", + "enabled": true, + "publicClient": true, + "standardFlowEnabled": false, + "directAccessGrantsEnabled": true, + "defaultClientScopes": [ + "email", + "basic" + ], + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": false, + "config": { + "claim.name": "groups", + "full.path": "false", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + }, + { + "name": "litellm-audience", + "protocol": "openid-connect", + "protocolMapper": "oidc-audience-mapper", + "consentRequired": false, + "config": { + "included.custom.audience": "litellm-e2e-other-app", + "access.token.claim": "true", + "id.token.claim": "false" + } + } + ] + } + ], + "clientScopes": [ + { + "name": "litellm_proxy_admin", + "protocol": "openid-connect", + "attributes": { + "include.in.token.scope": "true", + "display.on.consent.screen": "false" + } + }, + { + "name": "email", + "protocol": "openid-connect", + "attributes": { + "include.in.token.scope": "true" + }, + "protocolMappers": [ + { + "name": "email", + "protocol": "openid-connect", + "protocolMapper": "oidc-usermodel-property-mapper", + "config": { + "user.attribute": "email", + "claim.name": "email", + "jsonType.label": "String", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + } + ] + }, + { + "name": "basic", + "protocol": "openid-connect", + "attributes": { + "include.in.token.scope": "false" + }, + "protocolMappers": [ + { + "name": "sub", + "protocol": "openid-connect", + "protocolMapper": "oidc-sub-mapper", + "config": { + "access.token.claim": "true" + } + } + ] + } + ] +} diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 1ef0d89a8f9..e17b92a13ed 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -10,9 +10,7 @@ import time from dataclasses import dataclass import jwt - from e2e_config import MASTER_KEY -from proxy_client import ProxyClient from e2e_http import ( AuthHeaders, NetworkError, @@ -37,6 +35,8 @@ from models import ( KeyDeleteBody, KeyGenerateBody, KeyGenerateResponse, + KeyInfoParams, + KeyInfoResponse, KeyListParams, KeyListResponse, KeyRegenerateBody, @@ -81,6 +81,7 @@ from models import ( UserNewResponse, UserUpdateBody, ) +from proxy_client import ProxyClient MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" @@ -152,13 +153,21 @@ class ManagementClient: def update_key_models(self, key: str, models: list[str]) -> None: _ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models))) - def delete_key_strict(self, key: str) -> None: + def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]: + return self.proxy.transport.get( + "/key/info", + headers=self.proxy.transport.bearer(caller_key), + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + + def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> None: """Strict delete for the act phase of a test: a failed delete is a hard failure, unlike the warn-only ProxyClient.delete_key used at teardown.""" _ = unwrap( self.proxy.transport.post( "/key/delete", - headers=self.proxy.transport.master, + headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py new file mode 100644 index 00000000000..22306da8eb8 --- /dev/null +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -0,0 +1,91 @@ +"""Management writes and tenant isolation under credentials issued by Keycloak.""" + +from __future__ import annotations + +from typing import Final + +import pytest +from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import UnauthorizedError, UnknownApiError, unwrap +from idp import ADMIN_CLIENT_ID, Identity, Keycloak +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody + +pytestmark = pytest.mark.e2e + + +class TestJwtManagement: + @pytest.mark.covers("mgmt.key.jwt.lifecycle") + def test_admin_creates_reads_updates_clears_and_deletes_a_key( + self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager + ) -> None: + admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + alias: Final = f"e2e-jwt-key-{unique_marker()}" + created: Final = unwrap( + client.generate_key( + KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), + caller_key=admin, + ) + ) + resources.defer(lambda: client.proxy.delete_key(created.key)) + + original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + assert original.key_alias == alias and original.team_id == jwt_identity.group + assert original.models == [CHEAP_OPENAI_MODEL] + + updated_alias: Final = f"{alias}-updated" + unwrap( + client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin) + ) + updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + assert updated.key_alias == updated_alias and updated.rpm_limit == 120 + assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction" + + unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin)) + cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + assert cleared.models == [] and cleared.rpm_limit == 120 + + assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1 + client.delete_key_strict(created.key, caller_key=admin) + assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0 + + @pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied") + def test_member_cannot_write_and_another_team_cannot_read_the_key( + self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager + ) -> None: + admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + member: Final = idp.access_token(jwt_identity) + alias: Final = f"e2e-jwt-owned-{unique_marker()}" + created: Final = unwrap( + client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin) + ) + resources.defer(lambda: client.proxy.delete_key(created.key)) + + client.add_team_member(jwt_identity.group, jwt_identity.user_id) + assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias + + refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member) + assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}" + assert "does not have permissions for endpoint" in refused.body.lower(), ( + f"expected a permission denial: {refused}" + ) + assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias + + marker: Final = unique_marker() + outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer) + resources.defer(lambda: client.proxy.delete_user(outsider.user_id)) + client.create_user( + UserNewBody( + user_id=outsider.user_id, user_email=f"{outsider.username}@example.com", user_role="internal_user" + ) + ) + team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=marker, team_id=outsider.group)) + resources.defer(lambda: client.proxy.delete_team(team_id)) + client.add_team_member(outsider.group, outsider.user_id) + outsider_token: Final = idp.access_token(outsider) + hidden: Final = client.key_info_as(created.key, caller_key=outsider_token) + assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, ( + f"another team must not read this key: {hidden}" + ) + assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index 8b7d5f0eb6f..353b0f7cf09 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -12,8 +12,8 @@ asserting once. from __future__ import annotations import time -from collections.abc import Callable -from typing import Literal +from collections.abc import Callable, Iterator +from typing import Final, Literal import pytest @@ -21,8 +21,11 @@ from e2e_config import unique_marker from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient -from models import KeyDeleteBody, KeyGenerateBody, KeyUpdateBody -from pydantic import BaseModel +from models import ( + CLEAR, ChatResponse, KeyDeleteBody, KeyGenerateBody, KeyInfo, KeyUpdateBody, + LiteLLMParamsBody, OrgNewBody, TeamNewBody, +) +from pydantic import BaseModel, RootModel pytestmark = pytest.mark.e2e @@ -131,7 +134,93 @@ def _unblock(client: ManagementClient, key: str) -> None: ) +class ProjectIdentity(BaseModel): + project_id: str + + +class ProjectCreateBody(BaseModel): + team_id: str + project_alias: str + models: list[str] + + +class ProjectBlockBody(ProjectIdentity): + blocked: bool + + +class ProjectDeleteBody(BaseModel): + project_ids: list[str] + + +@pytest.fixture +def project_resources(client: ManagementClient) -> Iterator[ResourceManager]: + manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True) + yield manager + manager.teardown() + + class TestKeyManagementRoutes: + @pytest.mark.covers("mgmt.key.update.persists") + def test_project_detachment_preserves_key_scope_and_refreshes_auth( + self, client: ManagementClient, project_resources: ResourceManager + ) -> None: + resources: Final = project_resources + name: Final = f"e2e-detach-{unique_marker()}" + model_id: Final = client.proxy.create_model( + name, LiteLLMParamsBody(model="openai/synthetic-detachment", api_key="synthetic", mock_response="orbit") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + org_id: Final = client.create_org(OrgNewBody(organization_alias=name, models=[name])) + resources.defer(lambda: client.delete_org(org_id)) + team_id: Final = client.create_team(TeamNewBody(team_alias=name, organization_id=org_id, models=[name])) + resources.defer(lambda: client.delete_team(team_id)) + project: Final = unwrap(client.proxy.transport.post( + "/project/new", headers=client.proxy.transport.master, + json=ProjectCreateBody(team_id=team_id, project_alias=name, models=[name]), + response_type=ProjectIdentity, + )) + resources.defer(lambda: unwrap(client.proxy.transport.delete( + "/project/delete", headers=client.proxy.transport.master, + json=ProjectDeleteBody(project_ids=[project.project_id]), response_type=RootModel[list[ProjectIdentity]], + ))) + key: Final = _generate_key(client, resources, KeyGenerateBody( + key_alias=name, team_id=team_id, organization_id=org_id, project_id=project.project_id, + models=[name], max_budget=5, tpm_limit=12345, rpm_limit=97, + )) + initial: Final = client.chat_status(key, name, "project attached") + assert initial.ok, initial.body + _ = unwrap(client.update_key(KeyUpdateBody(key=key, key_alias=f"{name}-saved"))) + assert client.proxy.key_info(key).project_id == project.project_id + _ = unwrap(client.update_key(KeyUpdateBody(key=key, project_id=project.project_id))) + rejected: Final = client.proxy.transport.send( + "/key/update", headers=client.proxy.transport.master, + json=KeyUpdateBody(key=key, project_id=f"{name}-different"), + ) + assert rejected.status_code == 400 and "reassignment" in rejected.body + assert client.proxy.key_info(key).project_id == project.project_id + _ = unwrap(client.proxy.transport.post( + "/project/update", headers=client.proxy.transport.master, + json=ProjectBlockBody(project_id=project.project_id, blocked=True), response_type=NoBody, + )) + blocked: Final = client.chat_status(key, name, "project blocked") + assert not blocked.ok and "is blocked" in blocked.body + detached: Final = unwrap(client.proxy.transport.post( + "/key/update", headers=client.proxy.transport.master, + json=KeyUpdateBody(key=key, project_id=CLEAR), response_type=KeyInfo, + )) + assert detached.project_id is None + saved: Final = client.proxy.key_info(key) + assert (saved.project_id, saved.team_id, saved.organization_id) == (None, team_id, org_id) + assert (saved.models, saved.max_budget, saved.tpm_limit, saved.rpm_limit) == ([name], 5, 12345, 97) + allowed: Final = client.chat_status(key, name, "project detached") + assert allowed.ok, allowed.body + message: Final = ChatResponse.model_validate_json(allowed.body).choices[0].message + assert message is not None and message.content == "orbit" + _ = unwrap(client.update_key(KeyUpdateBody(key=key, project_id=CLEAR))) + assert client.proxy.key_info(key).project_id is None + denied: Final = client.chat_status(key, f"{name}-outside", "outside key scope") + assert denied.status_code in (401, 403), denied.body + @pytest.mark.covers("mgmt.key.info.persists") def test_info_reflects_the_fields_the_key_was_created_with( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2a728d0ef40..3cab0334dea 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -11,7 +11,16 @@ from datetime import datetime from typing import Final, Literal from e2e_http import PartialBody -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator +from pydantic import ( + AliasChoices, + BaseModel, + ConfigDict, + Field, + JsonValue, + RootModel, + model_serializer, + model_validator, +) # ---------- keys ---------- @@ -67,6 +76,7 @@ class KeyGenerateBody(BaseModel): budget_duration: str | None = None user_id: str | None = None team_id: str | None = None + project_id: str | None = None organization_id: str | None = None budget_id: str | None = None key_alias: str | None = None @@ -130,6 +140,8 @@ class KeyInfo(BaseModel): models: list[str] = [] tpm_limit: int | None = None rpm_limit: int | None = None + project_id: str | None = None + organization_id: str | None = None team_id: str | None = None blocked: bool | None = None spend: float | None = None @@ -695,6 +707,7 @@ class SpendLogRow(BaseModel): total_tokens: int | None = None request_tags: list[str] | None = None metadata: SpendLogMetadata | None = None + proxy_server_request: JsonValue = None class SpendLogs(RootModel[list[SpendLogRow]]): @@ -926,6 +939,7 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None + cooldown_time: float | None = None order: int | None = None @@ -1046,6 +1060,7 @@ class KeyUpdateBody(BaseModel): clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale.""" key: str + project_id: str | Cleared | None = None models: list[str] | None = None key_alias: str | None = None tpm_limit: int | None = None @@ -1261,6 +1276,19 @@ class TagListResponse(RootModel[list[TagListEntry]]): # ---------- health / lifecycle ---------- +class ProcessMemory(BaseModel): + ram_usage_mb: float | None = None + system_memory_percent: float | None = None + error: str | None = None + + +class MemorySummaryResponse(BaseModel): + worker_pid: int + hostname: str | None = None + status: str + memory: ProcessMemory + + class ReadinessResponse(BaseModel): """GET /health/readiness (public probe). The low-detail payload a load balancer sees: `status` plus the resolved DB state (`connected`, diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 1aa83ac42c7..4313bbe4068 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -1,11 +1,14 @@ """Client for the `other` holding-pen suite: the auth gate (master key vs an -invalid key on an admin route) and the process-lifecycle health probes -(liveness, public readiness, authenticated readiness diagnostics). +invalid key on an admin route), JWT auth against the suite's Keycloak realm +(idp.py), and the process-lifecycle health probes (liveness, public readiness, +authenticated readiness diagnostics). Holds the shared ProxyClient so `resources` / `scoped_key` still clean up, and adds only the routes these behaviors need. The health probes deliberately send no auth header (public routes), so they go through the transport with an empty -headers model rather than a bearer. +headers model rather than a bearer. JWT tests reach the identity provider +through `idp`, which provisions identities and mints tokens through Keycloak's +own endpoints, so no test ever holds a signing key. """ from __future__ import annotations @@ -13,6 +16,7 @@ from __future__ import annotations from dataclasses import dataclass from e2e_http import NoBody, ProbeResult, Result +from idp import Keycloak, keycloak_from_env from models import ( ReadinessDetailsResponse, ReadinessResponse, @@ -26,6 +30,11 @@ from proxy_client import ProxyClient class OtherClient: proxy: ProxyClient + @property + def idp(self) -> Keycloak: + """Resolved per use, so the suite's non-JWT tests never need the IdP env.""" + return keycloak_from_env() + def liveness(self) -> ProbeResult: """GET /health/liveliness. Unauthenticated; the probe returns status + raw body so the test can assert the worker reports itself alive.""" diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py new file mode 100644 index 00000000000..2bed40f4d69 --- /dev/null +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -0,0 +1,157 @@ +"""Real Keycloak tokens exercise verification, attribution and virtual-key coexistence.""" + +from __future__ import annotations + +import base64 +import time +from typing import Final + +import pytest +from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import UnauthorizedError, UnknownApiError, unwrap +from idp import SHORT_LIVED_CLIENT_ID, WRONG_AUDIENCE_CLIENT_ID, Identity +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, TeamNewBody +from other_client import OtherClient +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + + +class IssuedClaims(BaseModel): + """Read the IdP's signed payload only to check the test precondition.""" + + exp: int + sub: str + iss: str + aud: str | list[str] + + +def _claims(token: str) -> IssuedClaims: + payload: Final = token.split(".")[1] + return IssuedClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) + + +def _provision(client: OtherClient, resources: ResourceManager, *, marker: str) -> Identity: + """A Keycloak group and a user in it, torn down with the test. The group name + is what the token's `groups` claim carries, which is what the proxy resolves + as a litellm team id.""" + identity: Final = client.idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer) + resources.defer(lambda: client.proxy.delete_user(identity.user_id)) + return identity + + +@pytest.fixture +def identity(client: OtherClient, resources: ResourceManager) -> Identity: + """An IdP identity whose group is also a real litellm team, so anything the + proxy rejects is about the token and never about an unresolvable team.""" + marker: Final = unique_marker() + provisioned: Final = _provision(client, resources, marker=marker) + team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=provisioned.group)) + resources.defer(lambda: client.proxy.delete_team(team_id)) + return provisioned + + +def _ping() -> ChatBody: + return ChatBody( + model=CHEAP_OPENAI_MODEL, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=16, + ) + + +def _corrupt_signature(token: str) -> str: + header, payload, signature = token.split(".") + flipped: Final = "A" if signature[10] != "A" else "B" + return f"{header}.{payload}.{signature[:10]}{flipped}{signature[11:]}" + + +class TestJwtAuth: + @pytest.mark.covers("other.auth.jwt.valid_token_allows", "other.auth.jwt.spend_attributed_to_claims") + def test_valid_token_for_an_existing_team_is_accepted_and_attributed( + self, client: OtherClient, identity: Identity + ) -> None: + token: Final = client.idp.access_token(identity) + + assert _claims(token).sub == identity.user_id, "IdP must emit the provisioned user as sub" + response: Final = unwrap(client.proxy.chat(token, _ping())) + assert response.id is not None and response.choices, ( + f"chat under a valid JWT returned no completion: {response}" + ) + + rows: Final = client.proxy.poll_logs_for_request_id(response.id) + assert rows, f"no spend log row for request {response.id} within the poll deadline" + row: Final = rows[0] + assert row.team_id == identity.group, ( + f"spend row must carry the team from the JWT groups claim {identity.group!r}, got {row.team_id!r}" + ) + assert row.user == identity.user_id, ( + f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}" + ) + + @pytest.mark.covers("other.auth.jwt.invalid_signature_denied") + def test_tampered_signature_is_rejected(self, client: OtherClient, identity: Identity) -> None: + tampered: Final = _corrupt_signature(client.idp.access_token(identity)) + + result: Final = client.proxy.chat(tampered, _ping()) + assert isinstance(result, UnauthorizedError), ( + f"a JWT whose signature does not verify must be rejected with 401, got {result}" + ) + assert "signature verification failed" in result.body.lower(), ( + f"the 401 must come from signature verification, not another auth failure, got {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.jwt.expired_denied") + def test_expired_token_is_rejected(self, client: OtherClient, identity: Identity) -> None: + expiring: Final = client.idp.access_token(identity, client_id=SHORT_LIVED_CLIENT_ID) + delay: Final = _claims(expiring).exp - time.time() + 1 + assert delay <= 5, f"short-lived client expiry or IdP clock drifted: wait would be {delay}s" + time.sleep(max(0, delay)) + + result: Final = client.proxy.chat(expiring, _ping()) + assert isinstance(result, UnauthorizedError), ( + f"an expired JWT must be rejected with 401 even though its signature verifies, got {result}" + ) + assert "expired" in result.body.lower(), f"the 401 must say the token expired, got {result.body[:300]}" + + @pytest.mark.covers("other.auth.jwt.wrong_issuer_denied") + def test_signed_token_from_the_wrong_issuer_is_rejected(self, client: OtherClient, identity: Identity) -> None: + token: Final = client.idp.access_token(identity, issuer_host="unexpected-issuer.invalid") + claims: Final = _claims(token) + assert claims.iss != client.idp.issuer and "litellm-e2e" in claims.aud + + result: Final = client.proxy.chat(token, _ping()) + assert isinstance(result, UnauthorizedError), f"wrong issuer must be rejected: {result}" + assert "issuer" in result.body.lower(), f"expected issuer validation to reject the token: {result}" + + @pytest.mark.covers("other.auth.jwt.wrong_audience_denied") + def test_signed_token_for_another_application_is_rejected(self, client: OtherClient, identity: Identity) -> None: + token: Final = client.idp.access_token(identity, client_id=WRONG_AUDIENCE_CLIENT_ID) + claims: Final = _claims(token) + assert claims.iss == client.idp.issuer and "litellm-e2e" not in ( + [claims.aud] if isinstance(claims.aud, str) else claims.aud + ) + + result: Final = client.proxy.chat(token, _ping()) + assert isinstance(result, UnauthorizedError), f"wrong audience must be rejected: {result}" + assert "audience" in result.body.lower(), f"expected audience validation to reject the token: {result}" + + @pytest.mark.covers("other.auth.jwt.unknown_team_denied") + def test_token_naming_a_team_that_does_not_exist_is_rejected( + self, client: OtherClient, resources: ResourceManager + ) -> None: + stranger: Final = _provision(client, resources, marker=unique_marker()) + token: Final = client.idp.access_token(stranger) + + result: Final = client.proxy.chat(token, _ping()) + assert isinstance(result, UnknownApiError) and result.status_code == 403, ( + f"a valid JWT whose groups name no existing team must be rejected with 403, got {result}" + ) + assert stranger.group in result.body, ( + f"the 403 must name the team it could not resolve ({stranger.group}), got {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.jwt.virtual_key_unaffected") + def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None: + response: Final = unwrap(client.proxy.chat(scoped_key, _ping())) + assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}" diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index ceb695ffcd6..6c87c7ef7ac 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -45,18 +45,16 @@ import hashlib import re import threading from collections import deque -from collections.abc import Mapping, Sequence -from contextlib import closing +from collections.abc import Generator, Mapping, Sequence +from contextlib import closing, contextmanager from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import islice from pathlib import Path from types import MappingProxyType -from typing import Final, Generator, Literal, assert_never +from typing import Final, Literal, assert_never from urllib.parse import parse_qsl, urlsplit -from pydantic import JsonValue, TypeAdapter - from e2e_http import ( NetworkError, StreamChunk, @@ -94,6 +92,7 @@ from fixture_mode import ( current_test_key, parse_fixture_mode, ) +from pydantic import JsonValue, TypeAdapter EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { @@ -495,7 +494,29 @@ class ReplayEdge: source: ReplaySource -type EdgeBackend = RecordEdge | ReplayEdge +@dataclass(frozen=True, slots=True) +class LiveEdge: + pass + + +type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge + + +@dataclass(slots=True) +class ProviderRequestObservation: + marker: str + _count: int = field(default=0, init=False) + _lock: threading.Lock = field(default_factory=threading.Lock, init=False) + + def observe(self, body: bytes | None) -> None: + if body is not None and self.marker.encode() in body: + with self._lock: + self._count += 1 + + @property + def count(self) -> int: + with self._lock: + return self._count @dataclass(frozen=True, slots=True) @@ -721,6 +742,24 @@ def _handle_record( assert_never(head) +def _handle_live( + method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float +) -> EdgeOutcome: + forwarded: Final = { + name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS + } + head: Final = forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) + match head: + case NetworkError(message=message): + return _recorded_outcome(_network_error_response(message)) + case StreamHead() if _is_streamed(head.headers): + return EdgeStream(head.status_code, _filtered_response_headers(head.headers), head.steps) + case StreamHead(): + return _recorded_outcome(_drain_to_response(head)) + case _: + assert_never(head) + + def _handle_replay(source: ReplaySource, request: RecordedRequest) -> EdgeOutcome: try: interaction: Final = source.next_interaction(request) @@ -753,6 +792,10 @@ def handle_edge_request( method, split.path, split.query, body, _header_value(headers, "content-type") ) match backend: + case LiveEdge(): + return _handle_live( + method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout + ) case RecordEdge(): return _handle_record( backend, @@ -792,6 +835,8 @@ class _EdgeHandler(BaseHTTPRequestHandler): assert isinstance(edge_server, _EdgeHTTPServer) length: Final = int(self.headers.get("content-length") or "0") body: Final = self.rfile.read(length) if length else None + if edge_server.observation is not None: + edge_server.observation.observe(body) outcome: Final = handle_edge_request( edge_server.backend, edge_server.mounts, @@ -857,11 +902,13 @@ class _EdgeHTTPServer(ThreadingHTTPServer): backend: EdgeBackend, mounts: Mapping[str, str], forward_timeout: float, + observation: ProviderRequestObservation | None, ) -> None: super().__init__(bind, _EdgeHandler) self.backend: Final = backend self.mounts: Final = mounts self.forward_timeout: Final = forward_timeout + self.observation: Final = observation @dataclass(frozen=True, slots=True) @@ -890,13 +937,14 @@ def start_provider_edge( bind_host: str = "127.0.0.1", advertise_host: str | None = None, forward_timeout: float = 60.0, + observation: ProviderRequestObservation | None = None, ) -> RunningEdge: """Boot an edge server on an OS-assigned port in a daemon thread. ``advertise_host`` is what api_base URLs name (it differs from the bind host when the proxy runs in a container and reaches the host machine via a gateway address like host.docker.internal).""" server: Final = _EdgeHTTPServer( - (bind_host, 0), backend=backend, mounts=mounts, forward_timeout=forward_timeout + (bind_host, 0), backend=backend, mounts=mounts, forward_timeout=forward_timeout, observation=observation ) thread: Final = threading.Thread(target=server.serve_forever, name="e2e-provider-edge", daemon=True) thread.start() @@ -979,3 +1027,40 @@ def provider_edge_api_base( return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout).api_base(mount) case _: assert_never(mode) + + +def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend: + mode: Final = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode(value=value): + raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}") + case "live": + return LiveEdge() + case "record": + return RecordEdge(_shared_recorder(bundle_dir), threading.Lock()) + case "replay": + return ReplayEdge(_shared_replay_source(bundle_dir)) + case _: + assert_never(mode) + + +@contextmanager +def observed_provider_edge( + observation: ProviderRequestObservation, + *, + mode_raw: str, + bundle_dir: Path, + bind_host: str, + advertise_host: str, + forward_timeout: float = 60.0, + mounts: Mapping[str, str] = EDGE_MOUNTS, +) -> Generator[ProviderEdge, None, None]: + running: Final = start_provider_edge( + _observed_backend(mode_raw, bundle_dir), mounts=mounts, + bind_host=bind_host, advertise_host=advertise_host, + forward_timeout=forward_timeout, observation=observation, + ) + try: + yield running.edge + finally: + running.shutdown() diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1bac5116a9d..1fe2ec905ef 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -63,6 +63,7 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListParams, + MemorySummaryResponse, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -72,6 +73,11 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + TeamDeleteBody, + TeamNewBody, + TeamNewResponse, + UserDeleteBody, + UserDeleteResponse, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -467,6 +473,17 @@ class ProxyClient: ) ).info + def memory_summary_everywhere(self) -> Mapping[str, Result[MemorySummaryResponse]]: + return { + url: transport.get( + "/debug/memory/summary", + headers=transport.master, + params=NoBody(), + response_type=MemorySummaryResponse, + ) + for url, transport in self.replicas.items() + } + def read_back_everywhere[R: BaseModel]( self, path: str, @@ -794,6 +811,41 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_credential({credential_name!r}) failed: {result}", stacklevel=2) + def create_team(self, body: TeamNewBody) -> str: + return unwrap( + self.transport.post( + "/team/new", + headers=self.transport.master, + json=body, + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + result = self.transport.post( + "/team/delete", + headers=self.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_team({team_id!r}) failed: {result}", stacklevel=2) + + def delete_user(self, user_id: str) -> None: + """Best-effort teardown; a 404 is not a leak, since JWT tests defer this for + a user the proxy only upserts after a successful auth.""" + result = self.transport.post( + "/user/delete", + headers=self.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + match result: + case Success() | UnknownApiError(status_code=404): + return + case _: + warnings.warn(f"delete_user({user_id!r}) failed: {result}", stacklevel=2) + # ---- LLM calls ------------------------------------------------------ def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]: diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 1efcb1a045b..df2ff03aa4a 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -11,6 +11,8 @@ body, so a single long-lived proxy serves every reliability behavior. from __future__ import annotations +from collections.abc import Sequence + from pydantic import ValidationError from proxy_client import ProxyClient @@ -49,6 +51,13 @@ def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str: + return proxy.create_model( + name, + LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, api_base="http://127.0.0.1:9/v1", cooldown_time=0), + ) + + def create_timeout_deployment(proxy: ProxyClient, name: str) -> str: """Register a deployment with a 1ms deadline the real backend always exceeds.""" return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001)) @@ -109,6 +118,7 @@ def chat_override( override: RouterSettingsOverride | None = None, stream: bool = False, cache: dict[str, bool] | None = {"no-cache": True}, + history: Sequence[ChatMessage] = (), ) -> StreamingResponse: """POST /chat/completions with an optional per-request router_settings_override, returning the raw outcome so tests read status, body, and reliability headers.""" @@ -117,7 +127,7 @@ def chat_override( headers=proxy.transport.bearer(key), json=ReliabilityChatBody( model=model, - messages=[ChatMessage(role="user", content=content)], + messages=[*history, ChatMessage(role="user", content=content)], max_tokens=512, stream=stream, router_settings_override=override, diff --git a/tests/e2e/router/test_reliability_cache_e2e.py b/tests/e2e/router/test_reliability_cache_e2e.py index 4ea05a1ecca..f7a2f2ffeb7 100644 --- a/tests/e2e/router/test_reliability_cache_e2e.py +++ b/tests/e2e/router/test_reliability_cache_e2e.py @@ -1,37 +1,98 @@ -"""Live e2e: the response cache returns a cached answer on an exact repeat. +"""An exact cache hit preserves the full choices and usage without another provider call. -The same unique prompt is sent twice to the real `gpt-5.5` deployment under the -same key: the first call is a cache miss (the proxy computes and stores the entry, -and returns no x-litellm-cache-key), the second is an exact hit (the proxy serves -from cache and returns x-litellm-cache-key). This relies on the standard Redis -response cache being enabled on the proxy under test. +Response IDs, creation timestamps and proxy headers are transport metadata; +compare every field within choices and usage, including provider extensions. """ from __future__ import annotations +from typing import Final + import pytest - from complexity_router_client import ComplexityRouterClient -from e2e_config import unique_marker -from reliability_support import chat_override +from e2e_config import ( + FIXTURE_DIR, + FIXTURE_MODE_RAW, + PROVIDER_EDGE_ADVERTISE_HOST, + PROVIDER_EDGE_BIND_HOST, + REQUEST_TIMEOUT, + unique_marker, +) +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from provider_edge import ProviderRequestObservation, observed_provider_edge +from pydantic import BaseModel, JsonValue -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.replayable] + + +class _CacheChatBody(ChatBody): + ttl: int = 600 + + +class _CachedAnswer(BaseModel): + model: str + choices: tuple[dict[str, JsonValue], ...] + usage: dict[str, JsonValue] class TestReliabilityCache: @pytest.mark.covers("reliability.cache.exact.returns_cached") - def test_exact_cache_returns_cached(self, client: ComplexityRouterClient, scoped_key: str) -> None: - prompt = f"cache probe {unique_marker()}" + def test_exact_cache_returns_cached( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + marker: Final = unique_marker() + model: Final = f"e2e-cache-{marker}" + prompt: Final = f"Reply with a short sentence about a blue lantern. Request marker: {marker}" + observation: Final = ProviderRequestObservation(marker) - first = chat_override(client.proxy, scoped_key, "gpt-5.5", prompt, cache=None) - assert first.status_code == 200, f"first call should succeed, got {first.status_code}: {first.body[:300]}" - assert "x-litellm-cache-key" not in first.headers, ( - "first (uncached) call must not report a cache-key header" - ) + with observed_provider_edge( + observation, + mode_raw=FIXTURE_MODE_RAW, + bundle_dir=FIXTURE_DIR, + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + forward_timeout=REQUEST_TIMEOUT, + ) as edge: + model_id: Final = client.proxy.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-5.6", + api_key="os.environ/OPENAI_API_KEY", + api_base=f"{edge.api_base('openai')}/v1", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + body: Final = _CacheChatBody( + model=model, + messages=[ChatMessage(role="user", content=prompt)], + max_completion_tokens=512, + reasoning_effort="none", + cache=None, + ) + first: Final = client.proxy.transport.send( + "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=body + ) + assert first.status_code == 200, f"first call should succeed, got {first.status_code}: {first.body[:300]}" + assert "x-litellm-cache-key" not in first.headers, "first call must be a cache miss" + answer: Final = ChatResponse.model_validate_json(first.body) + assert len(answer.choices) == 1 + choice: Final = answer.choices[0] + assert choice.message is not None and choice.message.role == "assistant" + assert choice.message.content is not None and choice.message.content.strip(), "first answer is empty" + assert choice.finish_reason == "stop" + assert answer.usage is not None + assert answer.usage.prompt_tokens is not None and answer.usage.prompt_tokens > 0 + assert answer.usage.completion_tokens is not None and answer.usage.completion_tokens > 0 + assert answer.usage.total_tokens == answer.usage.prompt_tokens + answer.usage.completion_tokens + assert observation.count == 1, "first miss must invoke the provider exactly once" - second = chat_override(client.proxy, scoped_key, "gpt-5.5", prompt, cache=None) - assert second.status_code == 200, f"second call should succeed, got {second.status_code}: {second.body[:300]}" - assert "x-litellm-cache-key" in second.headers, ( - "second identical call should hit the response cache and report a cache-key header " - "(requires the proxy's Redis response cache to be enabled)" - ) + second: Final = client.proxy.transport.send( + "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=body + ) + assert second.status_code == 200, f"second call should succeed, got {second.status_code}: {second.body[:300]}" + assert second.headers.get("x-litellm-cache-key"), "identical request must hit the response cache" + assert _CachedAnswer.model_validate_json(second.body) == _CachedAnswer.model_validate_json(first.body), ( + "cache hit changed the answer, finish reason or usage" + ) + assert observation.count == 1, "two successful requests must invoke the provider exactly once" diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/router/test_reliability_memory_e2e.py new file mode 100644 index 00000000000..17e3e1a1996 --- /dev/null +++ b/tests/e2e/router/test_reliability_memory_e2e.py @@ -0,0 +1,263 @@ +"""Live e2e: a few hundred requests that fail before any provider answers must not +grow the proxy's resident memory past a fixed budget once the proxy is warm. + +The regression this guards shipped in v1.100.0: every retry breadcrumb copied the +whole request and the copies nested into one router-global list, so a proxy under +retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that +shape: a model group whose deployments refuse at the socket (an unreachable base +URL) with cooldown_time 0 so the router keeps retrying them, per-request retries, +and a fallback group that refuses the same way, each request carrying a long chat +transcript so every whole-request copy costs hundreds of containers instead of a +handful. Under the stack's cooldown policy a +deployment that fails a handful of times in a row is benched (a bad-credential 401 +included), the router answers "No deployments available" without retrying, and the +retry loop that leaks stops running; cooldown_time 0 keeps it running. + +Two identical phases run back to back. The first is the warmup that grows the +proxy's caches and allocator arenas to their steady state, the second is the one +the budget applies to, so a healthy proxy shows the second phase adding roughly +nothing while a leaking one adds a fixed amount per request. RSS is read through +/debug/memory/summary on every configured replica; a burst of failing calls leaves +a transient bulge of garbage that gc reclaims within seconds, so each checkpoint +samples until no new worker has answered for a settle window and keeps the lowest +reading per worker. The growth is judged per worker (by replica address, hostname +and pid, since pods in their own pid namespaces report the same pids) so each +worker is compared with itself, and the two checkpoints must see the same workers: +a single load-balanced address reaches the workers behind it one answer at a time, +and a worker that answered only one checkpoint would otherwise drop out of the +comparison, which is where a leaking worker could hide. + +RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, +json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS +by only about 15 MB per 300 failing requests, while every failing request's stored +request snapshot carried a copy of the request per failed attempt, over 100 KB on +the first call and a couple of MB once the copies nested, against tens of KB with +the fix. So the first check sends one failing request before the phases, reads its +spend log back through /spend/logs, and holds the stored request body to a fixed +size budget: the deterministic catch for a breadcrumb that copies the whole +request. It runs before the phases because the leaking writer drops its own rows +under the phases' traffic (a queue budget hit, a recursion limit on the nested +copies), which would turn the size check into a missing-row check. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import ( + MEMORY_CONCURRENCY, + MEMORY_REQUESTS_PER_PHASE, + MEMORY_RETRIES_PER_REQUEST, + MEMORY_RSS_BUDGET_MB, + MEMORY_RSS_SAMPLE_INTERVAL_SECONDS, + MEMORY_RSS_SETTLE_SAMPLES, + MEMORY_STORED_REQUEST_BUDGET_KB, + MEMORY_TRANSCRIPT_TURNS, + unique_marker, +) +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ChatMessage, RouterSettingsOverride, SpendLogRow +from proxy_client import ProxyClient +from reliability_support import chat_override, create_never_benched_refusing_deployment + +pytestmark = pytest.mark.e2e + +DEPLOYMENTS_PER_GROUP: Final = 2 +RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES + + +@dataclass(frozen=True, slots=True) +class FailedCall: + status_code: int + seconds: float + body_head: str + call_id: str | None + + +WorkerKey = tuple[str, str | None, int] + + +@dataclass(frozen=True, slots=True) +class RssReading: + replica: str + hostname: str | None + worker_pid: int + ram_usage_mb: float + + @property + def worker(self) -> WorkerKey: + return (self.replica, self.hostname, self.worker_pid) + + +@dataclass(frozen=True, slots=True) +class WorkerGrowth: + warm: RssReading + after: RssReading + + @property + def growth_mb(self) -> float: + return self.after.ram_usage_mb - self.warm.ram_usage_mb + + +def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None: + for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): + resources.defer(lambda model_id=model_id: proxy.delete_model(model_id)) + + +def _transcript(turns: int) -> tuple[ChatMessage, ...]: + return tuple( + ChatMessage(role=role, content=f"turn {turn} {role}") + for turn in range(turns) + for role in ("user", "assistant") + ) + + +TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS) + + +def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall: + started: Final = time.perf_counter() + resp: Final = chat_override( + proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT + ) + return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id) + + +def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]: + with ThreadPoolExecutor(max_workers=MEMORY_CONCURRENCY) as pool: + futures: Final = tuple( + pool.submit(_fail_once, proxy, key, model, override) for _ in range(MEMORY_REQUESTS_PER_PHASE) + ) + return tuple(future.result() for future in futures) + + +def _read_rss_everywhere_after_pause(proxy: ProxyClient) -> tuple[RssReading, ...]: + time.sleep(MEMORY_RSS_SAMPLE_INTERVAL_SECONDS) + return tuple( + RssReading(replica, body.hostname, body.worker_pid, body.memory.ram_usage_mb) + for replica, result in proxy.memory_summary_everywhere().items() + for body in (unwrap(result),) + if body.memory.ram_usage_mb is not None + ) + + +def _readings_until_no_new_worker( + proxy: ProxyClient, readings: tuple[RssReading, ...], samples: int, samples_since_new_worker: int +) -> tuple[RssReading, ...]: + if samples >= RSS_SAMPLE_CAP or samples_since_new_worker >= MEMORY_RSS_SETTLE_SAMPLES: + return readings + sample: Final = _read_rss_everywhere_after_pause(proxy) + known: Final = frozenset(reading.worker for reading in readings) + new_worker_answered: Final = any(reading.worker not in known for reading in sample) + return _readings_until_no_new_worker( + proxy, readings + sample, samples + 1, 0 if new_worker_answered else samples_since_new_worker + 1 + ) + + +def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[WorkerKey, RssReading]: + readings: Final = _readings_until_no_new_worker(proxy, (), 0, 0) + assert readings, "no /debug/memory/summary read carried ram_usage_mb, so the proxy cannot report its RSS" + return MappingProxyType( + { + worker: min((reading for reading in readings if reading.worker == worker), key=lambda r: r.ram_usage_mb) + for worker in {reading.worker for reading in readings} + } + ) + + +def _heaviest_worker_growth( + warm: Mapping[WorkerKey, RssReading], after: Mapping[WorkerKey, RssReading] +) -> WorkerGrowth: + assert warm.keys() == after.keys(), ( + f"the workers answering /debug/memory/summary changed between the checkpoints, so not every worker can " + f"be compared with itself: gone after the measured batch {sorted(warm.keys() - after.keys())} (a worker " + f"that died or was restarted under failing traffic, which is what an OOM kill looks like), first seen " + f"after it {sorted(after.keys() - warm.keys())} (the warm window never reached them, so they have no " + f"baseline; raise E2E_MEMORY_RSS_SETTLE_SAMPLES if the stack has more workers than the window covers)" + ) + return max((WorkerGrowth(warm[worker], after[worker]) for worker in warm), key=lambda growth: growth.growth_mb) + + +def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fallback: str) -> None: + served: Final = tuple(call for call in calls if call.status_code == 200) + assert not served, ( + f"{len(served)} of {len(calls)} calls came back 200, so they reached a provider and never " + f"exercised the retry loop: {served[0].body_head}" + ) + without_fallback: Final = tuple(call for call in calls if fallback not in call.body_head) + assert not without_fallback, ( + f"{len(without_fallback)} of {len(calls)} failures never named the fallback group {fallback}, " + f"so the request did not run through retries into the fallback: {without_fallback[0].body_head}" + ) + + +def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float: + assert call.call_id, ( + f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}" + ) + rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id) + assert rows, ( + f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack " + "writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored " + "request outgrew the writer's queue budget" + ) + snapshot: Final = rows[0].proxy_server_request + assert snapshot, ( + f"spend log {call.call_id} stored no request body, so the stack is not running with " + "general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously" + ) + return len(json.dumps(snapshot).encode()) / 1024 + + +class TestReliabilityMemory: + @pytest.mark.covers("reliability.perf.memory.under_slo") + def test_failing_requests_do_not_grow_rss_or_stored_request( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + marker: Final = unique_marker() + primary: Final = f"reliability-memory-{marker}" + fallback: Final = f"reliability-memory-fb-{marker}" + _register_refusing_group(client.proxy, resources, primary) + _register_refusing_group(client.proxy, resources, fallback) + override: Final = RouterSettingsOverride( + num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}] + ) + + probe: Final = _fail_once(client.proxy, scoped_key, primary, override) + _assert_every_call_failed_through_fallback((probe,), fallback) + stored_kb: Final = _stored_request_kb(client.proxy, probe) + assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, ( + f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the " + f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with " + f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole " + f"request into the stored snapshot the way the v1.100.0 ones did" + ) + + warmup: Final = _fail_many(client.proxy, scoped_key, primary, override) + _assert_every_call_failed_through_fallback(warmup, fallback) + warm: Final = _settled_rss_per_worker(client.proxy) + + measured: Final = _fail_many(client.proxy, scoped_key, primary, override) + _assert_every_call_failed_through_fallback(measured, fallback) + after: Final = _settled_rss_per_worker(client.proxy) + + heaviest: Final = _heaviest_worker_growth(warm, after) + assert heaviest.growth_mb <= MEMORY_RSS_BUDGET_MB, ( + f"proxy RSS grew {heaviest.growth_mb:.1f} MB over a second batch of {MEMORY_REQUESTS_PER_PHASE} failing " + f"requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) after an identical warmup batch, " + f"past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: worker pid {heaviest.warm.worker_pid} on " + f"{heaviest.warm.hostname or 'an unnamed host'} behind {heaviest.warm.replica} settled at " + f"{heaviest.warm.ram_usage_mb:.1f} MB warm and " + f"{heaviest.after.ram_usage_mb:.1f} MB after; failing requests are leaking memory the way the " + f"v1.100.0 retry breadcrumbs did" + ) diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py new file mode 100644 index 00000000000..33a09a0f13a --- /dev/null +++ b/tests/e2e/test_idp.py @@ -0,0 +1,179 @@ +"""Harness coverage for idp.py: the pure parts of the Keycloak client, which are +the ones a wrong value in silently mistargets. No proxy and no IdP needed, so +these carry no `e2e` marker and run everywhere.""" + +from __future__ import annotations + +from collections.abc import Callable, Generator +from contextlib import ExitStack, contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from queue import SimpleQueue +from threading import Thread +from typing import Final + +import pytest +from e2e_http import ExternalWrite +from idp import ( + KEYCLOAK_ADMIN_PASSWORD_ENV, + KEYCLOAK_ADMIN_USER_ENV, + KEYCLOAK_REALM_ENV, + KEYCLOAK_URL_ENV, + Keycloak, + PasswordCredential, + UserCreateBody, + created_id, + keycloak_from_env, +) + +_REALM: Final = Keycloak( + base_url="http://keycloak:8080", realm="litellm-e2e", admin_username="admin", admin_password="pw" +) + + +def test_realm_urls_match_keycloaks_own_layout() -> None: + assert _REALM.issuer == "http://keycloak:8080/realms/litellm-e2e" + assert _REALM.jwks_url == "http://keycloak:8080/realms/litellm-e2e/protocol/openid-connect/certs" + assert _REALM.token_url("master") == "http://keycloak:8080/realms/master/protocol/openid-connect/token" + + +def test_created_id_is_the_last_segment_of_the_location_header() -> None: + created: Final = ExternalWrite( + status_code=201, location="http://keycloak:8080/admin/realms/litellm-e2e/groups/abc-123" + ) + assert created_id(created, "a group") == "abc-123" + + +def test_a_refused_create_fails_the_test_with_the_idps_own_words() -> None: + with pytest.raises(BaseException, match=r"409.*already exists"): + created_id(ExternalWrite(status_code=409, body="Group already exists"), "a group") + + +@pytest.mark.parametrize("location", ["", "http://keycloak/groups/"]) +def test_create_without_a_resource_id_fails(location: str) -> None: + with pytest.raises(pytest.fail.Exception, match="resource id"): + created_id(ExternalWrite(status_code=201, location=location), "a group") + + +@contextmanager +def _idp_server( + *, user_status: int = 201, delete_status: int = 204, admin_status: int = 200 +) -> Generator[tuple[Keycloak, SimpleQueue[str]]]: + """Exercise provisioning failures through the same HTTP transport as live tests.""" + deletions: SimpleQueue[str] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if self.path.endswith("/token"): + self.send_response(admin_status) + self.end_headers() + self.wfile.write(b'{"access_token":"synthetic-harness-token"}') + else: + self.send_response(user_status if self.path.endswith("/users") else 201) + self.send_header("Location", f"{self.path}/resource-1") + self.end_headers() + if user_status != 201 and self.path.endswith("/users"): + self.wfile.write(b"injected create failure") + + def do_DELETE(self) -> None: + deletions.put(self.path) + self.send_response(delete_status) + self.end_headers() + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield ( + Keycloak( + base_url=f"http://127.0.0.1:{server.server_port}", + realm="test", + admin_username="admin", + admin_password="pw", + ), + deletions, + ) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> None: + with _idp_server(user_status=500) as (idp, deletions): + with ExitStack() as cleanup: + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + with pytest.raises(pytest.fail.Exception, match="injected create failure"): + idp.provision(marker="partial", group="team", defer=defer) + assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1" + assert deletions.empty() + + +def test_successful_provisioning_cleans_up_user_before_group() -> None: + with _idp_server() as (idp, deletions): + with ExitStack() as cleanup: + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + idp.provision(marker="complete", group="team", defer=defer) + assert deletions.get_nowait() == "/admin/realms/test/users/resource-1" + assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1" + assert deletions.empty() + + +def test_cleanup_failure_is_visible() -> None: + with _idp_server(delete_status=500) as (idp, _): + with pytest.warns(RuntimeWarning, match="cleanup failed.*HTTP 500"): + idp.delete_group("group") + + +def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None: + with _idp_server(admin_status=401) as (idp, _): + cleanup: Final = ExitStack() + cleanup.callback(idp.delete_group, "group") + cleanup.callback(idp.delete_user, "user") + with pytest.warns(RuntimeWarning, match="cleanup could not authenticate") as warnings: + cleanup.close() + assert len(warnings) == 2 + + +def test_new_users_are_born_fully_set_up() -> None: + """A user without a profile or with a pending required action authenticates + nowhere: Keycloak answers every grant with "Account is not fully set up".""" + body: Final = UserCreateBody( + username="e2e", email="e2e@example.com", groups=("team",), credentials=(PasswordCredential(value="pw"),) + ).model_dump(by_alias=True) + + assert body["requiredActions"] == () + assert body["firstName"] and body["lastName"] and body["emailVerified"] is True + assert body["credentials"][0]["temporary"] is False + + +def test_connection_details_come_from_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(KEYCLOAK_URL_ENV, "http://keycloak.litellm.svc.cluster.local:8080/") + monkeypatch.setenv(KEYCLOAK_REALM_ENV, "other-realm") + monkeypatch.setenv(KEYCLOAK_ADMIN_USER_ENV, "admin") + monkeypatch.setenv(KEYCLOAK_ADMIN_PASSWORD_ENV, "pw") + + resolved: Final = keycloak_from_env() + + assert resolved.issuer == "http://keycloak.litellm.svc.cluster.local:8080/realms/other-realm" + assert resolved.admin_username == "admin" and resolved.admin_password == "pw" + + +@pytest.mark.parametrize("blank", ["", " "]) +def test_a_missing_admin_credential_fails_loudly_instead_of_skipping( + monkeypatch: pytest.MonkeyPatch, blank: str +) -> None: + monkeypatch.setenv(KEYCLOAK_ADMIN_USER_ENV, "admin") + monkeypatch.setenv(KEYCLOAK_ADMIN_PASSWORD_ENV, blank) + + with pytest.raises(BaseException, match=KEYCLOAK_ADMIN_PASSWORD_ENV): + keycloak_from_env() diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 8ab389ee43c..18f72ac0e7a 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -34,10 +34,7 @@ from pathlib import Path from typing import Final import pytest -from pydantic import TypeAdapter - from e2e_http import RawResponse, StreamChunk, forward -from fixture_canonical import canonicalize from fixture_bundle import ( BundleRecorder, Interaction, @@ -49,6 +46,7 @@ from fixture_bundle import ( prepare_bundle, slug_for_test, ) +from fixture_canonical import canonicalize from fixture_mode import current_test_key from provider_edge import ( REPLAY_MISS_STATUS, @@ -56,15 +54,18 @@ from provider_edge import ( EdgeReply, EdgeStream, ProviderEdge, + ProviderRequestObservation, RecordEdge, ReplayEdge, ReplaySource, edge_request, handle_edge_request, + observed_provider_edge, provider_edge_api_base, replay_leftover_error, start_provider_edge, ) +from pydantic import TypeAdapter CHAT_PATH = "/openai/v1/chat/completions" UPLOAD_PATH = "/openai/v1/files" @@ -1290,3 +1291,62 @@ class TestApiBaseSeam: assert second.endswith("/anthropic") assert first.rsplit("/", 1)[0] == second.rsplit("/", 1)[0] assert (root / "manifest.json").is_file() + + +class TestProviderRequestObservation: + def test_live_counts_repeated_marker_calls_without_recording(self, tmp_path: Path) -> None: + observation: Final = ProviderRequestObservation("observed-lantern") + with fake_provider() as provider: + with observed_provider_edge( + observation, mode_raw="live", bundle_dir=tmp_path / "unused", + bind_host="127.0.0.1", advertise_host="127.0.0.1", + mounts={"openai": provider_url(provider)}, + ) as edge: + assert observation.count == 0 + unrelated: Final = call_edge(edge, "POST", CHAT_PATH, body=chat_body("other-lantern")) + assert unrelated.status_code == 200 + assert observation.count == 0 + first: Final = call_edge(edge, "POST", CHAT_PATH, body=chat_body("observed-lantern")) + assert first.status_code == 200 + assert json_object(first.body)["echo"] == chat_body("observed-lantern").decode() + assert observation.count == 1 + second: Final = call_edge(edge, "POST", CHAT_PATH, body=chat_body("observed-lantern")) + assert second.status_code == 200 + assert observation.count == 2 + assert len(provider.hits) == 3 + assert not (tmp_path / "unused").exists() + + def test_record_and_replay_count_each_matching_call(self, tmp_path: Path) -> None: + with fake_provider() as provider: + for mode, observation in ( + ("record", ProviderRequestObservation("observed-lantern")), + ("replay", ProviderRequestObservation("observed-lantern")), + ): + with observed_provider_edge( + observation, mode_raw=mode, bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", advertise_host="127.0.0.1", + mounts={"openai": provider_url(provider)}, + ) as edge: + assert observation.count == 0 + for expected, response in ( + (index, call_edge(edge, "POST", CHAT_PATH, body=chat_body("observed-lantern"))) + for index in (1, 2) + ): + assert response.status_code == 200 + assert json_object(response.body)["hit"] == expected + assert observation.count == expected + assert len(provider.hits) == 2 + assert replay_leftover_error( + mode_raw="replay", bundle_dir=tmp_path / "bundle", test_key=current_test_key() + ) is None + + def test_failed_provider_attempt_is_counted(self, tmp_path: Path) -> None: + observation: Final = ProviderRequestObservation("observed-lantern") + with observed_provider_edge( + observation, mode_raw="live", bundle_dir=tmp_path / "unused", + bind_host="127.0.0.1", advertise_host="127.0.0.1", + mounts={"openai": "http://127.0.0.1:9"}, + ) as edge: + response: Final = call_edge(edge, "POST", CHAT_PATH, body=chat_body("observed-lantern")) + assert response.status_code == 502 + assert observation.count == 1 diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 3b84a47e3cc..1b0133f12cb 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -245,6 +245,7 @@ class TestReplicasFor: replica_urls=("http://gateway-1", "http://gateway-2"), ) assert set(client.replicas_for("/key/info")) == {"http://backend"} + assert set(client.replicas_for("/project/info")) == {"http://backend"} assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"} def test_monolith_reads_management_routes_back_from_every_replica(self) -> None: diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 44fdbaa3e41..e8caa801467 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -295,6 +295,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/user", "/team", "/organization", + "/project", "/customer", "/end_user", "/tag", diff --git a/tests/ocr_tests/conftest.py b/tests/ocr_tests/conftest.py index 259aad5f782..bf03efca744 100644 --- a/tests/ocr_tests/conftest.py +++ b/tests/ocr_tests/conftest.py @@ -5,6 +5,7 @@ # Vertex AI OCR) are replayed for 24h. See tests/llm_translation/Readme.md # for the design overview. +from typing import Final import pytest @@ -23,7 +24,12 @@ from tests._vcr_conftest_common import ( # noqa: E402,F401 vcr_config_dict, ) -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () +_VCR_INCOMPATIBLE_NODEID_SUFFIXES: Final[tuple[str, ...]] = ( + "test_rust_bridge.py::test_native_public_ocr_matches_python[mistral/mistral-ocr-latest-False]", + "test_rust_bridge.py::test_native_public_ocr_matches_python[mistral/mistral-ocr-latest-True]", + "test_rust_bridge.py::test_native_public_ocr_matches_python[azure_ai/doc-intelligence/prebuilt-read-False]", + "test_rust_bridge.py::test_native_public_ocr_matches_python[azure_ai/doc-intelligence/prebuilt-read-True]", +) _verbose_state = VerboseReporterState() diff --git a/tests/proxy_behavior/auth/__init__.py b/tests/proxy_behavior/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_behavior/auth/conftest.py b/tests/proxy_behavior/auth/conftest.py new file mode 100644 index 00000000000..21982fa25cd --- /dev/null +++ b/tests/proxy_behavior/auth/conftest.py @@ -0,0 +1,20 @@ +"""Session-scoped PrismaClient for auth behavior tests that run raw SQL against a real Postgres.""" + +import os +from unittest.mock import MagicMock + +import pytest +import pytest_asyncio + +from litellm.proxy.utils import PrismaClient + + +@pytest_asyncio.fixture(scope="session", loop_scope="session") +async def prisma(): + database_url = os.environ.get("DATABASE_URL") + if not database_url: + pytest.skip("DATABASE_URL not set") # test-quality-ok: this suite exists to run SQL on a real Postgres + client = PrismaClient(database_url=database_url, proxy_logging_obj=MagicMock()) + await client.connect() + yield client + await client.disconnect() diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py new file mode 100644 index 00000000000..e2d947f4284 --- /dev/null +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -0,0 +1,174 @@ +"""Runs the auth prefetch's raw SQL against a real Postgres: the join must bind the membership to the requested +team and hand the getters rows they validate. The per-regime round-trip counts are unit-tested with fakes in +tests/test_litellm/proxy/auth/test_auth_object_prefetch.py.""" + +import json +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.auth.auth_checks import ( + get_org_object, + get_team_membership, + get_team_object, + get_user_object, +) +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects +from litellm.proxy.auth.team_grants import team_model_aliases +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +def _dead_db() -> MagicMock: + prisma = MagicMock(name="prisma_client") + prisma.db.query_first = AsyncMock(return_value=None) + return prisma + + +async def test_join_binds_the_membership_to_the_requested_team(prisma): + """A user in two teams with different member budgets must get the requested team's row.""" + run = uuid4().hex + user_id, team_a, team_b, org_id = (f"pf-user-{run}", f"pf-team-a-{run}", f"pf-team-b-{run}", f"pf-org-{run}") + try: + await prisma.db.litellm_budgettable.create( + data={"budget_id": f"a-{run}", "max_budget": 11.0, "created_by": "t", "updated_by": "t"} + ) + await prisma.db.litellm_budgettable.create( + data={"budget_id": f"b-{run}", "max_budget": 22.0, "created_by": "t", "updated_by": "t"} + ) + await prisma.db.litellm_organizationtable.create( + data={ + "organization_id": org_id, + "organization_alias": "pf", + "created_by": "t", + "updated_by": "t", + "litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}}, + } + ) + await prisma.db.litellm_usertable.create(data={"user_id": user_id, "max_budget": 33.0}) + await prisma.db.litellm_teamtable.create(data={"team_id": team_a, "organization_id": org_id, "max_budget": 1.0}) + await prisma.db.litellm_teamtable.create(data={"team_id": team_b, "max_budget": 2.0}) + await prisma.db.litellm_teammembership.create( + data={"user_id": user_id, "team_id": team_a, "litellm_budget_table": {"connect": {"budget_id": f"a-{run}"}}} + ) + await prisma.db.litellm_teammembership.create( + data={"user_id": user_id, "team_id": team_b, "litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}}} + ) + + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id) + await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) + + dead_db = _dead_db() + membership = await get_team_membership( + user_id=user_id, team_id=team_a, prisma_client=dead_db, user_api_key_cache=cache + ) + team = await get_team_object(team_id=team_a, prisma_client=dead_db, user_api_key_cache=cache) + user = await get_user_object( + user_id=user_id, prisma_client=dead_db, user_api_key_cache=cache, user_id_upsert=False + ) + org = await get_org_object(org_id=org_id, prisma_client=dead_db, user_api_key_cache=cache) + assert dead_db.db.mock_calls == [], "getters must be served from the prefetched cache" + + assert membership is not None and membership.litellm_budget_table is not None + assert (membership.team_id, membership.litellm_budget_table.max_budget) == (team_a, 11.0) + assert (team.team_id, team.max_budget, team.organization_id, team.models) == (team_a, 1.0, org_id, []) + assert user is not None and user.max_budget == 33.0 + assert org is not None and (org.organization_id, org.models) == (org_id, []) + finally: + await prisma.db.litellm_teammembership.delete_many(where={"user_id": user_id}) + await prisma.db.litellm_teamtable.delete_many(where={"team_id": {"in": [team_a, team_b]}}) + await prisma.db.litellm_usertable.delete_many(where={"user_id": user_id}) + await prisma.db.litellm_organizationtable.delete_many(where={"organization_id": org_id}) + await prisma.db.litellm_budgettable.delete_many(where={"budget_id": {"in": [f"a-{run}", f"b-{run}"]}}) + + +async def test_join_reads_team_model_aliases_from_the_mapped_column(prisma): + """The model table stores aliases in a column named ``aliases``; the cached team must expose ``model_aliases``.""" + run = uuid4().hex + team_id = f"pf-team-{run}" + aliases = {"gpt-4o": f"gpt-4o-{run}"} + model_table = await prisma.db.litellm_modeltable.create( + data={"model_aliases": json.dumps(aliases), "created_by": "t", "updated_by": "t"} + ) + try: + await prisma.db.litellm_teamtable.create(data={"team_id": team_id, "model_id": model_table.id}) + expected_team = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, include={"litellm_model_table": True} + ) + + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + refs = AuthObjectRefs(user_id=None, team_id=team_id, membership_user_id=None, organization_id=None) + await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) + + dead_db = _dead_db() + team = await get_team_object(team_id=team_id, prisma_client=dead_db, user_api_key_cache=cache) + assert dead_db.db.mock_calls == [], "getters must be served from the prefetched cache" + + assert expected_team is not None and expected_team.litellm_model_table is not None + assert team.litellm_model_table is not None + assert team.litellm_model_table.model_aliases == expected_team.litellm_model_table.model_aliases == aliases + assert team_model_aliases(team) == aliases + finally: + await prisma.db.litellm_teamtable.delete_many(where={"team_id": team_id}) + await prisma.db.litellm_modeltable.delete_many(where={"id": model_table.id}) + + +async def test_join_reads_null_nested_lists_the_way_prisma_does(prisma): + """Prisma reads a NULL scalar list as []; the nested permission and budget rows must match, not carry null.""" + run = uuid4().hex + user_id, team_id, permission_id, budget_id = (f"pf-user-{run}", f"pf-team-{run}", f"pf-perm-{run}", f"pf-bud-{run}") + try: + await prisma.db.litellm_objectpermissiontable.create(data={"object_permission_id": permission_id}) + await prisma.db.litellm_budgettable.create(data={"budget_id": budget_id, "created_by": "t", "updated_by": "t"}) + await prisma.db.execute_raw( + 'UPDATE "LiteLLM_ObjectPermissionTable" SET mcp_servers = NULL, models = NULL ' + "WHERE object_permission_id = $1", + permission_id, + ) + await prisma.db.execute_raw( + 'UPDATE "LiteLLM_BudgetTable" SET allowed_models = NULL WHERE budget_id = $1', budget_id + ) + await prisma.db.litellm_usertable.create(data={"user_id": user_id}) + await prisma.db.litellm_teamtable.create(data={"team_id": team_id, "object_permission_id": permission_id}) + await prisma.db.litellm_teammembership.create( + data={"user_id": user_id, "team_id": team_id, "litellm_budget_table": {"connect": {"budget_id": budget_id}}} + ) + expected_team = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, include={"object_permission": True} + ) + expected_membership = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True} + ) + + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + refs = AuthObjectRefs(user_id=user_id, team_id=team_id, membership_user_id=user_id, organization_id=None) + await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) + + dead_db = _dead_db() + team = await get_team_object(team_id=team_id, prisma_client=dead_db, user_api_key_cache=cache) + membership = await get_team_membership( + user_id=user_id, team_id=team_id, prisma_client=dead_db, user_api_key_cache=cache + ) + assert dead_db.db.mock_calls == [], "getters must be served from the prefetched cache" + + assert expected_team is not None and expected_team.object_permission is not None + assert team.object_permission is not None + assert team.object_permission.mcp_servers == expected_team.object_permission.mcp_servers == [] + assert team.object_permission.models == expected_team.object_permission.models == [] + assert expected_membership is not None and expected_membership.litellm_budget_table is not None + assert membership is not None and membership.litellm_budget_table is not None + assert ( + membership.litellm_budget_table.allowed_models + == expected_membership.litellm_budget_table.allowed_models + == [] + ) + finally: + await prisma.db.litellm_teammembership.delete_many(where={"user_id": user_id}) + await prisma.db.litellm_teamtable.delete_many(where={"team_id": team_id}) + await prisma.db.litellm_usertable.delete_many(where={"user_id": user_id}) + await prisma.db.litellm_objectpermissiontable.delete_many(where={"object_permission_id": permission_id}) + await prisma.db.litellm_budgettable.delete_many(where={"budget_id": budget_id}) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 2bd3a50f39f..860e872dd44 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -1,7 +1,8 @@ from __future__ import annotations -import asyncio -from collections.abc import Awaitable, Callable +import json +import subprocess +from functools import cache from pathlib import Path from typing import Final, Protocol, cast @@ -28,12 +29,12 @@ class _GatewayClient(Protocol): def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: - import litellm from fastapi.testclient import TestClient + import litellm + from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.endpoints import user_api_key_auth - from litellm.proxy import proxy_server provider_model: Final = cast(str, fixture.kwargs["provider_model"]) model_alias: Final = cast(str, fixture.kwargs["model_alias"]) @@ -76,24 +77,24 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: - from litellm.rust_bridge import get_native_bridge - - bridge: Final[object | None] = get_native_bridge() - trace: Final[object | None] = getattr(bridge, "_trace", None) if bridge is not None else None - gateway_messages: Final[object | None] = getattr(trace, "gateway_messages", None) - if gateway_messages is None or not callable(gateway_messages): - raise RuntimeError("native Rust trace bridge does not expose gateway_messages") - invoke_gateway: Final = cast(Callable[[str, str, str, object], Awaitable[object]], gateway_messages) - - async def invoke() -> object: - return await invoke_gateway( - cast(str, fixture.kwargs["model_alias"]), - cast(str, fixture.kwargs["provider_model"]), - cast(str, fixture.kwargs["api_base"]), - fixture.kwargs["body"], - ) - - result: Final = asyncio.run(invoke()) + payload: Final = json.dumps( + { + "model_alias": fixture.kwargs["model_alias"], + "provider_model": fixture.kwargs["provider_model"], + "api_base": fixture.kwargs["api_base"], + "body": fixture.kwargs["body"], + } + ) + completed: Final = subprocess.run( + (_gateway_trace_binary(),), + input=payload, + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError(f"Rust gateway trace failed: {completed.stderr.strip()}") + result: Final = json.loads(completed.stdout) payload: Final = TraceResponsePayload.model_validate(result) response: Final = _GatewayResponsePayload.model_validate(payload.response) if response.status != 200: @@ -101,6 +102,34 @@ def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: return native_trace_events(payload) +@cache +def _gateway_trace_binary() -> Path: + repo_root: Final = next(parent for parent in Path(__file__).resolve().parents if (parent / "litellm-rust").is_dir()) + rust_root: Final = repo_root / "litellm-rust" + completed: Final = subprocess.run( + ( + "cargo", + "build", + "--quiet", + "--package", + "litellm-ai-gateway", + "--features", + "trace-parity", + "--bin", + "trace-parity-gateway", + "--target-dir", + rust_root / "target", + ), + cwd=rust_root, + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError(f"Rust gateway trace build failed: {completed.stderr.strip()}") + return rust_root / "target" / "debug" / "trace-parity-gateway" + + def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: try: with replay_server() as provider: diff --git a/tests/store_model_in_db_tests/test_openai_error_handling.py b/tests/store_model_in_db_tests/test_openai_error_handling.py index 5de7c427d19..22707d522bc 100644 --- a/tests/store_model_in_db_tests/test_openai_error_handling.py +++ b/tests/store_model_in_db_tests/test_openai_error_handling.py @@ -210,7 +210,8 @@ async def test_chat_completion_bad_model_with_spend_logs(): assert "traceback" in error_info assert error_info["error_code"] == "400" assert error_info["error_class"] in ("ProxyModelNotFoundError", "BadRequestError") - assert "non-existent-model" in error_info["error_message"] + assert "non-existent-model" not in error_info["error_message"] + assert "/chat/completions: Invalid model name passed in" in error_info["error_message"] # Verify request details assert log_entry["cache_hit"] == "False" diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 85e8308ae91..40ad4f0c6f0 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -250,3 +250,27 @@ def test_in_memory_cache_prunes_expired_heap_entries_below_capacity(): assert len(in_memory_cache.cache_dict) == 5 assert len(in_memory_cache.ttl_dict) == 5 assert len(in_memory_cache.expiration_heap) == 5 + + +def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None: + class Clock: + now = 0.0 + + def __call__(self) -> float: + return self.now + + clock = Clock() + cache = InMemoryCache(max_size_in_memory=2, default_ttl=60, clock=clock) + cache.set_cache("first", "original", ttl=10) + clock.now = 9.0 + cache.set_cache("second", "survivor") + assert cache.get_cache("first") == "original" + clock.now = 10.001 + assert cache.get_cache("first") is None + cache.set_cache("third", "replacement") + assert cache.get_cache("second") == "survivor" + clock.now = 69.001 + cache.set_cache("fourth", "new") + assert cache.get_cache("second") is None + assert cache.get_cache("third") == "replacement" + assert cache.get_cache("fourth") == "new" diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index bcae33b976e..4bf7894829e 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1,10 +1,12 @@ import asyncio import time from collections.abc import Iterator +from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError @@ -679,7 +681,6 @@ def test_sync_batch_get_cache_survives_a_service_callback_that_raises( 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 @@ -1202,3 +1203,131 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new new_probe_release.set() assert await new_probe == "new probe" assert breaker._state == breaker.CLOSED + + +class _RoundTripCountingRedis: + """Fake redis.asyncio client: one round trip per awaited command or pipeline execute.""" + + def __init__(self, ttl: int) -> None: + self.values: dict[str, float] = {} + self.ttls: dict[str, int] = {} + self.round_trips = 0 + self._initial_ttl = ttl + + async def incrbyfloat(self, name: str, amount: float) -> float: + self.round_trips += 1 + return self._incr(name, amount) + + async def expire(self, name: str, time: int) -> bool: + self.round_trips += 1 + self.ttls[name] = time + return True + + def _incr(self, name: str, amount: float) -> float: + self.values[name] = self.values.get(name, 0.0) + amount + self.ttls.setdefault(name, self._initial_ttl) + return self.values[name] + + def pipeline(self, transaction: bool) -> "_RoundTripCountingRedis._Pipeline": + return _RoundTripCountingRedis._Pipeline(self) + + class _Pipeline: + def __init__(self, client: "_RoundTripCountingRedis") -> None: + self._client = client + self._commands: list[tuple[str, tuple[object, ...]]] = [] + + async def __aenter__(self) -> "_RoundTripCountingRedis._Pipeline": + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + def incrbyfloat(self, name: str, amount: float) -> None: + self._commands.append(("incrbyfloat", (name, amount))) + + def expire(self, name: str, time: int) -> None: + self._commands.append(("expire", (name, time))) + + def ttl(self, name: str) -> None: + self._commands.append(("ttl", (name,))) + + async def execute(self) -> list[object]: + self._client.round_trips += 1 + results: list[object] = [] + for command, args in self._commands: + if command == "incrbyfloat": + results.append(self._client._incr(str(args[0]), float(args[1]))) # pyright: ignore[reportArgumentType] # fake stores str/float + elif command == "expire": + self._client.ttls[str(args[0])] = int(args[1]) # pyright: ignore[reportArgumentType] # fake stores int + results.append(True) + else: + results.append(self._client.ttls.get(str(args[0]), -2)) + return results + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("refresh_ttl", "existing_ttl", "expected_round_trips", "expected_ttl"), + [ + pytest.param(True, 100, 2, 60, id="refresh_ttl: INCRBYFLOAT+EXPIRE in one round trip each"), + pytest.param(False, 100, 2, 100, id="keep ttl: INCRBYFLOAT+TTL in one round trip each, no EXPIRE"), + pytest.param(False, -1, 3, 60, id="unexpiring key: INCRBYFLOAT+TTL then EXPIRE once, 1 trip after"), + ], +) +async def test_async_increment_pipelines_the_ttl_command( + monkeypatch, redis_no_ping, refresh_ttl, existing_ttl, expected_round_trips, expected_ttl +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = _RoundTripCountingRedis(ttl=existing_ttl) + + with patch.object(redis_cache, "init_async_client", return_value=client): + first = await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=refresh_ttl) + second = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=refresh_ttl) + + assert (first, second) == (1.5, 3.5) + assert client.values == {"ns:spend:key:k": 3.5} + assert client.ttls == {"ns:spend:key:k": expected_ttl} + assert client.round_trips == expected_round_trips + + +class _SetRecordingPipeline: + def __init__(self) -> None: + self.sets: list[tuple[str, str, timedelta | None]] = [] + self.executes = 0 + + async def __aenter__(self) -> "_SetRecordingPipeline": + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + def set(self, name: str, value: str, ex: timedelta | None) -> None: + self.sets.append((name, value, ex)) + + async def execute(self) -> list[bool]: + self.executes += 1 + return [True] * len(self.sets) + + +@pytest.mark.asyncio +async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + monkeypatch.setattr(litellm, "default_redis_ttl", 300) + redis_cache = RedisCache(namespace="ns") + pipe = _SetRecordingPipeline() + client = MagicMock() + client.pipeline = MagicMock(return_value=pipe) + + with patch.object(redis_cache, "init_async_client", return_value=client): + await redis_cache.async_set_cache_pipeline_with_ttls( + (("team_id:t1", {"team_id": "t1"}, 60), ("u1", {"user_id": "u1"}, 7), ("org_id:o1", {"a": 1}, None)) + ) + + client.pipeline.assert_called_once_with(transaction=False) + assert pipe.executes == 1 + assert pipe.sets == [ + ("ns:team_id:t1", '{"team_id": "t1"}', timedelta(seconds=60)), + ("ns:u1", '{"user_id": "u1"}', timedelta(seconds=7)), + ("ns:org_id:o1", '{"a": 1}', timedelta(seconds=300)), + ] diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 3283af26cf7..f72316f5d5e 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1704,8 +1704,9 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N "initialize_not_found", ), ) +@pytest.mark.parametrize("raise_on_error", (False, True)) async def test_optional_discovery_capabilities_and_errors( - method: str, outcome: str, caplog: pytest.LogCaptureFixture + method: str, outcome: str, caplog: pytest.LogCaptureFixture, raise_on_error: bool ) -> None: import logging from unittest.mock import Mock @@ -1783,7 +1784,11 @@ async def test_optional_discovery_capabilities_and_errors( "resources/list": client.list_resources, "resources/templates/list": client.list_resource_templates, }[method] - result: Final = await operation() + if raise_on_error and outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"): + with pytest.raises((McpError, httpx.HTTPError)): + await operation(raise_on_error=True) + return + result: Final = await operation(raise_on_error=raise_on_error) requests: Final = tuple( JSONRPCMessage.model_validate_json(call.args[0].content).root @@ -1907,3 +1912,25 @@ def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): ) assert result.returncode == 0, result.stderr assert result.stdout.strip() == "MCPServerManager" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("resolved", (False, True)) +async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved: bool) -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth + + def client(token: str) -> MCPClient: + return MCPClient( + server_url="https://example.com/mcp", + auth_type=MCPAuth.api_key, + auth_value=None if resolved else token, + resolved_auth=StaticHeaderAuth(token) if resolved else None, + ) + + original: Final = await client("private-original-credential").discovery_auth_fingerprint() + repeated: Final = await client("private-original-credential").discovery_auth_fingerprint() + replaced: Final = await client("private-replaced-credential").discovery_auth_fingerprint() + assert original == repeated + assert original != replaced + assert len(original) == 64 + assert "private-original-credential" not in original diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 55e2dcdc270..44dda57dd27 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -45,6 +45,21 @@ class TestSlackAlerting(unittest.TestCase): result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, -0.2) + def test_get_user_info_str_omits_absent_token_for_user_alert(self): + user_info = CallInfo( + spend=85.0, + max_budget=100.0, + user_id="user-1", + user_email="person@example.com", + event_group=Litellm_EntityType.USER, + ) + + result = self.slack_alerting._get_user_info_str(user_info) + + self.assertIn("*user_id:* `user-1`", result) + self.assertIn("*user_email:* `person@example.com`", result) + self.assertNotIn("*token:*", result) + def test_get_event_and_event_message_max_budget(self): # Initial setup with no event event = None diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/test_litellm/integrations/otel/test_langfuse_logger.py index 8f93a9a564f..8db84b090a0 100644 --- a/tests/test_litellm/integrations/otel/test_langfuse_logger.py +++ b/tests/test_litellm/integrations/otel/test_langfuse_logger.py @@ -1,9 +1,9 @@ -"""Tests for ``LangfuseOpenTelemetryV2``: the root observation's input and output are stamped from the -request-task hooks, while the root span is still recording, so Langfuse can show them on the trace.""" +"""Tests for the Langfuse OTel v2 loggers: the trace name and the root observation's input and output are +stamped from the request task while the root span is still recording, so Langfuse can show them on the trace.""" import asyncio import json -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from typing import Final import pytest @@ -14,7 +14,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE import litellm # noqa: E402 from litellm.caching.dual_cache import DualCache # noqa: E402 -from litellm.integrations.otel.logger import build_otel_v2_logger # noqa: E402 +from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger # noqa: E402 from litellm.integrations.otel.model.config import OpenTelemetryV2Config, is_otel_v2_enabled # noqa: E402 from litellm.integrations.otel.model.spans import LITELLM_PROXY_REQUEST_SPAN_NAME, SpanRole # noqa: E402 from litellm.integrations.otel.plumbing import context as otel_context # noqa: E402 @@ -41,6 +41,7 @@ from litellm.types.utils import ( # noqa: E402 INPUT_ATTR: Final = "langfuse.observation.input" OUTPUT_ATTR: Final = "langfuse.observation.output" +TRACE_NAME_ATTR: Final = "langfuse.trace.name" CHAT_DATA: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "ping"}]} @@ -306,6 +307,73 @@ def test_unrenderable_output_never_raises_into_the_request(): assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs +def _run_named_request( + logger: OpenTelemetryV2, exporter: InMemorySpanExporter, litellm_params: Mapping[str, object] +) -> tuple[Mapping[str, object], Mapping[str, object]]: + response: Final = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + root: Final = _start_root(logger) + logger.log_pre_api_call( + model="gpt-5.4-mini", messages=[], kwargs={"litellm_call_id": "call_1", "litellm_params": litellm_params} + ) + root.end() + payload: Final = { + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-5.4-mini", + "messages": CHAT_DATA["messages"], + "response": response.model_dump(), + "status": "success", + "litellm_call_id": "call_1", + "metadata": {}, + "hidden_params": {}, + } + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": payload, "litellm_params": litellm_params}, response, None, None + ) + ) + generation: Final = next( + span for span in exporter.get_finished_spans() if span.name != LITELLM_PROXY_REQUEST_SPAN_NAME + ) + return _root_attrs(exporter), dict(generation.attributes or {}) + + +@pytest.mark.parametrize("capture", ["span_only", "no_content"]) +def test_langfuse_trace_name_header_names_the_root_and_the_generation_over_body_metadata(capture): + logger, exporter = _logger(capture=capture) + + root_attrs, generation_attrs = _run_named_request( + logger, + exporter, + { + "metadata": {"trace_name": "from-body"}, + "proxy_server_request": {"headers": {"langfuse_trace_name": "from-header"}}, + }, + ) + + assert root_attrs[TRACE_NAME_ATTR] == "from-header" + assert generation_attrs[TRACE_NAME_ATTR] == "from-header" + + +def test_body_metadata_trace_name_names_the_root_and_the_generation(): + logger, exporter = _logger() + + root_attrs, generation_attrs = _run_named_request( + logger, exporter, {"metadata": {"trace_name": "from-body"}, "proxy_server_request": {"headers": {}}} + ) + + assert root_attrs[TRACE_NAME_ATTR] == "from-body" + assert generation_attrs[TRACE_NAME_ATTR] == "from-body" + + +def test_unnamed_request_leaves_the_trace_name_off_both_spans(): + logger, exporter = _logger() + + root_attrs, generation_attrs = _run_named_request(logger, exporter, {"proxy_server_request": {"headers": {}}}) + + assert TRACE_NAME_ATTR not in root_attrs and TRACE_NAME_ATTR not in generation_attrs + + @pytest.mark.parametrize( ("capture", "mappers"), [("no_content", ("genai", "langfuse")), ("span_only", ("genai",))], diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index addadf8e598..8baf9310538 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -28,6 +28,7 @@ from litellm.integrations.otel import ( ) from litellm.integrations.otel.mappers.genai import GenAIMapper from litellm.integrations.otel.model import spans as spans_mod +from litellm.integrations.otel.model.metadata import LLMCallEvent, caller_trace_name from litellm.integrations.otel.model.payloads import ( LLMCallSpanData, RequestIdentity, @@ -722,6 +723,37 @@ def test_request_identity_falls_back_to_legacy_team_keys(): assert ident.team_alias == "legacy" +@pytest.mark.parametrize( + ("request_data", "expected"), + [ + ({"proxy_server_request": {"headers": {"langfuse_trace_name": "from-header"}}}, "from-header"), + ({"metadata": {"trace_name": "from-body"}}, "from-body"), + ({"litellm_metadata": {"trace_name": "from-anthropic-body"}}, "from-anthropic-body"), + ( + { + "proxy_server_request": {"headers": {"langfuse_trace_name": "from-header"}}, + "metadata": {"trace_name": "from-body"}, + }, + "from-header", + ), + ({"proxy_server_request": {"headers": {"langfuse_trace_name": ""}}, "metadata": {"trace_name": "body"}}, "body"), + ({"proxy_server_request": {"headers": {}}, "metadata": {"user_api_key_team_id": "t1"}}, None), + ({}, None), + ], + ids=["header", "body", "anthropic-body", "header-beats-body", "blank-header-falls-through", "neither", "empty"], +) +def test_caller_trace_name_prefers_the_langfuse_header_over_body_metadata(request_data, expected): + assert caller_trace_name({"litellm_params": request_data}) == expected + assert LLMCallEvent.from_dict({"litellm_params": request_data}).trace_name == expected + + +def test_llm_span_data_carries_the_caller_trace_name(): + data: Final = LLMCallSpanData.from_standard_logging_payload(_sample_payload(), trace_name="nightly-eval") + + assert data.trace_name == "nightly-eval" + assert LLMCallSpanData.from_standard_logging_payload(_sample_payload()).trace_name is None + + def test_llm_span_carries_proxy_request_route(): """The LLM span records the proxy route the request arrived on, so it can be filtered by endpoint (``/v1/responses`` vs ``/v1/chat/completions``) without diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index 94cb79f53b8..bcdda93383a 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -134,6 +134,11 @@ def test_langfuse_mapper_observation_attrs(): assert attrs["langfuse.trace.metadata.team_id"] == "t1" +def test_langfuse_mapper_names_the_trace_from_the_caller(): + assert LangfuseMapper().map(_llm_call(trace_name="nightly-eval"))["langfuse.trace.name"] == "nightly-eval" + assert "langfuse.trace.name" not in LangfuseMapper().map(_llm_call(trace_name=None)) + + def test_langfuse_mapper_skips_when_no_messages(): data = _llm_call(messages_in=(), choices_out=()) attrs = LangfuseMapper().map(data) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1644d78ae37..bb4822eae57 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,5 @@ import asyncio +import datetime as dt from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock @@ -9,9 +10,16 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks, Mode -from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail +from litellm.types.utils import ( + Choices, + GenericGuardrailAPIInputs, + GuardrailTracingDetail, + Message, + ModelResponse, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -2345,6 +2353,57 @@ class TestUndecoratedApplyGuardrailIsLogged: assert _Labelled.seen_label == "docs-style" + @pytest.mark.asyncio + async def test_post_call_recorded_outside_decorator_reaches_standard_logging_object(self): + """LIT-7608 regression: the auto-wrapped pre_call apply_guardrail copies the request bucket + into logging_obj.litellm_params["metadata"]. A post_call entry recorded later without the + decorator (the Bedrock streaming hook) must not be shadowed by that stale copy.""" + messages: Final = [{"role": "user", "content": "hello there"}] + litellm_metadata: Final[dict] = {"user_api_key_user_id": "u1"} + logging_obj: Final = Logging( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=messages, + stream=True, + call_type=CallTypes.acompletion.value, + start_time=dt.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + litellm_params={"litellm_metadata": litellm_metadata}, + optional_params={}, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + ) + request_data: Final = { + "model": "bedrock-haiku", + "messages": messages, + "litellm_metadata": litellm_metadata, + "litellm_logging_obj": logging_obj, + } + guardrail: Final = _UndecoratedGuardrail(guardrail_name="bedrock-pre", event_hook=GuardrailEventHooks.pre_call) + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello there"]), + request_data=request_data, + input_type="request", + logging_obj=logging_obj, + ) + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + event_type=GuardrailEventHooks.post_call, + ) + await logging_obj.async_success_handler( + result=ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))]), + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + entries: Final = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_mode"] for e in entries] == ["pre_call", "post_call"] + class _ApplyOnlyObserver(CustomGuardrail): """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" diff --git a/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py b/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py new file mode 100644 index 00000000000..e49e6baa8b0 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py @@ -0,0 +1,289 @@ +""" +Post-request budget gauges read the key/team/user/org state auth already resolved +from request metadata. get_*_object only runs when that state is missing (custom +auth, SDK callers, failure paths) +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from prometheus_client import REGISTRY + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_OrganizationTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.carried_budget_state import ( + carried_budget_metadata, + carry_organization_budget_state, + carry_team_and_user_budget_state, +) +from litellm.types.proxy.carried_budget_state import ( + KeyBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) + +TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) +USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc) +KEY_RESET_AT = datetime(2026, 12, 1, tzinfo=timezone.utc) + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + yield + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +@pytest.fixture +def prometheus_logger(): + return PrometheusLogger() + + +@pytest.fixture +def getters(): + """Every response-path object getter, patched where prometheus imports them from.""" + mocks = { + "get_key_object": AsyncMock(return_value=UserAPIKeyAuth(token="hashed", budget_reset_at=KEY_RESET_AT)), + "get_team_object": AsyncMock( + return_value=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT, max_budget=300.0) + ), + "get_user_object": AsyncMock( + return_value=LiteLLM_UserTable( + user_id="u1", + budget_reset_at=USER_RESET_AT, + user_email="alice@example.com", + user_alias="Alice", + max_budget=50.0, + ) + ), + "get_org_object": AsyncMock( + return_value=LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=40.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=500.0), + ) + ), + } + with ( + patch.multiple( # test-quality-ok: prometheus reads these proxy_server globals at call time, no injection seam + "litellm.proxy.proxy_server", prisma_client=MagicMock(), user_api_key_cache=MagicMock() + ), + patch.multiple( # test-quality-ok: the getters are the DB boundary this test counts calls to + "litellm.proxy.auth.auth_checks", **mocks + ), + ): + yield mocks + + +def _authed_token() -> UserAPIKeyAuth: + token = UserAPIKeyAuth( + token="hashed", + key_alias="key-alias", + team_id="t1", + team_alias="team-alias", + user_id="u1", + user_email="alice@example.com", + org_id="o1", + spend=1.0, + max_budget=10.0, + team_spend=20.0, + team_max_budget=300.0, + user_spend=5.0, + user_max_budget=50.0, + budget_reset_at=KEY_RESET_AT, + ) + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT, max_budget=300.0), + user_object=LiteLLM_UserTable(user_id="u1", budget_reset_at=USER_RESET_AT, user_alias="Alice", max_budget=50.0), + ) + carry_organization_budget_state( + valid_token=token, + org_table=LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=40.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=500.0), + ), + ) + return token + + +def _request_metadata(token: UserAPIKeyAuth) -> dict: + """What add_user_api_key_auth_to_request_metadata leaves in litellm_params["metadata"].""" + return { + **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token), + **carried_budget_metadata(token), + } + + +def _stub_gauges(prometheus_logger: PrometheusLogger) -> None: + for name in ( + "litellm_remaining_api_key_budget_metric", + "litellm_api_key_max_budget_metric", + "litellm_api_key_budget_remaining_hours_metric", + "litellm_remaining_team_budget_metric", + "litellm_team_max_budget_metric", + "litellm_team_budget_remaining_hours_metric", + "litellm_remaining_user_budget_metric", + "litellm_user_max_budget_metric", + "litellm_user_budget_remaining_hours_metric", + "litellm_remaining_org_budget_metric", + "litellm_org_max_budget_metric", + "litellm_org_budget_remaining_hours_metric", + ): + setattr(prometheus_logger, name, MagicMock()) + + +async def _emit(prometheus_logger: PrometheusLogger, metadata: dict) -> None: + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team="t1", + user_api_team_alias="team-alias", + user_api_key="hashed", + user_api_key_alias="key-alias", + litellm_params={"metadata": metadata}, + response_cost=2.0, + user_id="u1", + user_api_key_org_id="o1", + ) + + +@pytest.mark.asyncio +async def test_authed_request_sets_every_gauge_without_any_object_getter(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _request_metadata(_authed_token())) + + assert all(getter.await_count == 0 for getter in getters.values()), { + name: getter.await_count for name, getter in getters.items() + } + remaining = { + "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], + "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], + "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], + "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], + } + assert remaining == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(500.0) + prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.assert_called_once() + prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.assert_called_once() + prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() + + +@pytest.mark.asyncio +async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, {"user_api_key_team_spend": 20.0, "user_api_key_team_max_budget": 300.0}) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 1, + } + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0) + prometheus_logger.litellm_team_budget_remaining_hours_metric.labels().set.assert_called_once() + + +@pytest.mark.asyncio +async def test_partial_carried_state_only_skips_the_carried_objects(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1") + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=TEAM_RESET_AT), + user_object=None, + ) + + await _emit(prometheus_logger, dict(carried_budget_metadata(token))) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 0, + "get_user_object": 1, + "get_org_object": 1, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_max_budget", [300.0, None], ids=["metadata has max_budget", "filled from object"]) +async def test_carried_objects_match_what_the_getters_would_have_produced( + prometheus_logger, getters, metadata_max_budget +): + metadata = _request_metadata(_authed_token()) + user_max_budget = 50.0 if metadata_max_budget is not None else None + + carried_team = await prometheus_logger._assemble_team_object( + team_id="t1", + team_alias="team-alias", + spend=20.0, + max_budget=metadata_max_budget, + response_cost=2.0, + carried=TeamBudgetSnapshot.from_metadata(metadata), + ) + fetched_team = await prometheus_logger._assemble_team_object( + team_id="t1", team_alias="team-alias", spend=20.0, max_budget=metadata_max_budget, response_cost=2.0 + ) + carried_user = await prometheus_logger._assemble_user_object( + user_id="u1", + spend=5.0, + max_budget=user_max_budget, + response_cost=2.0, + carried=UserBudgetSnapshot.from_metadata(metadata), + user_email="alice@example.com", + ) + fetched_user = await prometheus_logger._assemble_user_object( + user_id="u1", spend=5.0, max_budget=user_max_budget, response_cost=2.0 + ) + carried_key = await prometheus_logger._assemble_key_object( + user_api_key="hashed", + user_api_key_alias="key-alias", + key_max_budget=10.0, + key_spend=1.0, + response_cost=2.0, + carried=KeyBudgetSnapshot.from_metadata(metadata), + ) + fetched_key = await prometheus_logger._assemble_key_object( + user_api_key="hashed", user_api_key_alias="key-alias", key_max_budget=10.0, key_spend=1.0, response_cost=2.0 + ) + + assert carried_team == fetched_team + assert carried_team.max_budget == 300.0 + assert carried_user == fetched_user + assert carried_user.max_budget == 50.0 + assert carried_key == fetched_key + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 0, + } diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index eecd876219e..76efe9c8576 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -2,6 +2,7 @@ the detached pipeline's single attempt-row write, and the cache-first job lookup.""" import asyncio +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -421,6 +422,183 @@ class TestSurfaceNormalization: assert "previous_response_id" not in shadow_call assert "instructions" not in shadow_call + @pytest.mark.parametrize( + "call_type,search_params,model", + [ + ("completion", {"web_search_options": {}}, "anthropic/claude-fable-5"), + ( + "acompletion", + {"web_search_options": {"search_context_size": "high"}}, + "anthropic/claude-fable-5", + ), + ( + "acompletion", + {"tools": [{"type": "web_search_20260209", "name": "web_search"}]}, + "anthropic/claude-fable-5", + ), + ( + "anthropic_messages", + {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + "anthropic/claude-fable-5", + ), + ( + "anthropic_messages", + {"tools": [{"type": "web_search_20260209", "name": "web_search"}]}, + "anthropic/claude-fable-5", + ), + ( + "anthropic_messages", + {"tools": [{"name": "web_search"}]}, + "anthropic/claude-fable-5", + ), + ("aresponses", {"tools": [{"type": "web_search"}]}, "anthropic/claude-fable-5"), + ("responses", {"tools": [{"type": "web_search_preview"}]}, "anthropic/claude-fable-5"), + ("aresponses", {"tools": [{"type": "web_search_2025_08_26"}]}, "anthropic/claude-fable-5"), + ( + "responses", + {"tools": [{"type": "web_search_preview_2025_03_11"}]}, + "anthropic/claude-fable-5", + ), + ("aresponses", {"tools": [{"type": "web_search"}]}, "bedrock/us.anthropic.claude-fable-5"), + ( + "responses", + {"tools": [{"type": "web_search_preview"}]}, + "bedrock/us.anthropic.claude-fable-5", + ), + ( + "acompletion", + { + "tools": [ + {"type": "function", "function": {"name": "WebSearch", "parameters": {"type": "object"}}}, + {"type": "web_search_20260209", "name": "web_search"}, + ] + }, + "anthropic/claude-fable-5", + ), + ], + ids=[ + "chat-empty-options", + "chat-configured-options", + "chat-provider-transformed-tools", + "messages-native-search", + "messages-dated-search", + "messages-legacy-search-normalized", + "responses-search", + "responses-preview", + "responses-dated-search", + "responses-dated-preview", + "responses-bedrock-erases-search", + "responses-bedrock-erases-preview", + "chat-mixed-client-and-hosted-tools", + ], + ) + async def test_hosted_web_search_skips_shadow_calls_and_spend( + self, call_type: str, search_params: Mapping[str, object], model: str + ) -> None: + base_kwargs: Final = _success_kwargs(call_type=call_type, model=model) + is_chat: Final = call_type in ("completion", "acompletion") + is_responses: Final = call_type in ("responses", "aresponses") + hook_kwargs: Final = { + **base_kwargs, + "model": model, + "messages": "what is new" if is_responses else base_kwargs["messages"], + "standard_logging_object": { + **base_kwargs["standard_logging_object"], + "model_parameters": search_params if is_chat else {}, + }, + "litellm_params": { + **base_kwargs["litellm_params"], + "proxy_server_request": {"body": {} if is_chat else search_params}, + }, + } + prisma: Final = _prisma() + router: Final = _router() + counter: Final = {"spend:shadow_eval:job-1": 0.1, "spend:shadow_eval:job-2": 0.1} + logger: Final = _logger( + router=router, + prisma=prisma, + jobs=(_job(max_budget=0.2), _job(id="job-2", max_budget=0.2)), + counter_store=counter, + ) + + await logger.async_log_success_event( + hook_kwargs, RESPONSES_API_RESPONSE if is_responses else RESPONSE, None, None + ) + await _drain(logger) + + router.acompletion.assert_not_called() + prisma.db.litellm_shadowevalattempt.create.assert_not_called() + assert logger._test_funnel == [("job-1", "unjudgeable"), ("job-2", "unjudgeable")] + assert logger._job_starts == {} + assert logger._test_counter == {"spend:shadow_eval:job-1": 0.1, "spend:shadow_eval:job-2": 0.1} + + @pytest.mark.parametrize( + "call_type,tool_name", + [ + (call_type, tool_name) + for call_type in ("completion", "acompletion", "anthropic_messages", "responses", "aresponses") + for tool_name in ("WebSearch", "litellm_web_search", "web_search") + ], + ) + async def test_client_web_search_tools_remain_sampled(self, call_type: str, tool_name: str) -> None: + is_chat: Final = call_type in ("completion", "acompletion") + is_responses: Final = call_type in ("responses", "aresponses") + tool: Final = ( + {"type": "function", "function": {"name": tool_name, "parameters": {"type": "object"}}} + if is_chat + else {"type": "function", "name": tool_name, "parameters": {"type": "object"}} + if is_responses + else {"name": tool_name, "input_schema": {"type": "object", "properties": {}}} + ) + source: Final = {"tools": [tool], "web_search_options": None} + base_kwargs: Final = _success_kwargs(call_type=call_type) + hook_kwargs: Final = { + **base_kwargs, + "messages": "search for current news" if is_responses else base_kwargs["messages"], + "standard_logging_object": { + **base_kwargs["standard_logging_object"], + "model_parameters": source if is_chat else {}, + }, + "litellm_params": { + **base_kwargs["litellm_params"], + "proxy_server_request": {"body": {} if is_chat else source}, + }, + } + + prisma, router = await self._drive(hook_kwargs, RESPONSES_API_RESPONSE if is_responses else RESPONSE) + + assert router.acompletion.call_count == 2 + shadow_call: Final = router.acompletion.call_args_list[0].kwargs + assert shadow_call["tools"][0]["function"]["name"] == tool_name + assert "web_search_options" not in shadow_call + prisma.db.litellm_shadowevalattempt.create.assert_called_once() + + @pytest.mark.parametrize("call_type", ["completion", "acompletion"]) + async def test_chat_search_removed_by_guardrail_still_samples(self, call_type: str) -> None: + base_kwargs: Final = _success_kwargs( + call_type=call_type, + request_metadata={ + "standard_logging_guardrail_information": [{"guardrail_name": "g", "guardrail_mode": "pre_call"}] + }, + ) + hook_kwargs: Final = { + **base_kwargs, + "litellm_params": { + **base_kwargs["litellm_params"], + "proxy_server_request": { + "body": {"web_search_options": {}, "tools": [{"type": "web_search_20260209"}]} + }, + }, + } + + prisma, router = await self._drive(hook_kwargs, RESPONSE) + + shadow_call: Final = router.acompletion.call_args_list[0].kwargs + assert "web_search_options" not in shadow_call + assert "tools" not in shadow_call + assert router.acompletion.call_count == 2 + prisma.db.litellm_shadowevalattempt.create.assert_called_once() + @pytest.mark.parametrize("payload_shape", ["typed", "dict"]) @pytest.mark.parametrize("call_type", ["aresponses", "responses"]) async def test_responses_arms_normalize_bare_string_input_and_instructions(self, call_type, payload_shape): diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 18309595414..83201aef143 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -1382,6 +1382,52 @@ def test_current_content_block_type_tracking(): assert iterator.current_content_block_type is None +def test_web_search_calls_are_cumulative_through_incomplete_search(): + iterator = ModelResponseIterator(None, sync_stream=True) + first_start = iterator.chunk_parser( + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "server_tool_use", + "id": "srvtoolu_A", + "name": "web_search", + "input": {"query": "a"}, + }, + } + ) + first_result = iterator.chunk_parser( + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_A", + "content": [], + }, + } + ) + second_start = iterator.chunk_parser( + { + "type": "content_block_start", + "index": 2, + "content_block": { + "type": "server_tool_use", + "id": "srvtoolu_B", + "name": "web_search", + "input": {"query": "b"}, + }, + } + ) + + assert list(first_start.choices[0].delta.provider_specific_fields["web_search_calls"]) == ["srvtoolu_A"] + assert first_result.choices[0].delta.provider_specific_fields["web_search_calls"]["srvtoolu_A"].status == "completed" + calls = second_start.choices[0].delta.provider_specific_fields["web_search_calls"] + assert list(calls) == ["srvtoolu_A", "srvtoolu_B"] + assert calls["srvtoolu_A"].status == "completed" + assert calls["srvtoolu_B"].status == "in_progress" + + def test_web_search_tool_result_captured_in_provider_specific_fields(): """ Test that web_search_tool_result content is captured in provider_specific_fields. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index f6ee1cd71c0..03b9840b1c3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -3358,6 +3358,27 @@ def test_is_web_search_tool(): assert adapter._is_web_search_tool(regular_tool) is False +@pytest.mark.parametrize("schema", [{}, {"type": "object", "properties": {"query": {"type": "string"}}}]) +def test_translate_anthropic_client_web_search_preserves_schema_and_choice(schema: dict[str, object]) -> None: + from litellm.types.llms.anthropic import AnthropicMessagesRequest + + request: Final = AnthropicMessagesRequest( + model="gpt-5.4-mini", + max_tokens=128, + messages=[{"role": "user", "content": "Search for current news"}], + tools=[{"name": "web_search", "input_schema": schema}], + tool_choice={"type": "tool", "name": "web_search"}, + ) + + translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(request) + + assert "web_search_options" not in translated + assert translated["tools"] == [ + {"type": "function", "function": {"name": "web_search", "parameters": schema}} + ] + assert translated["tool_choice"] == {"type": "function", "function": {"name": "web_search"}} + + def test_translate_anthropic_to_openai_with_web_search_tool(): """ Test that Anthropic web search tools are converted to web_search_options parameter. diff --git a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py b/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py index 3f98e9b6a2d..2f457fcb25b 100644 --- a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py +++ b/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py @@ -1,9 +1,5 @@ -import base64 -import json - import pytest -import litellm from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.azure_ai.ocr.document_intelligence.transformation import AzureDocumentIntelligenceOCRConfig @@ -12,27 +8,6 @@ from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig MODEL = "azure_ai/Cohere-parse-v5" API_BASE = "https://resource.services.ai.azure.com" PARSE_URL = f"{API_BASE}/providers/cohere/v2/parse" -IMAGE_URL = "https://example.com/receipt.png" -PNG_BYTES = base64.b64decode( - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" -) -PNG_DATA_URI = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}" - - -def _parse_response() -> dict: - return { - "id": "882bf973-9dfa-4d02-9d30-709247008efd", - "pages": [{"index": 0, "type": "markdown", "markdown": {"content": "# Receipt\n\nTotal Due: $4.00"}}], - "meta": {"api_version": {"version": "2"}, "billed_units": {"pages": 1}}, - } - - -@pytest.fixture() -def disable_aiohttp_transport(monkeypatch): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - litellm.in_memory_llm_clients_cache.flush_cache() - yield - litellm.in_memory_llm_clients_cache.flush_cache() @pytest.mark.parametrize( @@ -95,90 +70,3 @@ def test_validate_environment_requires_api_base(monkeypatch) -> None: with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): AzureAICohereParseConfig().validate_environment(headers={}, model="Cohere-parse-v5", api_key="key") - - -@pytest.mark.asyncio -async def test_aocr_inlines_remote_image_and_posts_to_foundry(disable_aiohttp_transport, respx_mock): - respx_mock.get(IMAGE_URL).respond(content=PNG_BYTES, headers={"Content-Type": "image/png"}) - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - response = await litellm.aocr( - model=MODEL, - document={"type": "image_url", "image_url": IMAGE_URL}, - api_base=API_BASE, - api_key="azure-key", - ) - - request = route.calls.last.request - assert request.headers["Authorization"] == "Bearer azure-key" - assert json.loads(request.content) == { - "model": "Cohere-parse-v5", - "document": {"type": "image_url", "image_url": PNG_DATA_URI}, - "output_format": "markdown", - } - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - assert response.usage_info.pages_processed == 1 - - -@pytest.mark.asyncio -async def test_aocr_passes_data_uri_through_without_fetching(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - await litellm.aocr( - model=MODEL, - document={"type": "image_url", "image_url": PNG_DATA_URI}, - api_base=API_BASE, - api_key="azure-key", - output_format="blocks", - ) - - body = json.loads(route.calls.last.request.content) - assert body["document"]["image_url"] == PNG_DATA_URI - assert body["output_format"] == "blocks" - - -def test_ocr_sync_inlines_remote_image(respx_mock): - respx_mock.get(IMAGE_URL).respond(content=PNG_BYTES, headers={"Content-Type": "image/png"}) - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - response = litellm.ocr( - model=MODEL, - document={"type": "image_url", "image_url": IMAGE_URL}, - api_base=API_BASE, - api_key="azure-key", - ) - - assert json.loads(route.calls.last.request.content)["document"]["image_url"] == PNG_DATA_URI - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - - -@pytest.mark.asyncio -async def test_aocr_rejects_pdf_before_calling_foundry(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - with pytest.raises(litellm.BadRequestError, match="only accepts `image_url` documents") as exc_info: - await litellm.aocr( - model=MODEL, - document={"type": "document_url", "document_url": "https://example.com/doc.pdf"}, - api_base=API_BASE, - api_key="azure-key", - ) - - assert exc_info.value.llm_provider == "azure_ai" - assert not route.called - - -@pytest.mark.asyncio -async def test_ahealth_check_ocr_sends_an_image_to_the_foundry_cohere_parse_deployment( - disable_aiohttp_transport, respx_mock -): - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - result = await litellm.ahealth_check( - model_params={"model": MODEL, "api_base": API_BASE, "api_key": "test-key"}, mode="ocr" - ) - - document = json.loads(route.calls.last.request.content)["document"] - assert document["type"] == "image_url" - assert document["image_url"].startswith("data:image/png;base64,") - assert "error" not in result diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py index ef4c78553f1..be0dfb5724e 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -1,4 +1,5 @@ from unittest.mock import MagicMock +from typing import Final import httpx import pytest @@ -371,3 +372,35 @@ def test_validate_environment_falls_back_to_entra_token(monkeypatch): assert headers["Authorization"] == "Bearer entra-token" assert "Ocp-Apim-Subscription-Key" not in headers + + +@pytest.mark.parametrize( + ("request_headers", "expected_poll_headers"), + ( + ( + {"Ocp-Apim-Subscription-Key": "subscription-key"}, + {"Ocp-Apim-Subscription-Key": "subscription-key"}, + ), + ( + {"Authorization": "Bearer entra-token"}, + {"Authorization": "Bearer entra-token"}, + ), + ), +) +def test_get_polling_target_preserves_request_authentication( + request_headers: dict[str, str], expected_poll_headers: dict[str, str] +) -> None: + response: Final = httpx.Response( + status_code=202, + headers={"Operation-Location": "https://example.cognitiveservices.azure.com/operations/123"}, + request=httpx.Request( + "POST", + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze", + headers=request_headers, + ), + ) + + operation_url, poll_headers = AzureDocumentIntelligenceOCRConfig()._get_polling_target(response) + + assert operation_url == "https://example.cognitiveservices.azure.com/operations/123" + assert poll_headers == expected_poll_headers diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 9457d5faaff..40566261c84 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -539,6 +539,90 @@ class TestBedrockMantleServiceTier: assert "priority" in str(mock_warning.call_args) +class TestBedrockMantleReasoningSummary: + @pytest.mark.parametrize("summary", ["concise", "detailed"]) + def test_unsupported_reasoning_summary_dropped_when_drop_params_true(self, summary): + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": summary}}, + model="openai.gpt-5.6-sol", + drop_params=True, + ) + assert params["reasoning"] == {"effort": "medium"} + + def test_reasoning_summary_only_field_drops_reasoning(self): + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"reasoning": {"summary": "detailed"}}, + model="openai.gpt-5.6-sol", + drop_params=True, + ) + assert "reasoning" not in params + + @pytest.mark.parametrize("summary", ["concise", "detailed"]) + def test_unsupported_reasoning_summary_raises_when_drop_params_false(self, summary): + cfg = BedrockMantleResponsesAPIConfig() + with pytest.raises(litellm.UnsupportedParamsError) as excinfo: + cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": summary}}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert summary in str(excinfo.value) + assert "reasoning.summary" in str(excinfo.value) + assert "drop_params" in str(excinfo.value) + + def test_unhashable_reasoning_summary_raises_unsupported_params_error(self): + cfg = BedrockMantleResponsesAPIConfig() + with pytest.raises(litellm.UnsupportedParamsError) as excinfo: + cfg.map_openai_params( + response_api_optional_params={"reasoning": {"summary": ["detailed"]}}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert "reasoning.summary" in str(excinfo.value) + + @pytest.mark.parametrize("drop_params", [True, False]) + def test_supported_reasoning_summary_kept(self, drop_params): + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "auto"}}, + model="openai.gpt-5.6-sol", + drop_params=drop_params, + ) + assert params["reasoning"] == {"effort": "medium", "summary": "auto"} + + def test_reasoning_summary_kept_on_standard_path(self): + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + params = cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "detailed"}}, + model="openai.gpt-oss-120b", + drop_params=False, + ) + assert params["reasoning"] == {"effort": "medium", "summary": "detailed"} + + def test_absent_reasoning_untouched(self): + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"stream": True}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert params == {"stream": True} + + def test_drop_logged_at_warning_level(self, caplog): + cfg = BedrockMantleResponsesAPIConfig() + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "detailed"}}, + model="openai.gpt-5.6-sol", + drop_params=True, + ) + warnings = [record for record in caplog.records if "dropping unsupported reasoning.summary" in record.getMessage()] + assert len(warnings) == 1 + assert "detailed" in warnings[0].getMessage() + + class TestBedrockMantleCodexRequestEndToEnd: def test_codex_priority_tier_request_becomes_mantle_acceptable(self): cfg = BedrockMantleResponsesAPIConfig() diff --git a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py index cb9af56f5e0..1f120be6ffa 100644 --- a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py +++ b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py @@ -1,8 +1,11 @@ -import json +from typing import Final +from unittest.mock import Mock +import httpx import pytest import litellm +from litellm.llms.cohere.ocr.transformation import CohereParseConfig PARSE_URL = "https://api.cohere.com/v2/parse" MODEL = "cohere/parse-v5.0" @@ -57,173 +60,38 @@ def _blocks_response() -> dict: } -@pytest.fixture() -def disable_aiohttp_transport(monkeypatch): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - litellm.in_memory_llm_clients_cache.flush_cache() - yield - litellm.in_memory_llm_clients_cache.flush_cache() - - -@pytest.mark.asyncio -async def test_aocr_sends_markdown_parse_request_and_normalizes_pages(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - request = route.calls.last.request - assert request.headers["Authorization"] == "Bearer test-key" - assert json.loads(request.content) == { - "model": "parse-v5.0", - "document": IMAGE_DOCUMENT, - "output_format": "markdown", - } - assert response.object == "ocr" - assert [page.index for page in response.pages] == [0, 1] - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - assert response.pages[1].markdown == "Page two" - assert response.pages[1].images is None - image = response.pages[0].images[0] - assert image.bbox == BOUNDING_BOX - assert image.model_extra["description"] == "A parking receipt" - assert image.model_extra["bounding_box_normalized"]["bottom_right_x"] == 1 - assert response.usage_info.pages_processed == 2 - assert response.get_provider_native_response() is None - - -@pytest.mark.asyncio -async def test_aocr_usage_prefers_billed_units_over_page_count(disable_aiohttp_transport, respx_mock): - respx_mock.post(PARSE_URL).respond(json=_markdown_response(billed_pages=3)) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - assert response.usage_info.pages_processed == 3 - - -@pytest.mark.asyncio -async def test_aocr_usage_falls_back_to_page_count_without_meta(disable_aiohttp_transport, respx_mock): - respx_mock.post(PARSE_URL).respond(json=_markdown_response(billed_pages=None)) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - assert response.usage_info.pages_processed == 2 - - -@pytest.mark.asyncio -async def test_aocr_blocks_output_format_forwards_param_and_keeps_blocks(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_blocks_response()) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", output_format="blocks") - - assert json.loads(route.calls.last.request.content)["output_format"] == "blocks" - assert response.pages[0].markdown == "" - assert response.pages[0].model_extra["blocks"] == [{"type": "text", "text": "Total Due: $4.00"}] - assert response.usage_info.pages_processed == 1 - - -@pytest.mark.asyncio -async def test_aocr_native_format_carries_provider_payload(disable_aiohttp_transport, respx_mock): - payload = _markdown_response() - route = respx_mock.post(PARSE_URL).respond(json=payload) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", req_format="native") - - assert "req_format" not in json.loads(route.calls.last.request.content) - assert response.get_provider_native_response() == payload - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - - -@pytest.mark.asyncio -async def test_aocr_rejects_unknown_output_format_before_calling_provider(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - with pytest.raises(litellm.BadRequestError, match="Invalid `output_format`: 'html'") as exc_info: - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", output_format="html") - - assert exc_info.value.status_code == 400 - assert not route.called - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "document", - [ - {"type": "document_url", "document_url": "https://example.com/doc.pdf"}, - {"type": "image_url", "image_url": "data:application/pdf;base64,JVBERi0="}, - {"type": "image_url", "image_url": ""}, - ], -) -async def test_aocr_rejects_non_image_documents_before_calling_provider( - disable_aiohttp_transport, respx_mock, document -): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - with pytest.raises(litellm.BadRequestError, match="only accepts `image_url` documents") as exc_info: - await litellm.aocr(model=MODEL, document=document, api_key="test-key") - - assert exc_info.value.status_code == 400 - assert not route.called - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "api_base, expected_url", - [ - ("https://gateway.example.com", "https://gateway.example.com/v2/parse"), - ("https://gateway.example.com/cohere/", "https://gateway.example.com/cohere/v2/parse"), - ("https://gateway.example.com/v2", "https://gateway.example.com/v2/parse"), - ("https://gateway.example.com/v2/parse", "https://gateway.example.com/v2/parse"), - ], -) -async def test_aocr_posts_to_api_base_variants(disable_aiohttp_transport, respx_mock, api_base, expected_url): - route = respx_mock.post(expected_url).respond(json=_markdown_response()) - - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", api_base=api_base) - - assert route.called - - -@pytest.mark.asyncio -async def test_aocr_surfaces_provider_error_with_its_status_and_message(disable_aiohttp_transport, respx_mock): - respx_mock.post(PARSE_URL).respond( - status_code=400, json={"id": "83b0d95e", "message": "output_format must be `blocks` or `markdown`"} +@pytest.mark.parametrize("output_format", ["markdown", "blocks"]) +def test_transform_cohere_request_filters_options(output_format: str) -> None: + config: Final = CohereParseConfig() + params: Final = config.map_ocr_params( + {"output_format": output_format, "req_format": "native", "unknown": True}, {}, "parse-v5.0" ) - - with pytest.raises(litellm.BadRequestError, match="output_format must be") as exc_info: - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - assert exc_info.value.status_code == 400 + request: Final = config.transform_ocr_request("parse-v5.0", IMAGE_DOCUMENT, params, {}) + assert request.data == {"model": "parse-v5.0", "document": IMAGE_DOCUMENT, "output_format": output_format} -@pytest.mark.asyncio -async def test_aocr_reads_api_key_from_environment(disable_aiohttp_transport, respx_mock, monkeypatch): - monkeypatch.setenv("COHERE_API_KEY", "env-key") - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT) - - assert route.calls.last.request.headers["Authorization"] == "Bearer env-key" +@pytest.mark.parametrize("native", [False, True]) +def test_transform_cohere_response_keeps_images_and_native_payload(native: bool) -> None: + payload: Final = _markdown_response(3) + response: Final = CohereParseConfig().transform_ocr_response( + "parse-v5.0", httpx.Response(200, json=payload), Mock(), {"req_format": "native" if native else "litellm"} + ) + assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" + assert response.pages[0].images[0].bbox == BOUNDING_BOX + assert response.pages[0].images[0].model_extra["description"] == "A parking receipt" + assert response.pages[1].images is None + assert response.usage_info.pages_processed == 3 + assert response.get_provider_native_response() == (payload if native else None) -@pytest.mark.asyncio -async def test_aocr_without_api_key_names_the_env_var(disable_aiohttp_transport, respx_mock, monkeypatch): - monkeypatch.delenv("COHERE_API_KEY", raising=False) - monkeypatch.setattr(litellm, "cohere_key", None) - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - with pytest.raises(Exception, match="Missing COHERE_API_KEY"): - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT) - - assert not route.called +def test_transform_cohere_blocks() -> None: + response: Final = CohereParseConfig().transform_ocr_response( + "parse-v5.0", httpx.Response(200, json=_blocks_response()), Mock() + ) + assert response.pages[0].model_extra["blocks"] == [{"type": "text", "text": "Total Due: $4.00"}] + assert response.pages[0].markdown == "" -@pytest.mark.asyncio -async def test_ahealth_check_ocr_sends_an_image_cohere_parse_accepts(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - result = await litellm.ahealth_check(model_params={"model": MODEL, "api_key": "test-key"}, mode="ocr") - - document = json.loads(route.calls.last.request.content)["document"] - assert document["type"] == "image_url" - assert document["image_url"].startswith("data:image/png;base64,") - assert "error" not in result +def test_transform_cohere_rejects_unsupported_output_format() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="output_format"): + CohereParseConfig().map_ocr_params({"output_format": "html"}, {}, "parse-v5.0") diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d4ef4282b27..8479397efa7 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -501,8 +501,8 @@ def test_unmapped_model_fallback_function_calling(): assert info["supports_function_calling"] is True -def test_transform_messages_helper_strips_thinking_blocks(): - """thinking_blocks must not be forwarded to Fireworks chat completions.""" +def test_transform_messages_helper_strips_thinking_blocks_but_keeps_reasoning_content(): + """Fireworks rejects thinking_blocks but requires reasoning_content to be replayed for reasoning_history.""" config = FireworksAIConfig() messages = [ {"role": "user", "content": "Translate a poem."}, @@ -519,7 +519,7 @@ def test_transform_messages_helper_strips_thinking_blocks(): messages, model="accounts/fireworks/models/glm-5p1", litellm_params={} ) assert "thinking_blocks" not in out[1] - assert "reasoning_content" not in out[1] + assert out[1]["reasoning_content"] == "internal" assert out[1]["content"] == "I can help." diff --git a/tests/test_litellm/llms/reducto/conftest.py b/tests/test_litellm/llms/reducto/conftest.py new file mode 100644 index 00000000000..4ff3ab43006 --- /dev/null +++ b/tests/test_litellm/llms/reducto/conftest.py @@ -0,0 +1,11 @@ +from collections.abc import Generator + +import pytest + +from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service + + +@pytest.fixture +def reducto_server() -> Generator[RecordingServer]: + with recording_service() as server: + yield server diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/test_litellm/llms/reducto/test_parse_legacy.py index db19460baa3..252369cbd3d 100644 --- a/tests/test_litellm/llms/reducto/test_parse_legacy.py +++ b/tests/test_litellm/llms/reducto/test_parse_legacy.py @@ -1,7 +1,7 @@ -import json +import pytest import litellm -import pytest +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec @pytest.fixture() @@ -17,24 +17,28 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_legacy_wraps_enhance_under_options( - disable_aiohttp_transport, respx_mock -): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://legacy.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json={ - "usage": {"num_pages": 1, "credits": 1}, - "result": { - "chunks": [ - { - "content": "Legacy parse", - "blocks": [{"content": "Legacy parse", "bbox": {"page": 1}}], - } - ] - }, - } +async def test_parse_legacy_wraps_enhance_under_options(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 2 + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://legacy.pdf"})) + reducto_server.enqueue( + ResponseSpec( + body={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Legacy parse", + "blocks": [ + { + "content": "Legacy parse", + "bbox": {"page": 1}, + } + ], + } + ] + }, + } + ) ) response = await litellm.aocr( @@ -45,13 +49,15 @@ async def test_parse_legacy_wraps_enhance_under_options( "mime_type": "application/pdf", }, api_key="legacy-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, enhance={"agentic": [{"type": "table"}]}, ) - assert upload_route.called - assert parse_route.called - request_body = json.loads(parse_route.calls[0].request.read()) + upload_request, parse_request = reducto_server.requests + assert upload_request.path == "/upload" + assert parse_request.path == "/parse" + assert isinstance(parse_request.body, dict) + request_body = parse_request.body assert request_body == { "document_url": "reducto://legacy.pdf", "options": {"enhance": {"agentic": [{"type": "table"}]}}, diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/test_litellm/llms/reducto/test_parse_v3.py index 1d0c826ef8b..0ebc0d926c4 100644 --- a/tests/test_litellm/llms/reducto/test_parse_v3.py +++ b/tests/test_litellm/llms/reducto/test_parse_v3.py @@ -1,8 +1,7 @@ -import json - import pytest import litellm +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec def _reducto_parse_response() -> dict: @@ -69,11 +68,11 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transport, respx_mock): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://uploaded.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) +async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 2 + provider_response = _reducto_parse_response() + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded.pdf"})) + reducto_server.enqueue(ResponseSpec(body=provider_response)) response = await litellm.aocr( model="reducto/parse-v3", @@ -83,25 +82,24 @@ async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transpo "mime_type": "application/pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, formatting={"table_output_format": "html"}, retrieval={"chunk_mode": "section"}, settings={"ocr_system": "standard"}, + req_format="native", ) - assert upload_route.called - assert parse_route.called - assert len(upload_route.calls) == 1 - assert len(parse_route.calls) == 1 - - upload_request = upload_route.calls[0].request + upload_request, parse_request = reducto_server.requests + assert upload_request.path == "/upload" + assert parse_request.path == "/parse" assert upload_request.headers["authorization"] == "Bearer test-key" assert "application/json" not in upload_request.headers["content-type"] - upload_body = upload_request.read() + upload_body = upload_request.raw_body assert b'filename="document"' in upload_body assert b"application/pdf" in upload_body - parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert isinstance(parse_request.body, dict) + parse_request_body = parse_request.body assert parse_request_body["input"] == "reducto://uploaded.pdf" assert parse_request_body["formatting"] == {"table_output_format": "html"} assert parse_request_body["retrieval"] == {"chunk_mode": "section"} @@ -116,15 +114,12 @@ async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transpo assert getattr(response.pages[0], "blocks")[0]["bbox"]["page"] == 1 assert response.pages[1].markdown == "Page 2 block A" assert response.pages[2].markdown == "Page 3 block A" - assert response._hidden_params["reducto_raw"]["usage"]["credits"] == 3 + assert response.get_provider_native_response() == provider_response @pytest.mark.asyncio -async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_transport, respx_mock): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://should-not-upload.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) +async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.enqueue(ResponseSpec(body=_reducto_parse_response())) response = await litellm.aocr( model="reducto/parse-v3", @@ -133,13 +128,15 @@ async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_tran "document_url": "reducto://already-uploaded.pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, retrieval={"chunk_mode": "section"}, ) - assert not upload_route.called - assert parse_route.called - parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert len(reducto_server.requests) == 1 + parse_request = reducto_server.requests[0] + assert parse_request.path == "/parse" + assert isinstance(parse_request.body, dict) + parse_request_body = parse_request.body assert parse_request_body["input"] == "reducto://already-uploaded.pdf" assert parse_request_body["retrieval"]["chunk_mode"] == "section" assert response.pages[0].markdown.startswith("Page 1 block A") @@ -147,11 +144,9 @@ async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_tran @pytest.mark.asyncio async def test_unknown_model_uses_current_protocol_without_local_rejection( - disable_aiohttp_transport, respx_mock + disable_aiohttp_transport, reducto_server: RecordingServer ): - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json=_reducto_parse_response() - ) + reducto_server.enqueue(ResponseSpec(body=_reducto_parse_response())) response = await litellm.aocr( model="reducto/future-parse-model", @@ -160,11 +155,9 @@ async def test_unknown_model_uses_current_protocol_without_local_rejection( "document_url": "reducto://already-uploaded.pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, ) - assert parse_route.called - assert json.loads(parse_route.calls[0].request.read()) == { - "input": "reducto://already-uploaded.pdf" - } + assert reducto_server.requests[0].path == "/parse" + assert reducto_server.requests[0].body == {"input": "reducto://already-uploaded.pdf"} assert response.model == "future-parse-model" diff --git a/tests/test_litellm/llms/reducto/test_upload.py b/tests/test_litellm/llms/reducto/test_upload.py index 4fae90436bb..adfc2663fb0 100644 --- a/tests/test_litellm/llms/reducto/test_upload.py +++ b/tests/test_litellm/llms/reducto/test_upload.py @@ -1,16 +1,16 @@ -import json import os from unittest.mock import AsyncMock, Mock import httpx -import litellm import pytest +import litellm from litellm.llms.reducto.common import ( extract_file_id_or_bytes, upload_bytes_async, upload_bytes_sync, ) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec @pytest.fixture() @@ -28,7 +28,8 @@ def disable_aiohttp_transport(monkeypatch): @pytest.mark.asyncio -async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport): +async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 0 with pytest.raises(litellm.BadRequestError, match="upload the file first"): await litellm.aocr( model="reducto/parse-v3", @@ -37,29 +38,30 @@ async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport): "document_url": "https://example.com/document.pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, ) @pytest.mark.asyncio async def test_parse_v3_image_data_uri_upload_uses_image_mime( - disable_aiohttp_transport, respx_mock + disable_aiohttp_transport, reducto_server: RecordingServer ): - upload_route = respx_mock.post("https://custom.reducto.test/upload").respond( - json={"file_id": "reducto://uploaded-image.png"} - ) - parse_route = respx_mock.post("https://custom.reducto.test/parse").respond( - json={ - "usage": {"num_pages": 1, "credits": 1}, - "result": { - "chunks": [ - { - "content": "Image OCR", - "blocks": [{"content": "Image OCR", "bbox": {"page": 1}}], - } - ] - }, - } + reducto_server.expected_requests = 2 + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded-image.png"})) + reducto_server.enqueue( + ResponseSpec( + body={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Image OCR", + "blocks": [{"content": "Image OCR", "bbox": {"page": 1}}], + } + ] + }, + } + ) ) response = await litellm.aocr( @@ -70,41 +72,43 @@ async def test_parse_v3_image_data_uri_upload_uses_image_mime( "mime_type": "image/png", }, api_key="programmatic-key", - api_base="https://custom.reducto.test/", + api_base=f"{reducto_server.base_url}/", ) - assert upload_route.called - assert parse_route.called - upload_request = upload_route.calls[0].request + upload_request, parse_request = reducto_server.requests + assert upload_request.path == "/upload" + assert parse_request.path == "/parse" assert upload_request.headers["authorization"] == "Bearer programmatic-key" - assert b"image/png" in upload_request.read() + assert b"image/png" in upload_request.raw_body - parse_request_body = json.loads(parse_route.calls[0].request.read()) - assert parse_request_body["input"] == "reducto://uploaded-image.png" + assert isinstance(parse_request.body, dict) + assert parse_request.body["input"] == "reducto://uploaded-image.png" assert response.pages[0].markdown == "Image OCR" @pytest.mark.asyncio -async def test_parse_v3_uses_programmatic_api_key_over_env( - disable_aiohttp_transport, respx_mock -): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://uploaded.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json={ - "usage": {"num_pages": 1, "credits": 1}, - "result": { - "chunks": [ - { - "content": "Programmatic auth", - "blocks": [ - {"content": "Programmatic auth", "bbox": {"page": 1}} - ], - } - ] - }, - } +async def test_parse_v3_uses_programmatic_api_key_over_env(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 2 + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded.pdf"})) + reducto_server.enqueue( + ResponseSpec( + body={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Programmatic auth", + "blocks": [ + { + "content": "Programmatic auth", + "bbox": {"page": 1}, + } + ], + } + ] + }, + } + ) ) await litellm.aocr( @@ -115,11 +119,11 @@ async def test_parse_v3_uses_programmatic_api_key_over_env( "mime_type": "application/pdf", }, api_key="passed-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, ) - assert upload_route.calls[0].request.headers["authorization"] == "Bearer passed-key" - assert parse_route.calls[0].request.headers["authorization"] == "Bearer passed-key" + assert reducto_server.requests[0].headers["authorization"] == "Bearer passed-key" + assert reducto_server.requests[1].headers["authorization"] == "Bearer passed-key" def test_upload_bytes_sync_uses_shared_client(monkeypatch): diff --git a/tests/test_litellm/ocr/test_legacy.py b/tests/test_litellm/ocr/test_legacy.py new file mode 100644 index 00000000000..a30976f89af --- /dev/null +++ b/tests/test_litellm/ocr/test_legacy.py @@ -0,0 +1,200 @@ +import importlib +from collections.abc import AsyncGenerator +from datetime import datetime +from io import BytesIO +from typing import Final +from unittest.mock import Mock + +import httpx +import orjson +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.llms.custom_httpx import llm_http_handler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.ocr.legacy import _prepare_ocr_request +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE + + +@pytest.fixture +async def provider(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[Mock]: + configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust must not load"))) + handler: Final = Mock( + return_value=httpx.Response( + 200, + json={ + "pages": [{"index": 0, "markdown": "parsed document"}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, + }, + ) + ) + transport: Final = httpx.MockTransport(handler) + with httpx.Client(transport=transport) as sync_client: + async with httpx.AsyncClient(transport=transport) as async_client: + sync_handler: Final = HTTPHandler(client=sync_client) + async_handler: Final = AsyncHTTPHandler() + await async_handler.client.aclose() + async_handler.client = async_client + monkeypatch.setattr(llm_http_handler, "_get_httpx_client", lambda: sync_handler) + monkeypatch.setattr(llm_http_handler, "get_async_httpx_client", lambda llm_provider: async_handler) + yield handler + NATIVE_OCR_LIFECYCLE.reset() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async", "sync_async"]) +@pytest.mark.parametrize("dispatch", ["disabled", "declined", "unavailable"]) +async def test_python_request_response_and_callbacks( + provider: Mock, monkeypatch: pytest.MonkeyPatch, mode: str, dispatch: str +) -> None: + class Declined(Exception): + pass + + if dispatch != "disabled": + monkeypatch.setenv("LITELLM_RUST", "1") + NATIVE_OCR_LIFECYCLE.override(Mock(side_effect=Declined()) if dispatch == "declined" else None) + main: Final = importlib.import_module("litellm.ocr.main") + monkeypatch.setattr(main, "native_exception_types", lambda: (Declined, RuntimeError)) + logger: Final = Mock(spec=CustomLogger) + monkeypatch.setattr(litellm, "input_callback", [logger]) + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "file", "file": BytesIO(b"pdf"), "mime_type": "application/pdf"}, + "api_key": "test-key", + "api_base": "https://ocr.test/v1", + "timeout": 7.0, + "pages": [0, 2], + "include_image_base64": True, + "extra_headers": {"x-test-header": "preserved"}, + } + + async def call() -> OCRResponse: + if mode == "async": + return await litellm.aocr(**arguments) + if mode == "sync_async": + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = Logging( + model=arguments["model"], + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.now(), + litellm_call_id="test-call", + function_id="test-function", + ) + return await litellm.ocr(**arguments, aocr=True, litellm_logging_obj=logging_obj) + return litellm.ocr(**arguments) + + response: Final = await call() + assert response.pages[0].markdown == "parsed document" + assert response.usage_info.pages_processed == 1 + assert provider.call_count == 1 + request: Final = provider.call_args.args[0] + assert str(request.url) == "https://ocr.test/v1/ocr" + assert request.headers["authorization"] == "Bearer test-key" + assert request.headers["x-test-header"] == "preserved" + assert request.extensions["timeout"] == {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0} + assert orjson.loads(request.content) == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"}, + "pages": [0, 2], + "include_image_base64": True, + } + assert logger.log_pre_api_call.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_python_provider_errors_keep_public_exception(provider: Mock, asynchronous: bool) -> None: + provider.return_value = httpx.Response(429, json={"error": "rate limited"}) + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/file.pdf"}, + "api_key": "test-key", + "api_base": "https://ocr.test/v1", + "num_retries": 0, + } + + async def call() -> object: + if asynchronous: + return await litellm.aocr(**arguments) + return litellm.ocr(**arguments) + + with pytest.raises(litellm.RateLimitError) as error: + await call() + assert error.value.status_code == 429 + assert error.value.model == "mistral-ocr-latest" + assert error.value.llm_provider == "mistral" + assert provider.call_count == 1 + + +def test_document_intelligence_environment_key_is_not_replaced_by_generic_azure_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key") + monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", "document-key") + monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", "https://document.example.com") + prepared: Final = _prepare_ocr_request( + model="azure_ai/doc-intelligence/prebuilt-layout", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key=None, + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": Mock()}, + ) + + assert prepared.api_key is None + headers: Final = prepared.provider_config.validate_environment( + headers={}, + model=prepared.model, + api_key=prepared.api_key, + api_base=prepared.api_base, + litellm_params=prepared.litellm_params, + ) + assert headers["Ocp-Apim-Subscription-Key"] == "document-key" + + +def test_document_intelligence_explicit_connection_is_preserved(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://generic.example.com") + prepared: Final = _prepare_ocr_request( + model="azure_ai/doc-intelligence/prebuilt-layout", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key="explicit-key", + api_base="https://document.example.com", + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": Mock()}, + ) + + assert prepared.api_key == "explicit-key" + assert prepared.api_base == "https://document.example.com" + + +def test_generic_azure_connection_still_applies_to_foundry_ocr(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://generic.example.com") + prepared: Final = _prepare_ocr_request( + model="azure_ai/mistral-document-ai-2505", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key=None, + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": Mock()}, + ) + + assert prepared.api_key == "generic-key" + assert prepared.api_base == "https://generic.example.com" diff --git a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py deleted file mode 100644 index 460aff3e8d1..00000000000 --- a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py +++ /dev/null @@ -1,73 +0,0 @@ -""" -Regression tests for Azure Document Intelligence api_base ownership in OCR. - -`azure_ai` exposes two OCR services on one provider; the `doc-intelligence` -sub-route must defer environment resolution to Rust, not accept the generic -`AZURE_AI_API_BASE` fallback that `get_llm_provider` injects. An explicitly -supplied api_base is still always honoured. -""" - -from litellm.llms.azure_ai.ocr.common_utils import ( - is_azure_document_intelligence_model, -) -from litellm.ocr.main import _prepare_ocr_request - -_DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -_AZURE_AI_API_BASE = "https://generic-azure-ai.example.com" - - -class _FakeLogging: - def update_from_kwargs(self, **kwargs: object) -> None: - return None - - -def _prepare(model: str, api_base: str | None): - return _prepare_ocr_request( - model=model, - document=dict(_DOC), - api_key="test-key", - api_base=api_base, - timeout=None, - custom_llm_provider=None, - extra_headers=None, - kwargs={"litellm_logging_obj": _FakeLogging()}, - ) - - -class TestIsAzureDocumentIntelligenceModel: - def test_matches_doc_intelligence_route(self): - assert is_azure_document_intelligence_model("doc-intelligence/prebuilt-layout") - - def test_matches_documentintelligence_and_is_case_insensitive(self): - assert is_azure_document_intelligence_model("azure_ai/DocumentIntelligence/x") - - def test_does_not_match_mistral_route(self): - assert not is_azure_document_intelligence_model("mistral-document-ai-2505") - - -class TestDocIntelligenceApiBaseResolution: - def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch): - """The generic Azure base must not overwrite Rust-owned DI resolution.""" - monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False) - - prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None) - - assert prepared.api_base is None - - def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch): - """A caller-supplied api_base must always win, even for doc-intelligence.""" - monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - - custom = "https://my-di.cognitiveservices.azure.com" - prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom) - - assert prepared.api_base == custom - - def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch): - """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE.""" - monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - - prepared = _prepare("azure_ai/mistral-document-ai-2505", None) - - assert prepared.api_base == _AZURE_AI_API_BASE diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index feb98d14c03..3526d8c00d6 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -12,15 +12,32 @@ Tests that: import base64 import os import tempfile +from collections.abc import Generator from io import BytesIO from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from typing import Final +from unittest.mock import AsyncMock, MagicMock, Mock import orjson import pytest from starlette.datastructures import FormData -from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type +from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type + + +@pytest.fixture(autouse=True, params=["native", "disabled", "unavailable"]) +def document_runtime(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + from litellm.rust_bridge import bindings, configuration + + configuration.reset_rust_configuration() + monkeypatch.delenv("LITELLM_RUST", raising=False) + if request.param == "disabled": + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust is disabled"))) + elif request.param == "unavailable": + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + yield + configuration.reset_rust_configuration() class TestGetMimeType: @@ -480,3 +497,37 @@ class TestProxySecurityGuard: "data:application/pdf;base64," ) assert result["model"] == "mistral/mistral-ocr-latest" + + +@pytest.mark.asyncio +async def test_proxy_upload_stops_reading_at_size_limit() -> None: + from starlette.datastructures import UploadFile + + from litellm.ocr.input import get_max_file_bytes + from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form + + limit: Final = get_max_file_bytes() + with tempfile.TemporaryFile() as stream: + stream.truncate(limit * 2) + upload: Final = UploadFile(file=stream, filename="large.pdf") + request: Final = MagicMock(form=AsyncMock(return_value=FormData({"file": upload}))) + with pytest.raises(ValueError, match="exceeds the size limit"): + await _parse_multipart_form(request) + assert stream.tell() == limit + 1 + + +@pytest.mark.asyncio +async def test_proxy_upload_filename_is_only_metadata(tmp_path: Path) -> None: + from starlette.datastructures import UploadFile + + from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form + + secret: Final = tmp_path / "secret.pdf" + secret.write_bytes(b"server secret") + upload: Final = UploadFile(file=BytesIO(b"uploaded bytes"), filename=str(secret)) + request: Final = MagicMock(form=AsyncMock(return_value=FormData({"file": upload}))) + result: Final = await _parse_multipart_form(request) + assert result["document"] == { + "type": "document_url", + "document_url": "data:application/pdf;base64,dXBsb2FkZWQgYnl0ZXM=", + } diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 46e9a4d3729..4ad556f6941 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -2,37 +2,7 @@ Tests for the OCR `req_format` option in the SDK request path. """ -import pytest - -import litellm from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.ocr import LiteLLMOcrRequest - -DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} - - -def _request( - optional_params: dict[str, object], model: str = "azure_ai/doc-intelligence/prebuilt-layout" -) -> LiteLLMOcrRequest: - return LiteLLMOcrRequest( - model=model, - document=DOCUMENT, - api_key="fake-key", - api_base=None, - custom_llm_provider=None, - extra_headers=None, - timeout=60.0, - kwargs=optional_params, - ) - - -@pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) -def test_rust_ocr_serves_default_format(optional_params): - assert rust_ocr_bridge.supported(_request(optional_params)) is True - - -def test_rust_ocr_serves_native_format_for_document_intelligence(): - assert rust_ocr_bridge.supported(_request({"req_format": "native"})) is True def test_rust_ocr_response_retains_provider_native_response(): @@ -50,34 +20,3 @@ def test_rust_ocr_response_retains_provider_native_response(): assert response.get_provider_native_response() == provider_response assert response.model_dump().get("provider_native_response") is None - - -@pytest.mark.parametrize("model", ["cohere/cohere-parse", "azure_ai/cohere-parse"]) -def test_rust_ocr_skipped_for_unsupported_models(model): - assert rust_ocr_bridge.supported(_request({}, model)) is False - - -@pytest.mark.asyncio -async def test_native_format_rejected_for_provider_without_support_as_bad_request(): - with pytest.raises(litellm.BadRequestError, match="not supported for provider") as exc_info: - await litellm.aocr( - model="mistral/mistral-ocr-latest", - document=DOCUMENT, - api_key="fake-key", - req_format="native", - ) - - assert exc_info.value.status_code == 400 - - -@pytest.mark.asyncio -async def test_unknown_format_rejected_for_provider_without_support_as_bad_request(): - with pytest.raises(litellm.BadRequestError, match="Invalid `req_format`") as exc_info: - await litellm.aocr( - model="mistral/mistral-ocr-latest", - document=DOCUMENT, - api_key="fake-key", - req_format="raw", - ) - - assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py deleted file mode 100644 index dbb4f822d0b..00000000000 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ /dev/null @@ -1,1161 +0,0 @@ -"""Tests for the optional Rust-backed OCR path.""" - -import builtins -import importlib -import types - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import configuration - -# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` -# function onto `litellm.ocr` and shadows the submodule, so import the modules -# explicitly via importlib rather than attribute traversal. -ocr_main = importlib.import_module("litellm.ocr.main") -rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") -rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") -rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") - -MODEL = "mistral/mistral-ocr-latest" -DOCUMENT: dict[str, object] = { - "type": "document_url", - "document_url": "https://example.com/doc.pdf", -} - -FAKE_OCR_RESPONSE: dict[str, object] = { - "pages": [{"index": 0, "markdown": "hello world"}], - "model": "mistral-ocr-2505-completion", - "document_annotation": None, - "usage_info": {"pages_processed": 1}, - "object": "ocr", -} - - -class CapturedException(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - -class RecordingBridge: - """A fake ``RustOcr`` callable that records the args it was handed.""" - - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - self.calls.append( - { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "optional_params": optional_params, - "input_sources": input_sources, - "timeout_seconds": timeout_seconds, - } - ) - return dict(FAKE_OCR_RESPONSE) - - -class RecordingAsyncBridge: - """A fake async ``RustAocr`` callable that records the args it was handed.""" - - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - self.calls.append( - { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "optional_params": optional_params, - "input_sources": input_sources, - "timeout_seconds": timeout_seconds, - } - ) - return dict(FAKE_OCR_RESPONSE) - - -class RaisingBridge: - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise RuntimeError("bridge failed") - - -class RaisingAsyncBridge: - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise RuntimeError("bridge failed") - - -class RecordingLogging: - """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" - - def __init__(self) -> None: - self.pre_call_kwargs: dict[str, object] | None = None - - def update_from_kwargs(self, **kwargs: object) -> None: - self.update_kwargs = kwargs - - def pre_call( - self, - *, - input: str, - api_key: str | None, - additional_args: dict[str, object], - ) -> None: - self.pre_call_kwargs = { - "input": input, - "api_key": api_key, - "additional_args": additional_args, - } - - -def build_request( - *, - logging_obj: RecordingLogging | None = None, - model: str = "mistral-ocr-latest", - document: dict[str, object] = DOCUMENT, - api_key: str | None = "sk-test", - api_base: str | None = None, - custom_llm_provider: str | None = "mistral", - extra_headers: dict[str, object] | None = None, - optional_params: dict[str, object] | None = None, - litellm_params: dict[str, object] | None = None, - timeout: float | httpx.Timeout | None = 12.5, -) -> rust_bridge.LiteLLMOcrRequest: - return rust_bridge.LiteLLMOcrRequest( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout=timeout, - kwargs={ - **(optional_params or {}), - **(litellm_params or {}), - "litellm_logging_obj": logging_obj or RecordingLogging(), - }, - ) - - -@pytest.fixture(autouse=True) -def _reset_rust_flag(): - """Keep the global toggle isolated between tests.""" - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() - configuration.reset_rust_configuration() - rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - yield - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() - configuration.reset_rust_configuration() - rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - - -@pytest.fixture -def fake_bridge(): - """Enable the Rust path with an injected recording bridge (no native wheel).""" - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - return bridge - - -@pytest.fixture -def fake_async_bridge(): - """Enable the async Rust path with an injected recording bridge.""" - bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._AOCR.override(bridge) - return bridge - - -def test_load_rust_ocr_returns_injected_impl(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - assert rust_bridge.load_rust_ocr() is bridge - - -def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch): - real_import = builtins.__import__ - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - if name == "litellm.rust_bridge" and "_native" in fromlist: - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - - -def test_native_bridge_loader_caches_absent_extension(monkeypatch): - real_import = builtins.__import__ - attempts = 0 - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal attempts - if name == "litellm.rust_bridge" and "_native" in fromlist: - attempts += 1 - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - assert rust_bridge_loader.get_native_bridge() is None - assert attempts == 1 - - -def test_native_bridge_loader_reset_forces_relookup(monkeypatch): - real_import = builtins.__import__ - attempts = 0 - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal attempts - if name == "litellm.rust_bridge" and "_native" in fromlist: - attempts += 1 - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - rust_bridge_loader.reset_native_bridge_cache() - assert rust_bridge_loader.get_native_bridge() is None - assert attempts == 2 - - -def test_native_bridge_available_reflects_loader(monkeypatch): - fake_module = types.ModuleType("litellm.rust_bridge._native") - monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module) - - assert rust_bridge_loader.native_bridge_available() is True - - -def test_load_rust_aocr_returns_injected_impl(): - bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._AOCR.override(bridge) - assert rust_bridge.load_rust_aocr() is bridge - - -def test_toggle_without_ocr_arg_preserves_injected_impl(): - """The public flag must not clobber an internal test binding.""" - bridge = RecordingBridge() - async_bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) - - litellm.rust(False) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge - litellm.rust(True) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge - - -def test_explicit_ocr_none_clears_injected_impl(monkeypatch): - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: None, - ) - bridge = RecordingBridge() - async_bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) - - rust_bridge._OCR.override(None) - rust_bridge._AOCR.override(None) - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None - - -def test_load_rust_ocr_none_when_extension_absent(monkeypatch): - """With no injected impl and no compiled wheel, the loader returns None so the - caller degrades to the Python path instead of raising ImportError.""" - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: None, - ) - litellm.rust(True) # no impl injected; extension isn't built in CI - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None - - -def test_load_rust_ocr_uses_compiled_extension(monkeypatch): - """With no injected impl but a packaged ``litellm.rust_bridge._native`` importable, - the loader returns the extension's ``ocr`` callable. The native wheel isn't - built in CI, so stand in a fake module via the bridge loader.""" - fake_module = types.ModuleType("litellm.rust_bridge._native") - fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: fake_module, - ) - - litellm.rust(True) # enabled, no impl injected -> import the extension - assert rust_bridge.load_rust_ocr() is fake_module.ocr - assert rust_bridge.load_rust_aocr() is fake_module.aocr - - -def test_timeout_to_seconds_handles_float_timeout_and_none(): - assert rust_bridge._timeout_to_seconds(12.5) == 12.5 - assert rust_bridge._timeout_to_seconds(None) is None - assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 - - -def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): - bridge = RecordingBridge() - - litellm.rust(True) - - rust_bridge._OCR.override(bridge) - response = rust_bridge.ocr( - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - custom_llm_provider="mistral", - extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True, "pages": [0]}, - timeout=12.5, - ) - - assert response == FAKE_OCR_RESPONSE - call = bridge.calls[0] - assert call == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - }, - "optional_params": {"include_image_base64": True, "pages": [0]}, - "input_sources": {}, - "timeout_seconds": 12.5, - } - - -@pytest.mark.asyncio -async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): - bridge = RecordingAsyncBridge() - - litellm.rust(True) - - rust_bridge._AOCR.override(bridge) - response = await rust_bridge.aocr( - model="mistral-ocr-maas", - document=DOCUMENT, - api_key=None, - api_base=None, - custom_llm_provider="vertex_ai", - extra_headers=None, - optional_params={"vertex_project": "project-1"}, - timeout=httpx.Timeout(30.0, read=42.0), - ) - - assert response == FAKE_OCR_RESPONSE - assert bridge.calls[0] == { - "model": "mistral-ocr-maas", - "document": DOCUMENT, - "api_key": None, - "api_base": None, - "custom_llm_provider": "vertex_ai", - "extra_headers": None, - "optional_params": {"vertex_project": "project-1"}, - "input_sources": {}, - "timeout_seconds": 42.0, - } - - -def test_run_rust_ocr_prepares_request_and_wraps_response(): - bridge = RecordingBridge() - logging_obj = RecordingLogging() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - response = ocr_main._run_rust_ocr( - request=build_request( - logging_obj=logging_obj, - api_base="https://proxy.internal", - extra_headers={"x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True}, - timeout=12.5, - ), - resolve_api_key=lambda _name: None, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert bridge.calls[0] == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": { - "x-trace-id": "trace-1", - }, - "optional_params": {"include_image_base64": True}, - "input_sources": {}, - "timeout_seconds": 12.5, - } - - -def test_rust_upstream_error_uses_ocr_provider_error_mapping(): - error = RustUpstreamError(400, '{"message":"invalid model"}') - - mapped = ocr_main._map_rust_ocr_error( - error, - build_request(), - (RuntimeError, RustUpstreamError), - ) - - assert isinstance(mapped, BaseLLMException) - assert mapped.status_code == 400 - assert mapped.message == '{"message":"invalid model"}' - - -def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request(api_key=None, timeout=None), - resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, - ) - - assert bridge.calls[0]["api_key"] == "sk-from-vault" - - -def test_run_rust_ocr_prefers_explicit_key_over_resolver(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name: str) -> str | None: - raise AssertionError(f"resolver should not be called for {name}") - - ocr_main._run_rust_ocr( - request=build_request( - api_key="sk-explicit", - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["api_key"] == "sk-explicit" - - -def test_run_rust_ocr_uses_mistral_secret_manager_without_provider_config(): - bridge = RecordingBridge() - resolver_calls = [] - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name): - resolver_calls.append(name) - return "sk-provider-env" - - ocr_main._run_rust_ocr( - request=build_request( - model="mistral-ocr-latest", - api_key=None, - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert resolver_calls == ["MISTRAL_API_KEY"] - assert bridge.calls[0]["api_key"] == "sk-provider-env" - - -def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="vertex_ai", - model="mistral-ocr-maas", - litellm_params={ - "vertex_project": "project-1", - "vertex_location": "us-central1", - "vertex_credentials": "redacted", - }, - optional_params={"include_image_base64": True}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["optional_params"] == { - "include_image_base64": True, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "vertex_credentials": "redacted", - } - - -def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name: str) -> str | None: - return { - "VERTEXAI_PROJECT": "project-from-secret", - "VERTEXAI_LOCATION": "us-east5", - "VERTEXAI_CREDENTIALS": "credentials-from-secret", - }.get(name) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="vertex_ai", - model="mistral-ocr-maas", - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" - assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" - assert bridge.calls[0]["optional_params"]["vertex_credentials"] == "credentials-from-secret" - - -def test_prepare_rust_ocr_call_defers_azure_environment_resolution_to_rust(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - api_base=None, - timeout=None, - ), - resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), - ) - - assert bridge.calls[0]["api_base"] is None - assert bridge.calls[0]["api_key"] is None - assert bridge.calls[0]["extra_headers"] is None - - -def test_prepare_rust_ocr_call_defers_document_intelligence_environment_to_rust(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="doc-intelligence/prebuilt-layout", - api_base=None, - timeout=None, - ), - resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), - ) - - assert bridge.calls[0]["api_base"] is None - - -def test_prepare_rust_ocr_call_forwards_raw_azure_auth_inputs(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - api_base="https://azure.example.com", - extra_headers={"x-trace-id": "trace-1"}, - litellm_params={ - "azure_ad_token": "entra-token", - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_scope": "scope", - "azure_authority_host": "https://login.example.com", - "azure_credential": "ClientSecretCredential", - "azure_federated_token_file": "/token", - }, - timeout=None, - ), - resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), - ) - - call = bridge.calls[0] - assert call["api_key"] is None - assert call["api_base"] == "https://azure.example.com" - assert call["extra_headers"] == {"x-trace-id": "trace-1"} - assert call["optional_params"] == { - "azure_ad_token": "entra-token", - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_scope": "scope", - "azure_authority_host": "https://login.example.com", - "azure_credential": "ClientSecretCredential", - "azure_federated_token_file": "/token", - } - assert call["input_sources"] == {} - - -def test_prepare_rust_ocr_call_preserves_proxy_input_sources(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - request_values = { - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_authority_host": "https://login.example.com", - "api_base": "https://azure.example.com", - } - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key="request-key", - api_base="https://azure.example.com", - litellm_params={ - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_authority_host": "https://login.example.com", - "proxy_server_request": {"body": request_values, "credential_fields": ("api_key",)}, - }, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["input_sources"] == { - **{name: "request" for name in request_values}, - "api_key": "request", - } - - marshaled = rust_bridge._marshal( - build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key="request-key", - api_base="https://azure.example.com", - litellm_params={ - "proxy_server_request": { - "body": {"api_base": "https://azure.example.com"}, - "credential_fields": ("api_key",), - } - }, - ), - lambda _name: None, - lambda document: document, - ) - assert marshaled.input_sources == {"api_base": "request", "api_key": "request"} - - -def test_rust_ocr_logging_redacts_azure_credentials(): - bridge = RecordingBridge() - logging_obj = RecordingLogging() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - logging_obj=logging_obj, - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - litellm_params={"azure_ad_token": "token", "client_secret": "secret"}, - ), - resolve_api_key=lambda _name: None, - ) - - assert logging_obj.update_kwargs["optional_params"] == { - "azure_ad_token": "****", - "client_secret": "****", - } - assert logging_obj.pre_call_kwargs is not None - additional_args = logging_obj.pre_call_kwargs["additional_args"] - assert isinstance(additional_args, dict) - complete_input = additional_args["complete_input_dict"] - assert isinstance(complete_input, dict) - assert complete_input["azure_ad_token"] == "****" - assert complete_input["client_secret"] == "****" - - -def test_rust_eligibility_rejects_python_only_azure_auth_modes(): - for params in ( - {"azure_ad_token_provider": lambda: "token"}, - {"azure_username": "user"}, - {"azure_password": "password"}, - ): - assert not ocr_main._rust_ocr_supported( - build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - litellm_params=params, - ) - ) - - -def test_prepare_rust_ocr_call_forwards_global_azure_refresh(monkeypatch: pytest.MonkeyPatch): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - api_base="https://azure.example.com", - litellm_params={"proxy_server_request": {"body": {"enable_azure_ad_token_refresh": True}}}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["optional_params"] == {"enable_azure_ad_token_refresh": True} - assert bridge.calls[0]["input_sources"] == {"enable_azure_ad_token_refresh": "deployment"} - - -def test_run_rust_ocr_runs_pre_call_logging(): - logging_obj = RecordingLogging() - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - logging_obj=logging_obj, - api_base="https://api.mistral.ai/v1", - extra_headers={"x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert logging_obj.pre_call_kwargs is not None - assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" - additional_args = logging_obj.pre_call_kwargs["additional_args"] - complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT - assert complete_input["include_image_base64"] is True - assert additional_args["api_base"] == "https://api.mistral.ai/v1" - assert additional_args["headers"] == { - "x-trace-id": "trace-1", - } - - -def test_ocr_routes_to_rust_when_enabled(fake_bridge): - response = litellm.ocr( - model=MODEL, - document=DOCUMENT, - api_key="sk-test", - extra_headers={"x-trace-id": "trace-1"}, - include_image_base64=True, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert len(fake_bridge.calls) == 1 - call = fake_bridge.calls[0] - assert call["model"] == MODEL - assert call["document"] == DOCUMENT - assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] is None - assert call["extra_headers"] == { - "x-trace-id": "trace-1", - } - assert call["optional_params"].get("include_image_base64") is True - - -def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): - response = litellm.ocr( - model="azure_ai/pixtral-12b-2409", - document=DOCUMENT, - api_key="sk-test", - api_base="https://example.services.ai.azure.com", - ) - - assert isinstance(response, OCRResponse) - assert len(fake_bridge.calls) == 1 - assert fake_bridge.calls[0]["model"] == "azure_ai/pixtral-12b-2409" - assert fake_bridge.calls[0]["custom_llm_provider"] is None - assert fake_bridge.calls[0]["extra_headers"] is None - - -def test_ocr_routes_azure_entra_inputs_to_rust_without_python_auth(fake_bridge): - response = litellm.ocr( - model="azure_ai/pixtral-12b-2409", - document=DOCUMENT, - api_base="https://example.services.ai.azure.com", - azure_ad_token="entra-token", - tenant_id="tenant", - client_id="client", - ) - - assert isinstance(response, OCRResponse) - assert fake_bridge.calls[0]["api_key"] is None - assert fake_bridge.calls[0]["extra_headers"] is None - assert fake_bridge.calls[0]["optional_params"] == { - "azure_ad_token": "entra-token", - "tenant_id": "tenant", - "client_id": "client", - } - - -def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): - response = litellm.ocr( - model=MODEL, - document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, - api_key="sk-test", - ) - - assert isinstance(response, OCRResponse) - document = fake_bridge.calls[0]["document"] - assert document["type"] == "document_url" - assert document["document_url"].startswith("data:application/pdf;base64,") - - -def test_ocr_exception_type_uses_resolved_provider_context( - monkeypatch: pytest.MonkeyPatch, -): - captured: dict[str, object] = {} - - def fake_exception_type(**kwargs: object) -> CapturedException: - captured.update(kwargs) - return CapturedException("wrapped") - - monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - litellm.rust(True) - rust_bridge._OCR.override(RaisingBridge()) - - with pytest.raises(CapturedException): - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured["model"] == "mistral-ocr-latest" - assert captured["custom_llm_provider"] == "mistral" - - -@pytest.mark.asyncio -async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): - response = await litellm.aocr( - model=MODEL, - document=DOCUMENT, - api_key="sk-test", - extra_headers={"x-trace-id": "trace-1"}, - include_image_base64=True, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert len(fake_async_bridge.calls) == 1 - call = fake_async_bridge.calls[0] - assert call["model"] == MODEL - assert call["document"] == DOCUMENT - assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] is None - assert call["extra_headers"] == { - "x-trace-id": "trace-1", - } - assert call["optional_params"].get("include_image_base64") is True - - -@pytest.mark.asyncio -async def test_aocr_exception_type_uses_resolved_provider_context( - monkeypatch: pytest.MonkeyPatch, -): - captured: dict[str, object] = {} - - def fake_exception_type(**kwargs: object) -> CapturedException: - captured.update(kwargs) - return CapturedException("wrapped") - - monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - litellm.rust(True) - rust_bridge._AOCR.override(RaisingAsyncBridge()) - - with pytest.raises(CapturedException): - await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured["model"] == "mistral-ocr-latest" - assert captured["custom_llm_provider"] == "mistral" - - -def test_ocr_forwards_timeout_to_rust(fake_bridge): - """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s - client ceiling doesn't silently override shorter deadlines.""" - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test", timeout=12.5) - - assert fake_bridge.calls[0]["timeout_seconds"] == 12.5 - - -def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - from litellm.constants import request_timeout - - assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) - - -def test_ocr_does_not_route_to_rust_when_disabled(): - """With the flag off, the bridge must not be consulted even if an impl exists.""" - bridge = RecordingBridge() - litellm.rust(False) - rust_bridge._OCR.override(bridge) - # The impl stays available for injection, but the disabled flag gates usage, - # so ocr() never reaches the Rust path (asserted via the enabled-path test). - assert bridge.calls == [] - - -def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): - """Rust enabled but no bridge available (no injected impl, no compiled wheel): - ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) - litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI - - captured = {} - - def fake_handler_ocr(**kwargs): - captured["called"] = True - return OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") - - monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) - - response = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured.get("called") is True # Python path was used - assert isinstance(response, OCRResponse) - - -def test_ocr_provider_configs_expose_api_key_env_vars(): - from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( - AzureDocumentIntelligenceOCRConfig, - ) - from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig - from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig - from litellm.llms.mistral.ocr.transformation import MistralOCRConfig - from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( - VertexAIDeepSeekOCRConfig, - ) - from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig - - assert BaseOCRConfig().get_api_key_env_var() is None - assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" - assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" - assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" - assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" - assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.asyncio -async def test_rust_receives_unmapped_azure_options(asynchronous, fake_bridge, fake_async_bridge): - from typing import Final - - arguments: Final = { - "model": "azure_ai/doc-intelligence/prebuilt-layout", - "document": DOCUMENT, - "api_key": "test-key", - "pages": [0, 2], - "features": ["languages", "style"], - "provider_extension": {"enabled": True}, - } - if asynchronous: - await litellm.aocr(**arguments) - else: - litellm.ocr(**arguments) - call: Final = (fake_async_bridge if asynchronous else fake_bridge).calls[0] - assert call["model"] == arguments["model"] - assert call["custom_llm_provider"] is None - assert call["extra_headers"] is None - assert call["optional_params"] == { - "pages": [0, 2], - "features": ["languages", "style"], - "provider_extension": {"enabled": True}, - } - - -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.asyncio -async def test_python_fallback_maps_original_options_once(enabled, monkeypatch): - from io import BytesIO - from typing import Final - - class PythonHandler: - def __init__(self): - self.calls = [] - - def ocr(self, **kwargs): - self.calls.append(kwargs) - return OCRResponse(pages=[], model=kwargs["model"]) - - handler: Final = PythonHandler() - monkeypatch.setattr(ocr_main, "base_llm_http_handler", handler) - litellm.rust(enabled) - rust_bridge._OCR.override(None) - rust_bridge._AOCR.override(None) - for asynchronous in (False, True): - file: Final = BytesIO(b"test document") - arguments: Final = { - "model": "azure_ai/doc-intelligence/prebuilt-layout", - "document": {"type": "file", "file": file}, - "api_key": "test-key", - "pages": [0, 2], - } - if asynchronous: - await litellm.aocr(**arguments) - else: - litellm.ocr(**arguments) - assert handler.calls[-1]["optional_params"]["pages"] == "1,3" - assert handler.calls[-1]["document"]["document_url"].endswith("dGVzdCBkb2N1bWVudA==") - assert len(handler.calls) == 2 - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) -@pytest.mark.asyncio -async def test_native_public_ocr_matches_python(model, asynchronous): - import json - from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer - from threading import Thread - from typing import Final - from urllib.parse import parse_qsl, urlsplit - - native: Final = rust_bridge_loader.get_native_bridge() - if native is None: - pytest.skip("requires the compiled Rust extension") - calls: Final = [] - - class Handler(BaseHTTPRequestHandler): - def do_POST(self): - body: Final = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) - target: Final = urlsplit(self.path) - calls.append( - ( - target.path, - parse_qsl(target.query), - self.headers.get("Authorization"), - self.headers.get("Ocp-Apim-Subscription-Key"), - body, - ) - ) - payload: Final = ( - {"status": "succeeded", "analyzeResult": {"pages": []}} - if "doc-intelligence" in model - else {"pages": [{"index": 0, "markdown": "hello"}]} - ) - encoded: Final = json.dumps(payload).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(encoded))) - self.end_headers() - self.wfile.write(encoded) - - def log_message(self, *_args): - pass - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread: Final = Thread(target=server.serve_forever, daemon=True) - thread.start() - responses: Final = [] - try: - for enabled in (False, True): - litellm.rust(enabled) - arguments: Final = { - "model": model, - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "pages": [0, 2], - "timeout": 3.0, - } - response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - responses.append(response.model_dump()) - assert len(calls) == 2 - assert calls[0] == calls[1] - for key in ("model", "pages", "object"): - assert responses[0][key] == responses[1][key] - finally: - server.shutdown() - server.server_close() - thread.join(timeout=3) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e6c8d4ee039..c2f4f7163a0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -9266,17 +9266,24 @@ def _agent_prisma(object_permission_id=None, side_effect=None): @contextlib.contextmanager -def _entitlement_fault_globals(prisma_client=None): +def _entitlement_fault_globals(prisma_client=None, user_api_key_cache=None): from litellm.caching.dual_cache import DualCache with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client or MagicMock()), - patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache or DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", _proxy_logging_with_awaitable_hooks()), ): yield +def _proxy_logging_with_awaitable_hooks(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + return proxy_logging_obj + + @pytest.mark.asyncio class TestEntitlementFaultSemantics: """Each entitlement level distinguishes two fault classes for a KEY-authenticated caller. @@ -9412,6 +9419,71 @@ class TestEntitlementFaultSemantics: assert set(allowed) == {"srv1"} +async def _cache_with_end_user(end_user_id, *, mcp_tool_permissions=None, object_permission_id=None): + """A real DualCache already holding the end user row, so ``get_end_user_object`` answers from + cache and no ``litellm.`` internal has to be patched. ``object_permission_id`` without a + permission body models a row that NAMES an entitlement the DB then fails to serve.""" + from litellm.caching.dual_cache import DualCache + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + cache = DualCache() + await cache.async_set_cache( + key=end_user_cache_key(end_user_id), + value=LiteLLM_EndUserTable( + user_id=end_user_id, + blocked=False, + object_permission_id=object_permission_id or ("op-eu" if mcp_tool_permissions else None), + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", mcp_tool_permissions=mcp_tool_permissions + ) + if mcp_tool_permissions + else None, + ), + ) + return cache + + +@pytest.mark.asyncio +class TestEndUserToolCeiling: + """The end user (customer) level narrows the TOOLS axis exactly as it narrows the servers axis, + so `object_permission.mcp_tool_permissions` on `/customer/new` is enforced, not just stored.""" + + async def test_end_user_tool_permissions_intersect_key_tools(self): + auth = _key_auth_reaching("srv1", tools=["tool_a", "tool_b"], end_user_id="eu-1") + cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv1": ["tool_a"]}) + with _entitlement_fault_globals(user_api_key_cache=cache): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == ["tool_a"] + + async def test_end_user_tool_permissions_become_allowlist_when_key_is_unrestricted(self): + auth = _key_auth_reaching("srv1", end_user_id="eu-1") + cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv1": ["tool_a"]}) + with _entitlement_fault_globals(user_api_key_cache=cache): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == ["tool_a"] + + async def test_end_user_tool_permissions_on_another_server_place_no_ceiling(self): + auth = _key_auth_reaching("srv1", tools=["tool_a", "tool_b"], end_user_id="eu-1") + cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv2": ["tool_z"]}) + with _entitlement_fault_globals(user_api_key_cache=cache): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert sorted(tools) == ["tool_a", "tool_b"] + + async def test_end_user_named_but_unloadable_permission_denies_tools(self): + auth = _key_auth_reaching("srv1", tools=["tool_a"], end_user_id="eu-1") + cache = await _cache_with_end_user("eu-1", object_permission_id="op-eu") + with _entitlement_fault_globals(user_api_key_cache=cache): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == [], "an end-user entitlement we know exists but cannot read must deny its tools" + + async def test_no_end_user_row_places_no_tool_ceiling(self): + auth = _key_auth_reaching("srv1", tools=["tool_a"], end_user_id="eu-1") + with _entitlement_fault_globals(user_api_key_cache=await _cache_with_end_user("someone-else")): + tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth) + assert tools == ["tool_a"] + + @pytest.mark.asyncio class TestScopedSessionAdmission: """LIT-4917: a session bearer sealed to one server (RFC 8707 resource at authorize) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 9f2feddb0e3..55accfb169d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -97,6 +97,88 @@ def unauthenticated_client(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("byok_first", [True, False]) +def test_byok_challenge_discovers_api_key_flow(monkeypatch, byok_first): + from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints + from litellm.proxy._experimental.mcp_server.oauth_utils import get_byok_www_authenticate + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + app = FastAPI() + routers = (byok_oauth_endpoints.router, discoverable_endpoints.router) + for item in routers if byok_first else reversed(routers): + app.include_router(item) + with TestClient(app) as session: + challenge = get_byok_www_authenticate() + assert challenge == 'Bearer resource_metadata="/v1/mcp/oauth/protected-resource"' + response = session.get(challenge.split('"')[1]) + assert response.status_code == 200 + assert response.json() == { + "resource": "http://testserver", + "authorization_servers": ["http://testserver/v1/mcp/oauth"], + } + authorization = session.get("/.well-known/oauth-authorization-server/v1/mcp/oauth") + assert authorization.status_code == 200 + metadata = authorization.json() + assert metadata["issuer"] == response.json()["authorization_servers"][0] + assert metadata["authorization_endpoint"] == "http://testserver/v1/mcp/oauth/authorize" + assert metadata["token_endpoint"] == "http://testserver/v1/mcp/oauth/token" + assert metadata["code_challenge_methods_supported"] == ["S256"] + + +@pytest.mark.parametrize( + ("base_url", "root_path", "expected"), + [ + ("", "", "/v1/mcp/oauth/protected-resource"), + ("", "/proxy", "/proxy/v1/mcp/oauth/protected-resource"), + ("https://gateway.example.com/proxy", "/proxy", "https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"), + ], +) +def test_byok_challenge_preserves_external_base(monkeypatch, base_url, root_path, expected): + from litellm.proxy._experimental.mcp_server.oauth_utils import get_byok_www_authenticate + + monkeypatch.setenv("PROXY_BASE_URL", base_url) + monkeypatch.setenv("SERVER_ROOT_PATH", root_path) + assert get_byok_www_authenticate() == f'Bearer resource_metadata="{expected}"' + + +def test_byok_discovery_preserves_per_request_prefixes(monkeypatch): + from fastapi import FastAPI + + from litellm.proxy._experimental.mcp_server.server import _check_byok_credential + from litellm.proxy.middleware.per_request_root_path_middleware import PerRequestRootPathMiddleware + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a,/tenant-b") + app = FastAPI() + app.include_router(router) + app.add_middleware(PerRequestRootPathMiddleware, root_paths=("/tenant-a", "/tenant-b")) + server = MCPServer(server_id="byok-prefix", name="byok-prefix", transport=MCPTransport.http, is_byok=True) + + @app.get("/challenge") + async def challenge(): + await _check_byok_credential(server, None) + + with TestClient(app) as client: + for prefix in ("/tenant-a", "/tenant-b", ""): + challenge_response = client.get(f"{prefix}/challenge") + assert challenge_response.status_code == 401 + metadata_path = f"{prefix}/v1/mcp/oauth/protected-resource" + assert challenge_response.headers["www-authenticate"] == f'Bearer resource_metadata="{metadata_path}"' + prm = client.get(metadata_path) + assert prm.status_code == 200 + issuer = f"http://testserver{prefix}/v1/mcp/oauth" + assert prm.json()["authorization_servers"] == [issuer] + asm = client.get(f"/.well-known/oauth-authorization-server{prefix}/v1/mcp/oauth") + assert asm.status_code == 200 + assert asm.json()["issuer"] == issuer + assert asm.json()["authorization_endpoint"] == f"{issuer}/authorize" + assert asm.json()["token_endpoint"] == f"{issuer}/token" + assert client.get("/.well-known/oauth-authorization-server/unknown/v1/mcp/oauth").status_code == 404 + + def test_oauth_authorization_server_metadata(client): resp = client.get("/.well-known/oauth-authorization-server") assert resp.status_code == 200 @@ -107,15 +189,6 @@ def test_oauth_authorization_server_metadata(client): assert "S256" in data["code_challenge_methods_supported"] -def test_oauth_protected_resource_metadata(client): - resp = client.get("/.well-known/oauth-protected-resource") - assert resp.status_code == 200 - data = resp.json() - assert "resource" in data - assert "authorization_servers" in data - assert len(data["authorization_servers"]) == 1 - - # --------------------------------------------------------------------------- # Authorization GET endpoint # --------------------------------------------------------------------------- @@ -501,7 +574,7 @@ async def test_check_byok_credential_no_user_id(): @pytest.mark.asyncio -async def test_check_byok_credential_missing_credential(): +async def test_check_byok_credential_missing_credential(monkeypatch): """BYOK server with a known user but no stored credential → 401.""" from litellm.proxy._experimental.mcp_server.server import _check_byok_credential from litellm.proxy._types import UserAPIKeyAuth @@ -515,6 +588,11 @@ async def test_check_byok_credential_missing_credential(): ) user_auth = UserAPIKeyAuth(user_id="user-99", api_key="sk-test") + from litellm.proxy._experimental.mcp_server import server as server_module + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(server_module, "_byok_cred_cache", {}) mock_prisma = MagicMock() with ( @@ -526,6 +604,10 @@ async def test_check_byok_credential_missing_credential(): ): with pytest.raises(HTTPException) as exc_info: await _check_byok_credential(server, user_auth) + with pytest.raises(HTTPException) as cached_exc: + await _check_byok_credential(server, user_auth) + assert cached_exc.value.status_code == 401 + assert cached_exc.value.headers == exc_info.value.headers assert exc_info.value.status_code == 401 detail: Any = exc_info.value.detail @@ -533,7 +615,38 @@ async def test_check_byok_credential_missing_credential(): assert detail["server_id"] == "byok-2" headers = exc_info.value.headers or {} assert "WWW-Authenticate" in headers # type: ignore[operator] - assert "oauth-protected-resource" in headers["WWW-Authenticate"] # type: ignore[index] + assert headers["WWW-Authenticate"] == 'Bearer resource_metadata="/v1/mcp/oauth/protected-resource"' + + +@pytest.mark.asyncio +async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monkeypatch): + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy") + monkeypatch.setattr(mcp_module, "_byok_cred_cache", {}) + server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True) + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + with pytest.raises(HTTPException) as exc_info: + await mcp_module.execute_mcp_tool( + name="list_regions", + arguments={}, + allowed_mcp_servers=[server], + requested_server_id=server.server_id, + start_time=datetime.now(timezone.utc), + user_api_key_auth=UserAPIKeyAuth(user_id="byok-discovery-user"), + ) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail["server_id"] == server.server_id + assert exc_info.value.headers == { + "WWW-Authenticate": 'Bearer resource_metadata="https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"' + } @pytest.mark.asyncio 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 1b7e3d6a1c3..9ea870d3210 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 @@ -3433,51 +3433,84 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa global_mcp_server_manager.registry.clear() +@pytest.mark.parametrize("server_count", [0, 1, 2]) +@pytest.mark.parametrize("byok_first", [True, False]) +@pytest.mark.parametrize( + ("base_url", "origin"), + [ + ("https://gateway.example.com", "https://gateway.example.com"), + ("https://gateway.example.com/proxy", "https://gateway.example.com"), + ("http://[::1]:4000/proxy", "http://[::1]:4000"), + ], +) +def test_root_protected_resource_discovers_gateway(monkeypatch, server_count, byok_first, base_url, origin): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setenv("PROXY_BASE_URL", base_url) + monkeypatch.setattr( + global_mcp_server_manager, + "registry", + { + f"oauth_{index}": MCPServer( + server_id=f"oauth_{index}", + name=f"oauth_{index}", + server_name=f"oauth_{index}", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + for index in range(server_count) + }, + ) + app = FastAPI() + routers = (byok_oauth_endpoints.router, discoverable_endpoints.router) + for router in routers if byok_first else reversed(routers): + app.include_router(router) + with TestClient(app) as client: + response = client.get("/.well-known/oauth-protected-resource", params={"mcp_server_name": "oauth_0"}) + assert response.status_code == 200 + assert response.json() == { + "resource": origin, + "authorization_servers": [f"{base_url}/mcp"], + "scopes_supported": [], + } + authorization = client.get("/.well-known/oauth-authorization-server/mcp") + assert authorization.status_code == 200 + metadata = authorization.json() + assert metadata["issuer"] == response.json()["authorization_servers"][0] + assert metadata["authorization_endpoint"] == f"{base_url}/authorize/mcp-session" + assert metadata["token_endpoint"] == f"{base_url}/token" + assert metadata["registration_endpoint"] == f"{base_url}/register" + aggregate = client.get("/.well-known/oauth-protected-resource/mcp") + assert aggregate.status_code == 200 + assert aggregate.json()["resource"] == f"{base_url}/mcp" + + @pytest.mark.asyncio -async def test_oauth_protected_resource_root_resolved_single_server_keeps_relay_as(): - """The unnamed (bare-root) legacy shape resolves the single configured oauth2 server and - must keep advertising the per-server relay authorization server: only an EXPLICITLY - named request opts into the gateway-as-AS flow (LIT-4864), so pre-existing single-server - deployments discovering through the root document are byte-identical.""" - try: - from fastapi import Request +async def test_unnamed_protected_resource_builder_uses_gateway_origin(monkeypatch): + from fastapi import Request - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_protected_resource_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - only_server = MCPServer( - server_id="solo_mcp", - name="solo_mcp", - server_name="solo_mcp", - alias="solo_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/oauth/token", + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, ) - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - global_mcp_server_manager.registry.clear() - try: - global_mcp_server_manager.registry[only_server.server_id] = only_server - response = await _build_oauth_protected_resource_response( - request=mock_request, mcp_server_name=None, use_standard_pattern=False - ) - assert response["authorization_servers"] == ["https://litellm.example.com/solo_mcp"] - finally: - global_mcp_server_manager.registry.clear() + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = Request( + {"type": "http", "scheme": "https", "server": ("gateway.example.com", 443), "path": "/", "headers": []} + ) + response = await _build_oauth_protected_resource_response(request, None, False) + assert response == { + "resource": "https://gateway.example.com", + "authorization_servers": ("https://gateway.example.com/mcp",), + "scopes_supported": (), + } @pytest.mark.asyncio @@ -4137,8 +4170,9 @@ async def test_discovery_root_does_not_expose_private_server_for_external_client assert "/test_oauth/" not in authorization_response["authorization_endpoint"] assert "/test_oauth/" not in authorization_response["token_endpoint"] assert authorization_response["scopes_supported"] == [] - assert resource_response["authorization_servers"] == ["https://llm.example.com"] - assert resource_response["scopes_supported"] == [] + assert tuple(resource_response["authorization_servers"]) == ("https://llm.example.com/mcp",) + assert resource_response["resource"] == "https://llm.example.com" + assert not resource_response["scopes_supported"] finally: global_mcp_server_manager.registry.clear() @@ -9014,7 +9048,7 @@ def test_aggregate_wellknown_routes_serve_gateway_metadata(): assert asm.status_code == 200 assert asm.json()["issuer"] == "http://testserver/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize/mcp-session" assert "none" in asm.json()["token_endpoint_auth_methods_supported"] @@ -9077,11 +9111,7 @@ def test_well_known_root_suffix_reflects_server_root_path(): @pytest.mark.asyncio -async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): - """The always-on aggregate front door must not change bare-origin discovery: with one - oauth2 server configured, the no-suffix /.well-known/oauth-{authorization-server, - protected-resource} still resolves THAT server, so an existing single-server deployment's - discovery is unchanged. The aggregate document lives only at the /mcp-suffixed routes.""" +async def test_root_resource_uses_gateway_without_changing_authorization_relay(): from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -9105,10 +9135,10 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): resource_response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name=None, use_standard_pattern=True ) - # per-server, not aggregate: the single server's name is in the endpoints assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] assert authorization_response["issuer"] == "https://llm.example.com" - assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] + assert tuple(resource_response["authorization_servers"]) == ("https://llm.example.com/mcp",) + assert resource_response["resource"] == "https://llm.example.com" finally: global_mcp_server_manager.registry.clear() @@ -10640,7 +10670,7 @@ class TestPerRequestRootPathDiscovery: assert asm.status_code == 200 assert asm.json()["issuer"] == "http://testserver/tenant-a/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize" + assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize/mcp-session" # The prefixed authorize URL routes to the real handler (not 404): # under per-request root_path the whole app is reachable per-prefix, @@ -10799,6 +10829,68 @@ def _consent_flow_handle(page: str) -> str: return match.group(1) +@pytest.mark.parametrize("redirect_uri", ["http://127.0.0.1:51234/callback", "https://client.example.com/callback"]) +@pytest.mark.parametrize("signed_in", [True, False]) +def test_root_discovery_origin_authorizes_mcp_session(monkeypatch, redirect_uri, signed_in): + from urllib.parse import parse_qs, urlparse + + client, session_cookie, minted = _native_client_app(monkeypatch) + root = client.get("/.well-known/oauth-protected-resource") + assert root.status_code == 200 + assert root.json()["resource"] == "http://testserver" + authorization = client.get("/.well-known/oauth-authorization-server/mcp") + assert authorization.status_code == 200 + metadata = authorization.json() + registered = client.post(metadata["registration_endpoint"], json={"redirect_uris": [redirect_uri]}) + assert registered.status_code == 201 + if signed_in: + client.cookies.set("token", session_cookie) + response = client.get( + metadata["authorization_endpoint"], + params={ + "response_type": "code", + "client_id": registered.json()["client_id"], + "redirect_uri": redirect_uri, + "state": "mcp-state", + "code_challenge": _s256("v" * 43), + "code_challenge_method": "S256", + "resource": root.json()["resource"], + }, + follow_redirects=False, + ) + assert response.status_code == 303 + target = urlparse(response.headers["location"]) + assert target.path == ("/ui/connect" if signed_in else "/sso/key/generate") + if signed_in: + flow = parse_qs(target.query)["connect_flow"][0] + described = client.get("/authorize/flow", params={"flow": flow}) + assert described.status_code == 200 + assert described.json()["state"] == "unscoped" + assert minted == [] + + +@pytest.mark.parametrize("valid_client", [True, False]) +def test_mcp_session_authorize_rejects_invalid_registration_or_pkce(monkeypatch, valid_client): + client, session_cookie, minted = _native_client_app(monkeypatch) + redirect_uri = "https://client.example.com/callback" + registered = client.post("/register", json={"redirect_uris": [redirect_uri]}) + assert registered.status_code == 201 + client.cookies.set("token", session_cookie) + response = client.get( + "/authorize/mcp-session", + params={ + "client_id": registered.json()["client_id"] if valid_client else "unknown-client", + "redirect_uri": redirect_uri, + "response_type": "code", + }, + follow_redirects=False, + ) + assert response.status_code == 400 + assert response.json()["error"] == ("invalid_request" if valid_client else "invalid_client") + assert "location" not in response.headers + assert minted == [] + + def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(monkeypatch): """The whole ``lite login --pkce`` server side over the real router: a Go CLI reads the versioned discovery document, registers a loopback public client, the signed-in user consents to a team, @@ -11186,3 +11278,99 @@ async def test_dcr_refusal_is_actionable_without_upstream_body( assert f"HTTP {upstream_status}" in str(exc.value.detail) assert "pre-registered OAuth client" in str(exc.value.detail) assert "private upstream details" not in str(exc.value.detail) + + +@pytest.mark.parametrize("prefix", ["", "/tenant-a", "/tenant-b"]) +@pytest.mark.parametrize( + ("server_name", "pattern"), + [ + ("issuer_test", "mcp/{server}"), + ("issuer_test", "{server}/mcp"), + ("issuer_test", "{server}"), + ("mcp", "mcp/{server}"), + ("mcp", "{server}/mcp"), + ], +) +def test_per_server_authorization_metadata_issuer_matches_discovery_path( + _no_proxy_base_url, _isolated_mcp_registry, prefix, server_name, pattern +): + server = _create_oauth2_server(server_id=server_name, name=server_name, server_name=server_name, alias=server_name) + _isolated_mcp_registry[server.server_id] = server + client = _prefixed_discovery_client(["/tenant-a", "/tenant-b"]) + path = pattern.format(server=server_name) + response = client.get(f"{prefix}/.well-known/oauth-authorization-server/{path}") + assert response.status_code == 200 + metadata = response.json() + assert metadata["issuer"] == f"http://testserver{prefix}/{path}" + assert metadata["authorization_endpoint"] == f"http://testserver{prefix}/{server_name}/authorize" + assert metadata["token_endpoint"] == f"http://testserver{prefix}/{server_name}/token" + assert metadata["registration_endpoint"] == f"http://testserver{prefix}/{server_name}/register" + + +@pytest.mark.parametrize("prefix", ["", "/tenant-a"]) +@pytest.mark.parametrize("relay", [False, True]) +@pytest.mark.parametrize("pattern", ["mcp/{server}", "{server}/mcp"]) +def test_named_resource_discovery_follows_matching_authorization_issuer( + _no_proxy_base_url, _isolated_mcp_registry, prefix, relay, pattern +): + server = _create_oauth2_server().model_copy(update={"per_server_oauth_discovery": relay}) + _isolated_mcp_registry[server.server_id] = server + client = _prefixed_discovery_client(["/tenant-a"]) + path = pattern.format(server=server.server_name) + response = client.get(f"{prefix}/.well-known/oauth-protected-resource/{path}") + assert response.status_code == 200 + resource = response.json() + issuer_path = server.server_name if relay else "mcp" + assert resource["resource"] == f"http://testserver{prefix}/{path}" + assert resource["authorization_servers"] == [f"http://testserver{prefix}/{issuer_path}"] + authorization = client.get(f"{prefix}/.well-known/oauth-authorization-server/{issuer_path}") + assert authorization.status_code == 200 + assert authorization.json()["issuer"] == resource["authorization_servers"][0] + + +def test_static_root_path_authorization_discovery_preserves_issuer(monkeypatch, tmp_path): + import subprocess + import sys + + monkeypatch.setenv("SERVER_ROOT_PATH", "/gateway") + monkeypatch.setenv("PROXY_BASE_URL", "http://testserver/gateway") + monkeypatch.setenv("LITELLM_UI_PATH", str(tmp_path / "ui")) + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import json +from fastapi import FastAPI +from fastapi.testclient import TestClient +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router +from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +global_mcp_server_manager.registry['example'] = MCPServer( + server_id='example', name='example', server_name='example', alias='example', + transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + authorization_url='https://idp.example.com/authorize', token_url='https://idp.example.com/token', +) +app = FastAPI(root_path='/gateway') +app.include_router(router) +with TestClient(app) as client: + responses = { + path: client.get('/.well-known/oauth-authorization-server/gateway/' + path) + for path in ('mcp/example', 'example/mcp', 'example', 'mcp') + } + print(json.dumps({path: {'status': response.status_code, 'body': response.json()} + for path, response in responses.items()})) +""", + ], + capture_output=True, + text=True, + check=True, + timeout=60, + ) + responses = json.loads(result.stdout) + for path in ("mcp/example", "example/mcp", "example", "mcp"): + assert responses[path]["status"] == 200, responses[path] + assert responses[path]["body"]["issuer"] == f"http://testserver/gateway/{path}" + assert responses["example/mcp"]["body"]["token_endpoint"] == "http://testserver/gateway/example/token" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 56851d31241..67948375403 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1233,13 +1233,14 @@ class TestResolveByokMcpAuthHeader: assert result == "stored-cred" @pytest.mark.asyncio - async def test_byok_server_raises_401_when_no_credential_stored(self): + async def test_byok_server_raises_401_when_no_credential_stored(self, monkeypatch): from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _resolve_byok_mcp_auth_header, ) + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy") server = self._server(is_byok=True) user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard") @@ -1252,6 +1253,9 @@ class TestResolveByokMcpAuthHeader: assert exc_info.value.status_code == 401 assert exc_info.value.detail["error"] == "byok_auth_required" + assert exc_info.value.headers == { + "WWW-Authenticate": 'Bearer resource_metadata="https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"' + } @pytest.mark.asyncio async def test_byok_server_checks_credential_and_keeps_caller_header_when_supplied(self): 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 b637586d7ff..d2987c5112e 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 @@ -30,7 +30,7 @@ from mcp.types import ( TextResourceContents, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl +from pydantic import AnyUrl, TypeAdapter from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -3708,6 +3708,7 @@ class TestMCPServerManager: mock_prompt = Prompt(name="hello", description="Say hi") mock_client = AsyncMock() mock_client.list_prompts = AsyncMock(return_value=[mock_prompt]) + mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash") with patch.object( manager, @@ -3779,6 +3780,7 @@ class TestMCPServerManager: mock_client = AsyncMock() mock_resources = [Resource(name="file", uri="https://example.com/file")] mock_client.list_resources = AsyncMock(return_value=mock_resources) + mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash") prefixed_resources = [Resource(name="alias-server-file", uri="https://example.com/file")] with ( @@ -3788,11 +3790,6 @@ class TestMCPServerManager: new_callable=AsyncMock, return_value=mock_client, ) as mock_create_client, - patch.object( - manager, - "_create_prefixed_resources", - return_value=prefixed_resources, - ) as mock_prefix, ): result = await manager.get_resources_from_server( server=server, @@ -3808,7 +3805,6 @@ class TestMCPServerManager: assert called_kwargs["mcp_auth_header"] == "auth" assert called_kwargs["extra_headers"] == {"X-Test": "1", "X-Static": "static"} mock_client.list_resources.assert_awaited_once() - mock_prefix.assert_called_once_with(mock_resources, server, add_prefix=True) assert result == prefixed_resources @pytest.mark.asyncio @@ -3832,9 +3828,10 @@ class TestMCPServerManager: ) ] mock_client.list_resource_templates = AsyncMock(return_value=mock_templates) - prefixed_templates = [ + mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash") + expected_templates = [ ResourceTemplate( - name="alias-server-template", + name="template", uriTemplate="https://example.com/{id}", ) ] @@ -3846,11 +3843,6 @@ class TestMCPServerManager: new_callable=AsyncMock, return_value=mock_client, ) as mock_create_client, - patch.object( - manager, - "_create_prefixed_resource_templates", - return_value=prefixed_templates, - ) as mock_prefix, ): result = await manager.get_resource_templates_from_server( server=server, @@ -3866,10 +3858,10 @@ class TestMCPServerManager: extra_headers=None, stdio_env=None, subject_token=None, + user_api_key_auth=None, ) mock_client.list_resource_templates.assert_awaited_once() - mock_prefix.assert_called_once_with(mock_templates, server, add_prefix=False) - assert result == prefixed_templates + assert result == expected_templates @pytest.mark.asyncio async def test_read_resource_from_server_success(self): @@ -13020,3 +13012,436 @@ async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, mon assert cached.status == "healthy" assert len(attempts) == 2 assert route.call_count == 1 + + +class _DiscoveryClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +class _DiscoveryUpstream: + def __init__(self) -> None: + self.requests: tuple[tuple[str, str], ...] = () + self.outcome = "supported" + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.release.set() + + async def respond(self, request: httpx.Request) -> httpx.Response: + from mcp.types import JSONRPCMessage, JSONRPCRequest + + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + self.requests = (*self.requests, (payload.method, request.headers.get("authorization", ""))) + if payload.method == "initialize": + return httpx.Response(200, json={ + "jsonrpc": "2.0", "id": payload.id, + "result": {"protocolVersion": "2025-03-26", "serverInfo": {"name": "discovery", "version": "1"}, + "capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}}}, + }) + self.entered.set() + await self.release.wait() + if self.outcome == "failure": + return httpx.Response(503) + if self.outcome == "cancelled": + raise asyncio.CancelledError() + if self.outcome == "rejected": + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, + "error": {"code": -32601, "message": "Unsupported"}}) + result: Final = { + "prompts/list": {"prompts": [{"name": "example", "description": "original"}]}, + "resources/list": {"resources": [{"name": "example", "uri": "test://example", "description": "original"}]}, + "resources/templates/list": {"resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}]}, + "tools/list": {"tools": []}, + }[payload.method] + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + @property + def initializes(self) -> int: + return sum(method == "initialize" for method, _auth in self.requests) + + +def _discovery_server() -> MCPServer: + return MCPServer(server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ("prompts", "resources", "templates")) +async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None: + import respx + + clock: Final = _DiscoveryClock() + manager: Final = MCPServerManager(discovery_clock=clock) + upstream: Final = _DiscoveryUpstream() + operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server}[kind] + server: Final = _discovery_server() + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + first: Final = await operation(server, None) + assert len(first) == 1 + assert first[0].name == "discovery-example" + first[0].description = "caller changed it" + second: Final = await operation(server, None, add_prefix=False) + assert second[0].name == "example" + assert second[0].description == "original" + assert upstream.initializes == 1 + clock.now = 59.999 + assert (await operation(server, None))[0].name == "discovery-example" + assert upstream.initializes == 1 + clock.now = 60.001 + assert (await operation(server, None))[0].name == "discovery-example" + assert upstream.initializes == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ("prompts", "resources", "templates")) +@pytest.mark.parametrize("outcome", ("unsupported", "rejected", "failure")) +async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: str) -> None: + import respx + + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + upstream.outcome = outcome + operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server}[kind] + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + assert await operation(_discovery_server(), None) == [] + assert await operation(_discovery_server(), None) == [] + assert upstream.initializes == (2 if outcome == "failure" else 1) + if outcome == "failure": + upstream.outcome = "supported" + assert (await operation(_discovery_server(), None))[0].name == "discovery-example" + assert upstream.initializes == 3 + + +@pytest.mark.asyncio +async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_auth() -> None: + import respx + + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + server: Final = _discovery_server() + first_user: Final = UserAPIKeyAuth(user_id="first") + second_user: Final = UserAPIKeyAuth(user_id="second") + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + for user in (first_user, second_user): + assert len(await manager.get_prompts_from_server(server, user)) == 1 + assert upstream.initializes == 1 + for credential in ("first-secret", "second-secret", "first-secret"): + assert len(await manager.get_prompts_from_server(server, first_user, extra_headers={"Authorization": credential})) == 1 + assert upstream.initializes == 3 + assert {auth for method, auth in upstream.requests if method == "prompts/list"} == {"", "first-secret", "second-secret"} + + +@pytest.mark.asyncio +async def test_discovery_cache_coalesces_and_survives_waiter_cancellation() -> None: + import respx + + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + upstream.release.clear() + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + tasks: Final = tuple(asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10)) + await asyncio.wait_for(upstream.entered.wait(), timeout=5) + tasks[0].cancel() + with pytest.raises(asyncio.CancelledError): + await tasks[0] + upstream.release.set() + results: Final = await asyncio.wait_for(asyncio.gather(*tasks[1:]), timeout=5) + assert all(result[0].name == "discovery-example" for result in results) + assert upstream.initializes == 1 + assert results[0][0] is not results[1][0] + assert (await manager.get_prompts_from_server(_discovery_server(), None))[0].name == "discovery-example" + assert upstream.initializes == 1 + + +@pytest.mark.asyncio +async def test_discovery_cache_invalidation_during_fetch_does_not_repopulate_old_results() -> None: + import respx + + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + upstream.release.clear() + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + task: Final = asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) + await asyncio.wait_for(upstream.entered.wait(), timeout=5) + manager._invalidate_discovery_lists("discovery") + upstream.release.set() + assert (await task)[0].name == "discovery-example" + assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 + assert upstream.initializes == 2 + manager._invalidate_discovery_lists("discovery") + assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 + assert upstream.initializes == 3 + + +@pytest.mark.asyncio +async def test_discovery_cache_can_be_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + import respx + + monkeypatch.setenv("LITELLM_MCP_DISCOVERY_CACHE_TTL", "0") + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 + assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 + assert upstream.initializes == 2 + + +@pytest.mark.parametrize("value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5))) +def test_discovery_cache_ttl_validation(value: str, expected: float, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _mcp_discovery_cache_ttl + + monkeypatch.setenv("LITELLM_MCP_DISCOVERY_CACHE_TTL", value) + assert _mcp_discovery_cache_ttl() == expected + + +@pytest.mark.parametrize("auth_type", (MCPAuth.oauth2, MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag)) +def test_discovery_cache_keys_isolate_user_dependent_auth(auth_type: MCPAuth) -> None: + manager: Final = MCPServerManager() + server: Final = _discovery_server().model_copy(update={"auth_type": auth_type}) + first: Final = manager._discovery_key(server, UserAPIKeyAuth(user_id="first"), None, None, None, None) + second: Final = manager._discovery_key(server, UserAPIKeyAuth(user_id="second"), None, None, None, None) + anonymous: Final = manager._discovery_key(server, None, None, None, None, None) + assert len({first, second, anonymous}) == 3 + assert "first" not in str(first) + assert "second" not in str(second) + + +@pytest.mark.asyncio +async def test_discovery_cache_retries_cancelled_fetches() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + + async def cancelled() -> list[Prompt]: + raise asyncio.CancelledError() + + async def supported() -> list[Prompt]: + return [Prompt(name="recovered")] + + with pytest.raises(asyncio.CancelledError): + await cache.get(("server", None), cancelled) + assert [item.name for item in await cache.get(("server", None), supported)] == ["recovered"] + + +@pytest.mark.asyncio +async def test_discovery_cache_cancels_fetch_when_last_waiter_leaves() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + entered: Final = asyncio.Event() + stopped: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def fetch() -> list[Prompt]: + entered.set() + try: + await release.wait() + return [Prompt(name="result")] + finally: + stopped.set() + + tasks: Final = tuple(asyncio.create_task(cache.get(("server", None), fetch)) for _ in range(3)) + await asyncio.wait_for(entered.wait(), timeout=5) + for task in tasks: + task.cancel() + outcomes: Final = await asyncio.gather(*tasks, return_exceptions=True) + assert all(isinstance(outcome, asyncio.CancelledError) for outcome in outcomes) + try: + await asyncio.wait_for(stopped.wait(), timeout=1) + finally: + release.set() + + +@pytest.mark.asyncio +async def test_discovery_cache_bounds_detached_fetches_without_dropping_results() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + entered: Final[asyncio.Queue[None]] = asyncio.Queue() + release: Final = asyncio.Event() + + async def blocked() -> list[Prompt]: + await entered.put(None) + await release.wait() + return [Prompt(name="blocked")] + + tasks: Final = tuple(asyncio.create_task(cache.get((str(index), None), blocked)) for index in range(1024)) + try: + for _ in tasks: + await asyncio.wait_for(entered.get(), timeout=5) + active_tasks: Final = frozenset(asyncio.all_tasks()) + + async def overflow() -> list[Prompt]: + assert frozenset(asyncio.all_tasks()) <= active_tasks + return [Prompt(name="overflow")] + + result: Final = await cache.get(("overflow", None), overflow) + assert [item.name for item in result] == ["overflow"] + finally: + release.set() + outcomes: Final = await asyncio.gather(*tasks) + assert all(result[0].name == "blocked" for result in outcomes) + + +@pytest.mark.asyncio +async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> None: + import respx + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import UpstreamCredentialProvider + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError, ServerSpec, Subject + + class CredentialSource(UpstreamCredentialProvider): + def __init__(self) -> None: + super().__init__() + self.token: str | None = "token-a" + + async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: + if self.token is None: + return Error(CredError.of_unauthorized("Credential revoked")) + return Ok(StaticHeaderAuth("Bearer " + self.token)) + + source: Final = CredentialSource() + managers: Final = (MCPServerManager(cred_provider=source), MCPServerManager(cred_provider=source)) + server: Final = MCPServer( + server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client", + authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token", + ) + user: Final = UserAPIKeyAuth(user_id="same-user", api_key="same-key") + upstream: Final = _DiscoveryUpstream() + + async def respond(request: httpx.Request) -> httpx.Response: + response: Final = await upstream.respond(request) + if '"prompts/list"' not in request.content.decode(): + return response + from mcp.types import JSONRPCMessage, JSONRPCRequest + + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + assert isinstance(payload, JSONRPCRequest) + name: Final = {"Bearer token-a": "account-a", "Bearer token-b": "account-b"}[request.headers["authorization"]] + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"prompts": [{"name": name}]}}) + + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=respond) + for manager in managers: + assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-a"] + assert upstream.initializes == 2 + source.token = "token-b" + for manager in managers: + assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-b"] + assert upstream.initializes == 4 + source.token = None + for manager in managers: + assert await manager.get_prompts_from_server(server, user) == [] + assert upstream.initializes == 4 + + +@pytest.mark.asyncio +async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None: + import respx + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken + + class TokenStore: + def __init__(self) -> None: + self.calls: tuple[tuple[str, str], ...] = () + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + self.calls = (*self.calls, (user_id, server_id)) + return OAuthToken(access_token="stored-token") + + async def invalidate(self, user_id: str, server_id: str) -> None: + return None + + store: Final = TokenStore() + manager: Final = MCPServerManager(per_user_oauth_token_store=store) + server: Final = MCPServer( + server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client", + authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token", + ) + user: Final = UserAPIKeyAuth(user_id="requesting-user") + upstream: Final = _DiscoveryUpstream() + with respx.mock(base_url="https://discovery.example") as router: + router.route().mock(side_effect=upstream.respond) + assert len(await manager.get_prompts_from_server(server, user)) == 1 + assert len(await manager.get_prompts_from_server(server, user)) == 1 + assert store.calls == (("requesting-user", "discovery"), ("requesting-user", "discovery")) + assert upstream.initializes == 1 + assert ("prompts/list", "Bearer stored-token") in upstream.requests + + +@pytest.mark.asyncio +async def test_discovery_cache_evicts_results_at_capacity() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + + async def original() -> list[Prompt]: + return [Prompt(name="original")] + + async def refetched() -> list[Prompt]: + return [Prompt(name="refetched")] + + for index in range(1025): + assert (await cache.get((f"server-{index:04}", None), original))[0].name == "original" + assert (await cache.get(("server-1024", None), refetched))[0].name == "original" + assert (await cache.get(("server-0000", None), refetched))[0].name == "refetched" + + +@pytest.mark.asyncio +async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_fetches() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def original() -> list[Prompt]: + return [Prompt(name="original")] + + async def blocked() -> list[Prompt]: + entered.set() + await release.wait() + return [Prompt(name="pending")] + + async def refetched() -> list[Prompt]: + return [Prompt(name="refetched")] + + assert (await cache.get(("server", None), original))[0].name == "original" + assert (await cache.get(("server-extra", None), original))[0].name == "original" + task: Final = asyncio.create_task(cache.get(("other", None), blocked)) + await asyncio.wait_for(entered.wait(), timeout=5) + cache.invalidate("server") + release.set() + assert (await asyncio.wait_for(task, timeout=5))[0].name == "pending" + assert (await cache.get(("other", None), refetched))[0].name == "pending" + assert (await cache.get(("server-extra", None), refetched))[0].name == "original" + assert (await cache.get(("server", None), refetched))[0].name == "refetched" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("description", ("x" * 96_000, "é" * 40_000), ids=("ascii", "unicode")) +async def test_discovery_cache_returns_oversized_results_without_retaining_them(description: str) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + fetch: Final = AsyncMock(return_value=[Prompt(name="large", description=description)]) + for _ in range(2): + result: Final = await cache.get(("server", None), fetch) + assert result[0].description == description + assert fetch.await_count == 2 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8777e24e209..d9a065db3cb 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -29,6 +29,7 @@ from litellm.proxy._types import ( ProxyException, SSOUserDefinedValues, UserAPIKeyAuth, + WebhookEvent, ) from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, @@ -53,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.constants import ( DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, @@ -5431,6 +5433,9 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5445,13 +5450,136 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): general_settings={}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) assert "User=u1" in str(over.value) +async def _run_internal_user_budget_alert( + *, + spend: float, +) -> tuple[AsyncMock, litellm.BudgetExceededError | None]: + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user: Final = LiteLLM_UserTable( + user_id="user-1", + user_email="person@example.com", + spend=0.0, + max_budget=100.0, + ) + token: Final = UserAPIKeyAuth(token="hashed-key-1", user_id="user-1") + slack_alerting: Final = SlackAlerting(alerting=["webhook"]) + send_alert: Final = AsyncMock() + alert_finished: Final = asyncio.Event() + + async def _get_spend( + counter_key: str, + fallback_spend: float, + max_budget: float | None = None, + **kwargs: object, + ) -> float: + assert counter_key == "spend:user:user-1" + assert fallback_spend == 0.0 + assert max_budget == 100.0 + return spend + + async def _budget_alerts( + *, + type: Literal["user_budget"], + user_info: CallInfo, + ) -> None: + assert type == "user_budget" + try: + await slack_alerting.budget_alerts(type=type, user_info=user_info) + finally: + alert_finished.set() + + proxy_logging_obj: Final = MagicMock(budget_alerts=_budget_alerts) + + async def _check() -> bool: + return await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=proxy_logging_obj, + valid_token=token, + request=MagicMock(spec=Request), + ) + + async def _check_for_error() -> litellm.BudgetExceededError | None: + if spend < 100.0: + assert await _check() is True + return None + + with pytest.raises(litellm.BudgetExceededError) as raised: + await _check() + return raised.value + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam + patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch.object(slack_alerting, "send_alert", send_alert), + ): + error: Final = await _check_for_error() + await asyncio.wait_for(alert_finished.wait(), timeout=1.0) + + return send_alert, error + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_below_threshold_does_not_emit_alert(): + send_alert, error = await _run_internal_user_budget_alert(spend=84.0) + + assert error is None + send_alert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_emits_user_threshold_event(): + send_alert, error = await _run_internal_user_budget_alert(spend=85.0) + + assert error is None + send_alert.assert_awaited_once() + event: Final = send_alert.await_args.kwargs["user_info"] + assert isinstance(event, WebhookEvent) + assert event.event == "threshold_crossed" + assert event.event_group == Litellm_EntityType.USER + assert event.user_id == "user-1" + assert event.user_email == "person@example.com" + assert event.spend == 85.0 + assert event.max_budget == 100.0 + assert event.token is None + assert event.key_alias is None + assert event.team_id is None + assert event.organization_id is None + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_emits_crossed_event_and_rejects(): + send_alert, error = await _run_internal_user_budget_alert(spend=100.0) + + assert error is not None + assert error.current_cost == 100.0 + assert error.max_budget == 100.0 + send_alert.assert_awaited_once() + event: Final = send_alert.await_args.kwargs["user_info"] + assert isinstance(event, WebhookEvent) + assert event.event == "budget_crossed" + assert event.event_group == Litellm_EntityType.USER + assert event.user_id == "user-1" + assert event.user_email == "person@example.com" + + @pytest.mark.asyncio async def test_common_checks_personal_user_budget_skipped_for_team_key(): """A user's personal max_budget does not apply to a team-scoped key. @@ -5475,6 +5603,9 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): async def _no_membership(*args, **kwargs): return None + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5489,11 +5620,12 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): general_settings={}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) assert result is True + proxy_logging_obj.budget_alerts.assert_not_awaited() @pytest.mark.asyncio @@ -5518,6 +5650,9 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag async def _no_membership(*args, **kwargs): return None + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5533,10 +5668,11 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag general_settings={"apply_user_budget_to_team_keys": True}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) assert "ExceededBudget: User=u1" in str(exc_info.value) @@ -5553,6 +5689,9 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5567,10 +5706,11 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key general_settings={"apply_user_budget_to_team_keys": True}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) @pytest.mark.parametrize( @@ -5651,6 +5791,42 @@ async def test_budget_checks_only_run_on_llm_api_routes(scope, route, expect_blo assert await _run() is True +@pytest.mark.asyncio +async def test_organization_budget_check_carries_org_state_on_the_token(): + """The org row auth already fetched is pinned on the token so the response path + (Prometheus org budget gauges) reads it from request metadata instead of calling + get_org_object again.""" + from litellm.proxy._types import LiteLLM_OrganizationTable + from litellm.proxy.auth.auth_checks import _organization_max_budget_check + from litellm.types.proxy.carried_budget_state import OrgBudgetSnapshot + + org_table = LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=12.5, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + token = UserAPIKeyAuth(token="k1", org_id="o1") + user_api_key_cache = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key="org_id:o1:with_budget", value=org_table, model_type=LiteLLM_OrganizationTable + ) + + await _organization_max_budget_check( + valid_token=token, + team_object=None, + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=MagicMock(), + ) + + assert token.organization_alias == "platform-org" + assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0) + + @pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"]) @pytest.mark.asyncio async def test_spend_capable_non_llm_routes_still_enforce_budget(route): diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py new file mode 100644 index 00000000000..0fd0dda3017 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -0,0 +1,338 @@ +"""Counts the Redis round trips and DB queries auth object reads cost per cache regime, and checks that the +per-object getters still enforce on their own when the prefetch cannot help.""" + +import json +from collections.abc import Sequence +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import ( + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + get_org_object, + get_team_membership, + get_team_object, + get_user_object, +) +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +USER_ID = "prefetch-user" +TEAM_ID = "prefetch-team" +ORG_ID = "prefetch-org" + +USER_ROW = { + "user_id": USER_ID, + "max_budget": 50.0, + "spend": 1.0, + "models": ["gpt-5.4-mini"], + "organization_memberships": [], +} +TEAM_ROW = { + "team_id": TEAM_ID, + "organization_id": ORG_ID, + "max_budget": 500.0, + "spend": 2.0, + "models": [], + "blocked": False, + "members_with_roles": {}, +} +MEMBERSHIP_ROW = { + "user_id": USER_ID, + "team_id": TEAM_ID, + "spend": 3.0, + "budget_id": "b1", + "litellm_budget_table": {"budget_id": "b1", "max_budget": 20.0}, +} +ORG_ROW = { + "organization_id": ORG_ID, + "organization_alias": "org", + "budget_id": "b2", + "created_by": "admin", + "updated_by": "admin", + "models": [], + "spend": 4.0, + "litellm_budget_table": {"budget_id": "b2", "max_budget": 1000.0}, +} +ALL_ROWS = { + "user_row": USER_ROW, + "team_row": TEAM_ROW, + "membership_row": MEMBERSHIP_ROW, + "organization_row": ORG_ROW, + "project_row": None, +} + + +class CountingRedis(RedisCache): + """Redis fake that counts commands and round trips (an MGET or a pipeline is one round trip).""" + + def __init__(self, store: dict[str, str] | None = None, fail: bool = False) -> None: + self.store: dict[str, str] = dict(store or {}) + self.fail = fail + self.round_trips = 0 + self.commands: list[str] = [] + + def _trip(self, *commands: str) -> None: + if self.fail: + raise ConnectionError("redis down") + self.round_trips += 1 + self.commands.extend(commands) + + async def async_get_cache(self, key: str, **kwargs: object) -> object: + self._trip(f"GET {key}") + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, object]: + self._trip(f"MGET {' '.join(key_list)}") + return {key: (json.loads(self.store[key]) if key in self.store else None) for key in key_list} + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + self._trip(f"SET {key}") + self.store[key] = json.dumps(value) + + async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> None: + self._trip(*(f"SET {key}" for key, _ in cache_list)) + for key, value in cache_list: + self.store[key] = json.dumps(value) + + async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: + self._trip(*(f"SET {key} ttl={ttl}" for key, _, ttl in cache_list)) + for key, value, _ in cache_list: + self.store[key] = json.dumps(value) + + async def async_delete_cache(self, key: str) -> None: + self._trip(f"DEL {key}") + self.store.pop(key, None) + + +def _prisma(rows: dict[str, object] | None = ALL_ROWS) -> MagicMock: + prisma = MagicMock(name="prisma_client") + prisma.db.query_first = AsyncMock(return_value=rows) + return prisma + + +def _non_prefetch_db_calls(prisma: MagicMock) -> list[str]: + return [str(call) for call in prisma.db.mock_calls if not str(call).startswith("call.query_first(")] + + +def _cache(redis: RedisCache | None) -> UserApiKeyCache: + return UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=redis) + + +def _refs() -> AuthObjectRefs: + return AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", user_id=USER_ID, team_id=TEAM_ID, org_id=ORG_ID)) + + +async def _read_all_through_getters( + cache: UserApiKeyCache, prisma: MagicMock +) -> tuple[ + LiteLLM_UserTable | None, + LiteLLM_TeamTableCachedObj, + LiteLLM_TeamMembership | None, + LiteLLM_OrganizationTable | None, +]: + return ( + await get_user_object(user_id=USER_ID, prisma_client=prisma, user_api_key_cache=cache, user_id_upsert=False), + await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache), + await get_team_membership(user_id=USER_ID, team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache), + await get_org_object(org_id=ORG_ID, prisma_client=prisma, user_api_key_cache=cache, include_budget_table=True), + ) + + +def test_refs_from_token_only_names_membership_when_both_ids_present(): + assert AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", team_id=TEAM_ID)).membership_user_id is None + assert AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", user_id=USER_ID)).membership_user_id is None + assert AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", user_id=USER_ID, team_id=TEAM_ID)) == AuthObjectRefs( + user_id=USER_ID, team_id=TEAM_ID, membership_user_id=USER_ID + ) + + +@pytest.mark.asyncio +async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_again(): + redis = CountingRedis() + prisma = _prisma() + cache = _cache(redis) + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + assert prisma.db.query_first.await_count == 1 + assert prisma.db.query_first.await_args.args[1:] == (USER_ID, TEAM_ID, USER_ID, ORG_ID, None) + mgets = [c for c in redis.commands if c.startswith("MGET")] + assert len(mgets) == 1 + assert set(mgets[0].split()[1:]) == { + USER_ID, + f"team_id:{TEAM_ID}", + f"{TEAM_ID}_{USER_ID}", + f"team_membership:{USER_ID}:{TEAM_ID}", + f"org_id:{ORG_ID}", + f"org_id:{ORG_ID}:with_budget", + } + sets = sorted(c for c in redis.commands if c.startswith("SET")) + assert sets == sorted( + [ + f"SET {TEAM_ID}_{USER_ID} ttl=5", + f"SET org_id:{ORG_ID} ttl=5", + f"SET org_id:{ORG_ID}:with_budget ttl=5", + f"SET {USER_ID} ttl=60", + f"SET team_id:{TEAM_ID} ttl=60", + f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", + ] + ) + assert redis.round_trips == 2, "one MGET, one pipeline" + + before = (redis.round_trips, prisma.db.query_first.await_count) + user, team, membership, org = await _read_all_through_getters(cache, prisma) + assert (redis.round_trips, prisma.db.query_first.await_count) == before + assert _non_prefetch_db_calls(prisma) == [] + + assert isinstance(user, LiteLLM_UserTable) and user.max_budget == 50.0 + assert isinstance(team, LiteLLM_TeamTableCachedObj) and team.organization_id == ORG_ID + assert team.last_refreshed_at is not None + assert isinstance(membership, LiteLLM_TeamMembership) and membership.litellm_budget_table is not None + assert membership.litellm_budget_table.max_budget == 20.0 + assert isinstance(org, LiteLLM_OrganizationTable) and org.litellm_budget_table is not None + assert org.litellm_budget_table.max_budget == 1000.0 + + +@pytest.mark.asyncio +async def test_redis_warm_regime_is_exactly_one_mget_and_zero_queries(): + seeded = CountingRedis() + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma()) + + redis = CountingRedis(store=seeded.store) + prisma = _prisma() + cache = _cache(redis) + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + assert redis.round_trips == 1 + assert redis.commands[0].startswith("MGET") + assert prisma.db.query_first.await_count == 0 + + user, team, membership, org = await _read_all_through_getters(cache, prisma) + assert redis.round_trips == 1 + assert prisma.db.mock_calls == [] + assert (user.user_id, team.team_id, membership.team_id, org.organization_id) == (USER_ID, TEAM_ID, TEAM_ID, ORG_ID) + + +@pytest.mark.asyncio +async def test_hot_regime_costs_nothing(): + redis = CountingRedis() + prisma = _prisma() + cache = _cache(redis) + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + redis.round_trips, redis.commands = 0, [] + prisma.db.reset_mock() + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + await _read_all_through_getters(cache, prisma) + + assert redis.round_trips == 0 + assert prisma.db.mock_calls == [] + + +@pytest.mark.asyncio +async def test_partial_redis_hit_queries_only_the_missing_objects(): + seeded = CountingRedis() + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma()) + for key in (f"team_id:{TEAM_ID}", f"{TEAM_ID}_{USER_ID}", f"team_membership:{USER_ID}:{TEAM_ID}"): + del seeded.store[key] + + prisma = _prisma() + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=prisma) + + assert prisma.db.query_first.await_count == 1 + assert prisma.db.query_first.await_args.args[1:] == (None, TEAM_ID, USER_ID, None, None) + + del seeded.store[f"org_id:{ORG_ID}:with_budget"] + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=prisma) + assert prisma.db.query_first.await_args.args[1:] == (None, None, None, ORG_ID, None) + + +@pytest.mark.asyncio +async def test_deleted_team_cache_entry_is_refetched_and_the_update_is_visible(): + redis = CountingRedis() + prisma = _prisma() + cache = _cache(redis) + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + await cache.async_delete_cache(f"team_id:{TEAM_ID}") + prisma.db.query_first.return_value = {**ALL_ROWS, "team_row": {**TEAM_ROW, "blocked": True}} + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + team = await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache) + assert team.blocked is True + assert prisma.db.query_first.await_count == 2 + assert prisma.db.query_first.await_args.args[1:] == (None, TEAM_ID, None, None, None) + + +@pytest.mark.asyncio +async def test_row_missing_a_required_column_is_not_cached_so_the_getter_still_fails_closed(): + prisma = _prisma({**ALL_ROWS, "team_row": {"max_budget": 1.0}}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + cache = _cache(CountingRedis()) + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + with pytest.raises(HTTPException) as exc: + await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache) + assert exc.value.status_code == 404 + assert prisma.db.litellm_teamtable.find_unique.await_count == 1 + assert cache.in_memory_cache.get_cache(USER_ID) is not None + + +@pytest.mark.asyncio +async def test_absent_rows_are_not_cached_as_present(): + redis = CountingRedis() + prisma = _prisma({key: None for key in ALL_ROWS}) + cache = _cache(redis) + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + assert [c for c in redis.commands if c.startswith("SET")] == [] + assert cache.in_memory_cache.get_cache(f"team_id:{TEAM_ID}") is None + + +@pytest.mark.asyncio +async def test_redis_failure_is_swallowed_and_getters_fall_back_to_their_own_reads(): + prisma = _prisma() + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + cache = _cache(CountingRedis(fail=True)) + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + assert prisma.db.query_first.await_count == 0 + assert cache.in_memory_cache.get_cache(USER_ID) is None + + +@pytest.mark.asyncio +async def test_no_prisma_still_uses_redis_but_never_queries(): + seeded = CountingRedis() + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma()) + redis = CountingRedis(store=seeded.store) + cache = _cache(redis) + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=None) + + assert redis.round_trips == 1 + assert cache.in_memory_cache.get_cache(f"org_id:{ORG_ID}") is not None + + +@pytest.mark.asyncio +async def test_no_redis_goes_straight_to_one_query(): + prisma = _prisma() + cache = _cache(None) + + await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) + + assert prisma.db.query_first.await_count == 1 + assert cache.in_memory_cache.get_cache(f"team_membership:{USER_ID}:{TEAM_ID}") is not None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6cce6d0316b..0cdbcde6abc 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5,7 +5,7 @@ import os import subprocess import sys from contextlib import contextmanager -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path from textwrap import dedent from types import SimpleNamespace @@ -23,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -48,6 +49,7 @@ from litellm.proxy.auth.user_api_key_auth import ( get_api_key, user_api_key_auth, ) +from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata class _RoutingRequest: @@ -1646,6 +1648,95 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.asyncio +@pytest.mark.parametrize("model_allowed", [True, False]) +async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_the_model(model_allowed): + """A request denied by the key's model list must not pay for the team/user/org MGET or DB join.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.proxy_server import hash_token + + api_key = "sk-prefetch-order-test" + valid_token = UserAPIKeyAuth(api_key=api_key, token=hash_token(api_key), user_id="u1", team_id="t1") + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.delete_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs_to_set = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} + denied = ProxyException( + message="Key not allowed to access model", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=401, + ) + try: + for attr, val in _attrs_to_set.items(): + setattr(_proxy_server_mod, attr, val) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with ( + patch( # test-quality-ok: the builder has no DI seam for the key lookup; stands in for the DB + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( # test-quality-ok: the observable is whether the prefetch runs before or after this check + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + side_effect=None if model_allowed else denied, + ), + patch( # test-quality-ok: counting prefetch calls on a denied request IS the regression being pinned + "litellm.proxy.auth.user_api_key_auth.prefetch_auth_objects", new_callable=AsyncMock + ) as mock_prefetch, + patch( # test-quality-ok: no DB in this test; the user lookup must not fail the allowed path + "litellm.proxy.auth.user_api_key_auth.get_user_object", new_callable=AsyncMock, return_value=None + ), + ): + call = _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o"}, + ) + if model_allowed: + assert isinstance(await call, UserAPIKeyAuth) + mock_prefetch.assert_awaited_once() + assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + else: + with pytest.raises(ProxyException) as exc: + await call + assert exc.value.type == ProxyErrorTypes.key_model_access_denied + mock_prefetch.assert_not_awaited() + finally: + for attr, val in _original_values.items(): + setattr(_proxy_server_mod, attr, val) + + @pytest.mark.asyncio async def test_return_user_api_key_auth_obj_user_spend_and_budget(): """ @@ -3930,6 +4021,65 @@ async def test_centralized_common_checks_routes_header_tags_to_litellm_metadata( assert "metadata" not in request_data +@pytest.mark.asyncio +async def test_centralized_common_checks_carries_team_and_user_budget_state_on_the_token(): + """The team and user objects auth resolves are pinned on the token so the + response path (Prometheus budget gauges) reads them from request metadata + instead of calling get_team_object / get_user_object again.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + reset_at = datetime(2026, 10, 1, tzinfo=timezone.utc) + token = UserAPIKeyAuth(api_key="sk-test", token="hashed", team_id="t1", user_id="u1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", budget_reset_at=reset_at, max_budget=300.0), + ) + await user_api_key_cache.async_set_cache( + key="u1", + value=LiteLLM_UserTable(user_id="u1", user_alias="Alice", budget_reset_at=None, max_budget=None), + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=None), + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "proxy_logging_obj": proxy_logging_obj, + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( # test-quality-ok: the authz gate has its own tests above; this one checks the carry step before it + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-5.4-mini"}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert dict(carried_budget_metadata(token)) == { + "user_api_key_team_budget_reset_at": "2026-10-01T00:00:00Z", + "user_api_key_team_table_max_budget": 300.0, + "user_api_key_user_budget_reset_at": None, + "user_api_key_user_table_max_budget": None, + "user_api_key_user_alias": "Alice", + } + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without diff --git a/tests/test_litellm/proxy/client/cli/conftest.py b/tests/test_litellm/proxy/client/cli/conftest.py index c77f516a768..e54dbeef875 100644 --- a/tests/test_litellm/proxy/client/cli/conftest.py +++ b/tests/test_litellm/proxy/client/cli/conftest.py @@ -1,5 +1,6 @@ import os -from collections.abc import Iterator +import shlex +from collections.abc import Callable, Iterator from pathlib import Path from typing import Final @@ -19,12 +20,52 @@ def _statusline_script_under_tmp(monkeypatch: pytest.MonkeyPatch, tmp_path: Path monkeypatch.setattr(claude_settings, "STATUSLINE_SCRIPT_PATH", tmp_path / "litellm-home" / "statusline.py") +@pytest.fixture +def fake_codex_version( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> Callable[[str | None, int], Path]: + directory: Final = tmp_path / "codex-bin" + directory.mkdir() + binary: Final = directory / ("codex.cmd" if os.name == "nt" else "codex") + version_output: Final = directory / "version-output.txt" + monkeypatch.setenv("PATH", str(directory)) + + def install(output: str | None, returncode: int = 0) -> Path: + if output is None: + binary.unlink(missing_ok=True) + return binary + version_output.write_text(output) + if os.name == "nt": + binary.write_text( + '@echo off\nif not "%~1"=="--version" exit /b 2\n' + 'if not "%~2"=="" exit /b 2\ntype "%~dp0version-output.txt"\n' + f'exit /b {returncode}\n' + ) + else: + binary.write_text( + '#!/bin/sh\nif [ "$#" -ne 1 ] || [ "$1" != "--version" ]; then\n exit 2\nfi\n' + f'/bin/cat {shlex.quote(str(version_output))}\nexit {returncode}\n' + ) + binary.chmod(0o700) + return binary + + install("codex-cli 0.129.0\n") + return install + + +@pytest.fixture(autouse=True) +def _isolated_codex_version_for_configure_tests(request: pytest.FixtureRequest) -> None: + if request.node.path.name in ("test_codex_settings.py", "test_configure_commands.py"): + request.getfixturevalue("fake_codex_version") + + @pytest.fixture(autouse=True) def isolated_claude_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]: before: Final = _current_bytes() monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(tmp_path / ".claude")) + monkeypatch.setenv("CODEX_HOME", str(tmp_path / ".codex")) yield tmp_path after: Final = _current_bytes() if after == before: diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index 3020b770e62..b122ea7b20e 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -172,7 +172,9 @@ class TestAgentLaunchArgs: assert 'model_providers.litellm.env_key="OPENAI_API_KEY"' in args assert 'model_providers.litellm.wire_api="responses"' in args assert "model_providers.litellm.supports_websockets=false" in args - assert joined.count("-c") == 6 + assert "model_providers.litellm.requires_openai_auth=false" in args + assert "model_providers.litellm.http_headers={}" in args + assert joined.count("-c") == 8 def test_codex_uses_basename(self): assert agent_launch_args("/usr/local/bin/codex", "http://localhost:4000") == ( diff --git a/tests/test_litellm/proxy/client/cli/test_codex_settings.py b/tests/test_litellm/proxy/client/cli/test_codex_settings.py new file mode 100644 index 00000000000..0d1f2b4056a --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_codex_settings.py @@ -0,0 +1,341 @@ +import json +import stat +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import pytest +import tomlkit + +from litellm.litellm_core_utils.private_json import commit_staged_json +from litellm.proxy.client.cli.commands import codex_settings as codex_settings_module +from litellm.proxy.client.cli.commands.agents import ( + agent_launch_args, + codex_config_path, +) +from litellm.proxy.client.cli.commands.codex_settings import ( + CodexSettingsError, + _snapshot, + _with, + codex_configure_state_path, + configure_codex_settings, + preflight_codex_settings, + unconfigure_codex_settings, +) + +GATEWAY: Final = "https://gateway.example.com/team" +KEY: Final = "sk-test-new-gateway-key" +MODEL: Final = "gateway-codex-model" + + +@pytest.mark.parametrize("path,existing,first_value,second_value", [ + ("model", 'model = "original" # starting model\n', 'value = "first"\n', 'value = "second"\n'), + ("model_providers.litellm", '', '[value]\nname = "first"\n', '[value]\nname = "second"\n'), + ("model_providers.litellm", '[model_providers.litellm]\nname = "original" # provider\n', + '[value]\nname = "first"\n', '[value]\nname = "second"\n'), +]) +def test_toml_transitions_leave_source_and_independent_results_unchanged( + path: str, existing: str, first_value: str, second_value: str +) -> None: + source: Final = tomlkit.parse('# user settings\n' + existing + '[profiles.work]\nmodel = "keep" # profile\n') + original_bytes: Final = source.as_string().encode() + first: Final = _with(source, path, first_value) + first_bytes: Final = first.as_string().encode() + second: Final = _with(source, path, second_value) + second_bytes: Final = second.as_string().encode() + removed: Final = _with(first, path, None) + first_snapshot: Final = _snapshot(first, path) + second_snapshot: Final = _snapshot(second, path) + assert source.as_string().encode() == original_bytes + assert first.as_string().encode() == first_bytes + assert second.as_string().encode() == second_bytes + assert first_snapshot is not None and tomlkit.parse(first_snapshot) == tomlkit.parse(first_value) + assert second_snapshot is not None and tomlkit.parse(second_snapshot) == tomlkit.parse(second_value) + assert _snapshot(removed, path) is None + for result in (first, second, removed): + assert result["profiles"] == source["profiles"] + assert '# user settings' in result.as_string() + assert '# profile' in result.as_string() + + +def test_persistent_provider_is_complete_and_preserves_unrelated_toml(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text( + '# user settings\nmodel = "old-model" # starting model\n' + 'model_provider = "openai"\nprofile = "work"\n' + '[model_providers.litellm]\nname = "old gateway"\n' + 'base_url = "https://old.example.com/v1"\nenv_key = "OLD_KEY"\n' + 'experimental_bearer_token = "sk-old"\nrequires_openai_auth = true\n' + '[model_providers.litellm.auth]\ncommand = "old-token-helper"\n' + '[model_providers.other]\nname = "Keep me" # other provider\n' + '[profiles.work]\nmodel = "work-model"\n' + '[[hooks.Stop]]\nhooks = [{type = "command", command = "echo done"}]\n' + ) + original: Final = tomlkit.parse(path.read_text()) + configure_codex_settings(GATEWAY, KEY, MODEL, path) + configured: Final = tomlkit.parse(path.read_text()) + assert configured["model"] == MODEL + assert configured["model_provider"] == "litellm" + assert "profile" not in configured + assert configured["model_providers"]["litellm"] == { + "name": "LiteLLM proxy", + "base_url": GATEWAY + "/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": False, + "http_headers": {"Authorization": "Bearer " + KEY}, + } + assert configured["model_providers"]["other"] == original["model_providers"]["other"] + assert configured["profiles"] == original["profiles"] + assert configured["hooks"] == original["hooks"] + assert "# user settings" in path.read_text() + assert "# other provider" in path.read_text() + assert KEY not in codex_configure_state_path(path).read_text() + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + assert stat.S_IMODE(codex_configure_state_path(path).stat().st_mode) == 0o600 + assert stat.S_IMODE(codex_configure_state_path(path).parent.stat().st_mode) == 0o700 + outcome: Final = unconfigure_codex_settings(path) + assert not outcome.kept and not outcome.file_removed + assert tomlkit.parse(path.read_text()) == original + assert "# starting model" in path.read_text() + assert "# other provider" in path.read_text() + assert not codex_configure_state_path(path).exists() + + +@pytest.mark.parametrize("original", [None, "", "# my preferences\n", '[model_providers]\n']) +def test_undo_distinguishes_missing_empty_and_existing_tables(tmp_path: Path, original: str | None) -> None: + path: Final = tmp_path / "config.toml" + if original is not None: + path.write_text(original) + configure_codex_settings(GATEWAY, KEY, MODEL, path) + outcome: Final = unconfigure_codex_settings(path) + assert outcome.file_removed == (original is None) + assert path.exists() == (original is not None) + if original is not None: + assert tomlkit.parse(path.read_text()) == tomlkit.parse(original) + assert original.strip() in path.read_text() + + +def test_repeat_setup_preserves_original_and_undo_keeps_user_edits(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('model = "original"\nmodel_provider = "openai"\n') + configure_codex_settings(GATEWAY, KEY, MODEL, path) + configure_codex_settings(GATEWAY + "/second", "sk-second", "second-model", path) + assert tomlkit.parse(path.read_text())["model"] == "second-model" + path.write_text(path.read_text().replace('model = "second-model"', 'model = "my-custom-model"')) + outcome: Final = unconfigure_codex_settings(path) + assert outcome.kept == ("model",) + assert tomlkit.parse(path.read_text()) == {"model": "my-custom-model", "model_provider": "openai"} + + +def test_repeat_setup_restores_the_user_value_it_displaced(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('model = "original"\n') + configure_codex_settings(GATEWAY, KEY, MODEL, path) + path.write_text(path.read_text().replace('model = "gateway-codex-model"', 'model = "user-edited"')) + configure_codex_settings(GATEWAY, "sk-rotated", "third-model", path) + unconfigure_codex_settings(path) + assert tomlkit.parse(path.read_text()) == {"model": "user-edited"} + + +def test_undo_keeps_provider_credentials_and_endpoint_together(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('[model_providers.litellm]\nbase_url = "https://old.example.com/v1"\n' + 'http_headers = { Authorization = "Bearer old-key" }\n') + configure_codex_settings(GATEWAY, KEY, MODEL, path) + path.write_text(path.read_text().replace(GATEWAY, "https://user.example.com")) + outcome: Final = unconfigure_codex_settings(path) + provider: Final = tomlkit.parse(path.read_text())["model_providers"]["litellm"] + assert outcome.kept == ("model_providers.litellm",) + assert provider["base_url"] == "https://user.example.com/v1" + assert provider["http_headers"] == {"Authorization": "Bearer " + KEY} + assert "old-key" not in path.read_text() + + +def test_user_deleted_config_is_not_recreated_to_restore_profile(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('profile = "old-profile"\n') + configure_codex_settings(GATEWAY, KEY, MODEL, path) + path.unlink() + outcome: Final = unconfigure_codex_settings(path) + assert outcome.file_removed and outcome.restored == () + assert not path.exists() + + +def test_user_comment_in_new_config_survives_undo(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + configure_codex_settings(GATEWAY, KEY, MODEL, path) + path.write_text("# keep my note\n" + path.read_text()) + assert not unconfigure_codex_settings(path).file_removed + assert "# keep my note" in path.read_text() + + +def test_code_home_and_symlink_aliases_share_receipt_and_write_target(tmp_path: Path) -> None: + target: Final = tmp_path / "real-config.toml" + target.write_text('model = "old"\n') + custom_home: Final = tmp_path / "codex-home" + custom_home.mkdir() + alias: Final = codex_config_path({"CODEX_HOME": str(custom_home)}) + alias.symlink_to(target) + configure_codex_settings(GATEWAY, KEY, MODEL, alias) + assert alias.is_symlink() + assert codex_configure_state_path(alias) == codex_configure_state_path(target) + assert tomlkit.parse(target.read_text())["model"] == MODEL + unconfigure_codex_settings(target) + assert alias.is_symlink() + assert tomlkit.parse(alias.read_text()) == {"model": "old"} + + +@pytest.mark.parametrize("invalid", [ + 'token = "sk-secret\n', + 'model_providers = "sk-secret"\n', + '[model_providers]\nlitellm = "sk-secret"\n', +]) +def test_invalid_settings_are_unchanged_and_errors_hide_content(tmp_path: Path, invalid: str) -> None: + path: Final = tmp_path / "config.toml" + path.write_text(invalid) + with pytest.raises(CodexSettingsError) as caught: + configure_codex_settings(GATEWAY, KEY, MODEL, path) + assert "sk-secret" not in str(caught.value) + assert path.read_text() == invalid + assert not codex_configure_state_path(path).exists() + + +def test_invalid_receipt_fails_preflight_before_settings_change(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('model = "keep"\n') + state: Final = codex_configure_state_path(path) + state.parent.mkdir() + state.write_text('{"previous": "sk-secret"}') + with pytest.raises(CodexSettingsError) as caught: + preflight_codex_settings(path) + assert "sk-secret" not in str(caught.value) + assert path.read_text() == 'model = "keep"\n' + + +@pytest.mark.parametrize("configured_before", [False, True]) +@pytest.mark.parametrize("failed_target", ["receipt", "settings"]) +def test_failed_commit_restores_receipt_and_cleans_private_staging( + tmp_path: Path, configured_before: bool, failed_target: str +) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('model = "old"\n') + state: Final = codex_configure_state_path(path) + if configured_before: + configure_codex_settings(GATEWAY, KEY, MODEL, path) + before: Final = path.read_bytes() + receipt_before: Final = state.read_bytes() if state.exists() else None + + def failing_commit(staged: str, destination: str) -> None: + if destination == str(state if failed_target == "receipt" else path.resolve()): + raise OSError("sk-secret OS error") + commit_staged_json(staged, destination) + + with pytest.raises(CodexSettingsError) as caught: + configure_codex_settings(GATEWAY, "sk-replacement", "new-model", path, commit=failing_commit) + assert "sk-secret" not in str(caught.value) + assert path.read_bytes() == before + assert (state.read_bytes() if state.exists() else None) == receipt_before + assert not tuple(tmp_path.rglob(".tmp-*")) + if configured_before: + unconfigure_codex_settings(path) + assert tomlkit.parse(path.read_text()) == {"model": "old"} + + +def test_failed_undo_keeps_the_settings_and_receipt_for_retry(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('model = "old"\n') + configure_codex_settings(GATEWAY, KEY, MODEL, path) + before: Final = path.read_bytes() + + def fail(staged: str, destination: str) -> None: + raise OSError("cannot replace") + + with pytest.raises(CodexSettingsError): + unconfigure_codex_settings(path, commit=fail) + assert path.read_bytes() == before + assert codex_configure_state_path(path).exists() + assert not tuple(tmp_path.rglob(".tmp-*")) + unconfigure_codex_settings(path) + assert tomlkit.parse(path.read_text()) == {"model": "old"} + + +def test_wrapper_and_persistent_provider_agree_except_credential_source(tmp_path: Path) -> None: + path: Final = tmp_path / "config.toml" + configure_codex_settings(GATEWAY, KEY, MODEL, path) + provider: Final = tomlkit.parse(path.read_text())["model_providers"]["litellm"] + args: Final = agent_launch_args("codex", GATEWAY) + overrides: Final = dict(argument.split("=", 1) for argument in args[1::2]) + for field in ("name", "base_url", "wire_api", "supports_websockets", "requires_openai_auth"): + assert json.loads(overrides[f"model_providers.litellm.{field}"]) == provider[field] + assert overrides["model_providers.litellm.http_headers"] == "{}" + assert overrides["model_providers.litellm.env_key"] == '"OPENAI_API_KEY"' + + +@pytest.mark.parametrize("version", [ + "codex-cli 0.129.0", "codex-cli 0.129.1", "codex-cli 0.130.0", "codex-cli 1.0.0", + " \ncodex-cli 0.129.0\n", +]) +def test_version_guard_accepts_the_fixed_release_and_newer_stable_versions(version: str) -> None: + assert codex_settings_module.require_safe_codex(version=lambda: version) is None + + +@pytest.mark.parametrize("version", [ + None, "", "codex-cli 0.99.0", "codex-cli 0.128.99", "codex-cli 0.129.0-alpha.1", + "codex-cli 1.0.0-beta.1", "0.129.0", "codex-cli 0.129.0 extra", "unparseable-sk-version-secret", +]) +def test_version_guard_refuses_missing_unsafe_or_unrecognized_versions(version: str | None) -> None: + with pytest.raises(CodexSettingsError) as caught: + codex_settings_module.require_safe_codex(version=lambda: version) + assert "0.129.0" in str(caught.value) + assert "sk-version-secret" not in str(caught.value) + + +@pytest.mark.parametrize("output,returncode", [ + (None, 0), + ("codex-cli 0.129.0\n", 7), +]) +def test_version_probe_handles_missing_or_failed_executable( + fake_codex_version: Callable[[str | None, int], Path], output: str | None, returncode: int +) -> None: + fake_codex_version(output, returncode) + assert codex_settings_module._codex_version() is None + + +@pytest.mark.parametrize("output,returncode", [ + (None, 0), + ("codex-cli 0.128.0\n", 0), + ("codex-cli 0.129.0-alpha.1\n", 0), + ("unparseable-sk-version-secret\n", 0), + ("codex-cli 0.129.0\n", 7), +]) +def test_writer_checks_the_installed_codex_before_replacing_a_key_or_receipt( + tmp_path: Path, fake_codex_version: Callable[[str | None, int], Path], + output: str | None, returncode: int, +) -> None: + path: Final = tmp_path / "config.toml" + path.write_text('model = "original"\n') + configure_codex_settings(GATEWAY, "sk-existing-gateway", MODEL, path) + state: Final = codex_configure_state_path(path) + before: Final = (path.read_bytes(), state.read_bytes()) + fake_codex_version(output, returncode) + with pytest.raises(CodexSettingsError) as caught: + configure_codex_settings(GATEWAY, KEY, "replacement-model", path) + assert "0.129.0" in str(caught.value) + assert KEY not in str(caught.value) and "sk-version-secret" not in str(caught.value) + assert (path.read_bytes(), state.read_bytes()) == before + assert not tuple(tmp_path.rglob(".tmp-*")) + + +def test_undo_does_not_require_codex_to_remain_installed( + tmp_path: Path, fake_codex_version: Callable[[str | None, int], Path] +) -> None: + path: Final = tmp_path / "config.toml" + original: Final = 'model = "original"\n' + path.write_text(original) + configure_codex_settings(GATEWAY, KEY, MODEL, path) + fake_codex_version(None, 0) + outcome: Final = unconfigure_codex_settings(path) + assert outcome.restored and not outcome.kept + assert tomlkit.parse(path.read_text()) == tomlkit.parse(original) + assert not codex_configure_state_path(path).exists() diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py index 8ed188af737..8f68bb1320b 100644 --- a/tests/test_litellm/proxy/client/cli/test_configure_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -1,11 +1,16 @@ +import io import json import os import stat +import time +from pathlib import Path +from types import SimpleNamespace import click import pytest import requests import responses +import tomlkit from click.testing import CliRunner from litellm.proxy.client.cli import cli @@ -36,6 +41,8 @@ def paths(monkeypatch, tmp_path): monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(settings_path.parent)) monkeypatch.setattr(claude_settings_module, "CLAUDE_SETTINGS_PATH", settings_path) monkeypatch.setattr(claude_settings_module, "CONFIGURE_STATE_PATH", state_path) + monkeypatch.delenv("LITELLM_PROXY_API_KEY", raising=False) + monkeypatch.delenv("LITELLM_PROXY_URL", raising=False) return settings_path, state_path @@ -56,6 +63,29 @@ def runner(): return CliRunner() +@pytest.fixture +def codex_path(): + return Path(os.environ["CODEX_HOME"]) / "config.toml" + + +class _TerminalInput(io.BytesIO): + def isatty(self): + return True + + +def _mock_agent_models(): + def listing(request): + assert request.headers["Authorization"] == f"Bearer {VALID_KEY}" + rows = ( + [{"id": "claude-router-6175746f", "source_model": "auto"}] + if request.headers.get("x-gateway-client") == "claude-code" + else [{"id": "auto"}] + ) + return 200, {"Content-Type": "application/json"}, json.dumps({"data": rows}) + + responses.add_callback(responses.GET, f"{PROXY}/v1/models", callback=listing) + + @pytest.fixture def lite_up_backup(monkeypatch, tmp_path): """A `lite up` session holding its backup, the local precondition every settings write refuses on.""" @@ -253,6 +283,259 @@ class TestInteractiveConfigure: assert "lite configure claude --api-key" in result.output +class TestConfigureAgents: + @responses.activate + @pytest.mark.parametrize("targets", [("claude",), ("codex",), ("claude", "codex")]) + def test_group_options_drive_the_agent_picker_and_write_only_selected_agents( + self, runner, paths, codex_path, monkeypatch, targets + ): + _mock_agent_models() + asked = [] + + def checkbox(**kwargs): + assert tuple(choice.value for choice in kwargs["choices"]) == ("claude", "codex") + return SimpleNamespace(execute=lambda: targets) + + def fuzzy(**kwargs): + assert "auto" in kwargs["choices"] + assert "claude-router-6175746f" not in kwargs["choices"] + assert not paths[0].exists() and not codex_path.exists() + asked.append(kwargs["message"]) + return SimpleNamespace(execute=lambda: "auto") + + monkeypatch.setattr(configure_module.inquirer, "checkbox", checkbox) + monkeypatch.setattr(configure_module.inquirer, "fuzzy", fuzzy) + result = runner.invoke( + cli, + ["configure", "--api-key", VALID_KEY, "--gateway-url", f"{PROXY}/v1/"], + input=_TerminalInput(), + ) + assert result.exit_code == 0, result.output + assert VALID_KEY not in result.output + assert len(asked) == len(targets) + assert paths[0].exists() == ("claude" in targets) + assert codex_path.exists() == ("codex" in targets) + if "claude" in targets: + claude = json.loads(paths[0].read_text()) + assert claude["model"] == "claude-router-6175746f" + assert claude["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert claude["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + if "codex" in targets: + codex = tomlkit.parse(codex_path.read_text()) + assert codex["model"] == "auto" + assert codex["model_provider"] == "litellm" + provider = codex["model_providers"]["litellm"] + assert provider["base_url"] == f"{PROXY}/v1" + assert provider["http_headers"]["Authorization"] == f"Bearer {VALID_KEY}" + assert "env_key" not in provider + assert [call.request.headers.get("x-gateway-client") for call in responses.calls] == [ + "claude-code" if target == "claude" else None for target in targets + ] + + @responses.activate + @pytest.mark.parametrize("target", ["claude", "codex"]) + @pytest.mark.parametrize("leaf_override", [False, True], ids=["inherit-group", "leaf-wins"]) + def test_group_connection_options_are_inherited_and_leaf_options_take_precedence( + self, runner, paths, codex_path, target, leaf_override + ): + _mock_agent_models() + group_url = "http://group.test" if leaf_override else PROXY + group_key = "sk-group" if leaf_override else VALID_KEY + args = [ + "--base-url", "http://global.test", "--api-key", "sk-global", "configure", + "--gateway-url", group_url, "--api-key", group_key, target, "--model", "auto", + ] + if leaf_override: + args.extend(["--base-url", f"{PROXY}/v1/", "--api-key", VALID_KEY]) + result = runner.invoke(cli, args) + assert result.exit_code == 0, result.output + assert all(key not in result.output for key in (VALID_KEY, group_key, "sk-global")) + if target == "claude": + written = json.loads(paths[0].read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert written["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert not codex_path.exists() + else: + provider = tomlkit.parse(codex_path.read_text())["model_providers"]["litellm"] + assert provider["http_headers"]["Authorization"] == f"Bearer {VALID_KEY}" + assert provider["base_url"] == f"{PROXY}/v1" + assert not paths[0].exists() + assert len(responses.calls) == 1 + + @responses.activate + @pytest.mark.parametrize("failure", ["invalid-model", "cancel"]) + def test_both_model_choices_complete_before_either_configuration_changes( + self, paths, codex_path, failure + ): + _mock_agent_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text('{"theme": "dark"}') + codex_path.parent.mkdir(parents=True) + codex_path.write_text('model = "original"\n') + before = (settings_path.read_bytes(), codex_path.read_bytes()) + + def pick_codex_model(listed): + assert listed == ("auto",) + assert (settings_path.read_bytes(), codex_path.read_bytes()) == before + if failure == "cancel": + raise KeyboardInterrupt() + return "not-listed" + + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY}) + expected = KeyboardInterrupt if failure == "cancel" else click.ClickException + with pytest.raises(expected): + interactive_configure( + ctx, + pick_targets=lambda: ("claude", "codex"), + pick_model=lambda listed: "auto", + pick_codex_model=pick_codex_model, + ) + assert (settings_path.read_bytes(), codex_path.read_bytes()) == before + assert not state_path.exists() + assert not (codex_path.parent / ".litellm").exists() + + @responses.activate + def test_both_configs_are_preflighted_before_fetching_models_or_writing( + self, paths, codex_path + ): + _mock_agent_models() + codex_path.parent.mkdir(parents=True) + codex_path.write_text("[invalid") + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY}) + with pytest.raises(click.ClickException, match="Could not read Codex settings"): + interactive_configure( + ctx, + pick_targets=lambda: ("claude", "codex"), + pick_model=lambda listed: "auto", + pick_codex_model=lambda listed: "auto", + ) + assert not paths[0].exists() and not paths[1].exists() + assert codex_path.read_text() == "[invalid" + assert len(responses.calls) == 0 + + @responses.activate + @pytest.mark.parametrize("targets", [("claude", "codex"), ("codex", "claude")]) + @pytest.mark.parametrize("version", [None, "codex-cli 0.128.0\n"]) + def test_unsafe_codex_blocks_both_targets_before_requests_or_writes( + self, paths, codex_path, fake_codex_version, targets, version + ): + _mock_agent_models() + fake_codex_version(version, 0) + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY}) + with pytest.raises(click.ClickException, match=r"0\.129\.0") as caught: + interactive_configure( + ctx, + pick_targets=lambda: targets, + pick_model=lambda listed: "auto", + pick_codex_model=lambda listed: "auto", + ) + assert VALID_KEY not in str(caught.value) + assert len(responses.calls) == 0 + assert not paths[0].exists() and not paths[1].exists() + assert not codex_path.exists() and not (codex_path.parent / ".litellm").exists() + + @responses.activate + def test_claude_only_configuration_does_not_require_codex( + self, runner, paths, codex_path, fake_codex_version + ): + _mock_agent_models() + fake_codex_version(None, 0) + result = runner.invoke( + cli, ["configure", "--api-key", VALID_KEY, "--gateway-url", PROXY, "claude", "--model", "auto"] + ) + assert result.exit_code == 0, result.output + assert json.loads(paths[0].read_text())["model"] == "claude-router-6175746f" + assert not codex_path.exists() + + @responses.activate + def test_codex_only_ignores_claudes_temporary_owner( + self, runner, paths, codex_path, lite_up_backup + ): + _mock_agent_models() + result = runner.invoke( + cli, ["configure", "--api-key", VALID_KEY, "--gateway-url", PROXY, "codex", "--model", "auto"] + ) + assert result.exit_code == 0, result.output + assert tomlkit.parse(codex_path.read_text())["model"] == "auto" + assert not paths[0].exists() and not paths[1].exists() + assert lite_up_backup.exists() + + def test_noninteractive_codex_requires_a_model(self, runner, paths, codex_path): + result = runner.invoke(cli, ["configure", "--api-key", VALID_KEY, "--gateway-url", PROXY, "codex"]) + assert result.exit_code != 0 and "Missing option '--model'" in result.output + assert not paths[0].exists() and not codex_path.exists() + + @responses.activate + @pytest.mark.parametrize("target", ["claude", "codex"]) + @pytest.mark.parametrize( + "option, value, expected", + [ + ("--api-key", "sk-secret\ninvalid", "must not be blank"), + ("--gateway-url", "https://user:sk-secret@proxy.test", "must not contain credentials"), + ("--gateway-url", "https://proxy.test?key=sk-secret", "must not include a query"), + ("--gateway-url", "file:///sk-secret", "must be a full http:// or https:// URL"), + ], + ) + def test_invalid_connection_input_never_writes_requests_or_echoes_secrets( + self, runner, paths, codex_path, target, option, value, expected + ): + result = runner.invoke( + cli, + [ + "configure", "--api-key", VALID_KEY, "--gateway-url", PROXY, + target, "--model", "auto", option, value, + ], + ) + assert result.exit_code != 0 and expected in result.output + assert "sk-secret" not in result.output and VALID_KEY not in result.output + assert not paths[0].exists() and not codex_path.exists() + assert len(responses.calls) == 0 + + @responses.activate + @pytest.mark.parametrize("failure", ["rejected", "connection", "response-body"]) + def test_gateway_failures_never_echo_the_key(self, runner, paths, codex_path, failure): + if failure == "rejected": + responses.get(f"{PROXY}/v1/models", status=401) + elif failure == "connection": + responses.get(f"{PROXY}/v1/models", body=requests.ConnectionError(VALID_KEY)) + else: + responses.get(f"{PROXY}/v1/models", json={"data": VALID_KEY}) + result = runner.invoke( + cli, ["configure", "--api-key", VALID_KEY, "--gateway-url", PROXY, "codex", "--model", "auto"] + ) + assert result.exit_code != 0 and "Error:" in result.output + assert VALID_KEY not in result.output + assert not paths[0].exists() and not codex_path.exists() + + @responses.activate + def test_configure_and_unconfigure_do_not_read_a_stored_login( + self, runner, paths, codex_path, tmp_path, secret_vault_factory, fake_codex_version + ): + _mock_agent_models() + token_path = tmp_path / ".litellm" / "token.json" + token_path.parent.mkdir() + token_path.write_text(json.dumps({"base_url": PROXY, "timestamp": time.time()})) + vault = secret_vault_factory(json.dumps({"base_url": PROXY, "key": "sk-login", "jwt_token": ""})) + missing = runner.invoke( + cli, ["configure", "--gateway-url", PROXY, "codex", "--model", "auto"], obj={"secret_vault": vault} + ) + assert missing.exit_code != 0 and "needs a long-lived virtual key" in missing.output + assert len(responses.calls) == 0 and not codex_path.exists() + configured = runner.invoke( + cli, + ["configure", "--api-key", VALID_KEY, "--gateway-url", PROXY, "codex", "--model", "auto"], + obj={"secret_vault": vault}, + ) + assert configured.exit_code == 0, configured.output + fake_codex_version(None, 0) + undone = runner.invoke(cli, ["unconfigure", "codex"], obj={"secret_vault": vault}) + assert undone.exit_code == 0, undone.output + assert vault.reads == 0 and vault.writes == [] and vault.erases == 0 + assert not codex_path.exists() and not paths[0].exists() + assert "Removed" in undone.output and "sk-login" not in missing.output + configured.output + undone.output + + class TestUnconfigureClaude: @responses.activate def test_restores_the_original_file_and_removes_the_receipt(self, runner, paths): diff --git a/tests/test_litellm/proxy/common_utils/test_debug_utils.py b/tests/test_litellm/proxy/common_utils/test_debug_utils.py new file mode 100644 index 00000000000..163ea530be9 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_debug_utils.py @@ -0,0 +1,69 @@ +import os +import socket +from pathlib import Path + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.debug_utils import ( + PSUTIL_MISSING_ERROR, + _ProcFilesystemProcess, + _summary_process_memory, + get_memory_summary, +) + +PAGE_SIZE = 4096 +STATM_SIZE_PAGES = 100_000 +STATM_RESIDENT_PAGES = 30_000 +MEMINFO_TOTAL_KB = 1_000_000 + + +@pytest.fixture +def proc_process(tmp_path: Path) -> _ProcFilesystemProcess: + statm = tmp_path / "statm" + statm.write_text(f"{STATM_SIZE_PAGES} {STATM_RESIDENT_PAGES} 5000 1 0 20000 0\n") + meminfo = tmp_path / "meminfo" + meminfo.write_text( + f"MemTotal: {MEMINFO_TOTAL_KB} kB\nMemFree: 400000 kB\nMemAvailable: 600000 kB\n" + ) + return _ProcFilesystemProcess(statm_path=str(statm), meminfo_path=str(meminfo), page_size=PAGE_SIZE) + + +def test_proc_filesystem_process_reads_resident_and_virtual_bytes_from_statm( + proc_process: _ProcFilesystemProcess, +) -> None: + memory_info = proc_process.memory_info() + + assert memory_info.rss == STATM_RESIDENT_PAGES * PAGE_SIZE + assert memory_info.vms == STATM_SIZE_PAGES * PAGE_SIZE + + +def test_proc_filesystem_process_reports_share_of_meminfo_total(proc_process: _ProcFilesystemProcess) -> None: + expected_percent = STATM_RESIDENT_PAGES * PAGE_SIZE / (MEMINFO_TOTAL_KB * 1024) * 100 + + assert proc_process.memory_percent() == pytest.approx(expected_percent) + + +def test_summary_reports_rss_from_the_proc_filesystem(proc_process: _ProcFilesystemProcess) -> None: + memory, health_status = _summary_process_memory(proc_process) + + assert memory["ram_usage_mb"] == round(STATM_RESIDENT_PAGES * PAGE_SIZE / (1024 * 1024), 2) + assert memory["system_memory_percent"] == pytest.approx(12.0) + assert health_status == "healthy" + assert "error" not in memory + + +def test_summary_without_any_memory_source_names_psutil_and_reports_no_rss() -> None: + memory, health_status = _summary_process_memory(None) + + assert memory == {"error": PSUTIL_MISSING_ERROR} + assert health_status == "healthy" + + +@pytest.mark.asyncio +async def test_memory_summary_names_the_host_and_worker_that_answered() -> None: + summary = await get_memory_summary(UserAPIKeyAuth()) + + assert summary["hostname"] == socket.gethostname() + assert summary["worker_pid"] == os.getpid() + assert summary["memory"]["ram_usage_mb"] > 0 diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index 875dca4bee3..f5fb1bda0c1 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -34,10 +34,12 @@ from pydantic import ValidationError from litellm.proxy.db.db_url_settings import ( PG_SSL_REQUEST, DatabaseURLSettings, + token_refresh_params_from_url, translate_libpq_ssl_params, unsupported_db_scheme, unsupported_db_scheme_message, ) +from litellm.proxy.db.pgbouncer import PgBouncerPlan, PgBouncerSettings, plan_pgbouncer from litellm.proxy.db.token_auth import AzureEntraTokenAuth, RdsIamTokenAuth @@ -51,6 +53,8 @@ _MANAGED_DB_ENV_VARS = ( "AZURE_POSTGRESQL_AUTH", "DATABASE_DISABLE_PREPARED_STATEMENTS", "DATABASE_MAX_IDLE_CONNECTION_LIFETIME", + "DATABASE_SSLMODE", + "DATABASE_SSLROOTCERT", "DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA", @@ -781,6 +785,109 @@ def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypa } +def _tls_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DATABASE_HOST", "writer.example.com") + monkeypatch.setenv("DATABASE_USER", "litellm") + monkeypatch.setenv("DATABASE_NAME", "litellm_db") + monkeypatch.setenv("DATABASE_SSLMODE", "verify-full") + monkeypatch.setenv("DATABASE_SSLROOTCERT", "/certs/rds-bundle.pem") + + +def test_tls_env_vars_make_the_minted_iam_writer_url_verify_the_server(monkeypatch: pytest.MonkeyPatch): + """The supervisor starts PgBouncer from the URL assembled here, before any + config.yaml is read, so an IAM URL with no TLS params leaves PgBouncer on + ``prefer`` (no SNI, no verification) and the RDS handshake fails.""" + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + _tls_env(monkeypatch) + + with _stub_iam_token("WRITER_TOKEN"): + assert _apply() is True + + url: Final = os.environ["DATABASE_URL"] + assert url.startswith("postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db?") + assert _query(url) == { + "sslmode": ["require"], + "sslcert": ["/certs/rds-bundle.pem"], + "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], + } + + +def test_tls_env_vars_apply_to_the_password_writer_and_the_assembled_reader(monkeypatch: pytest.MonkeyPatch): + _tls_env(monkeypatch) + monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t") + monkeypatch.setenv("DATABASE_SCHEMA", "public") + monkeypatch.setenv("DATABASE_HOST_READ_REPLICA", "reader.example.com") + + assert _apply() is True + + expected: Final = { + "schema": ["public"], + "sslmode": ["require"], + "sslcert": ["/certs/rds-bundle.pem"], + "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], + } + assert os.environ["DATABASE_URL"].startswith("postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?") + assert _query(os.environ["DATABASE_URL"]) == expected + assert os.environ["DATABASE_URL_READ_REPLICA"].startswith( + "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db?" + ) + assert _query(os.environ["DATABASE_URL_READ_REPLICA"]) == expected + + +def test_sslrootcert_env_var_alone_means_verify_full_for_prisma_and_pgbouncer(monkeypatch: pytest.MonkeyPatch): + """Under libpq's default ``prefer`` a root cert is never consulted, so a URL + carrying only ``sslrootcert`` would leave PgBouncer on ``prefer`` with the CA + loaded but unused. Supplying a CA and nothing else must verify.""" + _tls_env(monkeypatch) + monkeypatch.delenv("DATABASE_SSLMODE") + monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t") + + assert _apply() is True + + url: Final = os.environ["DATABASE_URL"] + assert _query(url) == { + "sslmode": ["require"], + "sslcert": ["/certs/rds-bundle.pem"], + "sslaccept": ["strict"], + "max_idle_connection_lifetime": ["60"], + } + plan: Final = plan_pgbouncer(url, PgBouncerSettings(enabled=True), Path("/run/pgb"), None) + assert isinstance(plan, PgBouncerPlan), plan + assert "server_tls_sslmode = verify-full" in plan.ini + assert "server_tls_ca_file = /run/pgb/server-ca.pem" in plan.ini + + +def test_tls_env_vars_never_override_a_pinned_database_url(monkeypatch: pytest.MonkeyPatch): + writer: Final = ( + "postgresql://pinned:url@db.example.com:5432/litellm_db?sslmode=disable&max_idle_connection_lifetime=60" + ) + reader: Final = "postgresql://pinned:url@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60" + monkeypatch.setenv("DATABASE_URL", writer) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", reader) + monkeypatch.setenv("DATABASE_HOST_READ_REPLICA", "reader.example.com") + _tls_env(monkeypatch) + + assert _apply() is False + + assert os.environ["DATABASE_URL"] == writer + assert os.environ["DATABASE_URL_READ_REPLICA"] == reader + + +def test_token_refresh_params_keep_the_prisma_tls_dialect_but_not_the_schema(): + kept: Final = token_refresh_params_from_url( + "postgresql://u:TOKEN@db.example.com:5432/litellm_db" + "?schema=tenant&connection_limit=5&sslmode=require&sslcert=/certs/root.pem&sslaccept=strict" + ) + assert dict(kept) == { + "connection_limit": "5", + "sslmode": "require", + "sslcert": "/certs/root.pem", + "sslaccept": "strict", + } + + def _issue_cert( subject: str, issuer: x509.Certificate | None, issuer_key: ec.EllipticCurvePrivateKey | None, ca: bool ) -> tuple[x509.Certificate, ec.EllipticCurvePrivateKey]: diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py index 7b5a0bf10f9..bf7df3077ea 100644 --- a/tests/test_litellm/proxy/db/test_pgbouncer.py +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -411,7 +411,7 @@ class TestPgBouncerProcess: port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1, - ready_timeout_seconds=0.3, + ready_timeout_seconds=2.0, ) assert pooler.start() is None first_pid: Final = pooler.pid @@ -421,8 +421,8 @@ class TestPgBouncerProcess: with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): os.kill(first_pid, signal.SIGKILL) assert _wait_until(lambda: _listening(wrong_port)) - assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records)) port_file.write_text(str(port)) + assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records)) assert _wait_until(lambda: _listening(port)) assert _wait_until(lambda: not _listening(wrong_port)) pooler.stop() diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index 963f6a5640f..99e494fccd5 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -299,6 +299,10 @@ def test_azure_entra_mint_writes_an_encoded_url_into_the_db_url_env_var(azure_en "connection_limit=20&pgbouncer=true&max_idle_connection_lifetime=45", {"connection_limit": ["20"], "pgbouncer": ["true"], "max_idle_connection_lifetime": ["45"]}, ), + ( + "sslmode=require&sslcert=/certs/root.pem&sslaccept=strict&schema=tenant", + {"sslmode": ["require"], "sslcert": ["/certs/root.pem"], "sslaccept": ["strict"]}, + ), ], ) def test_token_refresh_keeps_the_connection_params_of_the_url_it_replaces( diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index 966a638f6a4..3bc7e1f02f8 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -2,8 +2,9 @@ import asyncio import logging import os import sys -from typing import Any, Dict -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Any, Dict, Final +from unittest.mock import AsyncMock, MagicMock, call, patch +from urllib.parse import parse_qs, urlsplit import pytest @@ -927,6 +928,47 @@ def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails( ) +def test_prisma_client_init_keeps_reader_tls_params_on_the_minted_iam_url( + monkeypatch: pytest.MonkeyPatch, +): + """The initial reader mint rebuilds the URL from host/port/user/db, so the + Prisma TLS dialect on DATABASE_URL_READ_REPLICA must be carried over or + a verify-only database rejects the reader and reads fall to the writer.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", + "postgresql://reader_user@reader.aurora.local:5432/litellm" + "?schema=tenant&sslmode=require&sslcert=/certs/root.pem&sslaccept=strict", + ) + + prisma_factory: Final = MagicMock(name="Prisma") + fake_prisma_module: Final = MagicMock(Prisma=prisma_factory) + monkeypatch.setitem(sys.modules, "prisma", fake_prisma_module) + + fake_iam_module: Final = MagicMock(generate_iam_auth_token=MagicMock(return_value="READER-TOKEN")) + monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_iam_module) + + from litellm.proxy.utils import PrismaClient + + client: Final = PrismaClient( + database_url="postgresql://writer@writer.aurora.local:5432/litellm", + proxy_logging_obj=MagicMock(), + ) + + assert isinstance(client.db, RoutingPrismaWrapper) + reader_url: Final = os.environ["DATABASE_URL_READ_REPLICA"] + assert reader_url.startswith("postgresql://reader_user:READER-TOKEN@reader.aurora.local:5432/litellm?") + assert parse_qs(urlsplit(reader_url).query) == { + "schema": ["tenant"], + "sslmode": ["require"], + "sslcert": ["/certs/root.pem"], + "sslaccept": ["strict"], + } + assert prisma_factory.call_args_list == [call(), call(datasource={"url": reader_url})] + + @pytest.mark.asyncio async def test_connect_degrades_writer_when_reader_available(): """A writer connect failure with a healthy reader must NOT abort proxy diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 0361d5cfe8f..bca6344b3f7 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -8,6 +8,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations import asyncio +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -15,6 +16,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed @@ -78,6 +80,39 @@ def _row(window_start: datetime, spend: float) -> SimpleNamespace: return SimpleNamespace(window_start=window_start, spend=spend) +class _PausedSpendTable: + def __init__(self, spend: float) -> None: + self.spend: Final = spend + self.read_started: Final = asyncio.Event() + self.resume_read: Final = asyncio.Event() + + async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace: + self.read_started.set() + await self.resume_read.wait() + return _row(WINDOW_START, self.spend) + + +async def _reseed_with_paused_table( + table: _PausedSpendTable, cache: DualCache, counter_key: str, window: bool +) -> float | None: + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_usertable=table, litellm_budgetwindowspend=table)) + if window: + return await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + window_start=WINDOW_START, + ) + return await SpendCounterReseed.coalesced( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + ) + + @pytest.mark.asyncio async def test_window_from_table_reads_row_by_primary_key(): """The lookup must use the table's own entity_type values ("key"), not the @@ -270,6 +305,63 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): assert prisma.db.litellm_spendlogs.call_count == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0]) +async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache( + window: bool, + concurrent_spend: float, +) -> None: + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = "spend:team:team-1:window:1d" if window else "spend:user:user-1" + db_spend: Final = 989.01459411 + table: Final = _PausedSpendTable(db_spend) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) + + await asyncio.wait_for(table.read_started.wait(), timeout=5) + cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) + table.resume_read.set() + result: Final = await asyncio.wait_for(reseed_task, timeout=5) + + expected: Final = max(db_spend, concurrent_spend) + assert cache.in_memory_cache.get_cache(key=counter_key) == expected + assert result == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("batch", [False, True], ids=["single_increment", "batch_increment"]) +@pytest.mark.parametrize("increment", [5.0, -5.0], ids=["charge", "refund"]) +async def test_cold_reseed_preserves_concurrent_local_increment( + monkeypatch: pytest.MonkeyPatch, window: bool, batch: bool, increment: float +) -> None: + from litellm.proxy import proxy_server + + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = ( + f"spend:team:concurrent-{batch}-{increment}:window:1d" + if window + else f"spend:user:concurrent-{batch}-{increment}" + ) + table: Final = _PausedSpendTable(100.0) + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) + await asyncio.wait_for(table.read_started.wait(), timeout=5) + + increment_task: Final = asyncio.create_task( + proxy_server._apply_spend_counter_increments( + pending=(proxy_server.PendingSpendIncrement(counter_key=counter_key, increment=increment),) + ) + if batch + else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment) + ) + await asyncio.sleep(0) + table.resume_read.set() + await asyncio.wait_for(asyncio.gather(reseed_task, increment_task), timeout=5) + + assert cache.in_memory_cache.get_cache(key=counter_key) == 100.0 + increment + + @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) @@ -304,8 +396,7 @@ async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the @pytest.mark.asyncio async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error(): assert ( - await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") - is None + await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") is None ) assert ( await SpendCounterReseed.end_user_from_db( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_conduct.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_conduct.py new file mode 100644 index 00000000000..323756f8fa0 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_conduct.py @@ -0,0 +1,423 @@ +from __future__ import annotations + +import importlib.util +import json +import warnings +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Final, Literal + +import httpx +import pytest +import respx +from fastapi import HTTPException + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_endpoints import get_guardrail_ui_settings, get_provider_specific_params +from litellm.proxy.guardrails.guardrail_hooks.conduct import ( + DEFAULT_TIMEOUT_SECONDS, + ConductGuardrail, + initialize_guardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.conduct.conduct import ( + apply_conduct_guardrail, + binds_unreachable_fallback, + record_decision, + request_payload, +) +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams +from litellm.types.llms.openai import ChatCompletionAssistantMessage +from litellm.types.proxy.guardrails.guardrail_hooks.conduct import ( + ConductGuardrailConfigModel, + ConductGuardrailConfigModelOptionalParams, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +PACKAGE_INSTALLED: Final = importlib.util.find_spec("conduct_litellm_guard") is not None + + +class _RecordingGuardrail(CustomGuardrail): + """Stand-in with the ``conduct_litellm_guard.ConductGuard`` class contract.""" + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + return [GuardrailEventHooks.pre_call] + + def __init__( + self, + *, + api_url: str | None = None, + agent_token: str | None = None, + workspace_id: str | None = None, + unreachable_fallback: str | None = None, + tool_name: str = "llm_call", + timeout: float = 8.0, + guardrail_name: str | None = None, + event_hook: str | None = None, + default_on: bool = False, + supported_event_hooks: list[GuardrailEventHooks] | None = None, + ) -> None: + super().__init__( + guardrail_name=guardrail_name, + event_hook=event_hook, # pyright: ignore[reportArgumentType] # CustomGuardrail coerces the str at runtime + default_on=default_on, + supported_event_hooks=supported_event_hooks, + ) + self.api_url = api_url + self.agent_token = agent_token + self.workspace_id = workspace_id + self.unreachable_fallback = unreachable_fallback or "fail_closed" + self.tool_name = tool_name + self.timeout = timeout + + +@dataclass(frozen=True, slots=True) +class _Decision: + verdict: str + rule_id: str | None = None + + +class _Blocked(Exception): + def __init__(self, decision: _Decision) -> None: + super().__init__(decision.verdict) + self.decision = decision + + +@dataclass(slots=True) +class _RecordingCheck: + verdict: str + rule_id: str | None = None + calls: list[tuple[Mapping[str, object], str]] = field(default_factory=list) # mutable-ok: test spy + recorded: list[_Decision] = field(default_factory=list) # mutable-ok: test spy + + async def __call__(self, *, data: Mapping[str, object], call_type: str) -> _Decision: + self.calls.append((data, call_type)) + return _Decision(self.verdict, self.rule_id) + + def record(self, decision: _Decision) -> None: + self.recorded.append(decision) + + +async def _bridge( + check: _RecordingCheck, + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], +) -> GenericGuardrailAPIInputs: + return await apply_conduct_guardrail(inputs, request_data, input_type, check, _Blocked, check.record) + + +def _guardrail_records(request_data: Mapping[str, object]) -> list[tuple[str, object]]: + metadata: Final = request_data["metadata"] + assert isinstance(metadata, dict) + records: Final = metadata["standard_logging_guardrail_information"] + assert isinstance(records, list) + return [(record["guardrail_status"], record["guardrail_response"]) for record in records] + + +def _params(mode: str = "pre_call", **extras: object) -> LitellmParams: + return LitellmParams(guardrail="conduct", mode=mode, api_key="cond_agt_test", **extras) + + +def _guardrail(litellm_params: LitellmParams) -> Guardrail: + return Guardrail(guardrail_name="conduct-guard", litellm_params=litellm_params) + + +def _init(litellm_params: LitellmParams) -> _RecordingGuardrail: + callback: Final = initialize_guardrail( + litellm_params, _guardrail(litellm_params), guardrail_cls=_RecordingGuardrail + ) + assert isinstance(callback, _RecordingGuardrail) + return callback + + +@pytest.fixture(autouse=True) +def _isolate_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + +def test_maps_typed_fields_and_extras_onto_plugin_kwargs() -> None: + callback: Final = _init( + _params( + api_base="https://guard.example.test", + unreachable_fallback="fail_open", + timeout="3", + workspace_id="ws_123", + tool_name="workflow", + default_on=True, + ) + ) + + assert callback.api_url == "https://guard.example.test" + assert callback.agent_token == "cond_agt_test" + assert callback.unreachable_fallback == "fail_open" + assert callback.timeout == 3.0 + assert callback.workspace_id == "ws_123" + assert callback.tool_name == "workflow" + assert callback.guardrail_name == "conduct-guard" + assert callback.event_hook == "pre_call" + assert callback.default_on is True + assert litellm.callbacks == [callback] + + +def test_defaults_when_optional_config_is_omitted() -> None: + callback: Final = _init(_params()) + + assert callback.unreachable_fallback == "fail_closed" + assert callback.timeout == DEFAULT_TIMEOUT_SECONDS + assert callback.workspace_id is None + assert callback.tool_name == "llm_call" + + +def test_ui_form_defaults_match_what_the_initializer_forwards() -> None: + optional: Final = ConductGuardrailConfigModelOptionalParams() + model: Final = ConductGuardrailConfigModel(api_key="cond_agt_test") + callback: Final = _init( + _params(**{**model.model_dump(exclude={"api_key", "optional_params"}), **optional.model_dump()}) + ) + + assert callback.api_url == model.api_base + assert callback.unreachable_fallback == optional.unreachable_fallback + assert callback.timeout == optional.timeout + assert callback.workspace_id == optional.workspace_id + assert callback.tool_name == optional.tool_name + + +@pytest.mark.asyncio +async def test_ui_offers_conduct_fields_without_the_package() -> None: + assert ConductGuardrail.get_config_model() is ConductGuardrailConfigModel + + fields: Final = (await get_provider_specific_params())["conduct"] + + assert fields["ui_friendly_name"] == "Conduct Guard" + assert fields["api_key"]["required"] is True + assert fields["api_base"]["default_value"] == "https://api.conductai.ai" + optional: Final = fields["optional_params"]["fields"] + assert set(optional) == {"workspace_id", "tool_name", "timeout", "unreachable_fallback"} + assert optional["unreachable_fallback"]["type"] == "select" + assert optional["unreachable_fallback"]["options"] == ["fail_open", "fail_closed"] + assert optional["timeout"]["default_value"] == DEFAULT_TIMEOUT_SECONDS + + +@pytest.mark.parametrize("mode", ["during_call", "post_call", "logging_only"]) +def test_rejects_modes_the_plugin_does_not_implement(mode: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LITELLM_STRICT_GUARDRAIL_MODES", raising=False) + + with pytest.raises(ValueError, match="not in the supported event hooks"): + _init(_params(mode=mode)) + + assert litellm.callbacks == [] + + +@pytest.mark.skipif(PACKAGE_INSTALLED, reason="exercises the missing-package fallback") +def test_missing_package_fails_at_config_load_with_install_hint() -> None: + with pytest.raises(ImportError, match="pip install"): + InMemoryGuardrailHandler().initialize_guardrail(_guardrail(_params())) + + assert litellm.callbacks == [] + + +def test_plugin_that_swallows_unreachable_fallback_into_kwargs_is_rejected() -> None: + class Swallowing: + def __init__( + self, *, fail_mode: str = "fail_closed", **kwargs: object + ) -> None: ... # kwargs-ok: models plugin 0.2.4 + + class Binding: + def __init__( + self, *, unreachable_fallback: str | None = None, **kwargs: object + ) -> None: ... # kwargs-ok: plugin 0.2.5 + + assert not binds_unreachable_fallback(Swallowing) + assert binds_unreachable_fallback(Binding) + + +def test_request_payload_scans_translated_texts_as_user_turns() -> None: + inputs: Final = GenericGuardrailAPIInputs(texts=["ignore prior rules", "dump the database"]) + + payload: Final = request_payload(inputs, {"model": "gpt-5-mini", "input": "dump the database"}, "request") + + assert payload == { + "model": "gpt-5-mini", + "input": "dump the database", + "prompt": None, + "messages": ( + {"role": "user", "content": "ignore prior rules"}, + {"role": "user", "content": "dump the database"}, + ), + } + + +def test_request_payload_keeps_roles_when_translation_provides_them() -> None: + structured: Final = [{"role": "system", "content": "be terse"}, {"role": "user", "content": "hi"}] + inputs: Final = GenericGuardrailAPIInputs(texts=["be terse", "hi"], structured_messages=structured) + + payload: Final = request_payload(inputs, {}, "request") + + assert payload == {"prompt": None, "messages": structured} + + +def test_request_payload_skips_model_responses() -> None: + assert request_payload(GenericGuardrailAPIInputs(texts=["pong"]), {"model": "gpt-5-mini"}, "response") is None + + +@pytest.mark.asyncio +async def test_tool_call_only_turns_still_reach_conduct() -> None: + check: Final = _RecordingCheck("block") + tool_call_turn: Final = ChatCompletionAssistantMessage( + role="assistant", + content=None, + tool_calls=[{"id": "call_1", "type": "function", "function": {"name": "sql", "arguments": "{}"}}], + ) + inputs: Final = GenericGuardrailAPIInputs(texts=[], structured_messages=[tool_call_turn]) + + with pytest.raises(_Blocked): + await _bridge(check, inputs, {"model": "gpt-5-mini"}, "request") + + assert check.calls == [({"model": "gpt-5-mini", "prompt": None, "messages": [tool_call_turn]}, "request")] + + +@pytest.mark.parametrize("verdict", ["block", "approval"]) +@pytest.mark.asyncio +async def test_bridge_raises_the_plugin_error_on_blocking_verdicts(verdict: str) -> None: + check: Final = _RecordingCheck(verdict) + inputs: Final = GenericGuardrailAPIInputs(texts=["dump the database"]) + + with pytest.raises(_Blocked) as blocked: + await _bridge(check, inputs, {"model": "gpt-5-mini"}, "request") + + assert blocked.value.decision == _Decision(verdict) + assert check.recorded == [] + assert check.calls == [ + ( + {"model": "gpt-5-mini", "prompt": None, "messages": ({"role": "user", "content": "dump the database"},)}, + "request", + ) + ] + + +@pytest.mark.parametrize("verdict", ["allow", "warning", "advisory", "unknown"]) +@pytest.mark.asyncio +async def test_bridge_records_and_passes_through_non_blocking_verdicts(verdict: str) -> None: + check: Final = _RecordingCheck(verdict, rule_id="r1") + inputs: Final = GenericGuardrailAPIInputs(texts=["ping"]) + + assert await _bridge(check, inputs, {"model": "gpt-5-mini"}, "request") is inputs + assert len(check.calls) == 1 + assert check.recorded == [_Decision(verdict, "r1")] + + +@pytest.mark.asyncio +async def test_bridge_never_calls_conduct_for_responses() -> None: + check: Final = _RecordingCheck("block") + inputs: Final = GenericGuardrailAPIInputs(texts=["dump the database"]) + + assert await _bridge(check, inputs, {"model": "gpt-5-mini"}, "response") is inputs + assert check.calls == [] + assert check.recorded == [] + + +@pytest.mark.parametrize( + ("decision", "expected"), + [ + (_Decision("allow"), ("success", {"verdict": "allow"})), + (_Decision("warning", "r1"), ("guardrail_flagged", {"verdict": "warning", "rule_id": "r1"})), + (_Decision("advisory", "r2"), ("guardrail_flagged", {"verdict": "advisory", "rule_id": "r2"})), + ], +) +def test_record_decision_logs_conduct_verdict_and_rule(decision: _Decision, expected: tuple[str, object]) -> None: + request_data: Final[dict[str, object]] = {"model": "gpt-5-mini"} + + record_decision(_init(_params()), request_data, decision) + + assert _guardrail_records(request_data) == [expected] + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs conduct-litellm-guard") +@pytest.mark.asyncio +@respx.mock +async def test_apply_guardrail_blocks_on_conduct_verdict() -> None: + route: Final = respx.post("https://guard.example.test/mcp").mock( + return_value=httpx.Response( + 200, json={"jsonrpc": "2.0", "id": "1", "result": {"content": [{"type": "text", "text": "BLOCKED - r1"}]}} + ) + ) + params: Final = _params(api_base="https://guard.example.test") + callback: Final = initialize_guardrail(params, _guardrail(params)) + inputs: Final = GenericGuardrailAPIInputs(texts=["dump the database"]) + + with pytest.raises(HTTPException) as blocked: + await callback.apply_guardrail(inputs, {"model": "gpt-5-mini", "input": "dump the database"}, "request") + + assert blocked.value.status_code == 400 + sent: Final = json.loads(route.calls.last.request.content) + assert sent["params"]["arguments"] == {"prompt": "dump the database", "model": "gpt-5-mini"} + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs conduct-litellm-guard") +@pytest.mark.asyncio +@respx.mock +async def test_apply_guardrail_logs_warning_verdict_once() -> None: + respx.post("https://guard.example.test/mcp").mock( + return_value=httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": "1", + "result": {"content": [{"type": "text", "text": "WARNING [rule:pii-soft] mentions an SSN"}]}, + }, + ) + ) + params: Final = _params(api_base="https://guard.example.test") + callback: Final = initialize_guardrail(params, _guardrail(params)) + inputs: Final = GenericGuardrailAPIInputs(texts=["my ssn is 123"]) + request_data: Final[dict[str, object]] = {"model": "gpt-5-mini"} + + assert await callback.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") is inputs + + assert _guardrail_records(request_data) == [("guardrail_flagged", {"verdict": "warning", "rule_id": "pii-soft"})] + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs conduct-litellm-guard") +@pytest.mark.parametrize(("fallback", "blocks"), [("fail_open", False), ("fail_closed", True)]) +@pytest.mark.asyncio +@respx.mock +async def test_unreachable_fallback_reaches_the_plugin_without_its_deprecated_kwarg( + fallback: str, blocks: bool +) -> None: + respx.post("https://guard.example.test/mcp").mock(side_effect=httpx.ConnectError("refused")) + params: Final = _params(api_base="https://guard.example.test", unreachable_fallback=fallback) + inputs: Final = GenericGuardrailAPIInputs(texts=["ping"]) + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + callback: Final = initialize_guardrail(params, _guardrail(params)) + + if blocks: + with pytest.raises(HTTPException): + await callback.apply_guardrail(inputs, {"model": "gpt-5-mini"}, "request") + return + assert await callback.apply_guardrail(inputs, {"model": "gpt-5-mini"}, "request") is inputs + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs conduct-litellm-guard") +def test_config_loads_conduct_and_rejects_modes_the_plugin_lacks() -> None: + handler: Final = InMemoryGuardrailHandler() + + loaded: Final = handler.initialize_guardrail(_guardrail(_params())) + assert loaded is not None + assert loaded["litellm_params"].guardrail == "conduct" + assert [type(callback) for callback in litellm.callbacks] == [ConductGuardrail] + + with pytest.raises(ValueError, match="not in the supported event hooks"): + handler.initialize_guardrail(_guardrail(_params(mode="during_call"))) + + +@pytest.mark.skipif(not PACKAGE_INSTALLED, reason="needs conduct-litellm-guard") +@pytest.mark.asyncio +async def test_ui_only_offers_pre_call_for_conduct() -> None: + settings: Final = await get_guardrail_ui_settings() + + assert settings.supported_modes_by_provider["conduct"] == ["pre_call"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index a07157396df..a9ca13a463d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1,3 +1,7 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Final, cast +import json from unittest.mock import patch import httpx @@ -7,6 +11,9 @@ from pydantic import ValidationError import litellm from litellm.exceptions import Timeout +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import ( @@ -1719,3 +1726,165 @@ async def test_streaming_params_from_config_control_output_scan_cadence( handler = _initialize_from_config(mode="post_call", **configured) assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls + + +@asynccontextmanager +async def _guardrail_redacting(secret: str, replacement: str) -> AsyncIterator[CrowdStrikeAIDRHandler]: + def redacted(content: object) -> object: + if isinstance(content, str): + return content.replace(secret, replacement) + if isinstance(content, list): + return [ + {**part, "text": redacted(part["text"])} if isinstance(part, dict) and "text" in part else part + for part in content + ] + return content + + def respond(request: httpx.Request) -> httpx.Response: + sent: Final = json.loads(request.content)["guard_input"]["messages"] + return httpx.Response( + status_code=200, + json={ + "result": { + "blocked": False, + "transformed": True, + "guard_output": { + "messages": [{**message, "content": redacted(message.get("content"))} for message in sent] + }, + }, + }, + request=request, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler: Final = AsyncHTTPHandler() + handler.client = client + yield CrowdStrikeAIDRHandler( + mode="pre_call", + guardrail_name="crowdstrike-aidr-guard", + api_key="pts_crowdstrike_tokenid", + api_base="https://api.crowdstrike.com/aidr/aiguard", + async_handler=handler, + ) + + +class _MessageShapedGuardrail(CustomGuardrail): + """Returns one text per chat message and no ``structured_messages`` rewrite. + + Prompt Security and friends scan messages rather than Responses text parts, + which is the shape that outnumbers the endpoint's own bookkeeping. + """ + + def __init__(self, redacted: str) -> None: + super().__init__(guardrail_name="message-shaped") + self.redacted: Final = redacted + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: str, + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + messages: Final = inputs.get("structured_messages") or () + return {"texts": [self.redacted for _ in messages]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("case", "instructions", "responses_input"), + [ + ( + "instructions add a system message", + "be terse", + [{"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}], + ), + ( + "tool items add messages that carry no text", + None, + [ + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}, + {"type": "function_call", "call_id": "c1", "name": "get_x", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c1", "output": "42"}, + ], + ), + ], +) +async def test_unalignable_rewrite_is_rejected_never_sent_unredacted( + case: str, + instructions: str | None, + responses_input: list[dict[str, object]], +) -> None: + """An unalignable rewrite must fail the request, not forward the raw prompt. + + Skipping the write-back would hand the model the unredacted text, so a + guardrail could be bypassed by adding ``instructions`` or a tool call. + """ + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + data: dict[str, object] = {"model": "gpt-4o", "input": responses_input} + if instructions is not None: + data["instructions"] = instructions + + with pytest.raises(UnappliableRequestRewrite): + await OpenAIResponsesHandler().process_input_messages( + data=data, + guardrail_to_apply=_MessageShapedGuardrail("my ssn is "), + ) + + assert "078-05-1120" in str(responses_input), case + + +@pytest.mark.asyncio +async def test_aligned_rewrite_is_written_back() -> None: + """Matching counts must still redact the input in place.""" + responses_input: list[dict[str, object]] = [ + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]} + ] + + await OpenAIResponsesHandler().process_input_messages( + data={"model": "gpt-4o", "input": responses_input}, + guardrail_to_apply=_MessageShapedGuardrail("my ssn is "), + ) + + assert cast(list, responses_input[0]["content"])[0]["text"] == "my ssn is " + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("case", "responses_input", "redacted_input"), + [ + ( + "instructions add a system message", + [{"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}], + [{"role": "user", "content": [{"type": "input_text", "text": "my ssn is "}]}], + ), + ( + "tool items sit between two user turns", + [ + {"role": "user", "content": [{"type": "input_text", "text": "hello"}]}, + {"type": "function_call", "call_id": "c1", "name": "get_x", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c1", "output": "42"}, + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}, + ], + [ + {"role": "user", "content": [{"type": "input_text", "text": "hello"}]}, + {"type": "function_call", "call_id": "c1", "name": "get_x", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c1", "output": "42"}, + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is "}]}, + ], + ), + ], +) +async def test_structured_rewrite_lands_on_shapes_the_flat_path_cannot_align( + case: str, + responses_input: list[dict[str, object]], + redacted_input: list[dict[str, object]], +) -> None: + data: dict[str, object] = {"model": "gpt-5.6", "instructions": "be terse", "input": responses_input} + + async with _guardrail_redacting("078-05-1120", "") as guardrail: + await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["input"] == redacted_input, case + assert data["instructions"] == "be terse", case 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 3a7ae7aba61..2932373c77e 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 @@ -1,6 +1,8 @@ """Tests for unified guardrail.""" import logging +from types import SimpleNamespace +from typing import Final import pytest @@ -19,14 +21,14 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( openai_messages_without_system, openai_messages_without_tool, ) +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse +from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) -from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse -from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) @@ -644,6 +646,64 @@ class TestUnifiedLLMGuardrails: class TestOCRGuardrailE2E: """End-to-end tests: UnifiedLLMGuardrails -> OCRHandler.""" + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", [CallTypes.ocr, CallTypes.aocr, CallTypes.aresponses]) + async def test_post_call_logging_fallback_is_limited_to_ocr(self, call_type: CallTypes) -> None: + guardrail: Final = RecordingGuardrail() + response: Final = ( + TestUnifiedLLMGuardrails.TestResponsesRouteAliases._responses_api_response() + if call_type == CallTypes.aresponses + else OCRResponse(model="mistral-ocr-latest", pages=[OCRPage(index=0, markdown="Scan this page")]) + ) + + result: Final = await UnifiedLLMGuardrails().async_post_call_success_hook( + data={ + "guardrail_to_apply": guardrail, + "litellm_logging_obj": SimpleNamespace(call_type=call_type.value), + }, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert result is response + if call_type in (CallTypes.ocr, CallTypes.aocr): + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Scan this page"] + else: + assert guardrail.apply_calls == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"]) + async def test_ocr_logging_fallback_preserves_route_and_response_precedence( + self, request_route: str | None, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.types.utils import ModelResponse + + _patch_translation_mappings( + monkeypatch, + { + CallTypes.completion: OpenAIChatCompletionsHandler, + CallTypes.acompletion: OpenAIChatCompletionsHandler, + CallTypes.aocr: OCRHandler, + }, + ) + guardrail: Final = RecordingGuardrail() + response: Final = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Chat output"}}]) + + result: Final = await guardrail.async_post_call_success_deployment_hook( + request_data={ + "guardrails": [guardrail.guardrail_name], + "user_api_key_request_route": request_route, + "litellm_logging_obj": SimpleNamespace(call_type=CallTypes.aocr.value), + }, + response=response, + call_type=CallTypes.aocr, + ) + + assert result is response + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Chat output"] + @pytest.mark.asyncio async def test_pre_call_hook_invokes_ocr_handler_for_input(self): """ diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 8fde4cc9d5e..11c3d2f8b20 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -15,7 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. """ import asyncio -from typing import Any +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -297,6 +297,38 @@ async def test_no_flag_fires_create_task_normally(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("call_type", ["ocr", "aocr", "completion", "acompletion", "embedding", "responses"]) +@pytest.mark.parametrize("exception_raised", [False, True]) +def test_native_pending_logging_is_released_only_for_ocr(call_type: str, exception_raised: bool) -> None: + pending: Final = MagicMock() + enqueue: Final = MagicMock() + logger: Final = MagicMock( + call_type=call_type, + _native_pending_logging=pending, + _enqueue_deferred_logging=enqueue, + ) + + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( + logging_obj=logger, + exception_raised=exception_raised, + ) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( + logging_obj=logger, + exception_raised=exception_raised, + ) + + if call_type in ("ocr", "aocr"): + pending.release.assert_called_once_with(not exception_raised) + assert logger._native_pending_logging is None + else: + pending.release.assert_not_called() + assert logger._native_pending_logging is pending + if exception_raised: + enqueue.assert_not_called() + else: + enqueue.assert_called_once_with() + + def test_flush_deferred_async_logging_fires_on_success(): """ Happy path: with no exception, the production flush helper invokes the diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 527d46931fe..1ab50cc30de 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1,9 +1,12 @@ import asyncio +import copy import json import time -from typing import Final +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime, timedelta from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -19,6 +22,7 @@ from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64 from litellm.models.credentials import CredentialItem from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.router import Router from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, _show_no_redis_warning, @@ -1579,7 +1583,7 @@ async def test_health_endpoint_filters_model_list_by_user_access(): ): from fastapi import Response - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) assert "model_list" in captured, "health_endpoint did not call _perform_health_check_and_save" returned_names = {m["model_name"] for m in captured["model_list"]} @@ -1642,7 +1646,7 @@ async def test_health_endpoint_keeps_full_model_list_for_all_proxy_models(): ): from fastapi import Response - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) returned_names = {m["model_name"] for m in captured["model_list"]} assert returned_names == { @@ -1710,12 +1714,231 @@ async def test_health_endpoint_resolves_all_team_models_to_team_allowlist(): ): from fastapi import Response - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) returned_names = {m["model_name"] for m in captured["model_list"]} assert returned_names == {"model-b"}, f"all-team-models key should health-check the team's models: {returned_names}" +def _router_for(model_list: Sequence[Mapping[str, object]]) -> Router: + return Router(model_list=copy.deepcopy(list(model_list))) + + +_ACCESS_GROUP_MODEL_LIST = [ + { + "model_name": "bedrock-nova", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-bedrock", "access_groups": ["bedrock-group"]}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-openai"}, + }, +] +_ACCESS_GROUP_ROUTER = _router_for(_ACCESS_GROUP_MODEL_LIST) +_TEAM_MODEL_LIST = [ + _ACCESS_GROUP_MODEL_LIST[0], + { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": { + "id": "id-team-b", + "team_id": "team-b", + "team_public_model_name": "bedrock-nova", + "access_groups": ["bedrock-group"], + }, + }, +] +_TEAM_CACHED_RESULTS = { + "healthy_endpoints": [ + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-bedrock"}, + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, +} +_ACCESS_GROUP_CACHED_RESULTS = { + "healthy_endpoints": [ + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-bedrock"}, + {"model": "openai/gpt-5.4-mini", "model_id": "id-openai"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, +} + + +@contextmanager +def _proxy_health_globals( + llm_model_list: Sequence[Mapping[str, object]], + llm_router: object, + use_background_health_checks: bool = False, + health_check_results: Mapping[str, object] | None = None, +) -> Iterator[None]: + with ( + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.llm_model_list", list(llm_model_list) + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.llm_router", llm_router + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.prisma_client", None + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.use_background_health_checks", use_background_health_checks + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.user_model", None + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.health_check_results", dict(health_check_results or {}) + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.health_check_details", True + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.health_check_concurrency", 1 + ), + ): + yield + + +@pytest.mark.asyncio +async def test_health_endpoint_expands_access_group_on_live_path(): + """ + LIT-6907 / gh-28206: a key granted a model access group carries the group + name in user_api_key_dict.models. Matching it as a literal model_name + filtered every deployment out and /health answered 0/0 for a model the + same key could call. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} + + with ( + _proxy_health_globals(_ACCESS_GROUP_MODEL_LIST, _ACCESS_GROUP_ROUTER), + patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"]), + model=None, + model_id=None, + ) + + assert [m["model_name"] for m in captured["model_list"]] == ["bedrock-nova"] + + +@pytest.mark.asyncio +async def test_health_endpoint_expands_access_group_on_background_cache_path(): + """ + LIT-6907: the background-cache path scoped the cached entries through the + same literal model_name match, so an access-group key got an empty result + plus a warning blaming missing model_info.id. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _ACCESS_GROUP_MODEL_LIST, + _ACCESS_GROUP_ROUTER, + use_background_health_checks=True, + health_check_results=_ACCESS_GROUP_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"]), + model=None, + model_id=None, + ) + + assert [e["model_id"] for e in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + assert "warnings" not in result + + +@pytest.mark.asyncio +async def test_health_endpoint_treats_no_team_all_team_models_as_unrestricted(): + """ + A key granted "all-team-models" without a team resolves to an empty + allowlist in the auth layer, which means unrestricted. /health used to + keep the unresolved sentinel and filter every deployment out instead. + """ + from fastapi import Response + + from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} + + with ( + _proxy_health_globals(_ACCESS_GROUP_MODEL_LIST, None), + patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-test-key", models=[SpecialModelNames.all_team_models.value], team_id=None + ), + model=None, + model_id=None, + ) + + assert {m["model_name"] for m in captured["model_list"]} == {"bedrock-nova", "gpt-5.4-mini"} + + +@pytest.mark.asyncio +async def test_health_endpoint_omits_model_id_warning_when_no_deployment_matches(): + """ + The missing-model_info.id warning is only true when a matching deployment + exists without an id. A key whose grants match no deployment at all gets a + plain empty result, not advice to populate ids that are already there. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _ACCESS_GROUP_MODEL_LIST, + _ACCESS_GROUP_ROUTER, + use_background_health_checks=True, + health_check_results=_ACCESS_GROUP_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["no-such-model"]), + model=None, + model_id=None, + ) + + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 0 + assert "warnings" not in result + + @pytest.mark.asyncio async def test_health_endpoint_filters_background_cache_by_user_access(): """ @@ -1907,7 +2130,7 @@ async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not(): # withheld so clients that previously parsed them can detect the change. assert ( non_admin_response.headers.get("Litellm-Health-Field-Notice") - == "api_base and api_version are admin-only on this endpoint" + == "api_base, api_version, aws_bedrock_runtime_endpoint are admin-only on this endpoint" ) assert "Litellm-Health-Field-Notice" not in admin_response.headers @@ -1996,7 +2219,7 @@ async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cach cache filter was driven by an unvalidated ID and the global cache leaked id-b's entry to the caller. """ - from fastapi import Response + from fastapi import HTTPException, Response from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.health_endpoints._health_endpoints import health_endpoint @@ -2047,21 +2270,18 @@ async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cach ): # Calling with model="model-b" rather than model_id="id-b" because # the model_id branch raises 404 when llm_router is None. The bug - # being verified is the same: targeted resolver must drop entries - # not in the caller's scoped model_list. With the fix, the result - # has no leaked endpoints and the targeted-503 path fires. - result = await health_endpoint( - response=response, - user_api_key_dict=user_api_key_dict, - model="model-b", - model_id=None, - ) + # being verified is the same: a target outside the caller's scoped + # model_list is refused before the cache is read. + with pytest.raises(HTTPException) as refused: + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-b", + model_id=None, + ) - leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])} - leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])} - assert "id-b" not in leaked_ids, "background cache leaked an out-of-scope deployment to a scoped caller" - assert result["healthy_count"] == 0 - assert response.status_code == 503 + assert refused.value.status_code == 403 + assert "leaky-internal.test" not in str(refused.value.detail) @pytest.mark.asyncio @@ -2193,6 +2413,7 @@ async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_e response=response, user_api_key_dict=user_api_key_dict, model="model-a", + model_id=None, ) assert response.status_code == 503 @@ -2253,6 +2474,7 @@ async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endp response=response, user_api_key_dict=user_api_key_dict, model="model-a", + model_id=None, ) assert response.status_code == 200 @@ -2637,6 +2859,691 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, assert canary not in str(cleaned) +async def _live_probed_model_ids( + model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth, model: str | None = None +) -> set[str]: + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} + + with ( + _proxy_health_globals(model_list, _router_for(model_list)), + patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=model, model_id=None) + + return {m["model_info"]["id"] for m in captured["model_list"]} + + +@pytest.mark.asyncio +async def test_health_endpoint_hides_another_teams_deployment_behind_a_shared_access_group(): + """ + Expanding an access group must not reach past the team boundary: a + team-a key holding the group name may not probe team-b's deployment even + though that deployment sits in the same group. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"), + ) + + assert probed == {"id-bedrock"} + + +@pytest.mark.asyncio +async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team(): + """ + Routing never serves a team-owned deployment to a caller without a team + (``filter_team_based_models``), so a team-less access-group key must not + probe team-b's deployment with team-b's credentials either. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None), + ) + + assert probed == {"id-bedrock"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "expected_ids"), + [(None, {"id-bedrock"}), ("team-a", {"id-bedrock"}), ("team-b", {"id-bedrock", "id-team-b"})], +) +async def test_health_endpoint_keeps_an_unrestricted_non_admin_key_to_its_own_team(team_id, expected_ids): + """ + A key with no model restriction is still bound by routing's team rule: + it may probe global deployments and its own team's, never another team's. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=[], team_id=team_id), + ) + + assert probed == expected_ids + + +@pytest.mark.asyncio +async def test_health_endpoint_lets_a_proxy_admin_probe_every_teams_deployment(): + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert probed == {"id-bedrock", "id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_keeps_an_unrestricted_non_admin_key_to_its_own_team_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=[], team_id="team-a"), + model=None, + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + + +@pytest.mark.asyncio +async def test_health_endpoint_shows_a_teams_own_deployment_by_its_public_name(): + """ + A team key names its team deployment by ``team_public_model_name``, while + the proxy model list carries the internal ``__`` + name; the deployment must still be probed for its own team. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + ) + + assert probed == {"id-bedrock", "id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_hides_another_teams_deployment_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"), + model=None, + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + + +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_targeted_deployment_outside_the_callers_scope_on_live_path(): + """ + A scoped key asking for a deployment it may not see must get a 403 and no + probe at all: probing the rest of its scope instead would report another + deployment's health under the requested id and store it as such. + """ + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + fake_perform = AsyncMock() + + with ( + _proxy_health_globals(_TEAM_MODEL_LIST, _router_for(_TEAM_MODEL_LIST)), + patch( # test-quality-ok: the probe must never run; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + fake_perform, + ), + pytest.raises(HTTPException) as excinfo, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None), + model=None, + model_id="id-team-b", + ) + + assert excinfo.value.status_code == 403 + fake_perform.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_targeted_deployment_outside_the_callers_scope_on_background_cache_path(): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ), + pytest.raises(HTTPException) as excinfo, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"), + model="bedrock-nova_team-b_9f2c", + model_id=None, + ) + + assert excinfo.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None), + model=None, + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + + +_TEAM_ONLY_MODEL_LIST = [_TEAM_MODEL_LIST[1]] +_BARE_NAME_MODEL_LIST = [ + {"model_name": "gpt-5.4-nano", "litellm_params": {"model": "gpt-5.4-nano"}, "model_info": {"id": "id-nano"}}, + { + "model_name": "gpt-5.4-nano_team-b_7c3d", + "litellm_params": {"model": "gpt-5.4-nano"}, + "model_info": {"id": "id-nano-team-b", "team_id": "team-b", "team_public_model_name": "gpt-5.4-nano"}, + }, +] +_BARE_NAME_CACHED_RESULTS = { + "healthy_endpoints": [ + {"model": "gpt-5.4-nano", "model_id": "id-nano"}, + {"model": "gpt-5.4-nano", "model_id": "id-nano-team-b"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, +} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_a_team_only_deployment_by_its_public_name_on_live_path(): + """ + A team key targets its deployment by ``team_public_model_name``; when that + name resolves to nothing but the team deployment, the probe must run rather + than 403 as if the key were out of scope. + """ + probed = await _live_probed_model_ids( + _TEAM_ONLY_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + ) + + assert probed == {"id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_a_team_only_deployment_by_its_public_name_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_only_the_owning_teams_copy_behind_a_shared_public_name_on_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] + + +async def _live_narrowed_model_ids( + model_list: Sequence[Mapping[str, object]], + user_api_key_dict: UserAPIKeyAuth, + model: str | None = None, + model_id: str | None = None, +) -> set[str]: + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + probed = [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list] + return probed, [], {} + + with ( + _proxy_health_globals(model_list, _router_for(model_list)), + patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + ): + result = await health_endpoint( + response=Response(), user_api_key_dict=user_api_key_dict, model=model, model_id=model_id + ) + + return {ep["model_id"] for ep in result["healthy_endpoints"]} + + +_ADMIN_OUTSIDE_TEAM_B = UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_public_copies(): + """ + An admin outside team-b asks for ``bedrock-nova``. Team-b's copy answers to + that name only for team-b (routing keys public names by team), so probing + it too would spend team-b's credentials and let a healthy team copy mask a + down global deployment as 200. + """ + probed = await _live_narrowed_model_ids(_TEAM_MODEL_LIST, _ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova") + + assert probed == {"id-bedrock"} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_owning_teams_copy_behind_a_shared_public_name(): + """Team-b's requests for ``bedrock-nova`` route to its copy alone, so its health probe reaches only that copy.""" + probed = await _live_narrowed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + ) + + assert probed == {"id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_only_the_teams_copy_when_provider_model_equals_public_name(): + """A bare provider model equal to the public name must not pull the global copy into the team's probe.""" + probed = await _live_narrowed_model_ids( + _BARE_NAME_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-nano"], team_id="team-b"), + model="gpt-5.4-nano", + ) + + assert probed == {"id-nano-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_only_the_teams_copy_when_provider_model_equals_public_name_on_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _BARE_NAME_MODEL_LIST, + _router_for(_BARE_NAME_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_BARE_NAME_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["gpt-5.4-nano"], team_id="team-b"), + model="gpt-5.4-nano", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-nano-team-b"] + + +@pytest.mark.asyncio +async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_public_copies_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), user_api_key_dict=_ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_a_team_only_public_name_for_an_admin_on_live_path(): + """ + An admin's request for a public name only team-b's deployment carries routes + to that deployment, so the health probe for that name must reach it too + instead of answering an empty 503. + """ + probed = await _live_narrowed_model_ids(_TEAM_ONLY_MODEL_LIST, _ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova") + + assert probed == {"id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_a_team_only_public_name_for_an_admin_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), user_api_key_dict=_ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_keeps_a_team_only_public_name_off_a_team_less_key(use_background_health_checks): + """ + A key with no team holds the name ``bedrock-nova`` but never sees team-b's + deployment, so the public-name fallback an admin gets must not open it up. + """ + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_TEAM_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"]), + model="bedrock-nova", + model_id=None, + ) + + assert refused.value.status_code == 403 + assert "bedrock-nova" in str(refused.value.detail) + probe.assert_not_awaited() + + +def test_resolve_targeted_model_ids_lets_model_id_win_over_model(): + resolve = _health_endpoints_module._resolve_targeted_model_ids + + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", "id-team-b", None) == {"id-team-b"} + assert resolve([_TEAM_MODEL_LIST[0]], "bedrock-nova", "id-team-b", None) == set() + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, None) == {"id-bedrock"} + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, "team-b") == {"id-team-b"} + assert resolve(_TEAM_ONLY_MODEL_LIST, "bedrock-nova", None, None) == {"id-team-b"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_rejects_an_in_scope_model_paired_with_a_foreign_model_id(use_background_health_checks): + """ + A key scoped to ``bedrock-nova`` pairs that name with another team's + deployment id. The in-scope name must not carry the foreign id past the + 403: the live path narrows by id first, so the caller's own deployment + would be probed and its result stored under the foreign id. + """ + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_TEAM_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"]), + model="bedrock-nova", + model_id="id-team-b", + ) + + assert refused.value.status_code == 403 + assert "id-team-b" in str(refused.value.detail) + probe.assert_not_awaited() + + +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +@pytest.mark.asyncio +async def test_health_endpoint_returns_404_for_a_model_paired_with_an_unknown_model_id(use_background_health_checks): + """ + ``model_id`` wins over ``model``: pairing a known name with an id no + deployment carries gets the same 404 as the lone unknown id, before any + probe runs or a result is stored under the unknown id. + """ + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_TEAM_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=_ADMIN_OUTSIDE_TEAM_B, + model="bedrock-nova", + model_id="id-nobody-has", + ) + + assert refused.value.status_code == 404 + assert "id-nobody-has" in str(refused.value.detail) + probe.assert_not_awaited() + + +def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch): + """ + The dashboard's Test Connect button reads ``result.error`` and + ``result.raw_request_typed_dict`` from /health/test_connection, so the + allowlist must keep both while dropping the probe's own params. + """ + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + app = FastAPI() + app.include_router(_health_endpoints_module.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + client = TestClient(app) + + with ( + patch( # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + respx.mock(assert_all_called=True) as respx_mock, + ): + respx_mock.post(host="api.openai.com", path="/v1/chat/completions").respond( + status_code=401, json={"error": {"message": "Incorrect API key provided"}} + ) + response = client.post( + "/health/test_connection", + json={ + "mode": "chat", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test", "timeout": 7}, + }, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert body["status"] == "error" + assert "Incorrect API key provided" in body["result"]["error"] + assert "api.openai.com" in body["result"]["raw_request_typed_dict"]["raw_request_api_base"] + assert not {"api_key", "timeout", "exception"} & set(body["result"]) + + +def test_clean_endpoint_data_keeps_only_json_safe_diagnostics(): + """ + LIT-6907: _clean_endpoint_data used to copy every litellm_param not on a + deny list, so a nested mapping keyed by a tuple reached jsonable_encoder + and 500'd /health. Only the explicit allowlist survives now. + """ + from fastapi.encoders import jsonable_encoder + + from litellm.proxy.health_check import _clean_endpoint_data + + cleaned = _clean_endpoint_data( + { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + "metadata": {("us-east-1", "primary"): "canary-nested-mapping"}, + "allow_client_keepalive_override": False, + "api_key": "CANARY-API-KEY", + "x-ratelimit-remaining-requests": 99, + "raw_request_typed_dict": {"raw_request_api_base": "https://example.test"}, + "aws_bedrock_runtime_endpoint": "https://vpce-bedrock.example.test", + }, + details=True, + ) + + assert cleaned == { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + "x-ratelimit-remaining-requests": 99, + "raw_request_typed_dict": {"raw_request_api_base": "https://example.test"}, + "aws_bedrock_runtime_endpoint": "https://vpce-bedrock.example.test", + } + assert jsonable_encoder(cleaned) == cleaned + + +@pytest.mark.asyncio +async def test_health_endpoint_result_survives_non_json_safe_deployment_params(): + """ + LIT-6907: the full /health path with a deployment carrying a tuple-keyed + nested mapping must produce a response FastAPI can encode, with the + approved diagnostics intact and the offending param absent. + """ + from fastapi import Response + from fastapi.encoders import jsonable_encoder + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + model_list = [ + { + "model_name": "bedrock-nova", + "litellm_params": { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "aws_region_name": "us-east-1", + "aws_access_key_id": "CANARY-ACCESS-KEY", + "metadata": {("us-east-1", "primary"): "canary-nested-mapping"}, + }, + "model_info": {"id": "id-bedrock"}, + } + ] + + with ( + _proxy_health_globals(model_list, None), + patch( # test-quality-ok: the provider probe is faked; the assertion is the response shaping after it + "litellm.ahealth_check", AsyncMock(return_value={"x-ratelimit-remaining-requests": 99}) + ), + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-admin-key", user_role=LitellmUserRoles.PROXY_ADMIN), + model=None, + model_id=None, + ) + + encoded = jsonable_encoder(result) + assert encoded["healthy_count"] == 1 + entry = encoded["healthy_endpoints"][0] + assert entry["model_id"] == "id-bedrock" + assert entry["aws_region_name"] == "us-east-1" + assert entry["x-ratelimit-remaining-requests"] == 99 + assert "metadata" not in entry + assert "CANARY" not in str(encoded) + + class TestConfigBaseForHealthCheck: """A request that sets its own connection fields gets a base without the configuration's credentials; anything it leaves unset still comes from diff --git a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py index 860fb762450..1aa9382f3fe 100644 --- a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py @@ -5,6 +5,8 @@ Validates that email and secret manager operations are independent and non-block """ import asyncio +import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -475,7 +477,7 @@ class TestRotateVirtualKeyInSecretManager: class TestKeyUpdatedAuditLogObjectId: """Tests that /key/update audit logs never store the raw virtual key (issue #31620).""" - async def _run_updated_hook_and_capture_audit_log(self, request_key: str): + async def _run_updated_hook_and_capture_audit_log(self, request_key: str, detach_project: bool = False): import asyncio from litellm.proxy._types import ( @@ -493,6 +495,11 @@ class TestKeyUpdatedAuditLogObjectId: existing_key_row = LiteLLM_VerificationToken( token=hash_token("sk-raw-test-key-31620"), key_name="sk-...1620", + project_id="project-orbit", + ) + + data: Final = UpdateKeyRequest( + key=request_key, max_budget=2000.0, **({"project_id": None} if detach_project else {}) ) with ( @@ -503,7 +510,7 @@ class TestKeyUpdatedAuditLogObjectId: ), ): await KeyManagementEventHooks.async_key_updated_hook( - data=UpdateKeyRequest(key=request_key, max_budget=2000.0), + data=data, existing_key_row=existing_key_row, response=MagicMock(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin-key", user_id="admin"), @@ -530,13 +537,22 @@ class TestKeyUpdatedAuditLogObjectId: assert raw_key not in str(audit_row.updated_values) assert raw_key not in str(audit_row.before_value) + @pytest.mark.parametrize("detach_project", [False, True]) @pytest.mark.asyncio - async def test_update_audit_log_passes_through_hashed_key(self): + async def test_update_audit_log_passes_through_hashed_key(self, detach_project: bool): """An already-hashed token sent to /key/update is stored unchanged.""" from litellm.proxy.utils import hash_token hashed_key = hash_token("sk-raw-test-key-31620") - audit_row = await self._run_updated_hook_and_capture_audit_log(request_key=hashed_key) + audit_row: Final = await self._run_updated_hook_and_capture_audit_log( + request_key=hashed_key, detach_project=detach_project, + ) assert audit_row.object_id == hashed_key + updated_values: Final = json.loads(audit_row.updated_values) + assert ("project_id" in updated_values) is detach_project + if detach_project: + assert updated_values["project_id"] is None + assert json.loads(audit_row.before_value)["project_id"] == "project-orbit" + assert updated_values["max_budget"] == 2000.0 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index c96cfc3ee4a..fad137af66b 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,5 +1,7 @@ import asyncio +import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -7,6 +9,7 @@ import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( _get_budget_reservation_from_metadata, @@ -15,6 +18,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( _update_database_and_spend_counters, run_spend_event, ) +from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -2347,3 +2351,27 @@ async def test_sidecar_ignores_an_undecodable_event(): # test-quality-ok: a dis mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() await run_spend_event(b"garbage\n") mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_model_rejection(): + raw_model: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes" + writer: Final = MagicMock(spec=DBSpendUpdateWriter) + writer.update_database = AsyncMock() + logger: Final = _ProxyDBLogger(spend_writer=lambda: writer) + + await logger.async_post_call_failure_hook( + request_data={"model": raw_model, "messages": [{"role": "user", "content": "hi"}]}, + original_exception=ProxyModelNotFoundError(route="/chat/completions", model_name=raw_model), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + error_information: Final = writer.update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"][ + "error_information" + ] + assert "medical records" not in json.dumps(error_information) + assert ( + error_information["error_message"] + == "/chat/completions: Invalid model name passed in. Call `/v1/models` to view available models for your key." + ) + assert error_information["error_class"] == "ProxyModelNotFoundError" 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 d1d669cae38..92b1ab1586d 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 @@ -1,9 +1,12 @@ import json from datetime import datetime, timezone from types import SimpleNamespace +from typing import Final import pytest from fastapi.testclient import TestClient +from fastapi import HTTPException +from pytest_mock import MockerFixture from litellm.proxy._types import ( @@ -2128,6 +2131,128 @@ def test_update_internal_user_params_keeps_original_max_budget_when_not_provided assert "user_alias" in non_default_values +@pytest.mark.parametrize("cleared_budget", [{}, None], ids=["empty-map", "null"]) +def test_update_internal_user_params_clears_model_budget(cleared_budget: dict[str, object] | None) -> None: + request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=cleared_budget) + + update: Final = _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) + + assert update == {"user_id": "user-spruce", "model_max_budget": {}} + + +def test_update_internal_user_params_preserves_model_budget_presence_and_neighbors() -> None: + omitted: Final = UpdateUserRequest(user_id="user-spruce", user_alias="Spruce") + assert _update_internal_user_params(data_json=omitted.model_dump(), data=omitted) == { + "user_id": "user-spruce", + "user_alias": "Spruce", + } + + replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0} + request: Final = UpdateUserRequest( + user_id="user-spruce", + model_max_budget=replacement, + max_budget=50, + user_alias=None, + models=[], + allowed_cache_controls=[], + config={}, + ) + assert _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) == { + "user_id": "user-spruce", + "model_max_budget": replacement, + "max_budget": 50, + } + + +@pytest.mark.parametrize("invalid_budget", [{"model-spruce": "invalid"}, {"model-spruce": {"budget_limit": "invalid"}}]) +def test_update_internal_user_params_rejects_invalid_model_budget(invalid_budget: dict[str, object]) -> None: + request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=invalid_budget) + + with pytest.raises(HTTPException) as exc: + _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request) + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: MockerFixture) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper + + saved_user: Final = LiteLLM_UserTable( + user_id="user-spruce", + user_email="spruce@example.test", + model_max_budget={"model-spruce": {"budget_limit": 5, "time_period": "1d"}}, + max_budget=50, + ) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) + prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + await _update_single_user_helper( + user_request=UpdateUserRequest(user_email=saved_user.user_email, model_max_budget={}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert prisma_client.update_data.call_args.kwargs["data"]["model_max_budget"] == {} + assert "max_budget" not in prisma_client.update_data.call_args.kwargs["data"] + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) + + +@pytest.mark.asyncio +async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkUpdateUserRequest + + saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", model_max_budget={"model-spruce": {"budget_limit": 5}}) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) + prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + with pytest.raises(HTTPException) as exc: + await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": {"model-spruce": "invalid"}}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_changed_by=None, + ) + assert exc.value.status_code == 400 + prisma_client.db.litellm_usertable.update_many.assert_not_called() + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) == saved_user + + response: Final = await bulk_user_update( + data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": None}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_changed_by=None, + ) + + prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={"model_max_budget": "{}"}) + prisma_client.update_data.assert_not_called() + assert response.successful_updates == 1 + assert response.results[0].updated_user["model_max_budget"] == {} + assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) + + def test_generate_request_base_validator(): """ Test that GenerateRequestBase validator converts empty string to None for max_budget @@ -3498,7 +3623,11 @@ def test_enforce_user_info_access_blocks_cross_user_lookup(): @pytest.mark.asyncio -async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): +@pytest.mark.parametrize( + ("budget_field", "budget_value"), + [("max_budget", 999999), ("model_max_budget", {}), ("model_max_budget", None)], +) +async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budget_field, budget_value): """Non-admin updating their own record must be blocked from modifying max_budget (self-escalation).""" from fastapi import HTTPException @@ -3508,6 +3637,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): ) mock_prisma_client = mocker.MagicMock() + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-1", "data": {"user_id": "user-1"}}) existing_user = mocker.MagicMock() existing_user.model_dump.return_value = { "user_id": "user-1", @@ -3519,10 +3649,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - user_request = UpdateUserRequest( - user_id="user-1", - max_budget=999999, - ) + user_request = UpdateUserRequest.model_validate({"user_id": "user-1", budget_field: budget_value}) caller = UserAPIKeyAuth( user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER, @@ -3533,7 +3660,8 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker): user_request=user_request, user_api_key_dict=caller ) assert exc.value.status_code == 403 - assert "max_budget" in str(exc.value.detail) + assert budget_field in str(exc.value.detail) + mock_prisma_client.update_data.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 65cc23ea67f..646ae43f37a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1,3 +1,4 @@ +from typing import Final import json from datetime import datetime, timedelta, timezone @@ -18129,3 +18130,59 @@ def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatc ) is True ) + + +@pytest.mark.asyncio +async def test_project_detachment_preserves_omission_and_other_key_fields(): + existing: Final = LiteLLM_VerificationToken( + token="project-detach-token", project_id="project-orbit", team_id="team-orbit", + organization_id="org-orbit", models=["model-orbit"], max_budget=5, rpm_limit=97, + ) + omitted: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key=existing.token, key_alias="renamed"), existing_key_row=existing, + ) + assert "project_id" not in omitted + cleared: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key=existing.token, project_id=None), existing_key_row=existing, + ) + assert cleared == {"project_id": None, "metadata": {}} + assert existing.project_id == "project-orbit" + + +@pytest.mark.parametrize("project_id", [None, "project-orbit", "project-other", ""]) +@pytest.mark.asyncio +async def test_project_detachment_uses_effective_project_for_validation(project_id: str | None): + existing: Final = LiteLLM_VerificationToken(token="project-detach-token", project_id="project-orbit") + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + data: Final = UpdateKeyRequest(key=existing.token, project_id=project_id, models=["model-other"]) + if project_id is None: + await _validate_update_key_data( + data, existing, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, MagicMock(), cache, + ) + else: + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data, existing, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, MagicMock(), cache, + ) + assert exc.value.status_code == 400 + expected: Final = "not in project's allowed models" if project_id == "project-orbit" else "reassignment" + assert expected in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_key_creator_cannot_detach_project_without_admin_access(): + existing: Final = LiteLLM_VerificationToken( + token="project-detach-token", project_id="project-orbit", user_id="user-orbit", created_by="user-orbit", + ) + database: Final = MagicMock() + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id=None), existing, + UserAPIKeyAuth(user_id="user-orbit", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 403 + assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) 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 5325e069813..c66095dc1fd 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 @@ -3,7 +3,7 @@ import asyncio import contextlib import json from collections.abc import Mapping -from typing import Dict, Optional +from typing import Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3290,6 +3290,99 @@ def _build_db_model_with_pricing(): ) +class TestUpdateDBModelCompression: + @pytest.mark.parametrize( + "compression_patch, expected", + [ + ( + {}, + { + "auto_router_routing_compression": "routing-compressor", + "auto_router_model_compression": "model-compressor", + }, + ), + ({"auto_router_routing_compression": None}, {"auto_router_model_compression": "model-compressor"}), + ({"auto_router_model_compression": None}, {"auto_router_routing_compression": "routing-compressor"}), + ( + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"}, + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"}, + ), + ( + { + "auto_router_routing_compression": "new-compressor", + "auto_router_model_compression": "new-compressor", + }, + { + "auto_router_routing_compression": "new-compressor", + "auto_router_model_compression": "new-compressor", + }, + ), + ], + ) + def test_compression_patch_preserves_omissions_and_explicit_choices( + self, monkeypatch: pytest.MonkeyPatch, compression_patch: dict[str, str | None], expected: dict[str, str] + ): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + monkeypatch.setenv("LITELLM_SALT_KEY", "synthetic-compression-salt") + result: Final = update_db_model( + db_model=Deployment( + model_name="synthetic-router", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", + auto_router_routing_compression=encrypt_value_helper("routing-compressor"), + auto_router_model_compression=encrypt_value_helper("model-compressor"), + ), + model_info=ModelInfo(id="compression-router"), + ), + updated_patch=updateDeployment.model_validate({"litellm_params": compression_patch}), + ) + params: Final = json.loads(result["litellm_params"]) + assert { + key: decrypt_value_helper(value=val, key=key) + for key, val in params.items() + if key in ("auto_router_routing_compression", "auto_router_model_compression") + } == expected + + def test_explicit_compression_clear_removes_both_saved_overrides(self): + from litellm.proxy.guardrails.auto_router_compression import policy_from_litellm_params + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="synthetic-router", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", + auto_router_routing_compression="routing-compressor", + auto_router_model_compression="model-compressor", + api_base="http://127.0.0.1:9999/v1", + temperature=0, + ), + model_info=ModelInfo(id="compression-router", team_id="synthetic-team"), + ) + result: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment.model_validate( + { + "litellm_params": { + "auto_router_routing_compression": None, + "auto_router_model_compression": None, + "api_base": None, + }, + "model_info": {"team_id": None}, + } + ), + ) + + params: Final = json.loads(result["litellm_params"]) + assert "auto_router_routing_compression" not in params + assert "auto_router_model_compression" not in params + assert policy_from_litellm_params(params) is None + assert params["api_base"] == "http://127.0.0.1:9999/v1" + assert params["temperature"] == 0 + assert json.loads(result["model_info"])["team_id"] == "synthetic-team" + + class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored 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 2f6561046b1..2fb496d6231 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11220,13 +11220,14 @@ async def test_team_info_returns_model_aliases(): @pytest.mark.asyncio -async def test_team_info_hydrates_member_emails_from_the_user_table(): - """/team/info must fill in emails missing from the members_with_roles snapshot. +async def test_team_info_hydrates_member_names_and_emails_from_the_user_table(): + """/team/info must attach each member's display name and fill in emails missing + from the members_with_roles snapshot. - members_with_roles is written at add-time, so a member added by user_id alone - carries user_email=None forever. Without this join the Admin UI's member table - shows "-" for a user that has an email on their user row. A stored email is left - exactly as-is. + members_with_roles is written at add-time, so it never carries user_alias and a + member added by user_id alone carries user_email=None forever. Without this join + the Admin UI's member table can only show emails. A stored email is left exactly + as-is. """ from fastapi import Request @@ -11246,13 +11247,8 @@ async def test_team_info_hydrates_member_emails_from_the_user_table(): find_many = AsyncMock( return_value=[ - LiteLLM_UserTable( - user_id="no-email-on-roster", - user_email="real@example.com", - max_budget=None, - spend=0.0, - models=[], - ) + _user_row("no-email-on-roster", "real@example.com", "Real Person"), + _user_row("already-stored", "current@example.com", "Stored Person"), ] ) @@ -11270,12 +11266,12 @@ async def test_team_info_hydrates_member_emails_from_the_user_table(): ) members = response["team_info"].members_with_roles - assert [(m.user_id, m.user_email) for m in members] == [ - ("no-email-on-roster", "real@example.com"), - ("already-stored", "stored@example.com"), + assert [(m.user_id, m.user_email, m.user_alias) for m in members] == [ + ("no-email-on-roster", "real@example.com", "Real Person"), + ("already-stored", "stored@example.com", "Stored Person"), ] - # only the member actually missing an email is looked up - assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["no-email-on-roster"]}} + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["already-stored", "no-email-on-roster"]}} @pytest.mark.asyncio @@ -12472,89 +12468,93 @@ async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids repo.return_value.table.find_many.assert_not_awaited() -def _user_row(user_id: str, user_email: str | None) -> LiteLLM_UserTable: +def _user_row(user_id: str, user_email: str | None, user_alias: str | None = None) -> LiteLLM_UserTable: return LiteLLM_UserTable( - user_id=user_id, user_email=user_email, max_budget=None, spend=0.0, models=[] + user_id=user_id, user_email=user_email, user_alias=user_alias, max_budget=None, spend=0.0, models=[] ) @pytest.mark.asyncio -async def test_hydrate_member_emails_fills_in_emails_the_roster_snapshot_never_captured(): - """A member added by user_id alone has user_email=None on the stored roster entry. - - /team/info has to fill it in from the user row, or the UI renders "-" for a user - that plainly has an email. +async def test_hydrate_member_user_details_attaches_alias_and_fills_in_missing_email(): + """The stored roster never carries a display name, and a member added by user_id + alone has user_email=None. /team/info has to fill both in from the user row so the + UI can show and search by a human-readable name instead of only an email. """ - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details - find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com")]) + find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com", "Found Person")]) with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: repo.return_value.table.find_many = find_many - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), members=[Member(user_id="by-id", role="admin")], ) - assert [(m.user_id, m.user_email, m.role) for m in hydrated] == [("by-id", "found@example.com", "admin")] + assert [(m.user_id, m.user_email, m.user_alias, m.role) for m in hydrated] == [ + ("by-id", "found@example.com", "Found Person", "admin") + ] find_many.assert_awaited_once() assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id"]}} @pytest.mark.asyncio -async def test_hydrate_member_emails_never_overwrites_a_stored_email(): - """The snapshot wins wherever it has a value - hydration only fills blanks. - - Overwriting would be a real behavior change to /team/info; filling a null is not. - """ - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails - - find_many = AsyncMock(return_value=[_user_row("has-email", "current@example.com")]) +async def test_hydrate_member_user_details_never_overwrites_a_stored_email(): + """The snapshot wins wherever it has a value - hydration only fills blanks.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.table.find_many = find_many + repo.return_value.table.find_many = AsyncMock( + return_value=[_user_row("has-email", "current@example.com", "Current Name")] + ) - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), members=[Member(user_id="has-email", user_email="stored@example.com", role="user")], ) - assert hydrated[0].user_email == "stored@example.com" - # nothing was missing, so no round-trip either - find_many.assert_not_awaited() + assert (hydrated[0].user_email, hydrated[0].user_alias) == ("stored@example.com", "Current Name") @pytest.mark.asyncio -async def test_hydrate_member_emails_leaves_members_alone_when_the_user_row_has_no_email(): - """A user row with no email leaves the member as-is rather than inventing one.""" - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails +async def test_hydrate_member_user_details_leaves_blanks_when_the_user_row_is_bare_or_missing(): + """A user row with no email or alias, or no user row at all, must not invent values.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("no-email", None)]) + repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("bare", None)]) - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), - members=[Member(user_id="no-email", role="user"), Member(user_email="e@example.com", role="user")], + members=[ + Member(user_id="bare", role="user"), + Member(user_id="deleted", user_email="gone@example.com", role="user"), + Member(user_email="e@example.com", role="user"), + ], ) - assert [m.user_email for m in hydrated] == [None, "e@example.com"] + assert [(m.user_id, m.user_email, m.user_alias) for m in hydrated] == [ + ("bare", None, None), + ("deleted", "gone@example.com", None), + (None, "e@example.com", None), + ] @pytest.mark.asyncio -async def test_hydrate_member_emails_skips_the_query_when_every_member_has_one(): - """No blanks means /team/info pays for no extra query.""" - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails +async def test_hydrate_member_user_details_skips_the_query_when_no_member_has_a_user_id(): + """Email-only roster entries give nothing to look up, so /team/info pays for no query.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: repo.return_value.table.find_many = AsyncMock() - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), - members=[Member(user_id="a", user_email="a@example.com", role="user")], + members=[Member(user_email="a@example.com", role="user")], ) - assert hydrated[0].user_email == "a@example.com" + assert [(m.user_email, m.user_alias) for m in hydrated] == [("a@example.com", None)] repo.return_value.table.find_many.assert_not_awaited() diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 2f47736a398..6eac53df645 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1151,10 +1151,10 @@ async def test_prepare_window_spend_counter_increment_missing_window_start_inval # --------------------------------------------------------------------------- -def _two_pending_increments() -> tuple[ps._PendingSpendIncrement, ...]: +def _two_pending_increments() -> tuple[ps.PendingSpendIncrement, ...]: return ( - ps._PendingSpendIncrement(counter_key="spend:key:k", increment=1.5), - ps._PendingSpendIncrement(counter_key="spend:team:t", increment=1.5), + ps.PendingSpendIncrement(counter_key="spend:key:k", increment=1.5), + ps.PendingSpendIncrement(counter_key="spend:team:t", increment=1.5), ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py new file mode 100644 index 00000000000..fe852be775c --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py @@ -0,0 +1,119 @@ +"""Auth-resolved budget state rides on ``UserAPIKeyAuth`` and round-trips through request metadata.""" + +from datetime import datetime, timezone + +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_OrganizationTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.spend_tracking.carried_budget_state import ( + carried_budget_metadata, + carry_organization_budget_state, + carry_team_and_user_budget_state, +) +from litellm.types.proxy.carried_budget_state import ( + KeyBudgetSnapshot, + OrgBudgetSnapshot, + TeamBudgetSnapshot, + UserBudgetSnapshot, +) + +RESET_AT = datetime(2026, 10, 1, 12, 30, tzinfo=timezone.utc) + + +def test_team_and_user_state_round_trips_through_metadata(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=RESET_AT, max_budget=300.0), + user_object=LiteLLM_UserTable(user_id="u1", budget_reset_at=None, max_budget=None, user_alias="Alice"), + ) + + metadata = dict(carried_budget_metadata(token)) + + assert metadata == { + "user_api_key_team_budget_reset_at": RESET_AT.isoformat().replace("+00:00", "Z"), + "user_api_key_team_table_max_budget": 300.0, + "user_api_key_user_budget_reset_at": None, + "user_api_key_user_table_max_budget": None, + "user_api_key_user_alias": "Alice", + } + assert TeamBudgetSnapshot.from_metadata(metadata) == TeamBudgetSnapshot(budget_reset_at=RESET_AT, max_budget=300.0) + assert UserBudgetSnapshot.from_metadata(metadata) == UserBudgetSnapshot( + budget_reset_at=None, max_budget=None, user_alias="Alice" + ) + + +def test_missing_objects_leave_no_metadata_and_no_snapshot(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") + carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None) + + assert dict(carried_budget_metadata(token)) == {} + assert TeamBudgetSnapshot.from_metadata({}) is None + assert UserBudgetSnapshot.from_metadata({"user_api_key_user_alias": "Alice"}) is None + assert OrgBudgetSnapshot.from_metadata({"user_api_key_org_spend": 1.0}) is None + assert KeyBudgetSnapshot.from_metadata({}) is None + + +def test_organization_state_carries_alias_spend_and_max_budget(): + token = UserAPIKeyAuth(token="hashed", org_id="o1") + org = LiteLLM_OrganizationTable( + organization_id="o1", + organization_alias="platform-org", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=12.5, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + + carry_organization_budget_state(valid_token=token, org_table=org) + + assert token.organization_alias == "platform-org" + assert OrgBudgetSnapshot.from_metadata(carried_budget_metadata(token)) == OrgBudgetSnapshot( + spend=12.5, max_budget=100.0 + ) + + +def test_organization_without_budget_table_carries_no_cap(): + token = UserAPIKeyAuth(token="hashed", org_id="o1") + org = LiteLLM_OrganizationTable( + organization_id="o1", + budget_id="b1", + created_by="admin", + updated_by="admin", + spend=3.0, + ) + + carry_organization_budget_state(valid_token=token, org_table=org) + + assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=3.0, max_budget=None) + + +def test_key_snapshot_parses_the_iso_string_auth_metadata_writes(): + assert KeyBudgetSnapshot.from_metadata({"user_api_key_budget_reset_at": RESET_AT.isoformat()}) == KeyBudgetSnapshot( + budget_reset_at=RESET_AT + ) + assert KeyBudgetSnapshot.from_metadata({"user_api_key_budget_reset_at": None}) == KeyBudgetSnapshot( + budget_reset_at=None + ) + + +def test_snapshots_never_reach_the_serialized_token(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1") + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", budget_reset_at=RESET_AT), + user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), + ) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) + + dumped = token.model_dump() + + assert "team_budget_snapshot" not in dumped + assert "user_budget_snapshot" not in dumped + assert "org_budget_snapshot" not in dumped + assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py b/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py new file mode 100644 index 00000000000..3e4b817fab8 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_counter_batch.py @@ -0,0 +1,527 @@ +"""Exact Redis round-trip counts for the spend counters admission reads within one auth scope.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping, Sequence +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm.proxy.proxy_server as ps +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.spend_tracking.spend_counter_batch import ( + SpendCounterBatch, + active_spend_counter_batch, + admission_counter_keys, + bind_admission_counter_keys, + release_spend_counter_batch, + spend_counter_batch_scope, +) + +TOKEN = UserAPIKeyAuth(token="hashed", team_id="team", user_id="user", org_id="org") +TOKEN_KEYS = frozenset( + { + "spend:key:hashed", + "spend:team:team", + "spend:team_member:user:team", + "spend:user:user", + "spend:end_user:eu", + "spend:org:org", + } +) + + +class CountingRedis(RedisCache): + def __init__(self, store: dict[str, object] | None = None, fail: bool = False) -> None: + self.store: dict[str, object] = dict(store or {}) + self.fail = fail + self.commands: list[str] = [] + + async def async_get_cache(self, key: str, **kwargs: object) -> object: + if self.fail: + raise ConnectionError("redis down") + self.commands.append(f"GET {key}") + return self.store.get(key) + + async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, object]: + if self.fail: + raise ConnectionError("redis down") + self.commands.append(f"MGET {' '.join(key_list)}") + return {key: self.store.get(key) for key in key_list} + + def get_ttl(self, **kwargs: object) -> int | None: + return None + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self.commands.append(f"INCRBYFLOAT {key} {value}") + return self._incr(key, value) + + async def async_increment_pipeline( + self, increment_list: Sequence[Mapping[str, object]], **kwargs: object + ) -> list[float]: + self.commands.append(f"PIPELINE {' '.join(str(op['key']) for op in increment_list)}") + return [self._incr(str(op["key"]), float(str(op["increment_value"]))) for op in increment_list] + + def _incr(self, key: str, value: float) -> float: + total = float(str(self.store.get(key, 0.0))) + value + self.store[key] = total + return total + + +def _spend_counter_cache(redis: RedisCache | None, in_memory: dict[str, float] | None = None) -> MagicMock: + cache = MagicMock() + cache.redis_cache = redis + cache.in_memory_cache.get_cache = MagicMock(side_effect=lambda key: (in_memory or {}).get(key)) + return cache + + +def test_admission_counter_keys_cover_every_entity_the_checks_read(): + assert admission_counter_keys(TOKEN, end_user_id="eu") == TOKEN_KEYS + assert admission_counter_keys(UserAPIKeyAuth(token="hashed"), end_user_id=None) == {"spend:key:hashed"} + assert "spend:team_member:user:team" not in admission_counter_keys( + UserAPIKeyAuth(token="hashed", user_id="user"), end_user_id=None + ) + + +@pytest.mark.asyncio +async def test_bound_counters_share_one_mget_and_a_clean_miss_is_authoritative(): + redis = CountingRedis({"spend:key:hashed": 1.5, "spend:team:team": 2.5}) + batch = SpendCounterBatch(redis) + batch.bind(TOKEN_KEYS) + + reads = await asyncio.gather(*(batch.read(key) for key in sorted(TOKEN_KEYS))) + + assert len(redis.commands) == 1 + assert set(redis.commands[0].split()[1:]) == TOKEN_KEYS + assert dict(zip(sorted(TOKEN_KEYS), reads)) == { + "spend:end_user:eu": (None, True), + "spend:key:hashed": (1.5, True), + "spend:org:org": (None, True), + "spend:team:team": (2.5, True), + "spend:team_member:user:team": (None, True), + "spend:user:user": (None, True), + } + + +@pytest.mark.asyncio +async def test_unbound_counter_and_closed_batch_leave_the_read_to_the_caller(): + redis = CountingRedis({"spend:key:hashed": 1.0}) + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + + assert await batch.read("spend:tag:prod") is None + assert redis.commands == [] + assert await batch.read("spend:key:hashed") == (1.0, True) + + batch.close() + batch.bind(frozenset({"spend:org:org"})) + assert await batch.read("spend:key:hashed") is None + assert await batch.read("spend:org:org") is None + assert len(redis.commands) == 1 + + +@pytest.mark.asyncio +async def test_keys_bound_after_the_first_read_join_one_more_mget_for_only_the_new_keys(): + redis = CountingRedis({"spend:key:hashed": 1.0, "spend:org:org": 9.0}) + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + assert await batch.read("spend:key:hashed") == (1.0, True) + + batch.bind(frozenset({"spend:org:org", "spend:key:hashed"})) + assert await batch.read("spend:org:org") == (9.0, True) + assert await batch.read("spend:key:hashed") == (1.0, True) + + assert redis.commands == ["MGET spend:key:hashed", "MGET spend:org:org"] + + +@pytest.mark.asyncio +async def test_a_recorded_write_result_answers_later_reads_without_another_redis_read(): + redis = CountingRedis({"spend:key:hashed": 1.0}) + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + assert await batch.read("spend:key:hashed") == (1.0, True) + + batch.record("spend:key:hashed", 3.5) + batch.record("spend:org:org", 7.0) + + assert await batch.read("spend:key:hashed") == (3.5, True) + assert await batch.read("spend:org:org") == (7.0, True) + assert redis.commands == ["MGET spend:key:hashed"] + + +@pytest.mark.asyncio +async def test_a_forgotten_counter_is_read_fresh_from_redis_when_it_is_bound_again(): + redis = CountingRedis({"spend:key:hashed": 1.0}) + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + batch.record("spend:key:hashed", 3.5) + + batch.forget("spend:key:hashed") + assert await batch.read("spend:key:hashed") is None + + redis.store["spend:key:hashed"] = 9.0 + batch.bind(frozenset({"spend:key:hashed"})) + assert await batch.read("spend:key:hashed") == (9.0, True) + assert redis.commands == ["MGET spend:key:hashed"] + + +@pytest.mark.asyncio +async def test_failed_mget_hands_every_counter_back_to_the_caller(): + batch = SpendCounterBatch(CountingRedis(fail=True)) + batch.bind(TOKEN_KEYS) + + assert await batch.read("spend:key:hashed") is None + + +@pytest.mark.asyncio +async def test_non_numeric_counter_payload_hands_the_batch_back_to_the_caller(): + redis = CountingRedis() + redis.store["spend:key:hashed"] = "garbage" + batch = SpendCounterBatch(redis) + batch.bind(frozenset({"spend:key:hashed"})) + + assert await batch.read("spend:key:hashed") is None + + +def test_scope_installs_a_batch_only_when_redis_exists_and_release_closes_without_clearing(): + assert active_spend_counter_batch() is None + with spend_counter_batch_scope(None): + assert active_spend_counter_batch() is None + bind_admission_counter_keys(TOKEN, end_user_id=None) + + with spend_counter_batch_scope(CountingRedis()): + batch = active_spend_counter_batch() + assert batch is not None + bind_admission_counter_keys(TOKEN, end_user_id="eu") + assert batch.counter_keys == TOKEN_KEYS + release_spend_counter_batch() + assert active_spend_counter_batch() is batch + batch.bind(frozenset({"spend:tag:x"})) + assert batch.counter_keys == TOKEN_KEYS + assert active_spend_counter_batch() is None + + +@pytest.mark.asyncio +async def test_get_current_spend_inside_the_scope_costs_one_mget_for_all_admission_counters(monkeypatch): + redis = CountingRedis({"spend:key:hashed": 3.0, "spend:team:team": 4.0, "spend:org:org": 5.0}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id="eu") + key_spend = await ps.get_current_spend(counter_key="spend:key:hashed", fallback_spend=0.0) + team_spend = await ps.get_current_spend(counter_key="spend:team:team", fallback_spend=0.0) + org_spend = await ps.get_current_spend(counter_key="spend:org:org", fallback_spend=0.0) + user_spend = await ps.get_current_spend(counter_key="spend:user:user", fallback_spend=7.0) + + assert (key_spend, team_spend, org_spend, user_spend) == (3.0, 4.0, 5.0, 7.0) + assert [c for c in redis.commands if c.startswith("GET ")] == [], "the cold reseed reuses the MGET miss" + assert [c for c in redis.commands if c.startswith("MGET ")] == [f"MGET {' '.join(sorted(TOKEN_KEYS))}"] + + +@pytest.mark.asyncio +async def test_get_current_spend_outside_the_scope_still_reads_redis_per_counter(monkeypatch): + redis = CountingRedis({"spend:key:hashed": 3.0}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + + assert await ps.get_current_spend(counter_key="spend:key:hashed", fallback_spend=0.0) == 3.0 + assert redis.commands == ["GET spend:key:hashed"] + + +@pytest.mark.asyncio +async def test_after_release_a_read_goes_to_redis_directly_so_read_then_write_sees_fresh_values(monkeypatch): + redis = CountingRedis({"spend:key:hashed": 3.0}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id=None) + assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (3.0, True) + redis.store["spend:key:hashed"] = 8.0 + assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (3.0, True) + release_spend_counter_batch() + assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (8.0, True) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "GET"] + + +@pytest.mark.asyncio +async def test_batched_clean_miss_does_not_fall_back_to_the_per_pod_in_memory_copy(monkeypatch): + redis = CountingRedis() + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis, in_memory={"spend:key:hashed": 99.0})) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id=None) + assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (None, True) + + +@pytest.mark.asyncio +async def test_batched_redis_failure_falls_back_to_the_per_pod_in_memory_copy(monkeypatch): + redis = CountingRedis(fail=True) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis, in_memory={"spend:key:hashed": 99.0})) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id=None) + assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (99.0, False) + + +@pytest.mark.asyncio +async def test_scope_is_per_task_so_concurrent_requests_do_not_share_a_batch(): + redis = CountingRedis({"spend:key:a": 1.0, "spend:key:b": 2.0}) + + async def request(token: str) -> tuple[float | None, bool] | None: + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(UserAPIKeyAuth(token=token), end_user_id=None) + batch = active_spend_counter_batch() + assert batch is not None + await asyncio.sleep(0) + return await batch.read(f"spend:key:{token}") + + assert await asyncio.gather(request("a"), request("b")) == [(1.0, True), (2.0, True)] + assert sorted(redis.commands) == ["MGET spend:key:a", "MGET spend:key:b"] + + +@pytest.mark.asyncio +async def test_batch_reads_never_touch_a_prisma_client_when_redis_answers(monkeypatch): + redis = CountingRedis({"spend:key:hashed": 3.0}) + prisma = MagicMock() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock() + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", prisma) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id=None) + assert await ps.get_current_spend(counter_key="spend:key:hashed", fallback_spend=0.0) == 3.0 + + assert prisma.db.mock_calls == [] + + +def _reseed_prisma(spend: float) -> MagicMock: + prisma = MagicMock() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=MagicMock(spend=spend)) + return prisma + + +@pytest.mark.asyncio +async def test_reseed_reuses_the_admission_mget_instead_of_its_own_get(): + redis = CountingRedis({"spend:key:hashed": 7.5}) + prisma = _reseed_prisma(spend=1.0) + cache = _spend_counter_cache(redis) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id=None) + value = await SpendCounterReseed.coalesced(prisma, cache, counter_key="spend:key:hashed") + + assert value == 7.5 + assert redis.commands == [ + "MGET spend:key:hashed spend:org:org spend:team:team spend:team_member:user:team spend:user:user" + ] + assert prisma.db.mock_calls == [] + + +@pytest.mark.asyncio +async def test_reseed_treats_a_batched_clean_miss_as_authoritative_and_seeds_from_the_db(): + redis = CountingRedis() + redis.async_set_cache = AsyncMock(return_value=True) + prisma = _reseed_prisma(spend=2.25) + cache = _spend_counter_cache(redis, in_memory={"spend:key:hashed": 99.0}) + + with spend_counter_batch_scope(redis): + bind_admission_counter_keys(TOKEN, end_user_id=None) + value = await SpendCounterReseed.coalesced(prisma, cache, counter_key="spend:key:hashed") + + assert value == 2.25 + assert [c for c in redis.commands if c.startswith("GET")] == [] + redis.async_set_cache.assert_awaited_once_with(key="spend:key:hashed", value=2.25, nx=True) + + +@pytest.mark.asyncio +async def test_reseed_outside_the_scope_still_re_checks_redis_itself(): + redis = CountingRedis({"spend:key:hashed": 4.0}) + + value = await SpendCounterReseed.coalesced( + _reseed_prisma(spend=1.0), _spend_counter_cache(redis), "spend:key:hashed" + ) + + assert value == 4.0 + assert redis.commands == ["GET spend:key:hashed"] + + +POST_CALL_KEYS = TOKEN_KEYS | {"spend:tag:prod", "spend:model_access_group:premium"} + + +@pytest.mark.asyncio +async def test_post_call_increment_for_every_entity_costs_one_mget_and_one_pipeline(monkeypatch): + redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + + await ps.increment_spend_counters( + token="hashed", + team_id="team", + user_id="user", + org_id="org", + end_user_id="eu", + tags=["prod"], + model_access_groups=["premium"], + response_cost=0.5, + ) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands + assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS + assert set(redis.commands[1].split()[1:]) == POST_CALL_KEYS + assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.5 for key in POST_CALL_KEYS} + + +@pytest.mark.asyncio +async def test_post_call_cold_counters_seed_from_the_mget_miss_without_a_second_read(monkeypatch): + redis = CountingRedis({"spend:key:hashed": 1.0}) + redis.async_set_cache = AsyncMock(return_value=True) + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0)) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", prisma) + + await ps.increment_spend_counters(token="hashed", team_id="team", user_id=None, response_cost=0.5) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands + redis.async_set_cache.assert_awaited_once_with(key="spend:team:team", value=4.0, nx=True) + assert redis.store["spend:key:hashed"] == 1.5 + + +RESERVED_KEYS = frozenset( + {"spend:key:hashed", "spend:team:team", "spend:team_member:user:team", "spend:end_user:eu", "spend:org:org"} +) + + +def _reservation(reserved_cost: float, counter_keys: frozenset[str] = RESERVED_KEYS) -> dict[str, object]: + return { + "reserved_cost": reserved_cost, + "entries": [ + {"counter_key": key, "entity_type": "Key", "entity_id": key, "reserved_cost": reserved_cost} + for key in sorted(counter_keys) + ], + } + + +@pytest.mark.asyncio +async def test_post_call_with_a_reservation_costs_one_mget_one_reconcile_pipeline_one_increment_pipeline(monkeypatch): + redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + reservation = _reservation(reserved_cost=0.4) + + await ps.increment_spend_counters( + token="hashed", + team_id="team", + user_id="user", + org_id="org", + end_user_id="eu", + tags=["prod"], + model_access_groups=["premium"], + response_cost=0.5, + budget_reservation=reservation, + ) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "PIPELINE"], redis.commands + assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS, "reconcile and warm checks share the MGET" + assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS + assert set(redis.commands[2].split()[1:]) == POST_CALL_KEYS - RESERVED_KEYS + assert {key: round(redis.store[key], 6) for key in POST_CALL_KEYS} == { + key: (1.1 if key in RESERVED_KEYS else 1.5) for key in POST_CALL_KEYS + } + assert [round(entry["applied_adjustment"], 6) for entry in reservation["entries"]] == [0.1] * len(RESERVED_KEYS) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reconcile_settles_a_flushed_counter_on_its_own_after_the_shared_pipeline(monkeypatch): + from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation + + redis = CountingRedis({key: 1.0 for key in RESERVED_KEYS - {"spend:team:team"}}) + redis.async_set_max = AsyncMock(return_value=4.0) + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0)) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", prisma) + reservation = _reservation(reserved_cost=0.4) + + await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.5) + + assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "INCRBYFLOAT"], redis.commands + assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS - {"spend:team:team"} + assert redis.commands[2] == "INCRBYFLOAT spend:team:team 0.5" + redis.async_set_max.assert_awaited_once() + assert redis.async_set_max.await_args.kwargs["key"] == "spend:team:team" + assert all(round(entry["applied_adjustment"], 6) == 0.1 for entry in reservation["entries"]) + + +@pytest.mark.asyncio +async def test_pre_call_resize_against_an_inconsistent_counter_writes_nothing_and_denies(monkeypatch): + from litellm.proxy.spend_tracking.budget_reservation import _resize_applied_reservation + + redis = CountingRedis({"spend:key:hashed": 1.0}) + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + entries = _reservation(reserved_cost=0.4, counter_keys=frozenset({"spend:key:hashed", "spend:team:team"}))[ + "entries" + ] + + with pytest.raises(RuntimeError, match="spend:team:team"): + await _resize_applied_reservation(entries=entries, current_reserved_cost=0.4, new_reserved_cost=0.9) + + assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands + assert redis.store["spend:key:hashed"] == 1.0 + assert all("applied_adjustment" not in entry for entry in entries) + + +@pytest.mark.asyncio +async def test_a_failed_reconcile_pipeline_invalidates_every_reserved_counter_and_falls_back(monkeypatch): + redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS}) + redis.async_delete_cache = AsyncMock() + reconcile_pipeline_failed = False + + async def _pipeline(increment_list: Sequence[Mapping[str, object]], **kwargs: object) -> list[float]: + nonlocal reconcile_pipeline_failed + if not reconcile_pipeline_failed: + reconcile_pipeline_failed = True + raise ConnectionError("redis down") + return await CountingRedis.async_increment_pipeline(redis, increment_list, **kwargs) + + redis.async_increment_pipeline = _pipeline # pyright: ignore[reportAttributeAccessIssue] # instance override + monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis)) + monkeypatch.setattr(ps, "prisma_client", None) + reservation = _reservation(reserved_cost=0.4) + + await ps.increment_spend_counters( + token="hashed", + team_id="team", + user_id="user", + org_id="org", + end_user_id="eu", + response_cost=0.5, + budget_reservation=reservation, + ) + + assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == RESERVED_KEYS + assert all("applied_adjustment" not in entry for entry in reservation["entries"]) + assert redis.commands[-1].split()[0] == "PIPELINE" + assert set(redis.commands[-1].split()[1:]) == RESERVED_KEYS | {"spend:user:user"} + + +def test_a_scope_opened_inside_an_open_scope_joins_its_batch_and_a_closed_one_gets_its_own(): + redis = CountingRedis() + with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:a"})): + outer = active_spend_counter_batch() + assert outer is not None + with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:b"})): + assert active_spend_counter_batch() is outer + assert outer.counter_keys == {"spend:key:a", "spend:key:b"} + release_spend_counter_batch() + with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:c"})): + inner = active_spend_counter_batch() + assert inner is not outer + assert inner is not None and inner.counter_keys == {"spend:key:c"} + assert active_spend_counter_batch() is outer 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 dc8a87a97ad..e9bbaf0c96e 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 @@ -20,6 +20,7 @@ from litellm.constants import ( SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL, ) +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup @@ -2755,6 +2756,31 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( # ── _redact_logged_api_key unit tests ────────────────────────────────────── +@pytest.mark.parametrize( + ("original_exception", "expected_error_message"), + [ + ( + ProxyModelNotFoundError(route="/chat/completions", model_name=_RAW_MODEL_WITH_PROMPT), + "/chat/completions: Invalid model name passed in. Call `/v1/models` to view available models for your key.", + ), + (ValueError("provider timed out"), "provider timed out"), + ], +) +def test_sanitize_error_information_persists_no_raw_model_for_an_unknown_model_rejection( + original_exception: Exception, expected_error_message: str +): + error_information: Final = StandardLoggingPayloadSetup.get_error_information(original_exception=original_exception) + + sanitized: Final = _sanitize_error_information_for_spend_logs( + error_information, original_exception=original_exception + ) + + assert sanitized is not None + assert sanitized["error_message"] == expected_error_message + assert "medical records" not in json.dumps(sanitized) + assert sanitized["error_class"] == type(original_exception).__name__ + + def test_redact_logged_api_key_none_returns_none(): assert _redact_logged_api_key(None) is None diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index 0fdb43d60da..3073908fa54 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -196,6 +196,20 @@ def test_gateway_drops_ui_and_swagger_mounts(): f"Mount {path} must not be served by the gateway" +def test_gateway_keeps_memory_summary_and_trims_the_other_debug_routes(): + """The gateway serves /debug/memory/summary, since the RSS that matters is the + serving worker's and the memory regression e2e test reads it on every gateway + replica; the heavier and mutating /debug/memory routes stay on the backend.""" + debug_memory_routes = { + getattr(r, "path"): r for r in app.router.routes if str(getattr(r, "path", "")).startswith("/debug/memory/") + } + assert {"/debug/memory/summary", "/debug/memory/details", "/debug/memory/gc/configure"} <= set(debug_memory_routes) + assert _is_gateway_route(debug_memory_routes["/debug/memory/summary"]), \ + "/debug/memory/summary must survive the gateway route trim" + for path in ("/debug/memory/details", "/debug/memory/gc/configure"): + assert not _is_gateway_route(debug_memory_routes[path]), f"{path} must not be served by the gateway" + + def test_every_app_mount_is_assigned_to_a_component(): """Every Mount on the proxy app must be consciously assigned to a component. diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 5aa80213134..c0c853ae2c5 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -735,6 +735,121 @@ async def test_perform_health_check_and_save_forwards_skip_disabled_background_f assert call_kwargs["health_check_skip_disabled_background_models"] is True +@pytest.mark.asyncio +async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_name(): + """``/health?model=`` must probe the team deployment, not an empty list.""" + from litellm.proxy.health_check import perform_health_check + + team_deployment = { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-team-b", "team_id": "team-b", "team_public_model_name": "bedrock-nova"}, + } + other_deployment = { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-openai"}, + } + probe = AsyncMock(return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}], [], {})) + + with patch( # test-quality-ok: the deployments handed to the probe are the assertion; no injection seam + "litellm.proxy.health_check._perform_health_check", probe + ): + healthy, unhealthy, _ = await perform_health_check( + model_list=[team_deployment, other_deployment], model="bedrock-nova", team_id="team-b" + ) + + assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == ["id-team-b"] + assert [ep["model_id"] for ep in healthy] == ["id-team-b"] + assert unhealthy == [] + + +@pytest.mark.asyncio +async def test_perform_health_check_keeps_a_public_name_off_another_team(): + """A team's public model name is not a global alias: a caller from another team must not probe its deployment.""" + from litellm.proxy.health_check import perform_health_check + + team_deployment = { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-team-b", "team_id": "team-b", "team_public_model_name": "bedrock-nova"}, + } + probe = AsyncMock(return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}], [], {})) + + with patch( # test-quality-ok: the deployments handed to the probe are the assertion; no injection seam + "litellm.proxy.health_check._perform_health_check", probe + ): + healthy, unhealthy, _ = await perform_health_check( + model_list=[team_deployment], model="bedrock-nova", team_id="team-a" + ) + + probe.assert_not_awaited() + assert healthy == [] + assert unhealthy == [] + + +_GLOBAL_DEPLOYMENT = { + "model_name": "bedrock-nova", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-bedrock"}, +} +_TEAM_B_COPY = { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-team-b", "team_id": "team-b", "team_public_model_name": "bedrock-nova"}, +} +_GLOBAL_BARE_NAME = { + "model_name": "gpt-5.4-nano", + "litellm_params": {"model": "gpt-5.4-nano"}, + "model_info": {"id": "id-nano"}, +} +_TEAM_B_BARE_COPY = { + "model_name": "gpt-5.4-nano_team-b_7c3d", + "litellm_params": {"model": "gpt-5.4-nano"}, + "model_info": {"id": "id-nano-team-b", "team_id": "team-b", "team_public_model_name": "gpt-5.4-nano"}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "model", "model_list", "expected_ids"), + [ + (None, "bedrock-nova", [_TEAM_B_COPY], ["id-team-b"]), + (None, "bedrock-nova", [_GLOBAL_DEPLOYMENT, _TEAM_B_COPY], ["id-bedrock"]), + ("team-b", "bedrock-nova", [_GLOBAL_DEPLOYMENT, _TEAM_B_COPY], ["id-team-b"]), + ("team-b", "gpt-5.4-nano", [_GLOBAL_BARE_NAME, _TEAM_B_BARE_COPY], ["id-nano-team-b"]), + (None, "gpt-5.4-nano", [_GLOBAL_BARE_NAME, _TEAM_B_BARE_COPY], ["id-nano"]), + (None, "bedrock/us.amazon.nova-2-lite-v1:0", [_GLOBAL_DEPLOYMENT, _TEAM_B_COPY], ["id-bedrock", "id-team-b"]), + ], + ids=[ + "a team-less caller reaches a public name nothing else carries", + "model_name wins over a public name for a team-less caller", + "a team's own copy wins over the global model_name", + "a team's own copy wins over a litellm_params.model equal to the public name", + "model_name wins over a litellm_params.model equal to it for a team-less caller", + "a provider model string no name carries still matches litellm_params.model", + ], +) +async def test_perform_health_check_targets_a_name_the_way_a_request_for_it_routes( + team_id, model, model_list, expected_ids +): + """``/health?model=`` probes the deployments a request for that name from the same caller would route to.""" + from litellm.proxy.health_check import perform_health_check + + probe = AsyncMock( + return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": i} for i in expected_ids], [], {}) + ) + + with patch( # test-quality-ok: the deployments handed to the probe are the assertion; no injection seam + "litellm.proxy.health_check._perform_health_check", probe + ): + healthy, unhealthy, _ = await perform_health_check(model_list=model_list, model=model, team_id=team_id) + + assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == expected_ids + assert [ep["model_id"] for ep in healthy] == expected_ids + assert unhealthy == [] + + def test_parse_background_health_check_model_groups_unset_returns_none(): from litellm.proxy.health_check import parse_background_health_check_model_groups diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 97e308d7c3c..dd3669644af 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -872,9 +872,9 @@ def test_narrowing_by_an_id_that_matches_nothing_keeps_the_whole_list(): """Pinned because the disabled-dependency fix moved this filter into its own helper.""" deployments = [{"model_name": "a", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a-1"}}] - assert hc_module._narrow_to_target(deployments, None, "no-such-id") == tuple(deployments) - assert hc_module._narrow_to_target(deployments, None, "a-1") == tuple(deployments) - assert hc_module._narrow_to_target(deployments, "a", None) == tuple(deployments) + assert hc_module._narrow_to_target(deployments, None, "no-such-id", None) == tuple(deployments) + assert hc_module._narrow_to_target(deployments, None, "a-1", None) == tuple(deployments) + assert hc_module._narrow_to_target(deployments, "a", None, None) == tuple(deployments) def _nested_router_fixture(parent_tier: str): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 3b3031647dc..03a24ec5e98 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8419,7 +8419,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) - return ps._PendingSpendIncrement( + return ps.PendingSpendIncrement( counter_key=kwargs["counter_key"], increment=kwargs["increment"] ) diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 41ba57c4615..f7021763a4d 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -3,6 +3,7 @@ import pytest +from typing import Final from unittest.mock import MagicMock from fastapi import HTTPException @@ -1297,3 +1298,13 @@ async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through( assert agents_find_unique.await_count == 2 assert model_table.find_many_wheres == [] + + +def test_proxy_model_not_found_error_keeps_the_raw_model_only_in_the_client_response(): + raw_model: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes" + + error: Final = ProxyModelNotFoundError(route="/chat/completions", model_name=raw_model) + + assert raw_model in error.detail["error"] + assert raw_model not in error.spend_log_error_message + assert error.spend_log_error_message.startswith("/chat/completions: Invalid model name passed in") diff --git a/tests/test_litellm/proxy/utils/helpers/test_misc_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_misc_helpers.py index 7968fa40655..6267428e02e 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_misc_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_misc_helpers.py @@ -187,6 +187,22 @@ def test_construct_database_url_from_env_vars_with_schema(monkeypatch): } +def test_construct_database_url_from_env_vars_carries_tls_env(monkeypatch: pytest.MonkeyPatch): + """The CLI password path builds its URL here, so DATABASE_SSLMODE and + DATABASE_SSLROOTCERT must reach PgBouncer through it too.""" + monkeypatch.setenv("DATABASE_HOST", "db.example.com") + monkeypatch.setenv("DATABASE_USERNAME", "user") + monkeypatch.setenv("DATABASE_PASSWORD", "pass") + monkeypatch.setenv("DATABASE_NAME", "litellm") + monkeypatch.setenv("DATABASE_SCHEMA", "public") + monkeypatch.setenv("DATABASE_SSLMODE", "verify-full") + monkeypatch.setenv("DATABASE_SSLROOTCERT", "/etc/ssl/certs/ca-certificates.crt") + assert construct_database_url_from_env_vars() == ( + "postgresql://user:pass@db.example.com/litellm" + "?schema=public&sslmode=verify-full&sslrootcert=%2Fetc%2Fssl%2Fcerts%2Fca-certificates.crt" + ) + + def test_construct_database_url_from_env_vars_error_path_missing_host(monkeypatch): monkeypatch.delenv("DATABASE_HOST", raising=False) monkeypatch.setenv("DATABASE_USERNAME", "user") diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 761e87ac764..0827bbcdc38 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -266,7 +266,8 @@ def test_transcription_only_detection_rejects_speech_model(local_model_cost_map) @pytest.mark.asyncio -async def test_azure_health_check_keeps_beta_path_for_speech_model(): +async def test_azure_health_check_probes_the_ga_upstream_for_an_unconfigured_speech_model(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) connect = _CapturingConnect() with patch("websockets.connect", connect): assert await realtime_main._realtime_health_check( @@ -276,14 +277,18 @@ async def test_azure_health_check_keeps_beta_path_for_speech_model(): api_base="https://my-endpoint.openai.azure.com", api_version="2024-10-01-preview", ) - assert connect.url == ( - "wss://my-endpoint.openai.azure.com/openai/realtime" - "?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview" - ) + assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + + +_AZURE_BETA_HEALTH_URL: Final = ( + "wss://my-endpoint.openai.azure.com/openai/realtime" + "?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview" +) @pytest.mark.asyncio -async def test_azure_health_check_honors_deployment_realtime_protocol(): +async def test_azure_health_check_honors_deployment_realtime_protocol(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) connect = _CapturingConnect() with patch("websockets.connect", connect): assert await realtime_main._realtime_health_check( @@ -292,9 +297,24 @@ async def test_azure_health_check_honors_deployment_realtime_protocol(): api_key="fake-key", api_base="https://my-endpoint.openai.azure.com", api_version="2024-10-01-preview", - model_params={"realtime_protocol": "GA"}, + model_params={"realtime_protocol": "beta"}, ) - assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + assert connect.url == _AZURE_BETA_HEALTH_URL + + +@pytest.mark.asyncio +async def test_azure_health_check_honors_env_realtime_protocol(monkeypatch): + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") + connect = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main._realtime_health_check( + model="gpt-4o-realtime-preview", + custom_llm_provider="azure", + api_key="fake-key", + api_base="https://my-endpoint.openai.azure.com", + api_version="2024-10-01-preview", + ) + assert connect.url == _AZURE_BETA_HEALTH_URL class _ConnectThatStopsAfterCapturingTheUrl: @@ -327,7 +347,60 @@ async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai 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" + assert connect.url == "wss://my-project.services.ai.azure.com/openai/v1/realtime?model=gpt-realtime-mini" + + +class _ClientWebSocketWithHeaders: + def __init__(self, headers: tuple[tuple[bytes, bytes], ...]) -> None: + self.scope: Final = {"headers": headers} + + +_GA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=()) +_BETA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=((b"openai-beta", b"realtime=v1"),)) + + +async def _azure_backend_url_dialed_for(websocket: _ClientWebSocketWithHeaders, **kwargs: object) -> str | None: + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure/gpt-realtime", + websocket=websocket, + api_base="https://my-endpoint.openai.azure.com", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + **kwargs, + ) + return connect.url + + +@pytest.mark.asyncio +async def test_arealtime_azure_ga_client_without_beta_header_dials_the_ga_upstream(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert ( + await _azure_backend_url_dialed_for(_GA_CLIENT) + == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_beta_header_client_keeps_the_beta_upstream(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert await _azure_backend_url_dialed_for(_BETA_CLIENT) == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_explicit_beta_protocol_wins_over_a_ga_client(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert await _azure_backend_url_dialed_for(_GA_CLIENT, realtime_protocol="beta") == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypatch): + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") + assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 46249e50572..342ec4435a7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1,13 +1,23 @@ import json -from typing import Final +from copy import deepcopy +from typing import Final, Literal import pytest +from openai.types.responses.response_function_web_search import ( + ActionFind, + ActionOpenPage, + ActionSearch, + ActionSearchSource, + ResponseFunctionWebSearch, +) - +import litellm +from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, LiteLLMCompletionResponsesConfig, ) +from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( ChatCompletionMessageToolCall, Choices, @@ -3513,6 +3523,8 @@ class TestEnsureOutputItemContentPartAdded: iterator._custom_tool_names = set() iterator.responses_api_request = {} iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map(None) + iterator._web_search_calls = {} + iterator._queued_web_search_call_ids = set() return iterator def _make_text_chunk(self): @@ -4017,6 +4029,211 @@ def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): assert convert(openai)["id"] == "call_tokyo" +class TestHostedWebSearchReplay: + def test_emitted_hosted_search_output_round_trips_with_client_tool_result(self) -> None: + search_result: Final = { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_round_trip_search", + "content": [{"type": "web_search_result", "url": "https://example.com/forecast"}], + } + search: Final = build_web_search_call( + tool_id="srvtoolu_round_trip_search", tool_input={"query": "Paris forecast"}, result=search_result + ) + message: Final = Message( + role="assistant", + content="I found a forecast source.", + tool_calls=[ + ChatCompletionMessageToolCall( + id="srvtoolu_round_trip_search", + type="function", + function=Function(name="web_search", arguments='{"query":"Paris forecast"}'), + ), + ChatCompletionMessageToolCall( + id="call_round_trip_weather", + type="function", + function=Function(name="get_weather", arguments='{"city":"Paris"}'), + ), + ], + provider_specific_fields={"web_search_calls": [search], "web_search_results": [search_result]}, + ) + response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Find a forecast source and check the weather in Paris.", + responses_api_request={ + "tools": [ + {"type": "web_search"}, + {"type": "function", "name": "get_weather", "parameters": {"type": "object"}}, + ] + }, + chat_completion_response=_bridged_chat_completion_response( + choices=[Choices(index=0, finish_reason="tool_calls", message=message)] + ), + ) + assert [item for item in response.output if item.type == "web_search_call"] == [search] + assert [item.call_id for item in response.output if item.type == "function_call"] == ["call_round_trip_weather"] + history: Final = [ + {"role": "user", "content": "Find a forecast source and check the weather in Paris."}, + *(item.model_dump(exclude_none=True) for item in response.output), + {"type": "function_call_output", "call_id": "call_round_trip_weather", "output": "Paris is sunny."}, + ] + original: Final = deepcopy(history) + + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=history, responses_api_request={} + ) + + assert [item.get("role") for item in messages] == ["user", "assistant", "tool"] + assistant: Final = messages[1] + assert [call["id"] for call in assistant["tool_calls"]] == ["call_round_trip_weather"] + assert [call["function"]["name"] for call in assistant["tool_calls"]] == ["get_weather"] + content: Final = assistant["content"] + assert isinstance(content, list) + text_parts: Final = tuple(block["text"] for block in content if block.get("type") == "text") + assert text_parts[0] == "I found a forecast source." + replayed_searches: Final = tuple( + ResponseFunctionWebSearch.model_validate_json(text[text.index("{"):]) + for text in text_parts + if "web_search_call" in text + ) + assert replayed_searches == (search,) + assert messages[2]["tool_call_id"] == "call_round_trip_weather" + assert messages[2]["content"] == "Paris is sunny." + assert history == original + + @pytest.mark.parametrize( + "action", + ( + ActionSearch( + type="search", + query="hosted search history", + queries=["hosted search history", "search replay"], + sources=[ActionSearchSource(type="url", url="https://example.com/search-result")], + ), + ActionOpenPage(type="open_page", url="https://example.com/opened-page"), + ActionFind(type="find_in_page", url="https://example.com/find-page", pattern="search history"), + ), + ids=("search", "open_page", "find"), + ) + @pytest.mark.parametrize("status", ("completed", "failed")) + def test_replays_typed_search_action_without_client_tool_call( + self, + action: ActionSearch | ActionOpenPage | ActionFind, + status: Literal["completed", "failed"], + ) -> None: + search: Final = ResponseFunctionWebSearch( + id="ws_replayed_search", type="web_search_call", status=status, action=action + ) + input_item: Final = search.model_dump(exclude_none=True) + original: Final = deepcopy(input_item) + + messages: Final = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert not messages[0].get("tool_calls") + content: Final = messages[0].get("content") + assert isinstance(content, str) + replayed: Final = ResponseFunctionWebSearch.model_validate_json(content[content.index("{"):]) + assert replayed == search + assert input_item == original + + @pytest.mark.parametrize("order", ((0, 1, 2, 3), (1, 0, 3, 2), (1, 3, 0, 2))) + @pytest.mark.parametrize("modify_params", (False, True)) + @pytest.mark.parametrize("structured_content", (False, True)) + def test_search_replay_preserves_client_tool_result_adjacency( + self, + monkeypatch: pytest.MonkeyPatch, + order: tuple[int, int, int, int], + modify_params: bool, + structured_content: bool, + ) -> None: + monkeypatch.setattr(litellm, "modify_params", modify_params) + searches: Final = tuple( + ResponseFunctionWebSearch( + id=f"ws_search_{index}", + type="web_search_call", + status="completed", + action=ActionSearch( + type="search", + query=f"search query {index}", + queries=[f"search query {index}"], + sources=[ActionSearchSource(type="url", url=f"https://example.com/result-{index}")], + ), + ) + for index in (1, 2) + ) + replay_items: Final = ( + { + "type": "function_call", + "name": "get_weather", + "call_id": "call_weather", + "arguments": '{"city":"Paris"}', + }, + searches[0].model_dump(exclude_none=True), + {"type": "function_call", "name": "get_time", "call_id": "call_time", "arguments": "{}"}, + searches[1].model_dump(exclude_none=True), + ) + history: Final = [ + {"role": "user", "content": "Research the forecast and call get_weather."}, + { + "role": "assistant", + "content": [{"type": "output_text", "text": "I will check the forecast."}] + if structured_content + else "I will check the forecast.", + }, + *(replay_items[index] for index in order), + {"role": "assistant", "content": [{"type": "output_text", "text": "I found two sources."}]}, + {"type": "function_call_output", "call_id": "call_weather", "output": "Paris is sunny."}, + {"type": "function_call_output", "call_id": "call_time", "output": "12:00"}, + ] + original: Final = deepcopy(history) + + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=history, responses_api_request={} + ) + + assert [message.get("role") for message in messages] == ["user", "assistant", "tool", "tool"] + assistant: Final = messages[1] + assert [call["id"] for call in assistant["tool_calls"]] == ["call_weather", "call_time"] + assert [call["function"]["name"] for call in assistant["tool_calls"]] == ["get_weather", "get_time"] + assert messages[2]["tool_call_id"] == "call_weather" + assert messages[2]["content"] == "Paris is sunny." + assert messages[3]["tool_call_id"] == "call_time" + assert messages[3]["content"] == "12:00" + content: Final = assistant["content"] + assert isinstance(content, list) + text_parts: Final = tuple(block["text"] for block in content if block.get("type") == "text") + assert text_parts[0] == "I will check the forecast." + assert text_parts[-1] == "I found two sources." + replayed_searches: Final = tuple( + ResponseFunctionWebSearch.model_validate_json(text[text.index("{"):]) + for text in text_parts + if "web_search_call" in text + ) + assert replayed_searches == searches + assert history == original + + provider_messages: Final = anthropic_messages_pt( + messages=messages, model="claude-fable-5-1", llm_provider="anthropic" + ) + + assert [message["role"] for message in provider_messages] == ["user", "assistant", "user"] + assistant_blocks: Final = provider_messages[1]["content"] + result_blocks: Final = provider_messages[2]["content"] + assert [block["id"] for block in assistant_blocks if block.get("type") == "tool_use"] == [ + "call_weather", "call_time" + ] + assert [block["tool_use_id"] for block in result_blocks if block.get("type") == "tool_result"] == [ + "call_weather", "call_time" + ] + assert [block["content"] for block in result_blocks if block.get("type") == "tool_result"] == [ + "Paris is sunny.", "12:00" + ] + assert [block["text"] for block in assistant_blocks if block.get("type") == "text"] == list(text_parts) + assert history == original + + BRIDGED_CHAT_COMPLETION_ID = "chatcmpl-dfa2da3a-1586-4ff7-b64e-f59c692a5d11" @@ -4058,6 +4275,105 @@ class TestBridgedOutputItemIdPrefixes: chat_completion_response=chat_completion_response, ) + @pytest.mark.parametrize( + "tool_type,result_kind,expected_sources", + [ + ( + "web_search", + "valid", + {"srvtoolu_01Search": ["https://example.com/one"], "srvtoolu_02Search": ["https://example.com/two"]}, + ), + ( + "web_search_preview", + "valid", + {"srvtoolu_01Search": ["https://example.com/one"], "srvtoolu_02Search": ["https://example.com/two"]}, + ), + ("function", "valid", {}), + ("web_search", "unpaired", {"srvtoolu_01Search": ["https://example.com/one"]}), + ("web_search", "web_fetch", {"srvtoolu_02Search": ["https://example.com/two"]}), + ("web_search", "error", {"srvtoolu_01Search": [], "srvtoolu_02Search": ["https://example.com/two"]}), + ], + ) + def test_anthropic_web_search_output_mapping(self, tool_type, result_kind, expected_sources): + call_ids: Final = ("srvtoolu_01Search", "srvtoolu_02Search") + valid_results: Final = ( + { + "type": "web_search_tool_result", + "tool_use_id": call_ids[0], + "content": [{"type": "web_search_result", "url": "https://example.com/one"}], + }, + { + "type": "web_search_tool_result", + "tool_use_id": call_ids[1], + "content": [{"type": "web_search_result", "url": "https://example.com/two"}], + }, + ) + first_result: Final = ( + {**valid_results[0], "type": "web_fetch_tool_result"} + if result_kind == "web_fetch" + else {**valid_results[0], "content": {"type": "web_search_tool_result_error", "error_code": "unavailable"}} + if result_kind == "error" + else valid_results[0] + ) + results: Final = (first_result,) if result_kind == "unpaired" else (first_result, valid_results[1]) + message: Final = Message( + role="assistant", + content="answer", + tool_calls=[ + ChatCompletionMessageToolCall( + id=call_id, + type="function", + function=Function(name="web_search", arguments=json.dumps({"query": query})), + ) + for call_id, query in zip(call_ids, ("one", "two"), strict=True) + ] + + [ + ChatCompletionMessageToolCall( + id="toolu_regular", + type="function", + function=Function(name="get_weather", arguments='{"city":"Paris"}'), + ) + ], + provider_specific_fields={ + "web_search_results": results, + "web_search_calls": [ + build_web_search_call( + tool_id=result["tool_use_id"], + tool_input={"query": "one" if result["tool_use_id"].endswith("01Search") else "two"}, + result=result, + ) + for result in results + if tool_type != "function" and result["type"] == "web_search_tool_result" + ], + }, + ) + request_tools: Final = ( + [{"type": "function", "name": "web_search", "parameters": {"type": "object"}}] + if tool_type == "function" + else [{"type": tool_type}] + ) + response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="search", + responses_api_request={"tools": request_tools}, + chat_completion_response=_bridged_chat_completion_response( + choices=[Choices(index=0, finish_reason="stop", message=message)] + ), + ) + search_items: Final = { + item.id.removeprefix("ws_"): item for item in response.output if item.type == "web_search_call" + } + function_ids: Final = {item.call_id for item in response.output if item.type == "function_call"} + + assert set(search_items) == set(expected_sources) + assert function_ids == set(call_ids).difference(expected_sources) | {"toolu_regular"} + assert [item.content[0].text for item in response.output if item.type == "message"] == ["answer"] + for call_id, item in search_items.items(): + assert item.status == ("failed" if result_kind == "error" and call_id.endswith("01Search") else "completed") + assert item.action.type == "search" + assert item.action.query == ("one" if call_id.endswith("01Search") else "two") + assert item.action.queries == [item.action.query] + assert [source.url for source in item.action.sources] == expected_sources[call_id] + def test_message_item_id_uses_msg_prefix(self): response = self._transform(_bridged_chat_completion_response()) 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 850ee7ba623..5d97b0531d6 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 @@ -21,6 +21,7 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( Delta, ModelResponse, @@ -140,6 +141,214 @@ def test_tool_call_delta_is_emitted_as_responses_events(): assert len(evt2.delta) <= 10 # Chunks are max 10 characters +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "tool_type,result_kind,expected_sources", + [ + ( + "web_search", + "valid", + {"srvtoolu_01Search": ["https://example.com/one"], "srvtoolu_02Search": ["https://example.com/two"]}, + ), + ( + "web_search_preview", + "valid", + {"srvtoolu_01Search": ["https://example.com/one"], "srvtoolu_02Search": ["https://example.com/two"]}, + ), + ("function", "valid", {}), + ("web_search", "unpaired", {"srvtoolu_01Search": ["https://example.com/one"]}), + ("web_search", "web_fetch", {"srvtoolu_02Search": ["https://example.com/two"]}), + ("web_search", "error", {"srvtoolu_01Search": [], "srvtoolu_02Search": ["https://example.com/two"]}), + ], +) +async def test_web_search_stream_preserves_hosted_and_client_calls(sync_mode, tool_type, result_kind, expected_sources): + call_ids: Final = ("srvtoolu_01Search", "srvtoolu_02Search") + valid_results: Final = ( + { + "type": "web_search_tool_result", + "tool_use_id": call_ids[0], + "content": [{"type": "web_search_result", "url": "https://example.com/one"}], + }, + { + "type": "web_search_tool_result", + "tool_use_id": call_ids[1], + "content": [{"type": "web_search_result", "url": "https://example.com/two"}], + }, + ) + first_result: Final = ( + {**valid_results[0], "type": "web_fetch_tool_result"} + if result_kind == "web_fetch" + else {**valid_results[0], "content": {"type": "web_search_tool_result_error", "error_code": "unavailable"}} + if result_kind == "error" + else valid_results[0] + ) + results: Final = [first_result] if result_kind == "unpaired" else [first_result, valid_results[1]] + deltas: Final = ( + Delta( + role="assistant", + content=None, + tool_calls=[ + {"index": 0, "id": call_ids[0], "type": "function", "function": {"name": "web_search", "arguments": ""}} + ], + provider_specific_fields={ + "web_search_calls": [ + build_web_search_call( + call_ids[0], + {}, + {"content": []}, + status="in_progress", + ) + ] + if tool_type != "function" and result_kind != "web_fetch" + else [], + }, + ), + Delta( + content=None, + tool_calls=[ + { + "index": 1, + "id": "toolu_regular", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ], + ), + Delta(content=None, tool_calls=[{"index": 0, "function": {"arguments": '{"query":'}}]), + Delta(content=None, tool_calls=[{"index": 0, "function": {"arguments": '"one"}'}}]), + Delta( + content=None, + provider_specific_fields={ + "web_search_results": [first_result], + "web_search_calls": [ + build_web_search_call(call_ids[0], {"query": "one"}, first_result) + ] + if tool_type != "function" and first_result["type"] == "web_search_tool_result" + else [], + }, + ), + Delta( + content=None, + provider_specific_fields={ + "web_search_results": results, + "web_search_calls": [ + build_web_search_call( + result["tool_use_id"], + {"query": "one" if result["tool_use_id"].endswith("01Search") else "two"}, + result, + ) + for result in results + if tool_type != "function" and result["type"] == "web_search_tool_result" + ], + }, + ), + Delta( + content="answer", + tool_calls=[ + { + "index": 2, + "id": call_ids[1], + "type": "function", + "function": {"name": "web_search", "arguments": '{"query":"two"}'}, + } + ], + ), + ) + chunks: Final = tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-fable-5-1", + object="chat.completion.chunk", + choices=[ + StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == len(deltas) - 1 else None) + ], + ) + for index, delta in enumerate(deltas) + ) + request_tools: Final = ( + [{"type": "function", "name": "web_search", "parameters": {"type": "object"}}] + if tool_type == "function" + else [{"type": tool_type}] + ) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-fable-5-1", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="search", + responses_api_request={"tools": request_tools}, + custom_llm_provider="anthropic", + ) + events: Final = ( + [event.model_dump(exclude_none=True) for event in iterator] + if sync_mode + else [event.model_dump(exclude_none=True) async for event in iterator] + ) + completed: Final = events[-1] + search_items: Final = { + item["id"].removeprefix("ws_"): item + for item in completed["response"]["output"] + if item["type"] == "web_search_call" + } + function_items: Final = { + item["call_id"]: item for item in completed["response"]["output"] if item["type"] == "function_call" + } + function_events: Final = [event for event in events if "function_call_arguments" in event["type"]] + expected_functions: Final = set(call_ids).difference(expected_sources) | {"toolu_regular"} + search_indexes: Final = { + event["output_index"] for event in events if event["type"] == "response.web_search_call.completed" + } + completed_indexes: Final = {item["id"]: index for index, item in enumerate(completed["response"]["output"])} + + assert completed["type"] == "response.completed" + assert [item["content"][0]["text"] for item in completed["response"]["output"] if item["type"] == "message"] == [ + "answer" + ] + assert set(search_items) == set(expected_sources) + assert set(function_items) == expected_functions + assert {event["item_id"] for event in function_events} == {item["id"] for item in function_items.values()} + assert len(search_indexes) == len(expected_sources) + for call_id, item in search_items.items(): + search_events = [ + event for event in events if event.get("item_id", event.get("item", {}).get("id")) == item["id"] + ] + assert [event["type"] for event in search_events] == [ + "response.output_item.added", + "response.web_search_call.in_progress", + "response.web_search_call.searching", + "response.web_search_call.completed", + "response.output_item.done", + ] + assert {event["output_index"] for event in search_events} == {completed_indexes[item["id"]]} + assert search_events[0]["item"]["status"] == "in_progress" + assert search_events[-1]["item"] == item + assert item["status"] == ( + "failed" if result_kind == "error" and call_id.endswith("01Search") else "completed" + ) + assert item["action"]["type"] == "search" + assert item["action"]["query"] == ("one" if call_id.endswith("01Search") else "two") + assert item["action"]["queries"] == [item["action"]["query"]] + assert [source["url"] for source in item["action"]["sources"]] == expected_sources[call_id] + for call_id, item in function_items.items(): + argument_deltas = [ + event["delta"] + for event in function_events + if event["item_id"] == item["id"] and event["type"].endswith(".delta") + ] + assert json.loads("".join(argument_deltas)) == json.loads(item["arguments"]) + assert json.loads(item["arguments"]) == ( + {"city": "Paris"} + if call_id == "toolu_regular" + else {"query": "one" if call_id.endswith("01Search") else "two"} + ) + assert any( + event["type"] == "response.output_item.done" + and event.get("item") == item + and event["output_index"] == completed_indexes[item["id"]] + for event in events + ) + + def test_tool_calls_present_only_in_final_response_are_emitted_before_completed(): iterator = LiteLLMCompletionStreamingIterator( model="test-model", diff --git a/tests/test_litellm/responses/test_custom_tool_call.py b/tests/test_litellm/responses/test_custom_tool_call.py index 5122c1c1d67..e80301c3b2f 100644 --- a/tests/test_litellm/responses/test_custom_tool_call.py +++ b/tests/test_litellm/responses/test_custom_tool_call.py @@ -374,6 +374,8 @@ class TestTransformationCustomTools: "srvtoolu_01ServerCall", "toolu_01CustomCall", ] + assert result[1].type == "function_call" + assert result[1].name == "web_search" def test_transform_mixed_tool_calls(self): """Test transformation with both custom and regular tool calls.""" diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py index aff9d5acac1..08fa3bfc053 100644 --- a/tests/test_litellm/rust_bridge/test_configuration.py +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -52,6 +52,18 @@ def test_resolution_precedence( def test_release_default_remains_disabled() -> None: assert configuration.DEFAULT_RUST_ENABLED is False assert configuration.rust_enabled() is False + assert configuration.rust_ocr_enabled() is True + + +@pytest.mark.parametrize("process", [None, False, True]) +@pytest.mark.parametrize("environment", [None, "0", "1", "off"]) +def test_ocr_configuration(monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None) -> None: + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + if process is not None: + configuration.rust(process) + + assert configuration.rust_ocr_enabled() is (environment not in {"0", "off"} and process is not False) def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py b/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py new file mode 100644 index 00000000000..501a4e986c0 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py @@ -0,0 +1,230 @@ +from collections.abc import Generator, Mapping +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import legacy +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge.ocr import LiteLLMOcrRequest +from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE + + +@pytest.fixture(autouse=True) +def isolated_ocr_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + monkeypatch.delenv("LITELLM_RUST", raising=False) + configuration.reset_rust_configuration() + yield + NATIVE_OCR_LIFECYCLE.reset() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_unavailable_native_uses_legacy(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None: + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + NATIVE_OCR_LIFECYCLE.override(None) + document: Final = {"type": "document_url", "document_url": "https://example.com"} + + result: Final = ( + await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[0]) + if asynchronous + else litellm.ocr("mistral/mistral-ocr-latest", document, pages=[0]) + ) + + assert result is response + fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[0]) + + +def test_admitted_failure_is_returned_without_replay() -> None: + failure: Final = RuntimeError("admitted") + native: Final = Mock(side_effect=failure) + litellm.rust(True) + NATIVE_OCR_LIFECYCLE.override(native) + try: + with pytest.raises(RuntimeError) as caught: + litellm.ocr("mistral/mistral-ocr-latest", {"type": "document_url", "document_url": "https://example.com"}) + assert caught.value is failure + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + assert native.call_count == 1 + + +def test_public_binding_keeps_positional_fields_and_defaults_out_of_native_hook_kwargs() -> None: + document: Final = {"type": "document_url", "document_url": "https://example.com"} + captured: Final = [] + + def native( + request: LiteLLMOcrRequest, + args: tuple[object, ...], + kwargs: Mapping[str, object], + asynchronous: bool, + ) -> OCRResponse: + captured.append((request, args, kwargs, asynchronous)) + return OCRResponse(pages=[], model=request.model) + + litellm.rust(True) + NATIVE_OCR_LIFECYCLE.override(native) + try: + response: Final = litellm.ocr("mistral/mistral-ocr-latest", document) + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + + request, call_args, hook_kwargs, asynchronous = captured[0] + assert response.model == "mistral/mistral-ocr-latest" + assert request.model == "mistral/mistral-ocr-latest" + assert request.document is document + assert call_args == ("mistral/mistral-ocr-latest", document) + assert hook_kwargs == {} + assert asynchronous is False + + +def test_public_binding_keeps_keyword_model_and_document_in_native_hook_kwargs() -> None: + document: Final = {"type": "document_url", "document_url": "https://example.com"} + captured: Final = [] + + def native( + request: LiteLLMOcrRequest, + args: tuple[object, ...], + kwargs: Mapping[str, object], + asynchronous: bool, + ) -> OCRResponse: + assert args == () + captured.append(kwargs) + return OCRResponse(pages=[], model=request.model) + + litellm.rust(True) + NATIVE_OCR_LIFECYCLE.override(native) + try: + litellm.ocr(model="mistral/mistral-ocr-latest", document=document) + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + + assert captured[0]["model"] == "mistral/mistral-ocr-latest" + assert captured[0]["document"] is document + assert "timeout" not in captured[0] + + +@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"]) +def test_public_duplicate_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None: + native: Final = Mock(side_effect=AssertionError("binding errors precede admission")) + document: Final = {"type": "document_url", "document_url": "https://example.com"} + litellm.rust(enabled) + NATIVE_OCR_LIFECYCLE.override(native) + try: + with pytest.raises(TypeError, match=r"ocr\(\) got multiple values for argument 'model'"): + litellm.ocr("mistral/mistral-ocr-latest", document, model="duplicate") + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + assert native.call_count == 0 + + +@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"]) +def test_public_missing_required_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None: + native: Final = Mock(side_effect=AssertionError("binding errors precede admission")) + litellm.rust(enabled) + NATIVE_OCR_LIFECYCLE.override(native) + try: + with pytest.raises(TypeError, match=r"ocr\(\) missing 1 required positional argument: 'document'"): + litellm.ocr("mistral/mistral-ocr-latest") + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + assert native.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("enabled", [False, True, None]) +async def test_environment_opt_out_never_loads_native( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, enabled: bool | None +) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + load: Final = Mock(side_effect=AssertionError("native must not be loaded")) + monkeypatch.setattr(bindings, "get_native_bridge", load) + litellm.rust(enabled) + document: Final = {"type": "file", "file": b"pdf"} + + result: Final = ( + await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[1]) + if asynchronous + else litellm.ocr("mistral/mistral-ocr-latest", document, pages=[1]) + ) + + assert result is response + fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[1]) + load.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("environment", [None, "1"]) +async def test_native_is_enabled_by_default( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, environment: str | None +) -> None: + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + native: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + NATIVE_OCR_LIFECYCLE.override(native) + fallback: Final = Mock(side_effect=AssertionError("legacy must not run")) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + + result: Final = ( + await litellm.aocr("mistral/mistral-ocr-latest", {}) + if asynchronous + else litellm.ocr("mistral/mistral-ocr-latest", {}) + ) + + assert result is response + assert native.call_count == 1 + fallback.assert_not_called() + + +class Declined(Exception): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("declined", [False, True]) +async def test_only_native_declines_replay_on_legacy( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, declined: bool +) -> None: + failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called") + native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure) + NATIVE_OCR_LIFECYCLE.override(native) + import importlib + + main: Final = importlib.import_module("litellm.ocr.main") + monkeypatch.setattr(main, "native_exception_types", lambda: (Declined, RuntimeError)) + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + document: Final = {"type": "file", "file": b"pdf"} + + async def call() -> object: + if asynchronous: + return await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[0]) + return litellm.ocr("mistral/mistral-ocr-latest", document, pages=[0]) + + if declined: + assert await call() is response + fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[0]) + else: + with pytest.raises(RuntimeError) as caught: + await call() + assert caught.value is failure + fallback.assert_not_called() + assert native.call_count == 1 diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index 701938f5677..1303f46e8fa 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -14,6 +14,8 @@ import os import pytest +from litellm import completion_cost +from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import get_model_info @@ -56,17 +58,38 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["max_output_tokens"] == expected["max_output_tokens"] +def test_deepseek_v4p1_flash_twin_costs(local_model_cost_map): + for model in ( + "fireworks_ai/deepseek-v4p1-flash", + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", + ): + response = ModelResponse( + model=model, + choices=[Choices(index=0, message=Message(role="assistant", content="ok"))], + usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000), + ) + cost = completion_cost(completion_response=response, model=model) + assert cost == pytest.approx(8.8e-04) + + TWIN_PINNED_PRICES = { "deepseek-v4-flash-0731": { "input_cost_per_token": 2.2e-07, "cache_read_input_token_cost": 7e-09, "output_cost_per_token": 6.6e-07, }, + "deepseek-v4p1-flash": { + "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 7e-09, + "output_cost_per_token": 6.6e-07, + "supports_vision": True, + "max_output_tokens": 393216, + }, } -def test_deepseek_v4_flash_0731_twins_pin_published_pricing(model_data): - """Both 0731 entries carry the price published at docs.fireworks.ai/serverless/pricing.""" +def test_deepseek_v4_flash_twins_pin_published_pricing(model_data): + """Both entries of each Flash twin pair carry the price published at docs.fireworks.ai/serverless/pricing.""" for bare_suffix, expected in TWIN_PINNED_PRICES.items(): for key in ( f"fireworks_ai/{bare_suffix}", diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index c7e46829aba..19ed31c7b22 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import contextvars import json import logging import os @@ -7,6 +8,7 @@ import queue import threading from datetime import datetime, timedelta, timezone from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -60,6 +62,36 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_non_ocr_wrapper_preserves_logging_executor_and_context(monkeypatch: pytest.MonkeyPatch) -> None: + marker: Final = contextvars.ContextVar("non-ocr-logging-context", default="missing") + token: Final = marker.set("caller-context") + caller_thread: Final = threading.get_ident() + response: Final = object() + logger: Final = MagicMock() + observed: Final = queue.Queue[tuple[object, str, int]]() + + def record_success(result: object, start_time: datetime, end_time: datetime) -> None: + observed.put((result, marker.get(), threading.get_ident())) + + def embedding(**kwargs: object) -> object: + return response + + logger.success_handler.side_effect = record_success + monkeypatch.setattr("litellm.utils.function_setup", MagicMock(return_value=(logger, {}))) + try: + with ThreadPoolExecutor(max_workers=1) as executor: + monkeypatch.setattr("litellm.utils.executor", executor) + result: Final = client(embedding)() + logged_response, context, worker_thread = observed.get_nowait() + assert result is response + assert logged_response is response + assert context == "caller-context" + assert worker_thread != caller_thread + assert observed.empty() + finally: + marker.reset(token) + + def test_cloudflare_model_info_includes_rpm(local_model_cost_map: None) -> None: assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md deleted file mode 100644 index 4c117fb846b..00000000000 --- a/tests/test_litellm_rust/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# Rust OCR bridge tests - -This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR behavior tests live under `ocr/`; reusable OCR request, callback, and recording-server fixtures live under `support/` - -A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions - -`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport - -Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports - -Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs - -The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index b0c75d9d2f5..4387ea2e2fd 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -11,7 +11,7 @@ import pytest_asyncio import litellm from litellm import utils -from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils import litellm_logging, thread_pool_executor from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # preserve raw configuration state in test isolation _CONFIGURATION, @@ -29,11 +29,6 @@ CALLBACK_ATTRIBUTES: Final = ( "_async_success_callback", "_async_failure_callback", ) -EXPECTED_FAILURE_REASONS: Final = { - "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070", - "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070", - "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070", -} def _list_attribute(container: ModuleType, attribute: str) -> list[object]: @@ -76,7 +71,9 @@ async def isolate_ocr_test_state() -> AsyncIterator[None]: stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache stack.enter_context(_rebound(_CONFIGURATION, "override", None)) executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") + stack.enter_context(_rebound(litellm_logging, "executor", executor)) stack.enter_context(_rebound(utils, "executor", executor)) + stack.enter_context(_rebound(thread_pool_executor, "executor", executor)) try: yield finally: @@ -94,14 +91,6 @@ def recording_server() -> Generator[RecordingServer]: def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - for item in items: - if "test_litellm_rust" not in item.path.parts: - continue - relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :]) - reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path) - if reason is not None: - item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) - if not _parse_env_bool(os.environ.get("LITELLM_RUST")): skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") for item in items: diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index b08446412c0..1cfd04b1bff 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -41,7 +41,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_ observations: Final = [] class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observations.append((model, copy.deepcopy(kwargs["additional_args"]))) call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0]) @@ -64,13 +64,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["include_image_base64"] = True if raise_after_edit: raise RuntimeError("pre-call callback failed") class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(copy.deepcopy(request_body(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False) @@ -83,11 +83,11 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_headers(kwargs)["x-audit-tag"] = "reviewed" class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(dict(request_headers(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()]) @@ -96,6 +96,29 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" +def test_native_ocr_pre_call_header_rebinding_does_not_replace_execution_root(ocr_server: RecordingServer) -> None: + retained: Final = [] + observed: Final = [] + + class RetainMutateAndRebind(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + headers = request_headers(kwargs) + retained.append(headers) + kwargs["additional_args"]["headers"] = {"x-rebound": "not-sent"} + headers["x-retained"] = "sent" + + class ObserveRebinding(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(dict(request_headers(kwargs))) + + call_native_ocr_with_callbacks(ocr_server, [RetainMutateAndRebind(), ObserveRebinding()]) + + assert observed == [{"x-rebound": "not-sent"}] + assert retained[0]["x-retained"] == "sent" + assert ocr_server.requests[0].headers["x-retained"] == "sent" + assert "x-rebound" not in ocr_server.requests[0].headers + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references( @@ -107,12 +130,12 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ aliases: Final = [] class Retain(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): aliases.append(request_body(kwargs)["document"] is original) retained.append(request_body(kwargs)["document"]) class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): original["document_url"] = replacement_url arguments: Final = { @@ -143,7 +166,7 @@ def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_docum retained: Final = [] class RetainAndReplace(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): body = request_body(kwargs) retained.append(body["document"]) body["document"] = replacement @@ -165,11 +188,11 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov observed: Final = [] class Rebind(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(request_body(kwargs)) call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) @@ -182,11 +205,11 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_ queued: Final = [] class QueuePayload(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): queued.append(request_body(kwargs)) class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["queued-edit"] = True call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) @@ -200,7 +223,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o finished: Final = threading.Event() class Stash(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["test-token"] = token def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -280,7 +303,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal observed: Final = [] class TrackInFlightRequest(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["request-token"] = token def log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -364,7 +387,7 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context return "caller-token" class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): assert request_headers(kwargs)["Authorization"] == "Bearer caller-token" observations.append("pre_call") request_headers(kwargs)["Authorization"] = "Bearer edited" diff --git a/tests/test_litellm_rust/ocr/test_cohere.py b/tests/test_litellm_rust/ocr/test_cohere.py new file mode 100644 index 00000000000..2a35dc62bd1 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_cohere.py @@ -0,0 +1,141 @@ +from typing import Final + +import pytest + +import litellm +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension +MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0") +IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} +BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32} +PAYLOAD: Final = { + "pages": [ + { + "index": 4, + "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]}, + }, + {"markdown": {"content": "page two"}}, + ], + "meta": {"billed_units": {"pages": 3}}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_public_cohere_request_and_normalization( + recording_server: RecordingServer, model: str, asynchronous: bool +) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + args: Final = { + "model": model, + "document": IMAGE, + "api_base": recording_server.base_url, + "api_key": "test-key", + "req_format": "native", + "unrecognized": True, + } + response: Final = await litellm.aocr(**args) if asynchronous else litellm.ocr(**args) + request: Final = recording_server.requests[0] + assert request.path == ("/providers/cohere/v2/parse" if model.startswith("azure_ai/") else "/v2/parse") + assert request.headers["authorization"] == "Bearer test-key" + assert request.body == {"model": model.split("/", 1)[1], "document": IMAGE, "output_format": "markdown"} + assert [page.index for page in response.pages] == [4, 1] + assert response.pages[0].markdown == "receipt" + assert response.pages[0].images[0].bbox == BOX + assert response.pages[0].images[0].model_extra["description"] == "scan" + assert response.pages[1].images is None + assert response.usage_info.pages_processed == 3 + assert response.get_provider_native_response() == PAYLOAD + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_blocks_and_usage_fallback(recording_server: RecordingServer, model: str) -> None: + blocks: Final = [{"type": "text", "text": "total"}] + recording_server.enqueue(ResponseSpec(body={"pages": [{"blocks": blocks}]})) + response: Final = await litellm.aocr( + model=model, document=IMAGE, api_base=recording_server.base_url, api_key="test-key", output_format="blocks" + ) + assert recording_server.requests[0].body["output_format"] == "blocks" + assert response.pages[0].model_extra["blocks"] == blocks + assert response.pages[0].markdown == "" + assert response.usage_info.pages_processed == 1 + assert response.get_provider_native_response() is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize( + "document", + [ + {"type": "document_url", "document_url": "https://example.com/file.pdf"}, + {"type": "image_url", "image_url": "data:application/pdf;base64,YQ=="}, + {"type": "image_url", "image_url": ""}, + ], +) +async def test_public_cohere_rejects_non_images_before_network( + recording_server: RecordingServer, model: str, document: dict[str, str] +) -> None: + recording_server.expected_requests = 0 + with pytest.raises(litellm.BadRequestError, match="only accepts `image_url`"): + await litellm.aocr(model=model, document=document, api_base=recording_server.base_url, api_key="test-key") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_rejects_unknown_format(recording_server: RecordingServer, model: str) -> None: + recording_server.expected_requests = 0 + with pytest.raises(litellm.BadRequestError, match="output_format"): + await litellm.aocr( + model=model, document=IMAGE, api_base=recording_server.base_url, api_key="test-key", output_format="html" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_provider_failure(recording_server: RecordingServer, model: str) -> None: + recording_server.enqueue(ResponseSpec(status=400, body={"message": "output_format must be blocks or markdown"})) + with pytest.raises(litellm.BadRequestError, match="output_format must be") as caught: + await litellm.aocr(model=model, document=IMAGE, api_base=recording_server.base_url, api_key="test-key") + assert caught.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_health_check(recording_server: RecordingServer, model: str) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + response: Final = await litellm.ahealth_check( + model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr" + ) + assert "error" not in response + assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("suffix", ["", "/cohere/", "/v2", "/v2/parse"]) +async def test_public_cohere_url_variants(recording_server: RecordingServer, suffix: str) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + await litellm.aocr(model=MODELS[0], document=IMAGE, api_base=recording_server.base_url + suffix, api_key="test-key") + assert recording_server.requests[0].path == ("/cohere/v2/parse" if suffix == "/cohere/" else "/v2/parse") + + +@pytest.mark.asyncio +async def test_public_cohere_environment_key_and_remote_url( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("COHERE_API_KEY", "env-key") + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + document: Final = {"type": "image_url", "image_url": "https://example.com/receipt.png"} + await litellm.aocr(model=MODELS[0], document=document, api_base=recording_server.base_url) + assert recording_server.requests[0].headers["authorization"] == "Bearer env-key" + assert recording_server.requests[0].body["document"] == document + + +@pytest.mark.asyncio +async def test_public_cohere_missing_key(recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("COHERE_API_KEY", raising=False) + recording_server.expected_requests = 0 + with pytest.raises(Exception, match="Missing COHERE_API_KEY"): + await litellm.aocr(model=MODELS[0], document=IMAGE, api_base=recording_server.base_url) diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py index a6c76bc5d0e..7b4b9fab579 100644 --- a/tests/test_litellm_rust/ocr/test_dispatch.py +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -1,11 +1,9 @@ from typing import Final -from unittest.mock import Mock import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.ocr import main as ocr_main from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE @@ -18,18 +16,9 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer: return recording_server -@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"]) -def test_public_ocr_dispatches_according_to_rust_setting( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - rust_enabled: bool, -) -> None: - rust_call: Final = Mock(wraps=ocr_main.rust_ocr_bridge.ocr) - python_call: Final = Mock(wraps=ocr_main.base_llm_http_handler.ocr) - monkeypatch.setattr(ocr_main.rust_ocr_bridge, "ocr", rust_call) - monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_call) - litellm.rust(rust_enabled) - +@pytest.mark.parametrize("enabled", [False, True, None]) +def test_public_ocr_uses_native_route_independently_of_flag(ocr_server: RecordingServer, enabled: bool | None) -> None: + litellm.rust(enabled) response: Final = litellm.ocr( model=OCR_MODEL, document=OCR_DOCUMENT, @@ -39,6 +28,26 @@ def test_public_ocr_dispatches_according_to_rust_setting( assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "native OCR response" - assert rust_call.call_count == int(rust_enabled) - assert python_call.call_count == int(not rust_enabled) + assert len(ocr_server.requests) == 1 + assert not ocr_server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("caching", [None, False, True]) +async def test_ocr_does_not_depend_on_chat_cache( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None +) -> None: + from litellm.caching.caching import Cache + + monkeypatch.setattr(litellm, "cache", Cache(type="local", supported_call_types=["completion", "acompletion"])) + arguments: Final = { + "model": OCR_MODEL, + "document": OCR_DOCUMENT, + "api_key": "test-key", + "api_base": ocr_server.base_url, + "caching": caching, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert response.pages[0].markdown == "native OCR response" assert len(ocr_server.requests) == 1 diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py new file mode 100644 index 00000000000..1acad5527d8 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -0,0 +1,996 @@ +import asyncio +import datetime +import gc +import json +import sys +import threading +import weakref +from collections.abc import Coroutine +from contextvars import ContextVar +from typing import Final + +import pytest + +import litellm +from litellm._logging import trace_id_var +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_aocr, call_ocr + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["deployment", "failure"]) +async def test_cancellation_during_failure_obeys_phase_policy(ocr_server: RecordingServer, phase: str) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider failure"}, status=500)) + entered: Final = asyncio.Event() + observed: Final = [] + + class Observer(CustomLogger): + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, **kwargs): + if phase == "deployment": + entered.set() + await asyncio.Event().wait() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(kwargs["exception"]) + if phase == "failure": + entered.set() + await asyncio.Event().wait() + + observer: Final = Observer() + litellm.callbacks.append(observer) + task: Final = asyncio.create_task(call_aocr(ocr_server, callbacks=[observer])) + await asyncio.wait_for(entered.wait(), 5) + task.cancel() + if phase == "deployment": + with pytest.raises(litellm.InternalServerError) as caught: + await task + assert observed == [caught.value] + else: + with pytest.raises(asyncio.CancelledError): + await task + assert len(observed) == 1 + assert isinstance(observed[0], litellm.InternalServerError) + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +@pytest.mark.asyncio +async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + recorder: Final = RecordingLogger() + auth: Final = UserAPIKeyAuth(user_id="ocr-user") + response: Final = await call_aocr( + ocr_server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object() + ) + events: Final = await recorder.wait_for_async("async_log_success_event") + assert response.pages[0].markdown == "native OCR response" + assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user" + assert "metadata" not in ocr_server.requests[0].body + + +@pytest.mark.asyncio +async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None: + caller: Final = asyncio.current_task() + context: Final = ContextVar("lifecycle-test", default="before") + observations: Final = [] + recorder: Final = RecordingLogger() + + class Replace(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + context.set("pre") + observations.append(("pre", asyncio.current_task(), context.get())) + return {**kwargs, "pages": [2]} + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + observations.append(("post", asyncio.current_task(), context.get())) + return response.model_copy(update={"model": "replaced"}) + + litellm.callbacks.append(Replace()) + response: Final = await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="native-final") + events: Final = await recorder.wait_for_async("async_log_success_event") + assert observations == [("pre", caller, "pre"), ("post", caller, "pre")] + assert context.get() == "pre" + assert ocr_server.requests[0].body["pages"] == [2] + assert response.model == "replaced" + assert events[0].response is response + assert response._hidden_params["litellm_call_id"] == "native-final" + assert "response_cost" in response._hidden_params + + +@pytest.mark.asyncio +async def test_deployment_hook_replaces_complete_routing_request(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.05)) + original: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + observed: Final = [] + + class Replace(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + return { + **kwargs, + "model": "azure_ai/mistral-ocr-latest", + "custom_llm_provider": "azure_ai", + "document": replacement, + "api_key": "replacement-key", + "api_base": ocr_server.base_url, + "extra_headers": {"x-deployment": "replacement"}, + "timeout": 2, + "pages": [2], + } + + class Observe(Logging): + def pre_call(self, input, api_key, additional_args): + observed.append((additional_args["complete_input_dict"]["document"], api_key)) + + litellm.callbacks.append(Replace()) + logger: Final = Observe( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deployment-routing", + function_id="deployment-routing", + ) + response: Final = await call_aocr( + ocr_server, + document=original, + timeout=0.001, + litellm_logging_obj=logger, + ) + + assert response.pages[0].markdown == "native OCR response" + assert observed == [(replacement, "replacement-key")] + assert observed[0][0] is replacement + assert replacement == original + assert replacement is not original + assert original == {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr" + assert ocr_server.requests[0].headers["authorization"] == "Bearer replacement-key" + assert ocr_server.requests[0].headers["x-deployment"] == "replacement" + assert ocr_server.requests[0].body["document"] == replacement + assert ocr_server.requests[0].body["pages"] == [2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_metadata_failure_dispatches_only_failure_and_releases_logger( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + failure: Final = RuntimeError("metadata failed") + seen: Final = [] + + class FailingMetadata(Logging): + def _response_cost_calculator(self, *args, **kwargs): + raise failure + + def success_handler(self, *args, **kwargs): + seen.append("success") + + def failure_handler(self, exception, *args, **kwargs): + seen.append(("sync", exception)) + + async def async_failure_handler(self, exception, *args, **kwargs): + seen.append(("async", exception)) + + async def invoke(): + logger: Final = FailingMetadata( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr" if asynchronous else "ocr", + start_time=datetime.datetime.now(), + litellm_call_id="metadata", + function_id="metadata", + ) + reference: Final = weakref.ref(logger) + with pytest.raises(RuntimeError) as caught: + await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr( + ocr_server, litellm_logging_obj=logger + ) + assert caught.value is failure + failure.__traceback__ = None + return reference + + reference: Final = await invoke() + await drain_logging() + gc.collect() + assert seen == ([("sync", failure), ("async", failure)] if asynchronous else [("sync", failure)]) + assert reference() is None + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500)) + recorder: Final = RecordingLogger() + snapshots: Final = [] + + class Observe(CustomLogger): + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, **kwargs): + snapshots.append(exception) + exception.status_code = 418 + + litellm.callbacks.append(Observe()) + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, callbacks=[recorder]) + failures: Final = tuple(event for event in recorder.events if "failure" in event.name) + assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"] + assert all(event.kwargs["exception"] is caught.value for event in failures) + assert caught.value.status_code == 500 + assert snapshots[0] is not caught.value + assert snapshots[0].status_code == 418 + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["pre", "http", "post"]) +async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch( + ocr_server: RecordingServer, phase: str +) -> None: + entered: Final = asyncio.Event() + recorder: Final = RecordingLogger() + + class Pause(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + if phase == "pre": + entered.set() + await asyncio.Event().wait() + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + if phase == "post": + entered.set() + await asyncio.Event().wait() + + litellm.callbacks.append(Pause()) + if phase == "http": + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) + if phase == "pre": + ocr_server.expected_requests = 0 + restored: Final = [] + + async def invoke(): + trace_id_var.set("parent") + try: + await call_aocr(ocr_server, callbacks=[recorder], litellm_trace_id="native-call") + finally: + restored.append(trace_id_var.get()) + + task: Final = asyncio.create_task(invoke()) + if phase == "http": + await ocr_server.wait_for_requests(1) + else: + await asyncio.wait_for(entered.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await drain_logging() + assert restored == ["parent"] + assert not any("success" in name or "failure" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True]) +async def test_deferred_logging_requires_release_and_runs_at_most_once( + ocr_server: RecordingServer, blocked: bool +) -> None: + recorder: Final = RecordingLogger() + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deferred", + function_id="deferred", + dynamic_async_success_callbacks=[recorder], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger) + await drain_logging() + assert "async_log_success_event" not in recorder.names + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + await drain_logging() + events: Final = tuple(event for event in recorder.events if event.name == "async_log_success_event") + assert len(events) == int(not blocked) + if events: + assert events[0].response is response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [RuntimeError("native enqueue failed"), asyncio.CancelledError("cancelled")]) +async def test_deferred_release_handles_enqueue_failure_once_without_replay( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: BaseException +) -> None: + import inspect + + from litellm.litellm_core_utils import logging_worker + + attempts: Final[list[Coroutine[object, object, object]]] = [] + diagnostics: Final = [] + + class FailingWorker: + def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: + attempts.append(coroutine) + raise failure + + recorder: Final = RecordingLogger() + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="release-failure", + function_id="release-failure", + dynamic_async_success_callbacks=[recorder], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", FailingWorker()) + monkeypatch.setattr(sys, "unraisablehook", lambda event: diagnostics.append(event.exc_value)) + + if isinstance(failure, asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError, match="cancelled") as caught: + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + assert caught.value is failure + assert diagnostics == [] + else: + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + assert diagnostics == [failure] + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + + assert len(attempts) == 1 + assert inspect.getcoroutinestate(attempts[0]) == inspect.CORO_CLOSED + assert response.pages[0].markdown == "native OCR response" + assert len(ocr_server.requests) == 1 + assert not any("success" in name or "failure" in name for name in recorder.names) + + +@pytest.mark.asyncio +async def test_abandoned_deferred_logging_is_collectable(ocr_server: RecordingServer) -> None: + async def invoke(): + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="abandoned", + function_id="abandoned", + ) + logger._defer_async_logging = True + await call_aocr(ocr_server, litellm_logging_obj=logger) + return weakref.ref(logger) + + reference: Final = await invoke() + await drain_logging() + gc.collect() + assert reference() is None + + +def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: RecordingServer) -> None: + context: Final = ContextVar("sync-lifecycle", default="missing") + context.set("caller") + thread: Final = threading.current_thread() + finished: Final = threading.Event() + observations: Final = [] + + class Observe(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append((threading.current_thread(), context.get(), response_obj)) + finished.set() + + response: Final = call_ocr(ocr_server, callbacks=[Observe()]) + assert finished.wait(5) + assert observations[0][0] is not thread + assert observations[0][1] == "caller" + assert observations[0][2] is response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_invalid_response_runs_post_call_before_failure(ocr_server: RecordingServer, asynchronous: bool) -> None: + ocr_server.enqueue(ResponseSpec(body={"pages": "invalid"})) + events: Final = [] + + class Observe(Logging): + def pre_call(self, *args, **kwargs): + events.append("pre") + return super().pre_call(*args, **kwargs) + + def post_call(self, *args, **kwargs): + events.append(("post", kwargs["original_response"])) + return super().post_call(*args, **kwargs) + + def success_handler(self, *args, **kwargs): + events.append("success") + + def failure_handler(self, exception, *args, **kwargs): + events.append(("failure", exception)) + + async def async_failure_handler(self, exception, *args, **kwargs): + events.append(("async_failure", exception)) + + logger: Final = Observe( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr" if asynchronous else "ocr", + start_time=datetime.datetime.now(), + litellm_call_id="invalid", + function_id="invalid", + ) + with pytest.raises(litellm.APIConnectionError) as caught: + await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr( + ocr_server, litellm_logging_obj=logger + ) + assert events[0] == "pre" + assert events[1] == ("post", '{"pages": "invalid"}') + assert events[2] == ("failure", caught.value) + if asynchronous: + assert events[3] == ("async_failure", caught.value) + assert "success" not in events + + +@pytest.mark.asyncio +async def test_failing_terminal_handler_preserves_public_failure_and_runs_async_handler( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider failure"}, status=500)) + failures: Final = [] + + class BrokenHandler(Logging): + def failure_handler(self, exception, *args, **kwargs): + failures.append(exception) + raise RuntimeError("handler failed") + + async def async_failure_handler(self, exception, *args, **kwargs): + failures.append(exception) + + logger: Final = BrokenHandler( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="broken", + function_id="broken", + ) + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, litellm_logging_obj=logger) + assert failures == [caught.value, caught.value] + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 2 + recorder: Final = RecordingLogger() + outcomes: Final = [] + + class Nested(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + if kwargs.get("litellm_call_id") == "outer": + outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="inner")) + + litellm.callbacks.append(Nested()) + outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="outer")) + events: Final = await recorder.wait_for_async("async_log_success_event", count=2) + assert [event.kwargs["litellm_call_id"] for event in events] == ["inner", "outer"] + assert events[0].response is outcomes[0] + assert events[1].response is outcomes[1] + assert len(ocr_server.requests) == 2 + + +def test_sync_pre_call_can_make_nested_native_request(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 2 + observed: Final = [] + + class Nested(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + if kwargs["litellm_call_id"] == "outer-sync": + observed.append(call_ocr(ocr_server, litellm_call_id="inner-sync")) + + response: Final = call_ocr(ocr_server, callbacks=[Nested()], litellm_call_id="outer-sync") + assert observed[0].pages[0].markdown == response.pages[0].markdown + assert len(ocr_server.requests) == 2 + + +@pytest.mark.asyncio +async def test_retained_argument_aliases_and_body_roots_survive_envelope_replacement( + ocr_server: RecordingServer, +) -> None: + pages: Final = [0] + document: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + opaque: Final = object() + observed: Final = [] + + class Observe(Logging): + def pre_call(self, input, api_key, additional_args): + body: Final = additional_args["complete_input_dict"] + headers: Final = additional_args["headers"] + observed.append((body["document"] is document, body["pages"] is pages)) + pages.append(2) + headers["x-retained"] = "yes" + additional_args["complete_input_dict"] = {"discarded": True} + additional_args["headers"] = {} + observed.append((body, headers)) + + def post_call(self, original_response, additional_args): + observed.append( + (additional_args["complete_input_dict"] is observed[2][0], additional_args["headers"] is observed[2][1]) + ) + + class Deployment(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + observed.append(("model" in kwargs, "document" in kwargs, kwargs["opaque"] is opaque)) + + litellm.callbacks.append(Deployment()) + logger: Final = Observe( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="roots", + function_id="roots", + ) + response: Final = await litellm.aocr( + "mistral/mistral-ocr-latest", + document, + api_key="test-key", + api_base=ocr_server.base_url, + pages=pages, + opaque=opaque, + litellm_logging_obj=logger, + ) + assert response.pages[0].markdown == "native OCR response" + assert observed[0] == (False, False, True) + assert observed[1] == (True, True) + assert observed[3] == (True, True) + assert ocr_server.requests[0].body["pages"] == [0, 2] + assert ocr_server.requests[0].headers["x-retained"] == "yes" + + +def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None: + from litellm.ocr.main import _public_request + from litellm.rust_bridge import _native + + ocr_server.expected_requests = 0 + effects: Final = [] + + class File: + def read(self): + effects.append("read") + return b"abc" + + def create(): + file: Final = File() + kwargs: Final = {"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": file}} + coroutine: Final = _native._ocr_lifecycle(_public_request("aocr", (), kwargs), (), kwargs, True) + file.owner = coroutine + coroutine.close() + return weakref.ref(file) + + reference: Final = create() + gc.collect() + assert reference() is None + assert effects == [] + + +@pytest.mark.asyncio +async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server: RecordingServer) -> None: + effects: Final = [] + caller: Final = asyncio.current_task() + + class File: + def read(self): + effects.append(("read", asyncio.current_task())) + return b"abc" + + class Deployment(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + await asyncio.sleep(0) + effects.append(("hook", asyncio.current_task())) + + litellm.callbacks.append(Deployment()) + await call_aocr(ocr_server, document={"type": "file", "file": File()}) + assert effects == [("hook", caller), ("read", caller)] + + +@pytest.mark.asyncio +async def test_failure_callbacks_continue_within_both_families(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "failed"}, status=500)) + observed: Final = [] + + class Broken(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("broken-sync", kwargs["exception"])) + raise RuntimeError("sync observer") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("broken-async", kwargs["exception"])) + raise RuntimeError("async observer") + + class Following(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("following-sync", kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("following-async", kwargs["exception"])) + + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, callbacks=[Broken(), Following()]) + assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"] + assert all(error is caught.value for _, error in observed) + + +@pytest.mark.asyncio +async def test_cancelling_native_transport_closes_connection_before_return() -> None: + received: Final = asyncio.Event() + disconnected: Final = asyncio.Event() + + async def provider(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + headers: Final = await reader.readuntil(b"\r\n\r\n") + length: Final = next( + int(line.split(b":", 1)[1]) + for line in headers.split(b"\r\n") + if line.lower().startswith(b"content-length:") + ) + await reader.readexactly(length) + received.set() + assert await reader.read() == b"" + disconnected.set() + writer.close() + await writer.wait_closed() + + server: Final = await asyncio.start_server(provider, "127.0.0.1", 0) + async with server: + port: Final = server.sockets[0].getsockname()[1] + task: Final = asyncio.create_task( + litellm.aocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + api_key="test-key", + api_base=f"http://127.0.0.1:{port}", + ) + ) + await asyncio.wait_for(received.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(disconnected.wait(), 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["reducto/parse-v3", "reducto/parse-legacy"]) +async def test_reducto_lifecycle_retains_upload_parse_and_post_call_boundaries( + ocr_server: RecordingServer, model: str +) -> None: + ocr_server.expected_requests = 2 + ocr_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded.pdf"})) + ocr_server.enqueue(ResponseSpec(body={"result": {"chunks": [{"content": "parsed"}]}})) + boundaries: Final = [] + recorder: Final = RecordingLogger() + + class Observe(Logging): + def post_call(self, *args, **kwargs): + boundaries.append(tuple(request.path for request in ocr_server.requests)) + return super().post_call(*args, **kwargs) + + logger: Final = Observe( + model=model, + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="upload", + function_id="upload", + dynamic_async_success_callbacks=[recorder], + ) + response: Final = await call_aocr(ocr_server, model=model, litellm_logging_obj=logger) + events: Final = await recorder.wait_for_async("async_log_success_event") + assert boundaries == [("/upload", "/parse")] + assert b"abc" in ocr_server.requests[0].raw_body + assert "multipart/form-data" in ocr_server.requests[0].headers["content-type"] + assert ocr_server.requests[1].body["input" if model.endswith("v3") else "document_url"] == "reducto://uploaded.pdf" + assert response.pages[0].markdown == "parsed" + assert events[0].response is response + + +@pytest.mark.asyncio +async def test_document_intelligence_post_call_observes_submission_and_final_result( + ocr_server: RecordingServer, +) -> None: + ocr_server.expected_requests = 2 + ocr_server.enqueue( + ResponseSpec( + body={"status": "running"}, + status=202, + headers={"Operation-Location": f"{ocr_server.base_url}/operations/1", "Retry-After": "0"}, + ) + ) + ocr_server.enqueue(ResponseSpec(body={"status": "succeeded", "analyzeResult": {"pages": []}})) + boundaries: Final = [] + + class Observe(Logging): + def post_call(self, *args, **kwargs): + boundaries.append((tuple(request.method for request in ocr_server.requests), kwargs["original_response"])) + return super().post_call(*args, **kwargs) + + logger: Final = Observe( + model="azure_ai/doc-intelligence/prebuilt-read", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="poll", + function_id="poll", + ) + response: Final = await call_aocr( + ocr_server, model="azure_ai/doc-intelligence/prebuilt-read", litellm_logging_obj=logger + ) + assert [methods for methods, _ in boundaries] == [("POST",), ("POST", "GET")] + assert json.loads(boundaries[0][1])["status"] == "running" + assert json.loads(boundaries[1][1])["status"] == "succeeded" + assert [request.method for request in ocr_server.requests] == ["POST", "GET"] + assert ocr_server.requests[1].path == "/operations/1" + assert response.pages == [] + + +@pytest.mark.asyncio +async def test_vertex_deepseek_public_lifecycle_normalizes_before_success(ocr_server: RecordingServer) -> None: + ocr_server.enqueue( + ResponseSpec(body={"choices": [{"message": {"content": "recognized"}}], "usage": {"prompt_tokens": 1}}) + ) + recorder: Final = RecordingLogger() + response: Final = await call_aocr( + ocr_server, + model="vertex_ai/deepseek-ocr-maas", + document={"type": "document_url", "document_url": "gs://bucket/document.pdf"}, + vertex_project="project-1", + vertex_location="europe-west4", + callbacks=[recorder], + ) + events: Final = await recorder.wait_for_async("async_log_success_event") + assert response.pages[0].markdown == "recognized" + assert events[0].response is response + assert ( + ocr_server.requests[0].path + == "/v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("limit", ["budget", "retries"]) +async def test_shared_call_limits_still_reject_before_reading_ocr_file( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, limit: str +) -> None: + ocr_server.expected_requests = 0 + reads: Final = [] + + class File: + def read(self): + reads.append("read") + return b"abc" + + monkeypatch.setattr(litellm, "max_budget", 1 if limit == "budget" else None) + monkeypatch.setattr(litellm, "_current_cost", 2) + monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None) + expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"previous_models": ["earlier"]}} + with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"): + await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) + assert reads == [] + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("extra_bytes", [0, 1]) +async def test_response_limit_is_enforced_at_the_public_boundary( + ocr_server: RecordingServer, asynchronous: bool, extra_bytes: int +) -> None: + limit: Final = len(json.dumps(OCR_RESPONSE).encode()) - extra_bytes + if extra_bytes: + with pytest.raises(litellm.APIConnectionError, match="OCR response exceeds the size limit"): + await call_aocr(ocr_server, max_response_bytes=limit) if asynchronous else call_ocr( + ocr_server, max_response_bytes=limit + ) + else: + response: Final = ( + await call_aocr(ocr_server, max_response_bytes=limit) + if asynchronous + else call_ocr(ocr_server, max_response_bytes=limit) + ) + assert response.pages[0].markdown == "native OCR response" + assert len(ocr_server.requests) == 1 + body: Final = ocr_server.requests[0].body + assert isinstance(body, dict) + assert "max_response_bytes" not in body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("failure", [False, True]) +async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + failure: bool, + created_loggers: list[Logging], +) -> None: + from litellm import utils + from litellm.litellm_core_utils import litellm_logging, logging_worker + + class DispatchProbe: + deployments = 0 + submissions = 0 + enqueues = 0 + + def deployment(self, *args: object, **kwargs: object) -> None: + self.deployments += 1 + + def submit(self, *args: object, **kwargs: object) -> None: + self.submissions += 1 + + def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: + self.enqueues += 1 + coroutine.close() + + probe: Final = DispatchProbe() + for name in ( + "async_pre_call_deployment_hook", + "async_post_call_success_deployment_hook", + "async_post_call_failure_deployment_hook", + ): + monkeypatch.setattr(utils, name, probe.deployment) + monkeypatch.setattr(litellm_logging, "executor", probe) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) + if failure: + ocr_server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500)) + trace_id_var.set("callback-free-parent") + arguments: Final = {"litellm_trace_id": "callback-free-call", "litellm_call_id": "callback-free-id"} + if failure: + with pytest.raises(litellm.InternalServerError): + await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) + else: + response: Final = ( + await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert response._hidden_params["litellm_call_id"] == "callback-free-id" + assert response._hidden_params["response_cost"] is not None + assert response._hidden_params["_response_ms"] > 0 + assert trace_id_var.get() == "callback-free-parent" + assert probe.deployments == probe.submissions == probe.enqueues == 0 + assert len(created_loggers) == 1 + logger: Final = created_loggers[0] + assert not hasattr(logger, "_native_pending_logging") + assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] + assert "standard_logging_object" not in logger.model_call_details + assert ( + "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None + ) + assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) + assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "registration", ["success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"] +) +async def test_terminal_registration_added_during_http_is_observed( + ocr_server: RecordingServer, registration: str +) -> None: + failure: Final = "failure" in registration + observer: Final = RecordingLogger() + ocr_server.enqueue( + ResponseSpec( + body={"message": "provider failed"} if failure else OCR_RESPONSE, status=500 if failure else 200, delay=0.1 + ) + ) + task: Final = asyncio.create_task( + asyncio.to_thread(call_ocr, ocr_server) if registration == "success_callback" else call_aocr(ocr_server) + ) + await ocr_server.wait_for_requests(1) + getattr(litellm, registration).append(observer) + if failure: + with pytest.raises(litellm.InternalServerError): + await task + else: + await task + event: Final = ("async_" if registration.startswith("_async") else "") + ( + "log_failure_event" if failure else "log_success_event" + ) + await observer.wait_for_async(event) + assert event in observer.names + + +@pytest.fixture +def created_loggers(monkeypatch: pytest.MonkeyPatch) -> list[Logging]: + from litellm import utils + + original_setup: Final = utils.function_setup + loggers: Final[list[Logging]] = [] + + def setup( + call_type: str, + rules: utils.Rules, + start: datetime.datetime, + *args: object, + is_async_call: bool = True, + **kwargs: object, + ) -> tuple[Logging, dict[str, object]]: + logger, prepared = original_setup(call_type, rules, start, *args, is_async_call=is_async_call, **kwargs) + assert isinstance(logger, Logging) + setattr(logger, "_defer_async_logging", True) + loggers.append(logger) + return logger, prepared + + monkeypatch.setattr(utils, "function_setup", setup) + return loggers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("consumer", ["logger_fn", "raw_global", "request_debug"]) +async def test_explicit_logging_consumers_keep_request_and_response_payloads( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str +) -> None: + snapshots: Final[list[dict[str, object]]] = [] + if consumer == "raw_global": + monkeypatch.setattr(litellm, "log_raw_request_response", True) + arguments: Final = { + "logger_fn": {"logger_fn": lambda details: snapshots.append(dict(details))}, + "raw_global": {}, + "request_debug": {"litellm_request_debug": True}, + }[consumer] + response: Final = await call_aocr(ocr_server, **arguments) + details: Final = created_loggers[0].model_call_details + assert details["additional_args"]["complete_input_dict"]["model"] == "mistral-ocr-latest" + assert json.loads(details["original_response"])["pages"][0]["markdown"] == response.pages[0].markdown + if consumer.startswith("raw_"): + assert details["raw_request_typed_dict"]["raw_request_body"]["model"] == "mistral-ocr-latest" + if consumer == "logger_fn": + assert [item["log_event_type"] for item in snapshots] == ["pre_api_call", "post_api_call"] + + +@pytest.mark.asyncio +async def test_registration_removed_before_deferred_release_skips_queue( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging] +) -> None: + from litellm.litellm_core_utils import logging_worker + + class QueueProbe: + enqueues = 0 + + def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: + self.enqueues += 1 + coroutine.close() + + observer: Final = RecordingLogger() + litellm._async_success_callback.append(observer) + await call_aocr(ocr_server) + logger: Final = created_loggers[0] + assert hasattr(logger, "_native_pending_logging") + litellm._async_success_callback.clear() + probe: Final = QueueProbe() + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + assert probe.enqueues == 0 + assert not observer.names + assert logger.model_call_details["response_cost"] is not None diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index d241fe08fc8..4f4b39fa6c6 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -1,3 +1,4 @@ +from pathlib import Path from typing import Final import pytest @@ -5,13 +6,13 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, call_native_aocr, call_native_ocr, ) -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -79,6 +80,22 @@ def test_native_ocr_prepares_file_document_like_python(ocr_server: RecordingServ } +def test_native_ocr_reads_sdk_path_input(ocr_server: RecordingServer, tmp_path: Path) -> None: + document_path: Final = tmp_path / "document.pdf" + document_path.write_bytes(b"%PDF-1.4") + + response: Final = call_native_ocr( + ocr_server, + document={"type": "file", "file": document_path}, + ) + + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].body["document"] == { + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + } + + def test_native_ocr_sends_pages_and_image_options(ocr_server: RecordingServer) -> None: call_native_ocr(ocr_server, pages=[0, 2], include_image_base64=True) @@ -149,7 +166,7 @@ def test_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: Rec assert response.usage_info.pages_processed == 1 -def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: RecordingServer) -> None: +def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None: ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) with pytest.raises(litellm.BadRequestError) as caught: @@ -158,13 +175,23 @@ def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: assert caught.value.status_code == 400 assert caught.value.model == "mistral-ocr-latest" assert caught.value.llm_provider == "mistral" - assert "invalid OCR request" not in str(caught.value) + assert "invalid OCR request" in str(caught.value) -def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: +def test_native_ocr_rejects_unknown_response_format_before_provider_request(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 0 + + with pytest.raises(litellm.BadRequestError, match="Invalid `req_format`"): + call_native_ocr(ocr_server, req_format="raw") + + assert ocr_server.requests == [] + + +def test_ocr_raises_public_timeout_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: + litellm.rust(True) ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) - with pytest.raises(RuntimeError, match="OCR transport failed"): + with pytest.raises(litellm.Timeout): call_native_ocr(ocr_server, timeout=0.01) assert len(ocr_server.requests) == 1 @@ -301,13 +328,10 @@ async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callbac @pytest.mark.parametrize( "configuration", - [ - {"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}, - {"model": "azure_ai/doc-intelligence/prebuilt-read"}, - ], - ids=["oidc-assertion", "document-intelligence-model"], + [{"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}], + ids=["invalid-oidc-assertion"], ) -def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks( +def test_public_azure_ocr_maps_invalid_oidc_configuration_before_token_or_request( ocr_server: RecordingServer, isolated_azure_auth: None, configuration: dict[str, object], @@ -327,10 +351,10 @@ def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_call "callbacks": [recorder], **configuration, } - with pytest.raises(NotImplementedError): + with pytest.raises(litellm.APIConnectionError): call_native_ocr(ocr_server, **arguments) assert calls == [] - assert recorder.events == () + assert "log_pre_api_call" not in recorder.names assert ocr_server.requests == [] @@ -432,3 +456,170 @@ async def test_native_azure_ocr_rejects_coroutine_returned_by_sync_token_provide coroutine.close() assert calls == [] assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "override, expected_key", + [ + ({}, "credential-key"), + ({"api_key": "explicit-key"}, "explicit-key"), + ({"api_key": None}, "environment-key"), + ], + ids=["inherit", "explicit", "explicit-none"], +) +async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + override: dict[str, object], + expected_key: str, +) -> None: + from litellm.models.credentials import CredentialItem + + pages: Final = [0] + opaque: Final = object() + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}), + CredentialItem( + credential_name="ocr-test", + credential_info={}, + credential_values={ + "api_key": "credential-key", + "api_base": ocr_server.base_url, + "pages": pages, + "opaque": opaque, + }, + ), + CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}), + ], + ) + + class Observer(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + super().log_pre_api_call(model, messages, kwargs) + pages.append(2) + + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": OCR_DOCUMENT, + "litellm_credential_name": "ocr-test", + "callbacks": [Observer()], + **override, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}" + assert ocr_server.requests[0].body["pages"] == [0, 2] + + +@pytest.mark.parametrize("source", ["sdk", "proxy"]) +@pytest.mark.parametrize( + "filename,mime", [("scan.PNG", "image/png"), ("document.pdf", "application/pdf"), ("note.txt", "text/plain")] +) +def test_ocr_file_helpers_use_native_document_preparation(source: str, filename: str, mime: str) -> None: + from io import BytesIO + + from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type + from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload + + file: Final = BytesIO(b"abc") + file.name = filename + document: Final = ( + convert_file_document_to_url_document({"type": "file", "file": file}) + if source == "sdk" + else _build_document_from_upload(b"abc", filename, "application/octet-stream; charset=utf-8") + ) + field: Final = "image_url" if mime.startswith("image/") else "document_url" + assert get_mime_type(filename) == mime + assert document == {"type": field, field: f"data:{mime};base64,YWJj"} + + +@pytest.mark.parametrize("attribute", ["read", "name"]) +def test_native_file_preparation_preserves_property_errors(attribute: str) -> None: + from litellm.ocr.input import convert_file_document_to_url_document + + failure: Final = LookupError("file property failed") + + class File: + def __getattribute__(self, name: str): + if name == attribute: + raise failure + return super().__getattribute__(name) + + def read(self): + return b"abc" + + with pytest.raises(LookupError) as caught: + convert_file_document_to_url_document({"type": "file", "file": File()}) + assert caught.value is failure + + +@pytest.mark.parametrize("kind", ["bytes", "path", "reader"]) +def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Path) -> None: + from litellm.ocr.input import FileDocument, convert_file_document_to_url_document, get_max_file_bytes + + limit: Final = get_max_file_bytes() + path: Final = tmp_path / "large.pdf" + with path.open("wb") as stream: + stream.truncate(limit + 1) + + class Reader: + def read(self) -> bytes: + return b"a" * (limit + 1) + + document: Final[FileDocument] = { + "type": "file", + "file": path if kind == "path" else Reader() if kind == "reader" else b"a" * (limit + 1), + } + with pytest.raises(ValueError, match="exceeds the size limit"): + convert_file_document_to_url_document(document) + + +@pytest.mark.parametrize("kind", ["str", "path", "reader"]) +def test_native_upload_binding_rejects_filesystem_inputs(kind: str, tmp_path: Path) -> None: + from io import BytesIO + from typing import cast # noqa: TID251 # deliberately invalid inputs exercise the native runtime boundary + + from litellm.ocr.input import convert_upload_to_url_document + + path: Final = tmp_path / "secret.pdf" + path.write_bytes(b"server secret") + source: Final = str(path) if kind == "str" else path if kind == "path" else BytesIO(b"abc") + with pytest.raises(TypeError): + convert_upload_to_url_document(cast(bytes, source), "document.pdf", None) + + +@pytest.mark.parametrize("extra_bytes", [0, 1]) +def test_native_upload_enforces_file_size_limit(extra_bytes: int) -> None: + import base64 + + from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes + + content: Final = b"a" * (get_max_file_bytes() + extra_bytes) + if extra_bytes: + with pytest.raises(ValueError, match="exceeds the size limit"): + convert_upload_to_url_document(content, "scan.pdf", None) + return + document: Final = convert_upload_to_url_document(content, "scan.pdf", None) + assert document["type"] == "document_url" + assert base64.b64decode(document["document_url"].split(",", 1)[1]) == content + + +def test_native_file_preparation_preserves_reader_exception() -> None: + from litellm.ocr.input import convert_file_document_to_url_document + + failure: Final = RuntimeError("reader failed") + + class Reader: + def read(self) -> bytes: + raise failure + + with pytest.raises(RuntimeError) as caught: + convert_file_document_to_url_document({"type": "file", "file": Reader()}) + assert caught.value is failure diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py index 6de011b1414..d3749ccc095 100644 --- a/tests/test_litellm_rust/support/callback_recorder.py +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -81,7 +81,7 @@ class RecordingLogger(CustomLogger): await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout) return tuple(event for event in self.events if event.name == name) - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): self._record("log_pre_api_call", kwargs) def log_success_event(self, kwargs, response_obj, start_time, end_time): diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 5a9b9497c6e..228ed2cc454 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -58,7 +58,9 @@ def recording_service() -> Iterator[RecordingServer]: def _handle(self) -> None: content_length: Final = int(self.headers.get("Content-Length", "0")) raw_body: Final = self.rfile.read(content_length) if content_length else b"" - body: Final = json.loads(raw_body) if raw_body else None + body: Final = ( + json.loads(raw_body) if raw_body and self.headers.get_content_type() == "application/json" else None + ) requests.append( RecordedRequest( method=self.command, @@ -84,6 +86,7 @@ def recording_service() -> Iterator[RecordingServer]: pass do_POST = _handle + do_GET = _handle def log_message(self, format: str, *args: object) -> None: pass diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index d681d752ffd..7114e42a59e 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -2,7 +2,6 @@ from typing import Final import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import ocr as native_ocr from tests.test_litellm_rust.support.recording_server import RecordingServer OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} @@ -36,11 +35,11 @@ async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return native_ocr.ocr(ocr_arguments(server, **kwargs)) + return call_ocr(server, **kwargs) async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return await native_ocr.aocr(ocr_arguments(server, **kwargs)) + return await call_aocr(server, **kwargs) def request_body(kwargs: dict[str, object]) -> dict[str, object]: diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index ad1c8c652bb..e0e06d685b8 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -2,6 +2,7 @@ import json import threading from collections.abc import Generator from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from io import BytesIO from typing import Final import pytest @@ -99,6 +100,38 @@ def test_native_ocr_with_compiled_rust_extension( } +@pytest.mark.parametrize( + "file_input,mime_type,expected_type,expected_field,expected_uri", + [ + (b"abc", "application/pdf", "document_url", "document_url", "data:application/pdf;base64,YWJj"), + (BytesIO(b"abc"), "image/png", "image_url", "image_url", "data:image/png;base64,YWJj"), + ], +) +def test_native_lifecycle_core_encodes_python_file_input( + ocr_server, + file_input, + mime_type, + expected_type, + expected_field, + expected_uri, +): + server, requests = ocr_server + litellm.rust(True) + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "file", "file": file_input, "mime_type": mime_type}, + api_key="test-key", + api_base=f"http://127.0.0.1:{server.server_port}", + opaque_extension=object(), + ) + assert response.pages[0].markdown == "native OCR response" + assert requests[0]["body"]["document"] == { + "type": expected_type, + expected_field: expected_uri, + } + assert "opaque_extension" not in requests[0]["body"] + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) @pytest.mark.asyncio @@ -145,24 +178,20 @@ async def test_native_public_ocr_matches_python(model, asynchronous): server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) thread: Final = Thread(target=server.serve_forever, daemon=True) thread.start() - responses: Final = [] try: - for enabled in (False, True): - litellm.rust(enabled) - arguments: Final = { - "model": model, - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "pages": [0, 2], - "timeout": 3.0, - } - response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - responses.append(response.model_dump()) - assert len(calls) == 2 - assert calls[0] == calls[1] - for key in ("model", "pages", "object"): - assert responses[0][key] == responses[1][key] + litellm.rust(True) + arguments: Final = { + "model": model, + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "pages": [0, 2], + "timeout": 3.0, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + response_data: Final = response.model_dump() + assert len(calls) == 1 + assert response_data["object"] == "ocr" finally: server.shutdown() server.server_close() @@ -195,7 +224,7 @@ def test_native_ocr_rejects_invalid_input_before_network(ocr_server, custom_prov from litellm.rust_bridge import _native server, requests = ocr_server - with pytest.raises(ValueError, match=r"invalid (OCR request field|provider)|invalid request"): + with pytest.raises(ValueError, match="Document URL is required"): _native.ocr( model="mistral-ocr-latest", custom_llm_provider=custom_provider, @@ -224,7 +253,7 @@ async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, "num_retries": 0, } started = time.monotonic() - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.Timeout): await asyncio.wait_for( litellm.aocr(**arguments) if asynchronous else asyncio.to_thread(litellm.ocr, **arguments), timeout=3, diff --git a/ui/litellm-dashboard/public/assets/logos/conduct.png b/ui/litellm-dashboard/public/assets/logos/conduct.png new file mode 100644 index 00000000000..e68b32df916 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/conduct.png differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index dad8599e967..ebc97891744 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -1,11 +1,11 @@ import React from "react"; -import { render, screen, waitFor, within } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import AddAgentForm from "./add_agent_form"; import * as networking from "@/components/networking"; import type { AgentCreateInfo } from "@/components/networking"; -import { chooseSelectOption } from "../../../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders as render } from "../../../../../tests/test-utils"; vi.mock("@/components/networking", () => ({ createAgentCall: vi.fn(), @@ -351,4 +351,33 @@ describe("AddAgentForm submit payload", () => { expect(await screen.findByText("Agent Created!")).toBeInTheDocument(); expect(within(screen.getByText("Agent Created!").parentElement!).getByText("created-agent")).toBeInTheDocument(); }); + it("blocks creation after clearing the existing key and assigns the reselected key", async () => { + vi.mocked(networking.keyListCall).mockResolvedValue({ + keys: [{ token: "key-maple", key_alias: "Maple key" }], + }); + const user = userEvent.setup(); + renderForm(); + await user.type(await screen.findByLabelText("Agent Name"), "key-selection-agent"); + await user.type(screen.getByLabelText("Display Name"), "Key selection"); + await user.type(screen.getByPlaceholderText("Describe what this agent does..."), "d"); + for (let step = 0; step < 3; step++) { + await user.click(screen.getByRole("button", { name: /^Next/ })); + } + await user.click(screen.getByRole("radio", { name: "Assign an existing key" })); + const keySelector = await screen.findByPlaceholderText("Search by key name…"); + await chooseSelectOption(user, keySelector, "Maple key"); + await user.click(screen.getByRole("button", { name: "Clear" })); + await user.click(screen.getByRole("button", { name: /Create Agent/ })); + expect(networking.createAgentCall).not.toHaveBeenCalled(); + expect(networking.keyUpdateCall).not.toHaveBeenCalled(); + await chooseSelectOption(user, keySelector, "Maple key"); + await user.click(screen.getByRole("button", { name: /Create Agent/ })); + await waitFor(() => + expect(networking.keyUpdateCall).toHaveBeenCalledWith("tok", { + key: "key-maple", + agent_id: "agent-1", + }), + ); + expect(networking.createAgentCall).toHaveBeenCalledTimes(1); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index 108bae977e1..e71fed40209 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -338,6 +338,11 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok return; } + if (keyAssignOption === "existing_key" && !selectedExistingKey) { + toast.error("Please select an existing key to assign"); + return; + } + setIsSubmitting(true); try { const isValid = await form.trigger(); @@ -406,12 +411,7 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok selectedTeamId, ); setCreatedKeyValue(keyResponse.key || null); - } else if (keyAssignOption === "existing_key") { - if (!selectedExistingKey) { - toast.error("Please select an existing key to assign"); - setIsSubmitting(false); - return; - } + } else if (keyAssignOption === "existing_key" && selectedExistingKey) { await keyUpdateCall(accessToken, { key: selectedExistingKey, agent_id: agentId, @@ -963,8 +963,8 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok setSelectedExistingKey(value || null)} + value={selectedExistingKey} + onValueChange={setSelectedExistingKey} options={existingKeys.map((k) => ({ label: k.key_alias || k.token?.slice(0, 12) + "…", value: k.token, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx index 50f9c7cda3c..492a6b5c630 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx @@ -82,7 +82,7 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis control={form.control} name="tpm_limit" label="Max Tokens per minute" - description="Default is model limit." + description="Leave blank for no LiteLLM limit. Provider rate limits still apply." > {({ ref, value, onChange, ...field }) => ( = ({ isModalVisible, setIsModalVis control={form.control} name="rpm_limit" label="Max Requests per minute" - description="Default is model limit." + description="Leave blank for no LiteLLM limit. Provider rate limits still apply." > {({ ref, value, onChange, ...field }) => ( = ({ isModalVisible, setIs control={form.control} name="tpm_limit" label="Max Tokens per minute" - description="Default is model limit." + description="Leave blank for no LiteLLM limit. Provider rate limits still apply." > {({ ref, value, onChange, ...field }) => ( = ({ isModalVisible, setIs control={form.control} name="rpm_limit" label="Max Requests per minute" - description="Default is model limit." + description="Leave blank for no LiteLLM limit. Provider rate limits still apply." > {({ ref, value, onChange, ...field }) => ( = ({ field, embeddingModels, name={name} disabled={disabled} value={typeof value === "string" && value !== "" ? value : null} - onValueChange={(selected: string | null) => onChange(selected ?? "")} + onValueChange={onChange} > = ({ field, embeddingModels, onChange(model?.value ?? "")} + onValueChange={(model: EmbeddingModelOption | null) => onChange(model?.value ?? null)} itemToStringLabel={(model: EmbeddingModelOption) => model.label} isItemEqualToValue={(model: EmbeddingModelOption, other: EmbeddingModelOption) => model.value === other.value diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts index 7b9454a37c3..8da54e3ee78 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.ts @@ -1,6 +1,6 @@ import { CACHE_FIELDS, CacheField, CacheSection, REDACTED_VALUE, RedisType } from "./cacheSettingsFields"; -export type CacheFormValue = string | number | boolean | undefined; +export type CacheFormValue = string | number | boolean | null | undefined; export type CacheFormValues = Record; export type CacheSavePayloadValue = string | number | boolean | unknown[]; export type CacheSavePayload = Record; @@ -38,6 +38,9 @@ const initialValueForField = (field: CacheField, raw: unknown): CacheFormValue = return typeof source === "string" ? source : JSON.stringify(source, null, 2); } + if ((field.type === "select" || field.type === "model-select") && !hasValue(source)) { + return null; + } if (source === undefined || source === null) { return ""; } @@ -77,7 +80,7 @@ const saveValueForField = (field: CacheField, raw: CacheFormValue): CacheSavePay } if (typeof raw !== "string") { - return raw === undefined ? undefined : String(raw); + return raw == null ? undefined : String(raw); } const trimmed = raw.trim(); return trimmed === "" ? undefined : trimmed; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx index 918d8947151..0d1da4dc8ce 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx @@ -1,7 +1,8 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders, screen, waitFor } from "../../../../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import CacheSettings from "./index"; +import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; const { getCacheSettingsCall, testCacheConnectionCall, updateCacheSettingsCall } = vi.hoisted(() => ({ getCacheSettingsCall: vi.fn(), @@ -29,7 +30,7 @@ const LOADED_WITH_ADVANCED = { }, }; -const renderSettings = () => render(); +const renderSettings = () => renderWithProviders(); const save = async (user: ReturnType) => user.click(screen.getByRole("button", { name: /save changes/i })); @@ -174,4 +175,42 @@ describe("CacheSettings advanced settings round-trip", () => { await waitFor(() => expect(updateCacheSettingsCall).toHaveBeenCalledTimes(1)); expect(updateCacheSettingsCall.mock.calls[0][1]).not.toHaveProperty("ttl"); }); + + it("should omit a cleared cache model from save and test while retaining other settings", async () => { + vi.mocked(fetchAvailableModels).mockResolvedValue([ + { model_group: "synthetic-embedding", mode: "embedding" }, + ] as Awaited>); + getCacheSettingsCall.mockResolvedValue({ + current_values: { + redis_type: "semantic", + host: "localhost", + redis_semantic_cache_embedding_model: "synthetic-embedding", + password: "***REDACTED***", + ttl: 0, + ssl: false, + namespace: "synthetic-cache", + }, + }); + const user = userEvent.setup(); + renderSettings(); + await screen.findByRole("combobox", { name: "Embedding Model" }); + await user.click(screen.getByRole("button", { name: "Clear" })); + const expected = { + type: "redis", + host: "localhost", + port: "6379", + similarity_threshold: 0.8, + semantic_cache_scope: "key", + ssl: false, + ssl_check_hostname: false, + ttl: 0, + namespace: "synthetic-cache", + }; + await user.click(screen.getByRole("button", { name: "Test Connection" })); + await waitFor(() => expect(testCacheConnectionCall).toHaveBeenCalledWith("sk-test", expected)); + await save(user); + await waitFor(() => + expect(updateCacheSettingsCall).toHaveBeenCalledWith("sk-test", { ...expected, type: "redis-semantic" }), + ); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx index d85d26a21a8..06e332d2cfc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -170,8 +170,8 @@ interface StartFormValidityInputs { models: string[]; routerNames: string[]; direction: ShadowEvalDirection; - baselineModel: string; - judgeModel: string; + baselineModel: string | null; + judgeModel: string | null; percentage: string; maxBudget: string; } @@ -181,13 +181,13 @@ const startFormValidity = (inputs: StartFormValidityInputs) => { const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; const parsedMaxBudget = Number.parseFloat(inputs.maxBudget); const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; - const baselinePicked = inputs.direction === "forward" || inputs.baselineModel !== ""; + const baselinePicked = inputs.direction === "forward" || Boolean(inputs.baselineModel); const targetsPicked = inputs.apiKeyIds.length + inputs.teamIds.length + inputs.userIds.length > 0; const routerCountValid = inputs.routerNames.length >= 1 && inputs.routerNames.length <= MAX_ROUTERS; const routersMatchDirection = inputs.direction === "forward" || inputs.routerNames.length === 1; const routersValid = routerCountValid && routersMatchDirection; const scopeValid = routersValid && (inputs.direction === "reverse" || inputs.models.length <= MAX_MODELS); - const modelsPicked = scopeValid && inputs.judgeModel !== "" && baselinePicked; + const modelsPicked = scopeValid && Boolean(inputs.judgeModel) && baselinePicked; const filled = targetsPicked && modelsPicked; const boundsValid = percentageValid && maxBudgetValid; const valid = Boolean(inputs.accessToken) && filled && boundsValid; @@ -201,7 +201,7 @@ interface StartBodyInputs { models: string[]; routerNames: string[]; direction: ShadowEvalDirection; - baselineModel: string; + baselineModel: string | null; shadowPercentage: number; durationDays: number; maxBudget: number; @@ -215,7 +215,7 @@ const buildStartBody = (inputs: StartBodyInputs) => ({ models: inputs.direction === "forward" ? inputs.models : [], router_names: inputs.routerNames, direction: inputs.direction, - ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel } : {}), + ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel ?? undefined } : {}), shadow_percentage: inputs.shadowPercentage, duration_days: inputs.durationDays, max_budget: inputs.maxBudget, @@ -230,10 +230,10 @@ export const StartForm: React.FC = () => { const [models, setModels] = useState([]); const [routerNames, setRouterNames] = useState([]); const [direction, setDirection] = useState("forward"); - const [baselineModel, setBaselineModel] = useState(""); + const [baselineModel, setBaselineModel] = useState(null); const [percentage, setPercentage] = useState("10"); const [durationDays, setDurationDays] = useState("7"); - const [judgeModel, setJudgeModel] = useState(""); + const [judgeModel, setJudgeModel] = useState(null); const [maxBudget, setMaxBudget] = useState("10"); const { data: autoRouters } = useAutoRouters(); const configuredGroups = usePlainModelGroups(); @@ -286,6 +286,7 @@ export const StartForm: React.FC = () => { }; const { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid } = startFormValidity(validityInputs); const handleStart = () => { + if (!valid || !judgeModel) return; const bodyInputs: StartBodyInputs = { apiKeyIds, teamIds, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx index f3bd74260ad..b0612412059 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx @@ -15,7 +15,7 @@ const generateId = () => `entry-${Date.now()}-${Math.random().toString(36).subst const createDefaultEntry = (): ModelEntry => ({ id: generateId(), - model: "", + model: null, input_tokens: 1000, output_tokens: 500, num_requests_per_day: undefined, @@ -28,7 +28,7 @@ const PricingCalculator: React.FC = ({ accessToken, mode const { debouncedFetchForEntry, removeEntry, getMultiModelResult } = useMultiCostEstimate(accessToken); const handleEntryChange = useCallback( - (id: string, field: keyof ModelEntry, value: string | number | undefined) => { + (id: string, field: keyof ModelEntry, value: string | number | null | undefined) => { setEntries((prev) => { const updated = prev.map((entry) => (entry.id === id ? { ...entry, [field]: value } : entry)); const changedEntry = updated.find((e) => e.id === id); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts index 250857a74f5..859d2f3e8d9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts @@ -4,7 +4,7 @@ export interface PricingCalculatorProps { } export interface PricingFormValues { - model: string; + model: string | null; input_tokens: number; output_tokens: number; num_requests_per_day?: number; @@ -13,7 +13,7 @@ export interface PricingFormValues { export interface ModelEntry { id: string; - model: string; + model: string | null; input_tokens: number; output_tokens: number; num_requests_per_day?: number; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 1de4e697f64..d45cfc3fe7d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -48,7 +48,10 @@ const GUARDRAIL_MODES = [ ] as const; const submitGuardrailSchema = z.object({ - team_id: z.string().min(1, "Select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Select a team" }).min(1, "Select a team")), guardrail_name: z.string().min(1, "Enter a guardrail name"), mode: z.string().min(1, "Select a mode"), api_base: z.string().min(1, "Enter the API base URL").refine(isValidUrl, "Must be a valid URL"), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 7785a8e44ab..d0afc896260 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -318,4 +318,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + conduct: { + provider: "Conduct", + guardrailNameSuggestion: "Conduct Guard", + mode: "pre_call", + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index 1e486639840..9a9ab3a61d7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -28,6 +28,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { repelloai: "repelloai.png", straiker: "straiker.svg", alice: "alice.svg", + conduct: "conduct.png", }; describe("guardrail_garden_data logos", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index 931b3a111d8..165bd8f9967 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -474,6 +474,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Content Moderation", "Prompt Injection", "PII", "Policy"], providerKey: "Alice", }, + { + id: "conduct", + name: "Conduct Guard", + description: + "Conduct Guard evaluates prompts against workspace rules before the model call: prompt injection, PII, and custom policies, with block, warning, and approval verdicts.", + category: "partner", + logo: guardrailLogoMap["Conduct Guard"], + tags: ["Security", "Prompt Injection", "PII", "Policy"], + providerKey: "Conduct", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index f686ff5644a..fb3cf8f309a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -1,6 +1,7 @@ import aimSecurityLogo from "../../../../../public/assets/logos/aim_security.jpeg"; import aktoLogo from "../../../../../public/assets/logos/akto.svg"; import aliceLogo from "../../../../../public/assets/logos/alice.svg"; +import conductLogo from "../../../../../public/assets/logos/conduct.png"; import aporiaLogo from "../../../../../public/assets/logos/aporia.png"; import bedrockLogo from "../../../../../public/assets/logos/bedrock.svg"; import catoNetworksLogo from "../../../../../public/assets/logos/cato_networks.svg"; @@ -85,6 +86,7 @@ export const guardrail_provider_map: Record = { QostodianNexus: "qostodian_nexus", Repelloai: "repelloai", Alice: "alice", + Conduct: "conduct", }; // Function to populate provider map from API response - updates the original map @@ -208,6 +210,7 @@ export const guardrailLogoMap = { "RepelloAI Argus": repelloAiLogo.src, Straiker: straikerLogo.src, Alice: aliceLogo.src, + "Conduct Guard": conductLogo.src, } satisfies Record; export const getGuardrailLogo = (displayName: string): string | undefined => diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx index 90d38642a5d..48bac2a4dd9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx @@ -79,13 +79,16 @@ const ToolArgumentControl: React.FC<{ return ( @@ -108,8 +111,8 @@ const ToolArgumentControl: React.FC<{ if (prop.type === "boolean") { return ( onChange(setting.field_name, newValue ?? "")} - > + persist(ANTHROPIC_PROMPT_CACHING_TTL, newValue ?? "")} + value={ttlSetting.field_value ?? null} + onValueChange={(newValue) => persist(ANTHROPIC_PROMPT_CACHING_TTL, newValue)} > @@ -209,9 +206,11 @@ const GeneralSettings: React.FC = ({ accessToken, user return; } - let fieldValue = generalSettings.find((setting) => setting.field_name === fieldName)?.field_value; + const setting = generalSettings.find((setting) => setting.field_name === fieldName); + const fieldValue = setting?.field_value; - if (fieldValue == null || fieldValue == undefined) { + if (fieldValue == null) { + if (setting?.field_type === "Select") handleResetField(fieldName); return; } try { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.integration.test.tsx index b7962321333..dac43885da4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.integration.test.tsx @@ -1,5 +1,4 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "@/components/networking"; @@ -23,20 +22,16 @@ const providers = [ { provider_name: "tavily", ui_friendly_name: "Tavily Search" }, ]; -const renderModal = () => { - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); - return render( - - - , +const renderModal = () => + renderWithProviders( + , ); -}; const pickProvider = async (user: ReturnType, label: string) => { await user.click(screen.getAllByRole("combobox")[0]); @@ -45,6 +40,7 @@ const pickProvider = async (user: ReturnType, label: str describe("CreateSearchTools submit payload", () => { beforeEach(() => { + testQueryClient.clear(); vi.clearAllMocks(); vi.mocked(networking.fetchAvailableSearchProviders).mockResolvedValue({ providers }); vi.mocked(networking.createSearchTool).mockResolvedValue({ search_tool_id: "st-1" }); @@ -145,4 +141,24 @@ describe("CreateSearchTools submit payload", () => { ).toBeInTheDocument(); expect(networking.createSearchTool).not.toHaveBeenCalled(); }); + + it("should block creation after clearing the required provider and accept a restored choice", async () => { + const user = userEvent.setup(); + renderModal(); + fireEvent.change(await screen.findByLabelText(/Search Tool Name/), { target: { value: "synthetic-search" } }); + await pickProvider(user, "Perplexity AI"); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(networking.fetchAvailableSearchProviders).toHaveBeenCalledTimes(1); + await user.click(screen.getByRole("button", { name: "Add Search Tool" })); + expect(await screen.findByText("Please select a search provider")).toBeInTheDocument(); + expect(networking.createSearchTool).not.toHaveBeenCalled(); + await pickProvider(user, "Tavily Search"); + await user.click(screen.getByRole("button", { name: "Add Search Tool" })); + await waitFor(() => + expect(networking.createSearchTool).toHaveBeenCalledWith("test-token", { + search_tool_name: "synthetic-search", + litellm_params: { search_provider: "tavily" }, + }), + ); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx index a724d979af5..eb778726f22 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx @@ -65,7 +65,10 @@ const createSearchToolShape = { .string() .min(1, "Please enter a search tool name") .regex(/^[a-zA-Z0-9_-]+$/, "Name can only contain letters, numbers, hyphens, and underscores"), - search_provider: z.string().min(1, "Please select a search provider"), + search_provider: z + .string() + .nullable() + .pipe(z.string({ error: "Please select a search provider" }).min(1, "Please select a search provider")), api_key: z.string().optional(), description: z.string().optional(), }; @@ -74,7 +77,7 @@ const createSearchToolSchema = z.object(createSearchToolShape); type CreateSearchToolFormValues = z.infer; -const EMPTY_VALUES: CreateSearchToolFormValues = { search_tool_name: "", search_provider: "" }; +const EMPTY_VALUES: z.input = { search_tool_name: "", search_provider: null }; const labelWithHint = (label: string, hint: string): React.ReactNode => ( <> @@ -216,8 +219,8 @@ const CreateSearchTool: React.FC = ({ onChange(provider ?? "")} + value={value} + onValueChange={onChange} > = ({ placeholder="Select a search provider" className="h-10 w-full rounded-lg" disabled={isLoadingProviders} - showClear={value !== ""} + showClear={value != null && value !== ""} /> No matching search providers @@ -326,7 +329,7 @@ const CreateSearchTool: React.FC = ({ = ({ visible, onClose, accessT label={labelWithHint("Category (Optional)", "Select a category or enter a custom one")} > {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( - onChange(category ?? "")} - > + No matching categories diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx index 269f3b0af39..0239844851f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx @@ -20,7 +20,12 @@ import { useZodForm } from "@/lib/forms/useZodForm"; import { fetchClient } from "@/lib/http/api"; import { buildBody, settingsToForm, type DefaultInternalUserParams, type InternalUserSettings } from "./mapper"; -import { defaultUserSettingsSchema, EMPTY_TEAM_ROW, type DefaultUserSettingsFormValues } from "./schema"; +import { + defaultUserSettingsSchema, + EMPTY_TEAM_ROW, + type DefaultUserSettingsFormValues, + type DefaultUserSettingsSubmitValues, +} from "./schema"; const NO_RESET = "never"; @@ -63,7 +68,7 @@ interface RoleOption { description: string; } -type SettingsControl = Control; +type SettingsControl = Control; const TeamPickerField = ({ control, index }: { control: SettingsControl; index: number }) => { const [search, setSearch] = React.useState(""); @@ -233,7 +238,7 @@ const SettingsForm = ({ initialValues, roleOptions, updateSettings, onCancel, on const { isDirty } = form.formState; const mutation = useMutation({ - mutationFn: (values: DefaultUserSettingsFormValues) => updateSettings(buildBody(values)), + mutationFn: (values: DefaultUserSettingsSubmitValues) => updateSettings(buildBody(values)), onSuccess: (_result, values) => { toast.success("Default user settings updated successfully"); queryClient.invalidateQueries({ queryKey: SETTINGS_QUERY_KEY }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts index e8b350332c8..f4a7f001480 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import { buildBody, settingsToForm } from "./mapper"; -import type { DefaultUserSettingsFormValues } from "./schema"; +import type { DefaultUserSettingsSubmitValues } from "./schema"; const CONFIGURED_SETTINGS = { user_role: "internal_user", @@ -59,13 +59,13 @@ describe("settingsToForm", () => { it("degrades an unrecognisable team entry to a blank row instead of throwing", () => { expect(settingsToForm({ teams: [{ max_budget_in_team: 5 }, 7] }).teams).toStrictEqual([ - { team_id: "", max_budget_in_team: "", user_role: "user" }, - { team_id: "", max_budget_in_team: "", user_role: "user" }, + { team_id: null, max_budget_in_team: "", user_role: "user" }, + { team_id: null, max_budget_in_team: "", user_role: "user" }, ]); }); }); -const formValues = (overrides: Partial = {}): DefaultUserSettingsFormValues => ({ +const formValues = (overrides: Partial = {}): DefaultUserSettingsSubmitValues => ({ user_role: "internal_user", max_budget: "100", budget_duration: "30d", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts index 50365081afb..ac528542ae2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts @@ -2,7 +2,12 @@ import { z } from "zod/v4"; import type { components } from "@/lib/http/schema"; -import { EMPTY_TEAM_ROW, type DefaultTeamRowValues, type DefaultUserSettingsFormValues } from "./schema"; +import { + EMPTY_TEAM_ROW, + type DefaultTeamRowValues, + type DefaultUserSettingsFormValues, + type DefaultUserSettingsSubmitValues, +} from "./schema"; export type InternalUserSettings = components["schemas"]["InternalUserSettingsResponse"]; export type DefaultInternalUserParams = components["schemas"]["DefaultInternalUserParams"]; @@ -60,13 +65,13 @@ const textOrNull = (raw: string): string | null => (raw.trim() === "" ? null : r const listOrNull = (items: readonly T[]): T[] | null => (items.length === 0 ? null : [...items]); -const toTeamBody = (team: DefaultTeamRowValues): DefaultTeamBody => ({ +const toTeamBody = (team: DefaultUserSettingsSubmitValues["teams"][number]): DefaultTeamBody => ({ team_id: team.team_id, max_budget_in_team: numberOrNull(team.max_budget_in_team), user_role: team.user_role, }); -export const buildBody = (values: DefaultUserSettingsFormValues): DefaultInternalUserParams => ({ +export const buildBody = (values: DefaultUserSettingsSubmitValues): DefaultInternalUserParams => ({ user_role: asDefaultUserRole(values.user_role), max_budget: numberOrNull(values.max_budget), budget_duration: textOrNull(values.budget_duration), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts index 7309e3745da..cd66c3addc8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts @@ -10,14 +10,17 @@ const amountOrEmpty = z ); const defaultTeamRowSchema = z.object({ - team_id: z.string().min(1, "Select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Select a team" }).min(1, "Select a team")), max_budget_in_team: amountOrEmpty, user_role: z.enum(["user", "admin"]), }); -export type DefaultTeamRowValues = z.output; +export type DefaultTeamRowValues = z.input; -export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: "", max_budget_in_team: "", user_role: "user" }; +export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: null, max_budget_in_team: "", user_role: "user" }; const defaultUserSettingsShape = { user_role: z.string(), @@ -41,4 +44,5 @@ export const defaultUserSettingsSchema = z.object(defaultUserSettingsShape).supe ); }); -export type DefaultUserSettingsFormValues = z.output; +export type DefaultUserSettingsFormValues = z.input; +export type DefaultUserSettingsSubmitValues = z.output; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx index 875df74652c..8fa76a64af6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx @@ -192,7 +192,7 @@ export function UsersTable({ set("user_role", value)} + onValueChange={(value) => set("user_role", value ?? undefined)} placeholder="Select a role…" emptyText="No roles found" /> @@ -201,7 +201,7 @@ export function UsersTable({ set("team", value)} + onValueChange={(value) => set("team", value ?? undefined)} placeholder="Select a team…" emptyText="No teams found" /> diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index 0f7c356b8cc..83a72e50f5c 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -59,7 +59,7 @@ interface UISettings { interface CreateUserFormValues { user_email?: string; user_role: string; - team_id?: string; + team_id?: string | null; organization_ids?: string[]; metadata?: string; send_invite_email: boolean; diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx index ee2f42ad85b..31cd407a5e6 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../tests/test-utils"; import DeletedKeysPage from "./DeletedKeysPage"; import { useDeletedKeys, DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useDeletedKeys: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx index 7e30ef2c135..dd8007be264 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedKeysTable } from "./DeletedKeysTable"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const makeDeletedKey = (overrides: Partial = {}): DeletedKeyResponse => ({ token: "sk-1234567890abcdef", @@ -86,3 +88,19 @@ it("should show the empty state when there are no deleted keys", () => { expect(screen.getByText("No deleted keys found")).toBeInTheDocument(); }); + +it("links the owner, creator and deleter cells to their user detail pages", () => { + renderWithProviders(); + + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.getByRole("link", { name: "creator-1" })).toHaveAttribute("href", "/ui/users?user=creator-1"); + expect(screen.getByRole("link", { name: "deleter-1" })).toHaveAttribute("href", "/ui/users?user=deleter-1"); +}); + +it("leaves the default_user_id placeholder unlinked", () => { + const placeholderKey = makeDeletedKey({ user_id: "default_user_id", created_by: "default_user_id" }); + renderWithProviders(); + + expect(screen.getAllByText("default_user_id")).toHaveLength(2); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx index aa7d6380cc3..32185f867b4 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx @@ -3,8 +3,9 @@ import { ColumnDef } from "@tanstack/react-table"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; +import { DateCell, IdCell, IdentityCell, MoneyCell } from "@/components/shared/table_cells"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { userDetailHref } from "@/utils/entityLinks"; function TruncatedTextCell({ value }: { value: string | null | undefined }) { if (!value) { @@ -17,6 +18,17 @@ function TruncatedTextCell({ value }: { value: string | null | undefined }) { ); } +function UserLinkCell({ userId }: { userId: string | null | undefined }) { + if (!userId) { + return -; + } + return ( + + + + ); +} + export const getDeletedKeysTableColumns = (): ColumnDef[] => [ { id: "token", @@ -89,7 +101,7 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "User ID", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, { id: "created_at", @@ -107,7 +119,7 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "Created By", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, { id: "deleted_at", @@ -125,6 +137,6 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "Deleted By", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, ]; diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx index 952d8764463..d2d25f5e47e 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../tests/test-utils"; import DeletedTeamsPage from "./DeletedTeamsPage"; import { useDeletedTeams, DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useDeletedTeams: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx index e166f6b0d1b..837fb583f30 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx @@ -4,6 +4,8 @@ import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedTeamsTable } from "./DeletedTeamsTable"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const makeDeletedTeam = (overrides: Partial = {}): DeletedTeam => ({ team_id: "team-1", team_alias: "Test Team", @@ -81,3 +83,21 @@ it("renders the shared pagination footer with the server row count", () => { expect(screen.getByTestId("pagination-prev")).toBeEnabled(); expect(screen.getByTestId("pagination-next")).toBeDisabled(); }); + +it("links the organization and deleted by cells, leaving the deleted team id unlinked", () => { + renderWithProviders( + , + ); + + expect(screen.getByRole("link", { name: "org-1" })).toHaveAttribute("href", "/ui/organizations?org=org-1"); + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.queryByRole("link", { name: "team-1" })).not.toBeInTheDocument(); +}); + +it("leaves the default_user_id placeholder unlinked in the deleted by cell", () => { + const team = makeDeletedTeam({ deleted_by: "default_user_id", organization_id: null }); + renderWithProviders(); + + expect(screen.getByText("default_user_id")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx index e36077fd2c3..172f0417027 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx @@ -3,8 +3,20 @@ import { ColumnDef } from "@tanstack/react-table"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; +import { DateCell, IdCell, IdentityCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { orgDetailHref, userDetailHref } from "@/utils/entityLinks"; + +function EntityCell({ value, href }: { value: string | null | undefined; href: string | undefined }) { + if (!value) { + return -; + } + return ( + + + + ); +} export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ { @@ -78,7 +90,10 @@ export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ header: "Organization", size: 150, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => { + const orgId = row.original.organization_id; + return ; + }, }, { id: "deleted_at", @@ -97,15 +112,8 @@ export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ size: 120, enableSorting: false, cell: ({ row }) => { - const value = row.original.deleted_by; - if (!value) { - return -; - } - return ( - - {value} - - ); + const deletedBy = row.original.deleted_by; + return ; }, }, ]; diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx index 535694f14da..f4c9cc7c44c 100644 --- a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx @@ -1,4 +1,4 @@ -import { renderWithProviders, screen } from "../../tests/test-utils"; +import { fireEvent, renderWithProviders, screen } from "../../tests/test-utils"; import { vi } from "vitest"; import { EnvCredentialLoginWarningBanner } from "./EnvCredentialLoginWarningBanner"; import type { HealthReadinessDetailsResponse } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; @@ -23,6 +23,22 @@ const mockRole = (userRole: string) => { }; describe("EnvCredentialLoginWarningBanner", () => { + beforeEach(() => { + localStorage.clear(); + }); + + it("should hide the banner when dismissed and stay hidden on remount", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + const first = renderWithProviders(); + fireEvent.click(screen.getByRole("button", { name: "Dismiss banner" })); + expect(first.container).toBeEmptyDOMElement(); + + first.unmount(); + const second = renderWithProviders(); + expect(second.container).toBeEmptyDOMElement(); + }); + it("should warn an admin when env-credential login is enabled", () => { mockRole("Admin"); mockDetails({ status: "healthy", show_env_credential_login_warning: true }); diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx index 3a9d50011f1..a2cd531759e 100644 --- a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx @@ -1,26 +1,37 @@ "use client"; -import React from "react"; -import { TriangleAlert } from "lucide-react"; +import React, { useState } from "react"; +import { TriangleAlert, X } from "lucide-react"; +import { Button } from "@/components/ui/button"; import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; import { useAuth } from "@/contexts/AuthContext"; import { isAdminRole } from "@/utils/roles"; +const DISMISS_STORAGE_KEY = "litellm:envCredentialLoginWarningDismissed"; + export const EnvCredentialLoginWarningBanner: React.FC<{ accessToken: string | null }> = ({ accessToken }) => { const { userRole } = useAuth(); const { data: healthData } = useHealthReadinessDetails(accessToken); + const [dismissed, setDismissed] = useState( + () => typeof window !== "undefined" && localStorage.getItem(DISMISS_STORAGE_KEY) === "true", + ); - if (!isAdminRole(userRole) || !healthData?.show_env_credential_login_warning) { + if (dismissed || !isAdminRole(userRole) || !healthData?.show_env_credential_login_warning) { return null; } + const handleDismiss = () => { + localStorage.setItem(DISMISS_STORAGE_KEY, "true"); + setDismissed(true); + }; + return (
onRoutingChange(value === "" ? undefined : value)} + value={routing} + onValueChange={(value) => onRoutingChange(value ?? undefined)} placeholder="Inherit from the request's own compression guardrails" emptyText="No compression guardrails found" aria-label="Routing decision compression" @@ -74,8 +73,8 @@ const CompressionControls: React.FC = ({ value, onChan
onModelChange(value === "" ? undefined : value)} + value={model} + onValueChange={(value) => onModelChange(value ?? undefined)} placeholder="None (no compression)" emptyText="No compression guardrails found" aria-label="Model call compression" diff --git a/ui/litellm-dashboard/src/components/add_model/ModelChoiceCombobox.tsx b/ui/litellm-dashboard/src/components/add_model/ModelChoiceCombobox.tsx index 9021536f22a..e73c22189c6 100644 --- a/ui/litellm-dashboard/src/components/add_model/ModelChoiceCombobox.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ModelChoiceCombobox.tsx @@ -17,8 +17,8 @@ export interface ModelChoice { interface ModelChoiceComboboxProps { id: string; - value: string; - onChange: (value: string) => void; + value: string | null; + onChange: (value: string | null) => void; choices: ModelChoice[]; placeholder: string; ariaInvalid: true | undefined; @@ -40,7 +40,7 @@ const ModelChoiceCombobox: React.FC = ({ onChange(choice?.value ?? "")} + onValueChange={(choice: ModelChoice | null) => onChange(choice?.value ?? null)} itemToStringLabel={(choice: ModelChoice) => choice.label} isItemEqualToValue={(choice: ModelChoice, current: ModelChoice) => choice.value === current.value} > @@ -50,7 +50,7 @@ const ModelChoiceCombobox: React.FC = ({ aria-describedby={ariaDescribedBy} placeholder={placeholder} className="w-full" - showClear={value !== ""} + showClear={value != null && value !== ""} /> No models found diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx rename to ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx index ecd5abfa451..495fd207266 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx @@ -47,7 +47,7 @@ describe("RouterConfigBuilder", () => { expect(onChange).toHaveBeenCalledWith({ routes: [ expect.objectContaining({ - name: "", + name: null, utterances: [], description: "", score_threshold: 0.5, diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts new file mode 100644 index 00000000000..d681c911c46 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts @@ -0,0 +1,10 @@ +import { describe, expect, it } from "vitest"; +import { serializeRouterConfig } from "./RouterConfigBuilder"; + +describe("serializeRouterConfig", () => { + it("rejects a cleared route model and preserves selected route settings", () => { + expect(() => serializeRouterConfig({ routes: [{ name: null }] })).toThrow("Please select a model for every route"); + const config = { routes: [{ name: "model-silver", utterances: [], description: "", score_threshold: 0 }] }; + expect(JSON.parse(serializeRouterConfig(config))).toEqual(config); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx index 6dfb9f32091..9713d18872b 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx @@ -15,7 +15,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp interface Route { id: string; - model: string; + model: string | null; utterances: string[]; description: string; score_threshold: number; @@ -23,21 +23,28 @@ interface Route { interface SavedRoute { id?: string; - name?: string; - model?: string; + name?: string | null; + model?: string | null; utterances?: string[]; description?: string; score_threshold?: number; } -interface RouterConfig { +export interface RouterConfig { routes?: SavedRoute[]; } +export function serializeRouterConfig(config: RouterConfig | null): string { + if (config?.routes?.some((route) => !(route.name ?? route.model))) { + throw new Error("Please select a model for every route"); + } + return JSON.stringify(config); +} + interface RouterConfigBuilderProps { modelInfo: ModelGroup[]; - value?: RouterConfig; - onChange?: (config: any) => void; + value?: RouterConfig | null; + onChange?: (config: RouterConfig) => void; } interface UtteranceInputProps { @@ -136,7 +143,7 @@ const RouterConfigBuilder: React.FC = ({ modelInfo, va routeIds.push(id); return { id, - model: route.name || route.model || "", + model: route.name || route.model || null, utterances: route.utterances || [], description: route.description || "", score_threshold: route.score_threshold ?? 0.5, @@ -165,7 +172,7 @@ const RouterConfigBuilder: React.FC = ({ modelInfo, va const newRouteId = `route-${Date.now()}`; const updatedRoutes = [ ...routes, - { id: newRouteId, model: "", utterances: [], description: "", score_threshold: 0.5 }, + { id: newRouteId, model: null, utterances: [], description: "", score_threshold: 0.5 }, ]; setRoutes(updatedRoutes); updateConfig(updatedRoutes); diff --git a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx index 5684d32a319..0a5a9a5bfb6 100644 --- a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx +++ b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx @@ -61,7 +61,9 @@ const SemanticKeywordMatching: React.FC = ({ { + if (model !== null) onEmbeddingModelChange(model); + }} placeholder="Select an embedding model" emptyText="No embedding models found" aria-label="Embedding model" diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index d7c4a671317..632f4427a82 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -150,7 +150,10 @@ export const getSubmitBlockedReason = ( const autoRouterSchema = (requiresTeamScope: boolean) => z.object({ auto_router_name: z.string().min(1, "Auto router name is required"), - team_id: requiresTeamScope ? z.string().min(1, "Please select a team to continue") : z.string(), + team_id: z + .string() + .nullable() + .refine((teamId) => !requiresTeamScope || Boolean(teamId), "Please select a team to continue"), model_access_group: z.array(z.string()).optional(), }); @@ -158,12 +161,12 @@ type AddAutoRouterFormValues = z.infer>; const EMPTY_FORM_VALUES: AddAutoRouterFormValues = { auto_router_name: "", - team_id: "", + team_id: null, model_access_group: undefined, }; -const teamScopePayload = (requiresTeamScope: boolean, teamId: string): { team_id?: string } => - requiresTeamScope ? { team_id: teamId } : {}; +const teamScopePayload = (requiresTeamScope: boolean, teamId: string | null): { team_id?: string } => + requiresTeamScope && teamId ? { team_id: teamId } : {}; const BlockedReasonTooltip: React.FC<{ reason: string | null; children: React.ReactElement }> = ({ reason, @@ -458,7 +461,7 @@ const AddAutoRouterTab: React.FC = ({ const serverVerdict = await validateAutoRouterConfig( accessToken, complexityRouterConfigPayload as unknown as Record, - requiresTeamScope ? form.getValues("team_id") : undefined, + requiresTeamScope ? form.getValues("team_id") ?? undefined : undefined, ); const dryRunError = dryRunRejection(serverVerdict); if (dryRunError) { @@ -630,9 +633,7 @@ const AddAutoRouterTab: React.FC = ({ "Select the team this auto router belongs to. Only keys for this team will be able to call it.", )} > - {({ id, value, onChange }) => ( - onChange(next ?? "")} /> - )} + {({ id, value, onChange }) => } )} @@ -776,7 +777,7 @@ const AddAutoRouterTab: React.FC = ({ config={buildComplexityRouterConfig(complexityRouterConfigParams)} defaultModel={resolveComplexityDefaultModel(complexityRouterConfig, complexityRouterConfig.default_model)} routerName={watchedName} - teamId={requiresTeamScope ? watchedTeamId : undefined} + teamId={requiresTeamScope ? watchedTeamId ?? undefined : undefined} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts index 20d6af50d18..a3bd882243e 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -1,10 +1,23 @@ import { buildAutoRouterCompressionParams, + buildAutoRouterCompressionPatch, DEFAULT_AUTO_ROUTER_COMPRESSION, hydrateAutoRouterCompression, NO_COMPRESSION, } from "./buildAutoRouterCompression"; +describe("buildAutoRouterCompressionPatch", () => { + it.each([ + {}, + { auto_router_routing_compression: "routing-compressor" }, + { auto_router_model_compression: "model-compressor" }, + { auto_router_routing_compression: "none", auto_router_model_compression: "none" }, + { auto_router_routing_compression: "routing-compressor", auto_router_model_compression: "model-compressor" }, + ])("should preserve the exact stored fields on an untouched save: %j", (stored) => { + expect(buildAutoRouterCompressionPatch(hydrateAutoRouterCompression(stored), stored)).toEqual({}); + }); +}); + describe("buildAutoRouterCompressionParams", () => { it("omits both keys when routing was never configured", () => { expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({}); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 6f401a12865..c3503f78afb 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -30,6 +30,8 @@ export interface AutoRouterCompressionLitellmParams { auto_router_model_compression?: string; } +type AutoRouterCompressionPatch = Partial>; + export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = { routing: undefined, sameAsRouting: true, @@ -67,3 +69,18 @@ export const hydrateAutoRouterCompression = (litellmParams: { const sameAsRouting = model === routing; return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; }; + +export const buildAutoRouterCompressionPatch = ( + state: AutoRouterCompressionState, + stored: AutoRouterCompressionPatch, +): AutoRouterCompressionPatch => { + const initial = hydrateAutoRouterCompression(stored); + const modelUnchanged = state.sameAsRouting || state.model === initial.model; + if (state.routing === initial.routing && state.sameAsRouting === initial.sameAsRouting && modelUnchanged) { + return {}; + } + if (state.routing === undefined) { + return { auto_router_routing_compression: null, auto_router_model_compression: null }; + } + return buildAutoRouterCompressionParams(state); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx index bfa2df8f54d..8cf4b077f39 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx @@ -8,9 +8,9 @@ import { MountedFormField, type MountedFormValues } from "../common_components/M import { Providers } from "../provider_info_helpers"; interface LiteLLMModelNameFieldProps { - selectedProvider: Providers; + selectedProvider: string | null; providerModels: string[]; - getPlaceholder: (provider: Providers) => string; + getPlaceholder: (provider: string) => string; } const LiteLLMModelNameField: React.FC = ({ @@ -123,7 +123,7 @@ const LiteLLMModelNameField: React.FC = ({ id={control.id} value={(control.value as string | undefined) ?? ""} onBlur={control.onBlur} - placeholder={getPlaceholder(selectedProvider)} + placeholder={selectedProvider === null ? "Select a provider first" : getPlaceholder(selectedProvider)} onChange={(event) => { control.onChange(event); if (selectedProvider === Providers.Azure) { @@ -147,7 +147,7 @@ const LiteLLMModelNameField: React.FC = ({ value: "custom", }, { - label: `All ${selectedProvider} Models (Wildcard)`, + label: `All ${selectedProvider ?? "provider"} Models (Wildcard)`, value: "all-wildcard", }, ...providerModels.map((model) => ({ @@ -163,7 +163,7 @@ const LiteLLMModelNameField: React.FC = ({ value={(control.value as string | undefined) ?? ""} onChange={control.onChange} onBlur={control.onBlur} - placeholder={getPlaceholder(selectedProvider)} + placeholder={selectedProvider === null ? "Select a provider first" : getPlaceholder(selectedProvider)} /> ) } diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index e1a47e9792d..4add7b918f7 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -18,7 +18,7 @@ import { provider_map, Providers } from "../provider_info_helpers"; import { labelWithHint } from "@/components/shared/form/LabelWithHint"; interface ProviderSpecificFieldsProps { - selectedProvider: Providers; + selectedProvider: string | null; } const readTextFile = (file: File, onLoaded: (contents: string) => void) => { @@ -168,6 +168,7 @@ const ProviderSpecificFields: React.FC = ({ selecte }, [cacheEntries]); const allFields = React.useMemo(() => { + if (selectedProvider === null) return []; // First try to resolve from the in-memory cache. We support both the // enum/display-name form and the raw provider slug (e.g. "petals"). const cachedFields = diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx new file mode 100644 index 00000000000..a16f11eab60 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx @@ -0,0 +1,215 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import type { Member } from "@/components/networking"; + +import { renderWithProviders } from "../../../tests/test-utils"; +import MemberTable, { MemberTableColumn, memberRoleOptions } from "./MemberTable"; + +const MEMBERS: Member[] = [ + { user_id: "u-zed", user_email: "zed@example.com", user_alias: "Zed Ortiz", role: "user" }, + { user_id: "u-nameless", user_email: "mystery@example.com", user_alias: null, role: "user" }, + { user_id: "u-amy", user_email: "amy@example.com", user_alias: "amy chen", role: "admin" }, + { user_id: "u-bob", user_email: null, user_alias: "Bob Lee", role: "user" }, +]; + +const BUDGETS: Record = { "u-zed": 50, "u-nameless": null, "u-amy": 1000, "u-bob": 5 }; + +const budgetColumn: MemberTableColumn = { + title: "Budget", + key: "budget", + sortValue: (member) => BUDGETS[member.user_id ?? ""] ?? null, + render: (member) => {BUDGETS[member.user_id ?? ""] ?? "Unlimited"}, +}; + +const renderTable = (overrides: Partial> = {}) => { + const props = { + members: MEMBERS, + canEdit: true, + onEdit: vi.fn(), + onDelete: vi.fn(), + extraColumns: [budgetColumn], + ...overrides, + }; + renderWithProviders(); + return props; +}; + +const rowIds = (): (string | null)[] => + Array.from(document.querySelectorAll("tbody tr[data-row-id]")).map((row) => row.getAttribute("data-row-id")); + +const search = (value: string) => fireEvent.change(screen.getByTestId("datatable-search"), { target: { value } }); + +describe("MemberTable display", () => { + it("shows each member's name, falling back to a dash when there is none", () => { + renderTable(); + + expect(within(screen.getByRole("row", { name: /zed@example\.com/ })).getByText("Zed Ortiz")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /^name/i })).toBeInTheDocument(); + const cells = within(screen.getByRole("row", { name: /mystery@example\.com/ })).getAllByRole("cell"); + expect(cells[0]).toHaveTextContent("-"); + }); + + it("orders members by name with nameless members last by default", () => { + renderTable(); + + expect(rowIds()).toEqual(["u-amy", "u-bob", "u-zed", "u-nameless"]); + }); + + it("reports the full member count", () => { + renderTable(); + + expect(screen.getByText("4 Members")).toBeInTheDocument(); + }); +}); + +describe("MemberTable search", () => { + it("matches on name case-insensitively", async () => { + renderTable(); + + search("ZED"); + + await waitFor(() => expect(rowIds()).toEqual(["u-zed"])); + }); + + it("matches on email", async () => { + renderTable(); + + search("mystery@"); + + await waitFor(() => expect(rowIds()).toEqual(["u-nameless"])); + }); + + it("matches on user id", async () => { + renderTable(); + + search("u-bob"); + + await waitFor(() => expect(rowIds()).toEqual(["u-bob"])); + }); + + it("still searches names when the first member has no name", async () => { + renderTable({ members: [MEMBERS[1], MEMBERS[0]] }); + + search("ortiz"); + + await waitFor(() => expect(rowIds()).toEqual(["u-zed"])); + }); + + it("does not match on role", async () => { + renderTable(); + + search("admin"); + + await waitFor(() => expect(rowIds()).toEqual([])); + expect(screen.getByText("No members match your search or filters")).toBeInTheDocument(); + }); +}); + +describe("MemberTable sorting", () => { + it("sorts by email with missing emails last in both directions", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-user_email")); + expect(rowIds()).toEqual(["u-amy", "u-nameless", "u-zed", "u-bob"]); + + await user.click(screen.getByTestId("sort-header-user_email")); + expect(rowIds()).toEqual(["u-zed", "u-nameless", "u-amy", "u-bob"]); + }); + + it("sorts by role", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-role")); + expect(rowIds()[0]).toBe("u-amy"); + + await user.click(screen.getByTestId("sort-header-role")); + expect(rowIds()[3]).toBe("u-amy"); + }); + + it("sorts an extra column numerically by its sort value with blanks last", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-budget")); + expect(rowIds()).toEqual(["u-bob", "u-zed", "u-amy", "u-nameless"]); + + await user.click(screen.getByTestId("sort-header-budget")); + expect(rowIds()).toEqual(["u-amy", "u-zed", "u-bob", "u-nameless"]); + }); + + it("flips name order on the second click", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-user_alias")); + expect(rowIds()).toEqual(["u-zed", "u-bob", "u-amy", "u-nameless"]); + }); + + it("leaves extra columns without a sort value unsortable", () => { + renderTable({ + extraColumns: [{ title: "Rate Limits", key: "rate_limits", render: () => No Limits }], + }); + + expect(screen.getByRole("columnheader", { name: "Rate Limits" })).toBeInTheDocument(); + expect(screen.queryByTestId("sort-header-rate_limits")).not.toBeInTheDocument(); + }); +}); + +describe("MemberTable role filter", () => { + it("shows only members with the chosen role and clears on reset", async () => { + const user = userEvent.setup(); + renderTable({ roleColumnTitle: "Team Role" }); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-role")); + await user.click(await screen.findByRole("option", { name: "admin" })); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(rowIds()).toEqual(["u-amy"])); + expect(screen.getByTestId("filter-chip-role")).toHaveTextContent("Team Role"); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-drawer-reset")); + + await waitFor(() => expect(rowIds()).toHaveLength(4)); + }); + + it("offers the roles present in the roster", () => { + expect(memberRoleOptions(MEMBERS)).toEqual(["admin", "user"]); + expect(memberRoleOptions([{ user_id: "x", role: "" }])).toEqual([]); + }); +}); + +describe("MemberTable actions", () => { + it("passes the clicked member to onEdit and onDelete", async () => { + const user = userEvent.setup(); + const { onEdit, onDelete } = renderTable(); + const row = screen.getByRole("row", { name: /amy@example\.com/ }); + + await user.click(within(row).getByTestId("edit-member")); + await user.click(within(row).getByTestId("delete-member")); + + expect(onEdit).toHaveBeenCalledWith(MEMBERS[2]); + expect(onDelete).toHaveBeenCalledWith(MEMBERS[2]); + }); + + it("hides delete for members the caller excludes", () => { + renderTable({ showDeleteForMember: (member) => member.role !== "admin" }); + + expect(screen.getAllByTestId("delete-member")).toHaveLength(3); + expect( + within(screen.getByRole("row", { name: /amy@example\.com/ })).queryByTestId("delete-member"), + ).not.toBeInTheDocument(); + }); + + it("shows the empty text when there are no members at all", () => { + renderTable({ members: [], emptyText: "No members found" }); + + expect(screen.getByText("No members found")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx index 8a0b7671078..e54994c1ab7 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx @@ -1,17 +1,29 @@ -import { SimpleTooltip } from "@/components/ui/tooltip"; +import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table"; +import { Crown, Info, User, UserPlus } from "lucide-react"; +import React, { useState } from "react"; + import { Member } from "@/components/networking"; +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableSortHeader, + DataTableToolbar, +} from "@/components/shared/DataTable"; import { StatusBadge } from "@/components/shared/table_cells"; import { Button } from "@/components/ui/button"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; -import { Crown, Info, User, UserPlus } from "lucide-react"; -import React from "react"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { SimpleTooltip } from "@/components/ui/tooltip"; + import TableIconActionButton from "./IconActionButton/TableIconActionButtons/TableIconActionButton"; +export type MemberTableSortValue = string | number | null | undefined; + export interface MemberTableColumn { title: React.ReactNode; - key: React.Key; - dataIndex?: keyof Member; - render?: (value: Member[keyof Member], member: Member, index: number) => React.ReactNode; + key: string; + render: (member: Member) => React.ReactNode; + sortValue?: (member: Member) => MemberTableSortValue; } export interface MemberTableProps { @@ -27,12 +39,152 @@ export interface MemberTableProps { emptyText?: string; } -const extraColumnCell = (column: MemberTableColumn, member: Member, index: number): React.ReactNode => { - const value = column.dataIndex ? member[column.dataIndex] : undefined; - return column.render ? column.render(value, member, index) : value; +const ALL_ROLES = "all"; + +export const memberRowId = (member: Member): string => member.user_id ?? member.user_email ?? JSON.stringify(member); + +export const memberRoleOptions = (members: readonly Member[]): string[] => + Array.from(new Set(members.map((member) => member.role).filter((role) => role !== ""))).sort(); + +const isAdminRole = (role: string): boolean => { + const normalized = role.toLowerCase(); + return normalized === "admin" || normalized === "org_admin"; }; -const STICKY_ACTIONS_CLASS = "sticky right-0 w-[120px] bg-background"; +function RoleHeaderTitle({ title, tooltip }: { title: string; tooltip?: string }) { + if (tooltip === undefined) return <>{title}; + return ( + + {title} + + + + + ); +} + +const ACTIONS_COLUMN_WIDTH = 120; + +interface MemberColumnDeps { + canEdit: boolean; + onEdit: (member: Member) => void; + onDelete: (member: Member) => void; + roleColumnTitle: string; + roleTooltip?: string; + extraColumns: MemberTableColumn[]; + showDeleteForMember?: (member: Member) => boolean; +} + +const extraColumnDef = (column: MemberTableColumn): ColumnDef => { + const { sortValue } = column; + if (sortValue === undefined) { + return { + id: column.key, + header: () => {column.title}, + enableSorting: false, + enableGlobalFilter: false, + cell: ({ row }) => column.render(row.original), + }; + } + return { + id: column.key, + accessorFn: (member) => sortValue(member) ?? undefined, + header: ({ column: tableColumn }) => , + sortDescFirst: false, + sortUndefined: "last", + enableGlobalFilter: false, + cell: ({ row }) => column.render(row.original), + }; +}; + +const buildColumns = ({ + canEdit, + onEdit, + onDelete, + roleColumnTitle, + roleTooltip, + extraColumns, + showDeleteForMember, +}: MemberColumnDeps): ColumnDef[] => [ + { + id: "user_alias", + accessorFn: (member) => member.user_alias || undefined, + header: ({ column }) => , + sortingFn: "text", + sortUndefined: "last", + enableGlobalFilter: true, + meta: { title: "Name" }, + cell: ({ row }) => row.original.user_alias || -, + }, + { + id: "user_email", + accessorFn: (member) => member.user_email || undefined, + header: ({ column }) => , + sortingFn: "text", + sortUndefined: "last", + enableGlobalFilter: true, + meta: { title: "User Email" }, + cell: ({ row }) => row.original.user_email || "-", + }, + { + id: "user_id", + accessorFn: (member) => member.user_id ?? undefined, + header: "User ID", + enableSorting: false, + enableGlobalFilter: true, + cell: ({ row }) => + row.original.user_id === "default_user_id" ? ( + + ) : ( + row.original.user_id || "-" + ), + }, + { + id: "role", + accessorFn: (member) => member.role, + header: ({ column }) => ( + } /> + ), + sortingFn: "text", + filterFn: "equalsString", + enableGlobalFilter: false, + meta: { title: roleColumnTitle }, + cell: ({ row }) => ( + + {isAdminRole(row.original.role) ? : } + {row.original.role || "-"} + + ), + }, + ...extraColumns.map(extraColumnDef), + { + id: "actions", + header: "Actions", + size: ACTIONS_COLUMN_WIDTH, + enableSorting: false, + enableGlobalFilter: false, + meta: { pinned: "right" }, + cell: ({ row }) => + canEdit ? ( + + onEdit(row.original)} + /> + {(!showDeleteForMember || showDeleteForMember(row.original)) && ( + onDelete(row.original)} + /> + )} + + ) : null, + }, +]; export default function MemberTable({ members, @@ -46,90 +198,89 @@ export default function MemberTable({ showDeleteForMember, emptyText, }: MemberTableProps) { + const [globalFilter, setGlobalFilter] = useState(""); + const [columnFilters, setColumnFilters] = useState([]); + const [filtersOpen, setFiltersOpen] = useState(false); + + const columnDeps: MemberColumnDeps = { + canEdit, + onEdit, + onDelete, + roleColumnTitle, + roleTooltip, + extraColumns, + showDeleteForMember, + }; + const columns = buildColumns(columnDeps); + const roleFilterItems = [ + { value: ALL_ROLES, label: "All Roles" }, + ...memberRoleOptions(members).map((role) => ({ value: role, label: role })), + ]; + + const isNarrowed = globalFilter !== "" || columnFilters.length > 0; + return (
{members.length} Member{members.length !== 1 ? "s" : ""} - - - - User Email - User ID - - {roleTooltip ? ( - - {roleColumnTitle} - - - - - ) : ( - roleColumnTitle + + {isNarrowed ? "No members match your search or filters" : emptyText ?? "No data"} + + } + toolbar={(table) => ( + <> + setFiltersOpen(true)} + showViewOptions={false} + /> + + {({ get, set }) => ( + + + )} - - {extraColumns.map((column) => ( - {column.title} - ))} - Actions - - - - {members.length === 0 ? ( - - - {emptyText ?? "No data"} - - - ) : ( - members.map((member, memberIndex) => ( - - {member.user_email || "-"} - - {member.user_id === "default_user_id" ? ( - - ) : ( - member.user_id || "-" - )} - - - - {member.role?.toLowerCase() === "admin" || member.role?.toLowerCase() === "org_admin" ? ( - - ) : ( - - )} - {member.role || "-"} - - - {extraColumns.map((column) => ( - {extraColumnCell(column, member, memberIndex)} - ))} - - {canEdit ? ( - - onEdit(member)} - /> - {(!showDeleteForMember || showDeleteForMember(member)) && ( - onDelete(member)} - /> - )} - - ) : null} - - - )) - )} - -
+ + + )} + /> {onAddMember && canEdit && ( diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/editAutoRouterFormSchema.ts b/ui/litellm-dashboard/src/components/edit_auto_router/editAutoRouterFormSchema.ts new file mode 100644 index 00000000000..af84a0897c6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_auto_router/editAutoRouterFormSchema.ts @@ -0,0 +1,42 @@ +import { z } from "zod/v4"; + +const sharedShape = { + auto_router_name: z.string().min(1, "Auto router name is required"), + model_access_group: z.array(z.string()), +}; + +const complexityRouterShape = { + ...sharedShape, + auto_router_default_model: z + .string() + .nullable() + .transform((value) => value ?? ""), + auto_router_embedding_model: z + .string() + .nullable() + .transform((value) => value ?? ""), +}; + +const semanticRouterShape = { + ...sharedShape, + auto_router_default_model: z + .string() + .nullable() + .pipe(z.string({ error: "Default model is required" }).min(1, "Default model is required")), + auto_router_embedding_model: z + .string() + .nullable() + .pipe(z.string({ error: "Embedding model is required" }).min(1, "Embedding model is required")), +}; + +export const complexityRouterSchema = z.object(complexityRouterShape); +export const semanticRouterSchema = z.object(semanticRouterShape); + +export type EditAutoRouterFormValues = z.infer; + +export const EMPTY_FORM_VALUES: z.input = { + auto_router_name: "", + auto_router_default_model: null, + auto_router_embedding_model: null, + model_access_group: [], +}; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx rename to ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 3367c810991..277a3825d5d 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -1064,6 +1064,55 @@ describe("EditAutoRouterModal prompt compression", () => { />, ); + it("should clear both saved compression overrides when inheritance is selected", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "routing-compressor", + auto_router_model_compression: "model-compressor", + }); + + await user.click(await screen.findByText("Advanced: Compression")); + await user.click(screen.getAllByRole("button", { name: "Clear", exact: true })[0]); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(modelPatchUpdateCall).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + litellm_params: expect.objectContaining({ + model: "auto_router/complexity_router", + auto_router_routing_compression: null, + auto_router_model_compression: null, + }), + }), + "auto-1", + ), + ); + }); + + it("should discard a cancelled clear and preserve compression when the saved choice is restored", async () => { + const user = userEvent.setup(); + const stored = { auto_router_routing_compression: "none", auto_router_model_compression: "model-compressor" }; + const view = renderWithStoredCompression(stored); + + await user.click(await screen.findByText("Advanced: Compression")); + await user.click(screen.getAllByRole("button", { name: "Clear", exact: true })[0]); + await user.click(screen.getByRole("button", { name: "Cancel", exact: true })); + expect(modelPatchUpdateCall).not.toHaveBeenCalled(); + view.unmount(); + + renderWithStoredCompression(stored); + await user.click(await screen.findByText("Advanced: Compression")); + expect(screen.getByRole("combobox", { name: "Routing decision compression" })).toHaveValue("None (no compression)"); + await user.click(screen.getAllByRole("button", { name: "Clear", exact: true })[0]); + await user.click(screen.getByRole("combobox", { name: "Routing decision compression" })); + await user.click(screen.getByRole("option", { name: "None (no compression)" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()).toMatchObject(stored); + }); + it("leaves both compression keys out of an untouched save when none were stored", async () => { const user = userEvent.setup(); renderWithStoredCompression(); @@ -1075,18 +1124,24 @@ describe("EditAutoRouterModal prompt compression", () => { expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression"); }); - it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => { + it.each([ + { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "headroom-a" }, + { auto_router_routing_compression: "routing-compressor" }, + { auto_router_model_compression: "model-compressor" }, + ])("should preserve the exact stored compression fields through an untouched save: %j", async (stored) => { const user = userEvent.setup(); - renderWithStoredCompression({ - auto_router_routing_compression: "headroom-a", - auto_router_model_compression: "headroom-a", - }); + renderWithStoredCompression(stored); await user.click(await screen.findByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); - expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); - expect(savedLitellmParams()?.auto_router_model_compression).toBe("headroom-a"); + expect( + Object.fromEntries( + Object.entries(savedLitellmParams()).filter( + ([key]) => key === "auto_router_routing_compression" || key === "auto_router_model_compression", + ), + ), + ).toEqual(stored); }); it("shows a stored different-compression choice as Use a different compression, not Same", async () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 8a62e86e842..c4166692f78 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -1,5 +1,10 @@ import React, { useEffect, useMemo, useState } from "react"; -import { z } from "zod/v4"; +import { + complexityRouterSchema, + semanticRouterSchema, + EMPTY_FORM_VALUES, + type EditAutoRouterFormValues, +} from "./editAutoRouterFormSchema"; import { toast } from "@/lib/toast"; import { CircleHelp } from "lucide-react"; import { FieldGroup } from "@/components/ui/field"; @@ -13,7 +18,7 @@ import AccessGroupTagsCombobox from "../add_model/AccessGroupTagsCombobox"; import ModelChoiceCombobox, { type ModelChoice } from "../add_model/ModelChoiceCombobox"; import { modelAvailableCall, modelPatchUpdateCall, validateAutoRouterConfig } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; -import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; +import RouterConfigBuilder, { type RouterConfig, serializeRouterConfig } from "../add_model/RouterConfigBuilder"; import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; import { type ActiveTierSet, @@ -44,7 +49,7 @@ import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; import { type AutoRouterCompressionState, - buildAutoRouterCompressionParams, + buildAutoRouterCompressionPatch, DEFAULT_AUTO_ROUTER_COMPRESSION, hydrateAutoRouterCompression, } from "../add_model/buildAutoRouterCompression"; @@ -400,35 +405,6 @@ export const buildUpdatedComplexityRouterConfig = ( }; }; -const sharedShape = { - auto_router_name: z.string().min(1, "Auto router name is required"), - model_access_group: z.array(z.string()), -}; - -const complexityRouterShape = { - ...sharedShape, - auto_router_default_model: z.string(), - auto_router_embedding_model: z.string(), -}; - -const semanticRouterShape = { - ...sharedShape, - auto_router_default_model: z.string().min(1, "Default model is required"), - auto_router_embedding_model: z.string().min(1, "Embedding model is required"), -}; - -const complexityRouterSchema = z.object(complexityRouterShape); -const semanticRouterSchema = z.object(semanticRouterShape); - -type EditAutoRouterFormValues = z.infer; - -const EMPTY_FORM_VALUES: EditAutoRouterFormValues = { - auto_router_name: "", - auto_router_default_model: "", - auto_router_embedding_model: "", - model_access_group: [], -}; - const labelWithHint = (label: string, hint: string): React.ReactNode => ( <> {label} @@ -452,7 +428,7 @@ const EditAutoRouterModal: React.FC = ({ const [modelInfo, setModelInfo] = useState([]); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); - const [routerConfig, setRouterConfig] = useState(null); + const [routerConfig, setRouterConfig] = useState(null); const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [keywordTierRules, setKeywordTierRules] = useState([]); const [escalationKeywords, setEscalationKeywords] = useState([]); @@ -587,8 +563,8 @@ const EditAutoRouterModal: React.FC = ({ // Set form values form.reset({ auto_router_name: modelData.model_name, - auto_router_default_model: modelData.litellm_params?.auto_router_default_model || "", - auto_router_embedding_model: modelData.litellm_params?.auto_router_embedding_model || "", + auto_router_default_model: modelData.litellm_params?.auto_router_default_model || null, + auto_router_embedding_model: modelData.litellm_params?.auto_router_embedding_model || null, model_access_group: modelData.model_info?.access_groups || [], }); } catch (error) { @@ -679,7 +655,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, - ...buildAutoRouterCompressionParams(autoRouterCompression), + ...buildAutoRouterCompressionPatch(autoRouterCompression, modelData.litellm_params ?? {}), }; const updatedModelInfo = { ...modelData.model_info, @@ -706,7 +682,7 @@ const EditAutoRouterModal: React.FC = ({ // Prepare the updated litellm_params const updatedLitellmParams = { ...modelData.litellm_params, - auto_router_config: JSON.stringify(routerConfig), + auto_router_config: serializeRouterConfig(routerConfig), auto_router_default_model: values.auto_router_default_model, auto_router_embedding_model: values.auto_router_embedding_model || undefined, }; @@ -745,7 +721,7 @@ const EditAutoRouterModal: React.FC = ({ })(); } catch (error) { console.error("Error updating auto router:", error); - toast.fromError("Failed to update auto router configuration"); + toast.fromError(error); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx index c0c12356c6f..dbc270fa62e 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx @@ -96,10 +96,10 @@ export function BudgetFallbacksEditor({ value, onChange, availableModels }: Budg ({ label: m, value: m }))} - value={entry.primaryModel ?? ""} + value={entry.primaryModel} onValueChange={(v) => { const newFallbacks = entry.fallbackModels.filter((m) => m !== v); - updateEntry(entry.id, { primaryModel: v === "" ? null : v, fallbackModels: newFallbacks }); + updateEntry(entry.id, { primaryModel: v, fallbackModels: newFallbacks }); }} placeholder="Select model" emptyText="No models found" diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index ef5a3e15b52..0fab5555343 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -153,8 +153,8 @@ export function ModelMaxBudgetEditor({ ({ label: model, value: model }))} - value={entry.model ?? ""} - onValueChange={(model) => updateEntry(entry.id, { model: model === "" ? null : model })} + value={entry.model} + onValueChange={(model) => updateEntry(entry.id, { model })} placeholder="Select model" emptyText="No models found" disabled={!premiumUser} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx index b1174d1d37d..0da3628a76b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render, screen } from "@testing-library/react"; +import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { describe, it, expect } from "vitest"; import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "./MCPToolArgumentsForm"; @@ -10,7 +10,7 @@ const toolWith = (schema: InputSchema | string): MCPTool => const renderForm = (schema: InputSchema | string) => { const ref = React.createRef(); - render(); + renderWithProviders(); return ref; }; @@ -101,7 +101,7 @@ describe("MCPToolArgumentsForm", () => { it("resets dotted defaults and positional values when the selected tool changes", async () => { const ref = React.createRef(); - const { rerender } = render( + const { rerender } = renderWithProviders( { await expect(submit(ref)).resolves.toEqual({}); }); }); + +it("should distinguish an unset enum from empty string and retain explicit false", async () => { + const user = userEvent.setup(); + const ref = renderForm({ + type: "object", + properties: { + mode: { type: "string", enum: ["", "fast"], default: "fast" }, + active: { type: "boolean", default: true }, + }, + }); + await user.click(screen.getByRole("combobox", { name: "mode" })); + await user.click(await screen.findByRole("option", { name: "Select mode" })); + await user.click(screen.getByRole("combobox", { name: "active" })); + await user.click(await screen.findByRole("option", { name: "False" })); + await expect(submit(ref)).resolves.toEqual({ active: false }); + await user.click(screen.getByRole("combobox", { name: "mode" })); + await user.click(await screen.findByRole("option", { name: "Empty string" })); + expect(screen.getByRole("combobox", { name: "mode" })).toHaveTextContent("Empty string"); + await expect(submit(ref)).resolves.toEqual({ mode: "", active: false }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx index ab3213d9b37..7d0b7fc1bd7 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx @@ -23,6 +23,9 @@ const BOOLEAN_ITEMS = [ const isBlank = (value: unknown): boolean => value === undefined || value === null || value === ""; +const isUnsetArgument = (prop: InputSchemaProperty | undefined, value: unknown): boolean => + prop?.type === "string" && prop.enum ? value == null : isBlank(value); + const jsonErrorFor = (prop: InputSchemaProperty, value: unknown): string | null => { try { const parsed = typeof value === "string" ? JSON.parse(value) : value; @@ -45,10 +48,15 @@ const collectErrors = ( ): Record => { const entries = Object.entries(actualSchema.properties ?? {}).flatMap<[string, FieldError]>(([key, prop]) => { const value = values[key]; - const blank = isBlank(value); + const blank = isUnsetArgument(prop, value); if (actualSchema.required?.includes(key) && blank) { return [[key, { type: "required", message: requiredMessages[key] ?? `Please enter ${key}` }]]; } + if (prop.type === "string" && prop.enum) { + if (!blank && !prop.enum.includes(String(value))) { + return [[key, { type: "validate", message: `Please select a valid ${key}` }]]; + } + } if (prop.type !== "object" && prop.type !== "array") return []; if (blank) return []; const message = jsonErrorFor(prop, value); @@ -146,6 +154,7 @@ function buildDefaultValue(prop?: InputSchemaProperty, overrideDefault?: any): a } const getInitialValueForField = (prop: InputSchemaProperty): any => { + if (prop.type === "string" && prop.enum && prop.default === undefined) return null; const defaultValue = buildDefaultValue(prop); if (prop.type === "object" || prop.type === "array") { const fallback = prop.type === "array" ? [] : {}; @@ -164,7 +173,7 @@ function convertFormValues( Object.entries(values).forEach(([key, value]) => { const prop = schemaToUse.properties?.[key]; - if (prop && value !== null && value !== undefined && value !== "") { + if (prop && !isUnsetArgument(prop, value)) { switch (prop.type) { case "boolean": convertedValues[key] = value === "true" || value === true; @@ -202,7 +211,7 @@ function convertFormValues( default: convertedValues[key] = value; } - } else if (value !== null && value !== undefined && value !== "") { + } else if (!isUnsetArgument(prop, value)) { convertedValues[key] = value; } }); @@ -342,20 +351,22 @@ const MCPToolArgumentsForm = forwardRef { if (prop.type === "string" && prop.enum) { return ( - - + + {field.value === "" ? "Empty string" : undefined} + - {!required && Select {key}} + {!required && Select {key}} {prop.enum.map((v) => ( - {v} + {v === "" ? "Empty string" : v} ))} @@ -364,7 +375,11 @@ const MCPToolArgumentsForm = forwardRef + { control.onChange(value); - resetCredentialFormOnProviderChange(formAdapter, value as Providers, setSelectedProvider); + resetCredentialFormOnProviderChange(formAdapter, value, setSelectedProvider); }} /> )} diff --git a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts index 7190c539c81..e7b1e861811 100644 --- a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts +++ b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts @@ -1,5 +1,3 @@ -import { Providers } from "../provider_info_helpers"; - interface CredentialFormAdapter { getFieldValue: (field: string) => unknown; resetFields: () => void; @@ -25,8 +23,8 @@ interface CredentialFormAdapter { */ export function resetCredentialFormOnProviderChange( form: CredentialFormAdapter, - newProvider: Providers, - setSelectedProvider: (p: Providers) => void, + newProvider: string | null, + setSelectedProvider: (p: string | null) => void, ): void { const preservedName = form.getFieldValue("credential_name"); form.resetFields(); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index d7cf5b6110b..cab073dc808 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2946,6 +2946,7 @@ export interface Member { role: string; user_id: string | null; user_email?: string | null; + user_alias?: string | null; max_budget_in_team?: number | null; tpm_limit?: number | null; rpm_limit?: number | null; diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts index 37a973d5def..311e9357825 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts @@ -202,6 +202,8 @@ export const buildKeyCreatePayload = (input: KeyCreateInput): KeyPayloadResult = endpoint: input.keyOwner === "service_account" ? "service_account" : "standard", payload: { ...withoutKeys(values, dropped), + ...(values.organization_id === null && { organization_id: undefined }), + ...(values.project_id === null && { project_id: undefined }), ...(input.keyOwner === "you" && { user_id: input.userID }), ...(input.keyOwner === "agent" && { agent_id: input.selectedAgentId }), ...(input.autoRotationEnabled && { auto_rotate: true, rotation_interval: input.rotationInterval }), diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 2bf39bf4cd7..0d5d9f5ec8d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -16,7 +16,7 @@ const state = vi.hoisted(() => ({ can: {} as Record, uiSettings: {} as Record, tags: {} as Record, - teams: [] as { team_id: string; team_alias: string; models: string[] }[], + teams: [] as { team_id: string; team_alias: string; models: string[]; organization_id?: string }[], organizations: [] as { organization_id: string; organization_alias: string }[], accessGroups: [] as { access_group_id: string; access_group_name: string }[], projects: [] as { project_id: string; project_alias: string; team_id?: string; models?: string[] }[], @@ -806,6 +806,36 @@ describe("CreateKey", () => { expect((await createdPayload()).organization_id).toBe("org-1"); }); + it("discards the old project and team when the organization changes", async () => { + state.uiSettings = { enable_projects_ui: true }; + state.organizations = [ + { organization_id: "scope-silver", organization_alias: "Silver" }, + { organization_id: "scope-copper", organization_alias: "Copper" }, + ]; + state.teams = [{ team_id: "group-maple", team_alias: "Maple", organization_id: "scope-silver", models: [] }]; + state.projects = [{ project_id: "project-orbit", project_alias: "Orbit", team_id: "group-maple", models: [] }]; + await openModal({ teams: state.teams as Team[] }); + await nameTheKey(); + await userEvent.click(await screen.findByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Silver/ })); + await userEvent.click(await screen.findByLabelText("Project")); + await userEvent.click(await screen.findByRole("option", { name: /Orbit/ })); + await waitFor(() => expect(screen.getByLabelText("Team")).toHaveValue("Maple")); + expect(screen.getByLabelText("Team")).toBeDisabled(); + + await userEvent.click(screen.getByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Copper/ })); + expect(screen.getByLabelText("Project")).toHaveValue(""); + expect(screen.getByLabelText("Team")).toHaveValue(""); + expect(screen.getByLabelText("Team")).toBeEnabled(); + await submit(); + + const payload = JSON.parse(JSON.stringify(await createdPayload())); + expect(payload.organization_id).toBe("scope-copper"); + expect(payload.team_id).toBeNull(); + expect(payload).not.toHaveProperty("project_id"); + }); + it("drops organization_id when the chosen organization is cleared again", async () => { state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; await openModal(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 831cb0cf6a6..82580fd4667 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -592,35 +592,35 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; const changeOrganization = (write: FieldWrite) => (orgId: string | null) => { - write(orgId ?? undefined); + write(orgId); setSelectedOrganizationId(orgId); // Clear team and project when org changes setSelectedCreateKeyTeam(null); setSelectedProjectId(null); - form.setValue("team_id", undefined); - form.setValue("project_id", undefined); + form.setValue("team_id", null); + form.setValue("project_id", null); }; const selectTeam = (team: Team | null) => { setSelectedCreateKeyTeam(team); setSelectedProjectId(null); - form.setValue("project_id", undefined); + form.setValue("project_id", null); // Auto-populate org from team for non-admin users if (team?.organization_id) { setSelectedOrganizationId(team.organization_id); form.setValue("organization_id", team.organization_id); } else if (!team) { setSelectedOrganizationId(null); - form.setValue("organization_id", undefined); + form.setValue("organization_id", null); } }; - const changeProject = (write: FieldWrite) => (projectId: string) => { + const changeProject = (write: FieldWrite) => (projectId: string | null) => { write(projectId); if (!projectId) { setSelectedProjectId(null); setSelectedCreateKeyTeam(null); - form.setValue("team_id", undefined); + form.setValue("team_id", null); return; } setSelectedProjectId(projectId); @@ -756,8 +756,8 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp inputId="create-key-agent" placeholder="Select an agent" emptyText="No agents found" - value={selectedAgentId ?? undefined} - onValueChange={(value) => setSelectedAgentId(value === "" ? null : value)} + value={selectedAgentId} + onValueChange={setSelectedAgentId} options={agentsList.map((a) => ({ label: a.agent_name || a.agent_id, value: a.agent_id, @@ -783,7 +783,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ return
Organization not found
; } + const orgMemberById = new Map((orgData.members || []).map((m) => [m.user_id, m])); + const orgMemberFor = (record: Member) => (record.user_id != null ? orgMemberById.get(record.user_id) : undefined); + const orgExtraColumns: MemberTableColumn[] = [ { title: "Spend (USD)", key: "spend", - render: (_: unknown, record: Member) => { - const orgMember = - record.user_id != null ? (orgData.members || []).find((m) => m.user_id === record.user_id) : undefined; - return ; - }, + sortValue: (record: Member) => orgMemberFor(record)?.spend ?? null, + render: (record: Member) => , }, { title: "Created At", key: "created_at", - render: (_: unknown, record: Member) => { - const orgMember = - record.user_id != null ? (orgData.members || []).find((m) => m.user_id === record.user_id) : undefined; - return {orgMember?.created_at ? new Date(orgMember.created_at).toLocaleString() : "-"}; + sortValue: (record: Member) => orgMemberFor(record)?.created_at ?? null, + render: (record: Member) => { + const createdAt = orgMemberFor(record)?.created_at; + return {createdAt ? new Date(createdAt).toLocaleString() : "-"}; }, }, ]; @@ -252,10 +252,12 @@ const OrganizationInfoView: React.FC = ({
({ role: m.user_role || "", user_id: m.user_id, user_email: m.user_email, + user_alias: m.user?.user_alias ?? null, }))} canEdit={canEditOrg} onEdit={(member) => { diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 72ec5990557..de83cd00790 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -453,7 +453,7 @@ export const getPlaceholder = (selectedProvider: string): string => { return providerPlaceholderMap[resolvedProvider] ?? "gpt-3.5-turbo"; }; -export const getProviderModels = (provider: Providers, modelMap: any): Array => { +export const getProviderModels = (provider: string, modelMap: any): Array => { let providerKey = provider; let custom_llm_provider = provider_map[providerKey]; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 162fc39a632..8ed8e392ae1 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -390,6 +390,28 @@ describe("DataTable filtering", () => { expect(names()).toEqual(["Alice"]); }); + it("client global filter searches an opted-in column even when the first row has no value", () => { + const nicknameColumns: ColumnDef[] = [ + ...nameEmailColumns, + { + id: "nickname", + accessorFn: (row) => (row.id === "b" ? "Bobby" : undefined), + enableGlobalFilter: true, + header: "Nickname", + }, + ]; + render( + , + ); + expect(names()).toEqual(["Bob"]); + }); + it("server mode never filters locally even when columnFilters is set", () => { render( (columns: ColumnDef[]): Colu return { left: collect("left"), right: collect("right") }; } +function columnCanGlobalFilter(firstRow: TData | undefined, column: Column): boolean { + if (column.columnDef.enableGlobalFilter === true) return true; + if (firstRow === undefined || column.accessorFn === undefined) return false; + const firstValue: unknown = column.accessorFn(firstRow, 0); + return typeof firstValue === "string" || typeof firstValue === "number"; +} + function buildRowModels( sortingMode: SortingMode, paginationMode: PaginationMode, @@ -516,6 +523,7 @@ function useDataTableInstance( onRowSelectionChange: rowSelectionState.onChange, onColumnVisibilityChange: setColumnVisibility, onColumnSizingChange: setColumnSizing, + getColumnCanGlobalFilter: (column) => columnCanGlobalFilter(data[0], column), getCoreRowModel: getCoreRowModel(), ...buildRowModels(sortingMode, paginationMode, filterMode, expansionGuard), ...(getRowId !== undefined ? { getRowId } : {}), diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx similarity index 89% rename from ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx rename to ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx index 04c8d3e9012..2b948ca8420 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -51,7 +51,7 @@ describe("PaginatedSearchSelect", () => { const onSearchChange = vi.fn(); function Controlled() { - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( { expect(onSearchChange).not.toHaveBeenCalled(); }); - it("still reports a cleared input so the unfiltered page comes back", async () => { + it("should keep a cleared selection empty after a late page arrives and reset the query", async () => { const user = userEvent.setup(); const onSearchChange = vi.fn(); - renderSelect({ onSearchChange, value: "alias-alpha" }); - - await user.click(document.querySelector('[data-slot="combobox-clear"]') as HTMLElement); - - await waitFor(() => expect(onSearchChange).toHaveBeenCalledWith("")); + const onValueChange = vi.fn(); + function Controlled({ options }: { options: SearchSelectOption[] }) { + const [value, setValue] = useState("alias-alpha"); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + const { rerender } = render(); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(onValueChange).toHaveBeenLastCalledWith(null); + rerender( ({ ...option }))} />); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("")); + expect(screen.getByRole("combobox")).toHaveValue(""); + expect(screen.queryByRole("button", { name: "Clear" })).not.toBeInTheDocument(); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); + expect(onValueChange).toHaveBeenLastCalledWith("alias-beta"); }); it("requests the next page once the list is scrolled near the bottom", async () => { @@ -147,7 +166,7 @@ describe("PaginatedSearchSelect", () => { function ServerBacked() { const [search, setSearch] = useState(""); - const [value, setValue] = useState("alias-alpha"); + const [value, setValue] = useState("alias-alpha"); const freshlyBuiltOptions = OPTIONS.filter((option) => option.label.includes(search)).map((option) => ({ ...option, })); @@ -236,7 +255,7 @@ describe("PaginatedSearchSelect", () => { function Refetching() { const [options, setOptions] = useState([{ label: "Beta Team", value: "team-2" }]); - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( <> { function ServerBacked() { const [search, setSearch] = useState(""); - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( option.label.includes(search))} diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index 4b7ef7401c7..1314afdf2bf 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -17,8 +17,8 @@ import { usePaginatedCombobox } from "./usePaginatedCombobox"; interface PaginatedSearchSelectProps { options: SearchSelectOption[]; - value?: string; - onValueChange: (value: string) => void; + value?: string | null; + onValueChange: (value: string | null) => void; onSearchChange: (query: string) => void; onLoadMore?: () => void; hasNextPage?: boolean; @@ -86,7 +86,7 @@ export function PaginatedSearchSelect({ }; const selected = useMemo(() => { - if (value === undefined || value === "") return null; + if (value == null || value === "") return null; return ( options.find((option) => option.value === value) ?? (pickedOption?.value === value ? pickedOption : { label: value, value }) @@ -118,7 +118,7 @@ export function PaginatedSearchSelect({ inputValue={typedQuery ?? selected?.label ?? ""} onValueChange={(item: SearchSelectOption | null) => { setPickedOption(item); - onValueChange(item?.value ?? ""); + onValueChange(item?.value ?? null); }} onInputValueChange={(next, eventDetails) => handleTypedInput(next, eventDetails.reason)} onOpenChange={(nextOpen, eventDetails) => handleOpenChange(nextOpen, eventDetails.reason)} @@ -139,7 +139,7 @@ export function PaginatedSearchSelect({ onKeyDown={snapshotWholeSelection} onPaste={snapshotWholeSelection} placeholder={placeholder} - showClear={value !== undefined && value !== ""} + showClear={value != null && value !== ""} className={`w-full ${className ?? ""}`} /> diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx similarity index 76% rename from ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx rename to ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index 62a510268dd..c981010dff9 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -1,4 +1,5 @@ -import { fireEvent, render, screen } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen } from "../../../tests/test-utils"; +import { useState } from "react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; @@ -36,11 +37,30 @@ describe("SearchSelect", () => { expect(screen.getByRole("combobox")).toHaveValue("Growth"); }); - it("shows a clear control only when a value is selected", () => { - const { rerender } = render(); - expect(document.querySelector('[data-slot="combobox-clear"]')).toBeNull(); - rerender(); - expect(document.querySelector('[data-slot="combobox-clear"]')).not.toBeNull(); + it("should clear to null and allow selecting again through the real control", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + function Controlled() { + const [value, setValue] = useState(null); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + render(); + expect(screen.queryByRole("button", { name: "Clear" })).not.toBeInTheDocument(); + await chooseSelectOption(user, screen.getByRole("combobox"), "Growth"); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(onValueChange).toHaveBeenLastCalledWith(null); + expect(screen.getByRole("combobox")).toHaveValue(""); + await chooseSelectOption(user, screen.getByRole("combobox"), "Data Team"); + expect(onValueChange).toHaveBeenLastCalledWith("team-3"); }); it("filters the options client-side as you type", async () => { diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx index 53a5371d72d..21d38e9458c 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -21,8 +21,8 @@ export interface SearchSelectOption { interface SearchSelectProps { options: SearchSelectOption[]; - value?: string; - onValueChange: (value: string) => void; + value?: string | null; + onValueChange: (value: string | null) => void; placeholder?: string; emptyText?: string; disabled?: boolean; @@ -51,9 +51,7 @@ export function SearchSelect({ "aria-label": ariaLabel, }: SearchSelectProps) { const selected = - value === undefined || value === "" - ? null - : options.find((option) => option.value === value) ?? { label: value, value }; + value == null || value === "" ? null : options.find((option) => option.value === value) ?? { label: value, value }; const items = selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options; @@ -61,7 +59,7 @@ export function SearchSelect({ onValueChange(item?.value ?? "")} + onValueChange={(item: SearchSelectOption | null) => onValueChange(item?.value ?? null)} isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} itemToStringLabel={(item: SearchSelectOption) => item.label} filter={matchesQuery} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx index c715210fdde..c6720ecc7b1 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx @@ -4,6 +4,10 @@ import { describe, expect, it, vi } from "vitest"; import { IdCell } from "./id_cell"; +const { routerPushMock } = vi.hoisted(() => ({ routerPushMock: vi.fn() })); + +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: routerPushMock }) })); + const { copyToClipboardMock } = vi.hoisted(() => ({ copyToClipboardMock: vi.fn() })); vi.mock("@/utils/dataUtils", async (importOriginal) => ({ @@ -83,4 +87,23 @@ describe("IdCell", () => { render(); expect(screen.getByTestId("key-id-cell")).toHaveTextContent("k-1"); }); + + it("renders the id as a link and routes client side when href is set", async () => { + const user = userEvent.setup(); + render(); + + const link = screen.getByRole("link", { name: "user-42" }); + expect(link).toHaveAttribute("href", "/ui/users?user=user-42"); + expect(link).toHaveClass("cursor-pointer"); + + await user.click(link); + expect(routerPushMock).toHaveBeenCalledWith("/ui/users?user=user-42"); + }); + + it("stays plain text when href is undefined", () => { + render(); + + expect(screen.getByText("default_user_id").tagName).toBe("SPAN"); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx index 33c7f835e64..47d0d750223 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx @@ -3,6 +3,7 @@ import { Copy } from "lucide-react"; import * as React from "react"; +import { useEntityLinkClick } from "@/components/shared/EntityLink"; import { cn } from "@/lib/cva.config"; import { copyToClipboard } from "@/utils/dataUtils"; @@ -13,6 +14,7 @@ export type IdCellVariant = "pill" | "plain"; interface IdCellProps { value: string | null | undefined; variant?: IdCellVariant; + href?: string; onClick?: (value: string) => void; copyable?: boolean; copyLabel?: string; @@ -38,6 +40,7 @@ const VARIANT_CLASS: Record export function IdCell({ value, variant = "pill", + href, onClick, copyable = false, copyLabel = "Copy ID", @@ -52,16 +55,17 @@ export function IdCell({ return {fallback}; } + const linked = !!href && !disabled; const clickable = !!onClick && !disabled; const classes = cn( VARIANT_CLASS[variant].base, - clickable && VARIANT_CLASS[variant].clickable, + (linked || clickable) && VARIANT_CLASS[variant].clickable, truncate && "block max-w-[15ch] truncate", disabled && "opacity-50", className, ); - const idElement = clickable ? ( + const unlinkedElement = clickable ? ( @@ -71,6 +75,14 @@ export function IdCell({ ); + const idElement = linked ? ( + + {value} + + ) : ( + unlinkedElement + ); + const withTooltip = ; if (!copyable) { @@ -94,3 +106,21 @@ export function IdCell({ ); } + +interface IdLinkProps extends React.ComponentPropsWithoutRef<"a"> { + href: string; + dataTestId?: string; +} + +const IdLink = React.forwardRef(function IdLink( + { href, dataTestId, children, ...props }, + ref, +) { + const handleClick = useEntityLinkClick(href); + + return ( + + {children} + + ); +}); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 988e2e082aa..8e3a17c2622 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -333,7 +333,10 @@ const teamUpdateFieldsSchema = z.object({ modelLimits: z .array( z.object({ - model: z.string().min(1, "Missing model"), + model: z + .string() + .nullable() + .refine((model) => Boolean(model), "Missing model"), tpm: z.number().nullish(), rpm: z.number().nullish(), }), @@ -1879,7 +1882,7 @@ const TeamInfoView: React.FC = ({ onChange(next === "" ? null : next)} + onValueChange={onChange} options={userOrganizations.map((org) => ({ value: org.organization_id ?? "", label: org.organization_alias || org.organization_id || "", diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 6711514bfe9..2c119bb5848 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -1,4 +1,4 @@ -import { screen, within } from "@testing-library/react"; +import { fireEvent, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -168,6 +168,27 @@ describe("TeamMembersComponent", () => { expect(table).toHaveTextContent("admin"); }); + it("clears the member search when a different team is shown", () => { + const props = { + canEditTeam: false, + handleMemberDelete: mockHandleMemberDelete, + setSelectedEditMember: mockSetSelectedEditMember, + setIsEditMemberModalVisible: mockSetIsEditMemberModalVisible, + setIsAddMemberModalVisible: mockSetIsAddMemberModalVisible, + }; + const { rerender } = renderWithProviders(); + + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "user2" } }); + expect(screen.queryByText("user1@test.com")).not.toBeInTheDocument(); + + const otherTeam = createMockTeamData({ team_id: "team-456" }); + rerender(); + + expect(screen.getByTestId("datatable-search")).toHaveValue(""); + expect(screen.getAllByText("user1@test.com").length).toBeGreaterThanOrEqual(1); + expect(screen.getAllByText("user2@test.com").length).toBeGreaterThanOrEqual(1); + }); + it("should render Add Member button", () => { renderWithProviders( ), key: "model_scope", - render: (_: unknown, record: Member) => { + render: (record: Member) => { const models = getUserAllowedModels(record.user_id); if (!models) { return (all team models); @@ -141,9 +141,8 @@ export default function TeamMemberTab({ ), key: "spend", - render: (_: unknown, record: Member) => ( - - ), + sortValue: (record: Member) => getUserCurrentCycleSpend(record.user_id), + render: (record: Member) => , }, { title: ( @@ -155,19 +154,22 @@ export default function TeamMemberTab({ ), key: "total_spend", - render: (_: unknown, record: Member) => , + sortValue: (record: Member) => getUserTotalSpend(record.user_id), + render: (record: Member) => , }, { title: "Team Member Budget (USD)", key: "budget", - render: (_: unknown, record: Member) => ( + sortValue: (record: Member) => getUserBudget(record.user_id), + render: (record: Member) => ( ), }, { title: "Budget Reset", key: "budget_reset", - render: (_: unknown, record: Member) => , + sortValue: (record: Member) => getUserBudgetReset(record.user_id), + render: (record: Member) => , }, { title: ( @@ -179,12 +181,13 @@ export default function TeamMemberTab({ ), key: "rate_limits", - render: (_: unknown, record: Member) => {getUserRateLimits(record.user_id)}, + render: (record: Member) => {getUserRateLimits(record.user_id)}, }, ]; return ( { diff --git a/ui/litellm-dashboard/src/components/templates/KeyProjectField.tsx b/ui/litellm-dashboard/src/components/templates/KeyProjectField.tsx new file mode 100644 index 00000000000..4c523b0ae01 --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/KeyProjectField.tsx @@ -0,0 +1,58 @@ +import type { Organization, Team } from "@/components/networking"; +import { isOrgAdminForAnyOrg, isProxyAdminRole } from "@/utils/roles"; +import { useId } from "react"; +import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; +import { Button } from "@/components/ui/button"; +import { Field, FieldLabel } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; + +type KeyProjectFieldProps = { + projectId: string | null | undefined; + canDetach: boolean; + pending: boolean; + disabled: boolean; + onToggle: () => void; +}; + +export function KeyProjectField({ projectId, canDetach, pending, disabled, onToggle }: KeyProjectFieldProps) { + const id = useId(); + const { data: projects } = useProjects(); + const alias = projects?.find((project) => project.project_id === projectId)?.project_alias; + const display = alias ? `${alias} (${projectId})` : projectId; + return ( + + Project + + {canDetach && ( + <> + {pending && ( +

+ The project will be removed when you save. Team, organization, and key limits will stay the same. +

+ )} + + + )} + + ); +} + +type ProjectKeyTeam = Pick & { + team_member_permissions?: string[] | null; +}; + +export function canDetachKeyProject( + team: ProjectKeyTeam | undefined, + organizations: Organization[] | undefined, + userID: string | null, + userRole: string | null, +): boolean { + if (isProxyAdminRole(userRole ?? "")) return true; + const member = team?.members_with_roles?.find((candidate) => candidate.user_id === userID); + if (member?.role === "admin") return true; + const canUpdateKey = member != null && team?.team_member_permissions?.includes("/key/update"); + const keyOrganizations = organizations?.filter((org) => org.organization_id === team?.organization_id); + return Boolean(canUpdateKey && isOrgAdminForAnyOrg(keyOrganizations, userID)); +} diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index fec8e749143..233b58b48ab 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -49,6 +49,7 @@ export interface KeyEditFormValues { skills?: string[]; organization_id?: string | null; team_id?: string | null; + project_id?: string | null; logging_settings?: unknown[]; metadata?: string; duration?: string | null; @@ -106,6 +107,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => skills: keyData.object_permission?.skills || [], organization_id: keyData.organization_id, team_id: keyData.team_id, + project_id: keyData.project_id, logging_settings: extractLoggingSettings(keyData.metadata), metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)), duration: (keyData as { duration?: string }).duration ?? "", @@ -153,6 +155,7 @@ export const keyEditFormSchema = z.object({ skills: z.custom(), organization_id: z.custom(), team_id: z.custom(), + project_id: z.string().nullable().optional(), logging_settings: z.custom(), metadata: z.custom(), duration: z.custom(), diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx similarity index 93% rename from ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx rename to ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index b2c2a381b42..cbe17b67865 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -1,12 +1,13 @@ import { fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import { KeyResponse } from "../key_team_helpers/key_list"; import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "../key_team_helpers/ModelMaxBudgetEditor"; import { getPassThroughEndpointsCall, getPoliciesList, + getUiSettings, getPromptsList, modelAvailableCall, vectorStoreListCall, @@ -22,6 +23,7 @@ vi.mock("../networking", async () => { const actual = await vi.importActual("../networking"); return { ...actual, + getUiSettings: vi.fn().mockResolvedValue({ values: { enable_projects_ui: false } }), getPromptsList: vi.fn().mockResolvedValue({ prompts: [{ prompt_id: "prompt-1" }, { prompt_id: "prompt-2" }], }), @@ -88,7 +90,11 @@ vi.mock("../common_components/RouterSettingsAccordion", async () => { vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: vi.fn().mockReturnValue({ data: [ - { organization_id: "org-1", organization_alias: "Engineering" }, + { + organization_id: "org-1", + organization_alias: "Engineering", + members: [{ user_id: "user-orbit", user_role: "org_admin" }], + }, { organization_id: "org-2", organization_alias: "Sales" }, ], isLoading: false, @@ -366,6 +372,8 @@ describe("KeyEditView", () => { beforeEach(() => { vi.clearAllMocks(); can.mockReturnValue(true); + vi.mocked(getUiSettings).mockResolvedValue({ values: { enable_projects_ui: false } }); + testQueryClient.removeQueries({ queryKey: ["uiSettings"] }); }); describe("policy and prompt fields", () => { @@ -1483,11 +1491,11 @@ describe("KeyEditView", () => { }); }); - it("submits organization_id as null after the organization is cleared", async () => { + it("clears the organization and its dependent team in the update payload", async () => { const onSubmit = vi.fn().mockResolvedValue(undefined); renderWithProviders( {}} onSubmit={onSubmit} accessToken="" @@ -1504,9 +1512,87 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null })); + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null, team_id: null })); }); - expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toHaveProperty("organization_id", null); + expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toMatchObject({ + organization_id: null, + team_id: null, + }); + }); + + it("should save an explicit project detach while keeping parents locked until the saved key changes", async () => { + vi.mocked(getUiSettings).mockResolvedValue({ values: { enable_projects_ui: true } }); + const onSubmit = vi.fn().mockResolvedValue(undefined); + const onCancel = vi.fn(); + const key = { ...MOCK_KEY_DATA, organization_id: "org-1", team_id: "group-maple", project_id: "project-orbit" }; + const team = { + team_id: "group-maple", + organization_id: "org-1", + members_with_roles: [] as { user_id: string; role: string }[], + team_member_permissions: [] as string[], + }; + const renderEditor = (keyData: KeyResponse = key, role = "Admin", editorTeam = team) => ( + + ); + const view = renderWithProviders(renderEditor()); + await userEvent.click(await screen.findByRole("button", { name: "Detach from project" })); + expect(screen.getByRole("combobox", { name: "Organization" })).toBeDisabled(); + expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled(); + await userEvent.click(screen.getByRole("button", { name: "Cancel" })); + expect(onCancel).toHaveBeenCalledOnce(); + expect(onSubmit).not.toHaveBeenCalled(); + view.rerender(renderEditor({ ...key })); + await userEvent.click(await screen.findByRole("button", { name: "Detach from project" })); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + const expectedDetach = { project_id: null, organization_id: "org-1", team_id: "group-maple", models: key.models }; + await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expectedDetach))); + expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled(); + view.rerender(renderEditor({ ...key, project_id: null })); + expect(screen.getByRole("combobox", { name: "Team ID" })).toBeEnabled(); + expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); + view.rerender(renderEditor(key, "Internal User")); + expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); + view.rerender(renderEditor(key, "Org Admin")); + expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); + const memberTeam = { ...team, members_with_roles: [{ user_id: "user-orbit", role: "user" }] }; + view.rerender(renderEditor(key, "Org Admin", memberTeam)); + expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); + const permittedTeam = { ...memberTeam, team_member_permissions: ["/key/update"] }; + view.rerender(renderEditor(key, "Org Admin", permittedTeam)); + expect(await screen.findByRole("button", { name: "Detach from project" })).toBeInTheDocument(); + const adminTeam = { ...team, members_with_roles: [{ user_id: "user-orbit", role: "admin" }] }; + view.rerender(renderEditor(key, "Internal User", adminTeam)); + expect(await screen.findByRole("button", { name: "Detach from project" })).toBeInTheDocument(); + }); + + it("keeps project key relationships locked and omits project updates when the project UI is disabled", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + expect(await screen.findByRole("combobox", { name: "Organization" })).toBeDisabled(); + expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ organization_id: "org-1", team_id: "group-maple" }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("project_id"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index e464dfe5008..d327db21a9e 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -1,6 +1,6 @@ +import { canDetachKeyProject, KeyProjectField } from "./KeyProjectField"; import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import PolicySelector from "@/components/policies/PolicySelector"; import { Button } from "@/components/ui/button"; @@ -119,17 +119,12 @@ export function KeyEditView({ const modelBudget = useModelMaxBudgetField(keyData.token, keyData.model_max_budget); const routerSettingsRef = useRef(null); const keyTypeFieldId = React.useId(); - const projectFieldId = React.useId(); const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); - const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); const hasProject = Boolean(keyData.project_id); - const projectDisplay = (() => { - if (!keyData.project_id) return null; - const project = projects?.find((p) => p.project_id === keyData.project_id); - return project?.project_alias ? `${project.project_alias} (${keyData.project_id})` : keyData.project_id; - })(); + const detachProject = hasProject && form.watch("project_id") === null; + const canDetachProject = canDetachKeyProject(team, organizations, userID, userRole); const allowedRoutesValue = form.watch("allowed_routes"); const selectedModels = (form.watch("models") as string[] | undefined) ?? []; @@ -296,7 +291,12 @@ export function KeyEditView({ values.router_settings = routerSettings; } - await onSubmit(withNormalizedEstimates(values)); + await onSubmit( + withNormalizedEstimates({ + ...values, + ...(detachProject && enableProjectsUI && canDetachProject ? { project_id: null } : {}), + }), + ); } finally { setIsKeySaving(false); } @@ -305,7 +305,7 @@ export function KeyEditView({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); setSelectedOrganizationId(orgId); - form.setValue("team_id", undefined); + form.setValue("team_id", null); }; const handleTeamChange = (setField: (value: string | null) => void, teamId: string | null) => { @@ -316,7 +316,7 @@ export function KeyEditView({ form.setValue("organization_id", selectedTeam.organization_id); } else if (!teamId) { setSelectedOrganizationId(null); - form.setValue("organization_id", undefined); + form.setValue("organization_id", null); } }; @@ -769,14 +769,15 @@ export function KeyEditView({ "Organization", "The organization this key belongs to. Selecting an organization filters the available teams.", )} + description={hasProject ? "Organization is locked because this key belongs to a project" : undefined} > {({ value, onChange, id }) => ( handleOrganizationChange(onChange, orgId)} /> )} @@ -786,15 +787,13 @@ export function KeyEditView({ control={form.control} name="team_id" label="Team ID" - description={ - enableProjectsUI && hasProject ? "Team is locked because this key belongs to a project" : undefined - } + description={hasProject ? "Team is locked because this key belongs to a project" : undefined} > {({ value, onChange, id }) => ( - + form.setValue("project_id", detachProject ? keyData.project_id : null)} + /> )} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx index fc7e34bd5e1..e97552838da 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx @@ -68,7 +68,7 @@ function TeamFilterField({ onChange(emptyToUndefined(next))} + onValueChange={(next) => onChange(next ?? undefined)} placeholder="Search or select a team" emptyText="No teams found" /> @@ -108,7 +108,7 @@ function KeyAliasFilterField({ onChange(emptyToUndefined(next))} + onValueChange={(next) => onChange(next ?? undefined)} onSearchChange={setSearch} onLoadMore={() => void fetchNextPage()} hasNextPage={hasNextPage} @@ -146,7 +146,7 @@ function ModelFilterField({ value, onChange }: { value: string; onChange: (value onChange(emptyToUndefined(next))} + onValueChange={(next) => onChange(next ?? undefined)} onSearchChange={setSearch} onLoadMore={() => void fetchNextPage()} hasNextPage={hasNextPage} @@ -191,7 +191,7 @@ function UserIdFilterField({ onChange(emptyToUndefined(next))} + onValueChange={(next) => onChange(next ?? undefined)} onSearchChange={setSearch} onLoadMore={() => void fetchNextPage()} hasNextPage={hasNextPage} @@ -236,7 +236,7 @@ function EndUserFilterField({ onChange(emptyToUndefined(next))} + onValueChange={(next) => onChange(next ?? undefined)} onSearchChange={setSearch} onLoadMore={() => void fetchNextPage()} hasNextPage={hasNextPage} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index dd455cba44a..9e844b992b2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -207,17 +207,8 @@ export interface paths { path?: never; cookie?: never; }; - /** - * Oauth Protected Resource Mcp - * @description OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. - * - * Legacy pattern: /{server_name}/mcp - * Discovery path: /.well-known/oauth-protected-resource/{server_name}/mcp - * - * This endpoint is kept for backward compatibility. New integrations should - * use the standard MCP pattern (/mcp/{server_name}) instead. - */ - get: operations["oauth_protected_resource_mcp__well_known_oauth_protected_resource_get"]; + /** Oauth Protected Resource Root */ + get: operations["oauth_protected_resource_root__well_known_oauth_protected_resource_get"]; put?: never; post?: never; delete?: never; @@ -1186,6 +1177,23 @@ export interface paths { patch?: never; trace?: never; }; + "/authorize/mcp-session": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Authorize Mcp Session */ + get: operations["authorize_mcp_session_authorize_mcp_session_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/auto_router/benchmarks": { parameters: { query?: never; @@ -4161,6 +4169,7 @@ export interface paths { * * Returns: * - worker_pid: Process ID + * - hostname: Host (the pod on Kubernetes) the worker runs on * - status: Overall health based on memory usage * - memory: Process memory usage and RAM info * - caches: Cache item counts and descriptions @@ -8080,6 +8089,7 @@ export interface paths { * - user_id: Optional[str] - User ID associated with key * - team_id: Optional[str] - Team ID associated with key * - agent_id: Optional[str] - The agent id associated with the key. + * - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. * - organization_id: Optional[str] - The organization id of the key. * - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. * - models: Optional[list] - Model_name's a user is allowed to call @@ -37982,6 +37992,11 @@ export interface components { } | null; /** Policies */ policies?: string[] | null; + /** + * Project Id + * @description Omit to retain the project, or send null to detach. Assigning a different project is not supported. + */ + project_id?: string | null; /** Prompts */ prompts?: string[] | null; /** Rotation Interval */ @@ -40285,11 +40300,9 @@ export interface operations { }; }; }; - oauth_protected_resource_mcp__well_known_oauth_protected_resource_get: { + oauth_protected_resource_root__well_known_oauth_protected_resource_get: { parameters: { - query?: { - mcp_server_name?: string | null; - }; + query?: never; header?: never; path?: never; cookie?: never; @@ -40302,16 +40315,9 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; - }; - }; - /** @description Validation Error */ - 422: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["HTTPValidationError"]; + "application/json": { + [key: string]: string | string[]; + }; }; }; }; @@ -41666,6 +41672,43 @@ export interface operations { }; }; }; + authorize_mcp_session_authorize_mcp_session_get: { + parameters: { + query: { + redirect_uri: string; + client_id: string; + state?: string; + code_challenge?: string | null; + code_challenge_method?: string | null; + response_type?: string | null; + resource?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_auto_router_benchmarks_auto_router_benchmarks_get: { parameters: { query?: { diff --git a/uv.lock b/uv.lock index 88f558221e7..9c659be658f 100644 --- a/uv.lock +++ b/uv.lock @@ -4390,6 +4390,7 @@ cli = [ { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, + { name = "tomlkit" }, ] extra-proxy = [ { name = "a2a-sdk" }, @@ -4444,6 +4445,7 @@ proxy = [ { name = "rq" }, { name = "soundfile" }, { name = "starlette" }, + { name = "tomlkit" }, { name = "uvicorn" }, { name = "uvloop", marker = "sys_platform != 'win32'" }, { name = "websockets" }, @@ -4669,6 +4671,8 @@ requires-dist = [ { name = "starlette", marker = "extra == 'proxy'", specifier = ">=1.0.1,<2.0" }, { name = "tiktoken", specifier = ">=0.8.0,<1.0" }, { name = "tokenizers", specifier = ">=0.21.0,<1.0" }, + { name = "tomlkit", marker = "extra == 'cli'", specifier = ">=0.13.3,<1.0" }, + { name = "tomlkit", marker = "extra == 'proxy'", specifier = ">=0.13.3,<1.0" }, { name = "uvicorn", marker = "extra == 'proxy'", specifier = ">=0.33.0,<1.0" }, { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.22.1,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" },