diff --git a/.github/deploy-on-aws.png b/.github/deploy-on-aws.png new file mode 100644 index 00000000000..06d41f2a5e0 Binary files /dev/null and b/.github/deploy-on-aws.png differ diff --git a/.github/deploy-on-gcp.png b/.github/deploy-on-gcp.png new file mode 100644 index 00000000000..e831a8c2e4e Binary files /dev/null and b/.github/deploy-on-gcp.png differ diff --git a/.github/workflows/guard-fork-dependencies.yml b/.github/workflows/guard-fork-dependencies.yml index bf7282688ef..f4cbdd63cdf 100644 --- a/.github/workflows/guard-fork-dependencies.yml +++ b/.github/workflows/guard-fork-dependencies.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" paths: - "uv.lock" diff --git a/.github/workflows/guard-main-branch.yml b/.github/workflows/guard-main-branch.yml index 1c1ce0de079..21aad18d298 100644 --- a/.github/workflows/guard-main-branch.yml +++ b/.github/workflows/guard-main-branch.yml @@ -31,12 +31,12 @@ jobs: echo "PR head repo: $HEAD_REPO" echo "PR head branch: $HEAD_REF" if [ "$HEAD_REPO" != "$BASE_REPO" ]; then - echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the 'litellm_oss_branch' branch instead." + echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the 'litellm_oss_staging' branch instead." exit 1 fi if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then echo "Allowed source branch." exit 0 fi - echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_oss_branch' instead." + echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_oss_staging' instead." exit 1 diff --git a/.github/workflows/osv-scan.yml b/.github/workflows/osv-scan.yml index 9dd321f88db..31104002dab 100644 --- a/.github/workflows/osv-scan.yml +++ b/.github/workflows/osv-scan.yml @@ -5,13 +5,8 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" - paths: - - uv.lock - - ui/litellm-dashboard/package-lock.json - - osv-scanner.toml - - .github/workflows/osv-scan.yml schedule: - cron: "23 6 * * *" workflow_dispatch: diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 4f09857eb1b..872a1799d98 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index de7e1b68346..f7c53ff443a 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: @@ -14,7 +14,7 @@ permissions: jobs: lint: runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: 15 steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -50,10 +50,10 @@ jobs: run: | uv sync --frozen - - name: Check Black formatting + - name: Check ruff format run: | cd litellm - uv run --no-sync black --check --exclude '/enterprise/' . + uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' . cd .. - name: Debug - Check file state @@ -87,9 +87,11 @@ jobs: run: | uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')" - - name: Run basedpyright type checking + - name: Check basedpyright budget (delta vs base) + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} run: | - (uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py + (uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA" - name: Check for circular imports run: | diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index b83119712a7..ce8d8cb9c95 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -7,7 +7,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" jobs: @@ -111,4 +111,4 @@ jobs: if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }} run: | npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true - node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json + node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json --check eslint-metrics.json diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 2ae60951afc..cefb77980ac 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-model-map.yaml b/.github/workflows/test-model-map.yaml index 49821fca3a8..b2170d9f6a4 100644 --- a/.github/workflows/test-model-map.yaml +++ b/.github/workflows/test-model-map.yaml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml new file mode 100644 index 00000000000..13e1dc4ad5e --- /dev/null +++ b/.github/workflows/test-rust.yml @@ -0,0 +1,65 @@ +name: LiteLLM Rust + +on: + push: + paths: + - "litellm-rust/**" + - ".github/workflows/test-rust.yml" + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + paths: + - "litellm-rust/**" + - ".github/workflows/test-rust.yml" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + rust-checks: + name: rustfmt, clippy, test + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: litellm-rust + env: + CARGO_TERM_COLOR: always + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Rust + run: | + rustup toolchain install stable --profile minimal --component clippy,rustfmt + rustup default stable + + - name: Cache Cargo registry and target + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo- + + - name: Check Rust formatting + run: cargo fmt --check + + - name: Run Clippy + run: cargo clippy --workspace --all-targets --locked -- -D warnings + + - name: Run Rust tests + run: cargo test --workspace --locked diff --git a/.github/workflows/test-semgrep.yml b/.github/workflows/test-semgrep.yml index 2ba23e44da8..f0dcb9887be 100644 --- a/.github/workflows/test-semgrep.yml +++ b/.github/workflows/test-semgrep.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-core-utils.yml b/.github/workflows/test-unit-core-utils.yml index da1267756cd..d6d6353238f 100644 --- a/.github/workflows/test-unit-core-utils.yml +++ b/.github/workflows/test-unit-core-utils.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index b2a8640223a..8ad9fb6a73b 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-enterprise-routing.yml b/.github/workflows/test-unit-enterprise-routing.yml index ffc09dd8f94..13136c968d1 100644 --- a/.github/workflows/test-unit-enterprise-routing.yml +++ b/.github/workflows/test-unit-enterprise-routing.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-integrations.yml b/.github/workflows/test-unit-integrations.yml index b316ad5dfdf..c95ed4e7c24 100644 --- a/.github/workflows/test-unit-integrations.yml +++ b/.github/workflows/test-unit-integrations.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-llm-providers.yml b/.github/workflows/test-unit-llm-providers.yml index 2a1912ce92d..df78564ab0c 100644 --- a/.github/workflows/test-unit-llm-providers.yml +++ b/.github/workflows/test-unit-llm-providers.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index a7363ac3b43..133c135d97a 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: @@ -32,7 +32,9 @@ jobs: tests/test_litellm/repositories tests/test_litellm/images tests/test_litellm/interactions + tests/test_litellm/ocr tests/test_litellm/passthrough + tests/test_litellm/sandbox tests/test_litellm/vector_stores tests/test_litellm/test_*.py workers: 2 diff --git a/.github/workflows/test-unit-proxy-auth.yml b/.github/workflows/test-unit-proxy-auth.yml index 99882066a8e..97dfaed6e81 100644 --- a/.github/workflows/test-unit-proxy-auth.yml +++ b/.github/workflows/test-unit-proxy-auth.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index d9b6a348b60..d4f00050596 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" workflow_dispatch: diff --git a/.github/workflows/test-unit-proxy-infra.yml b/.github/workflows/test-unit-proxy-infra.yml index 336e53ee3d7..884d62289b9 100644 --- a/.github/workflows/test-unit-proxy-infra.yml +++ b/.github/workflows/test-unit-proxy-infra.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: @@ -29,6 +29,7 @@ jobs: tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/experimental tests/test_litellm/proxy/common_utils + tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py workers: 2 reruns: 2 diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml index 5768551f9b0..922e3c2eddf 100644 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ b/.github/workflows/test-unit-proxy-legacy.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test-unit-responses-caching-types.yml b/.github/workflows/test-unit-responses-caching-types.yml index 13069be9e3a..2f177587997 100644 --- a/.github/workflows/test-unit-responses-caching-types.yml +++ b/.github/workflows/test-unit-responses-caching-types.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" permissions: diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml index 985653796c2..ac363071d55 100644 --- a/.github/workflows/test_server_root_path.yml +++ b/.github/workflows/test_server_root_path.yml @@ -7,7 +7,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" jobs: diff --git a/.gitignore b/.gitignore index fda3311fe02..3563d7c8c2d 100644 --- a/.gitignore +++ b/.gitignore @@ -123,3 +123,6 @@ crash.*.log # and should be committed. .vscode .pin_list.txt + +# pytest coverage data +.coverage diff --git a/CLAUDE.md b/CLAUDE.md index 2070b6fcdd6..b721064aaa7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -29,7 +29,7 @@ If you ever make public-facing PR descriptions, comments, issues, commit message - don't use "—". Instead, reach for ";", ".", etc. - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. - don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose -- don't add a trailing "." at the end of paragraphs (just like this file) +- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: unless there's a sentence immediately after, don't add a "." - don't use →. Instead, prefer not to use arrows, and if need be, use -> instead Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs diff --git a/Dockerfile b/Dockerfile index 4d55148ff89..af49dc8d8cf 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/Makefile b/Makefile index 27150aec938..7c74526e130 100644 --- a/Makefile +++ b/Makefile @@ -20,13 +20,13 @@ help: @echo " make install-test-deps - Install the full local test environment" @echo " make install-helm-unittest - Install helm unittest plugin" @echo " make install-hooks - Install git hooks (Conventional Commits + Branches)" - @echo " make format - Apply Black code formatting" - @echo " make format-check - Check Black code formatting (matches CI)" - @echo " make lint - Run all linting (Ruff, basedpyright, Black check, circular imports, import safety)" + @echo " make format - Apply ruff format code formatting" + @echo " make format-check - Check ruff format code formatting (matches CI)" + @echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)" @echo " make lint-ruff - Run Ruff linting only" @echo " make lint-basedpyright - Run basedpyright strict, gated by per-rule error counts" @echo " make lint-basedpyright-budget-update - Re-capture the basedpyright per-rule budget (ratchet)" - @echo " make lint-black - Check Black formatting (matches CI)" + @echo " make lint-format - Check ruff format formatting (matches CI)" @echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its ceiling" @echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)" @echo " make lint-ruff-budget-update - Re-capture per-rule baselines in ruff-strict-budget.json (ratchet)" @@ -82,11 +82,13 @@ install-hooks: ./scripts/install_git_hooks.sh # Formatting +# 88-column wrap matches the Black width the whole repo is formatted to; ruff.toml's +# global line-length is 120 (for E501/isort), so 88 is forced here. format: install-dev - cd litellm && $(UV_RUN) black . && cd .. + cd litellm && $(UV_RUN) ruff format --line-length 88 --exclude '/enterprise/' . && cd .. format-check: install-dev - cd litellm && $(UV_RUN) black --check . && cd .. + cd litellm && $(UV_RUN) ruff format --check --line-length 88 --exclude '/enterprise/' . && cd .. # Linting targets lint-ruff: install-dev @@ -125,12 +127,13 @@ lint-ruff-FULL-dev: install-dev else echo "No changed .py files to check."; fi lint-basedpyright: install-dev - ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py + git fetch origin litellm_internal_staging + ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging lint-basedpyright-budget-update: install-dev ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update -lint-black: format-check +lint-format: format-check lint-ruff-budget: install-dev $(UV_RUN) python scripts/ruff_strict_gate.py diff --git a/README.md b/README.md index b26ad39eada..3d0f7282d7c 100644 --- a/README.md +++ b/README.md @@ -6,10 +6,10 @@

Open Source AI Gateway for 100+ LLMs. Self-hosted. Enterprise-ready. Call any LLM in OpenAI format.

- Deploy to Render - - Deploy on Railway - + Deploy to Render + Deploy on Railway + Deploy on AWS + Deploy on GCP

LiteLLM Proxy Server (AI Gateway) | Hosted Proxy | Enterprise Tier | Website

@@ -406,6 +406,140 @@ You can use LiteLLM through either the Proxy Server or Python SDK. Both give you Support for more providers. Missing a provider or LLM Platform, raise a [feature request](https://github.com/BerriAI/litellm/issues/new?assignees=&labels=enhancement&projects=&template=feature_request.yml&title=%5BFeature%5D%3A+). +### Deploy on AWS or GCP with Terraform + +Run the LiteLLM proxy as a production-ready componentized stack (gateway, backend, UI on separate services; managed Postgres + Redis + object store) using the published Terraform modules. Both modules are on the [public Terraform Registry](https://registry.terraform.io/namespaces/BerriAI) — no auth needed. + +#### AWS — ECS Fargate + Aurora + ElastiCache + ALB + +[![Launch in AWS CloudShell](https://img.shields.io/badge/Launch-AWS_CloudShell-FF9900?logo=amazon-aws&logoColor=white)](https://console.aws.amazon.com/cloudshell/home) — opens an in-browser shell, already authenticated to your AWS account. Once inside, run: + +```bash +git clone https://github.com/BerriAI/litellm.git +cd litellm/terraform/litellm/aws/examples/default +cp terraform.tfvars.example terraform.tfvars # edit region/tenant/env +terraform init && terraform apply +``` + +[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/aws/latest) + +Or call the module from your own root config: + +```hcl +# main.tf +terraform { + required_version = ">= 1.6.0" + required_providers { + aws = { source = "hashicorp/aws", version = "~> 5.60" } + } +} + +provider "aws" { + region = "us-west-2" +} + +module "litellm" { + source = "BerriAI/litellm/aws" + version = "~> 1.89" + + region = "us-west-2" + azs = ["us-west-2a", "us-west-2b"] + tenant = "acme" + env = "prod" + + # Production: provide an ACM cert. Without one, set allow_plaintext_alb = true + # (dev/trial only). + # acm_certificate_arn = "arn:aws:acm:us-west-2:111122223333:certificate/..." + allow_plaintext_alb = true +} + +output "litellm_url" { + value = module.litellm.alb_dns_name +} +``` + +```bash +terraform init +terraform apply +``` + +Provider API keys live in AWS Secrets Manager; reference ARNs via `gateway_extra_secrets`. Full input list and architecture diagram on the [registry page](https://registry.terraform.io/modules/BerriAI/litellm/aws/latest?tab=inputs). + +#### GCP — Cloud Run + Cloud SQL + Memorystore + HTTPS LB + +[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.png)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true) + +Real 1-click. Opens Cloud Shell, clones this repo, and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md) — pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply. + +[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/google/latest) + +To call the module from your own config instead, Cloud Run can't pull from `ghcr.io` directly, so first set up a one-time Artifact Registry remote repo backed by GHCR: + +```bash +gcloud artifacts repositories create litellm \ + --location=us-central1 \ + --repository-format=docker \ + --mode=remote-repository \ + --remote-docker-repo=https://ghcr.io \ + --project=my-gcp-project +``` + +Then: + +```hcl +# main.tf +terraform { + required_version = ">= 1.6.0" + required_providers { + google = { source = "hashicorp/google", version = "~> 6.10" } + google-beta = { source = "hashicorp/google-beta", version = "~> 6.10" } + } +} + +provider "google" { project = "my-gcp-project"; region = "us-central1" } +provider "google-beta" { project = "my-gcp-project"; region = "us-central1" } + +module "litellm" { + source = "BerriAI/litellm/google" + version = "~> 1.89" + + project_id = "my-gcp-project" + region = "us-central1" + tenant = "acme" + env = "prod" + + # Replace my-gcp-project with your GCP project ID (same value as project_id above). + image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai" + + # Production: provide DNS already pointing at the LB IP for Google-managed certs. + # Without one, set allow_plaintext_lb = true (dev/trial only). + # lb_domains = ["proxy.example.com"] + allow_plaintext_lb = true +} + +output "litellm_url" { + value = module.litellm.load_balancer_url +} +``` + +```bash +terraform init +terraform apply +``` + +Provider API keys live in Secret Manager; reference resource IDs (e.g. `projects/my-gcp-project/secrets/openai-api-key`) via `gateway_extra_secrets`. Full input list and architecture diagram on the [registry page](https://registry.terraform.io/modules/BerriAI/litellm/google/latest?tab=inputs). + +#### Both stacks include + +- The full componentized split (gateway / backend / UI as independent services) +- Managed Postgres (writer + reader) and Redis +- Versioned object store for proxy state + file uploads +- An auto-generated `LITELLM_MASTER_KEY` in your cloud's secret manager +- A one-off migration job that runs `prisma migrate deploy` before the proxy starts +- The same `proxy_config` surface as the [Helm chart](./helm/litellm/) — pass YAML as a typed map + +The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws/) and [`terraform/litellm/gcp/`](./terraform/litellm/gcp/) in this repo; the registry entries are read-only mirrors updated on each release. + ### Run in Developer Mode #### Services 1. Setup .env file in root diff --git a/backend/Dockerfile b/backend/Dockerfile index 2cfdde8a517..667bdb073eb 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 2f65f99c292..b67f7d42127 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -84,6 +84,8 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/active/callbacks", "/callbacks", "/team_callback", + # Rust data-plane gateway → proxy control-plane API (logging today, auth later) + "/v1/rust_control_plane/", # Alerting / email / IP allowlist "/alerting/", "/email/", diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 7ba7656e407..1af0148e452 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -69,7 +69,7 @@ }, "reportMatchNotExhaustive": { "baseline": 1, - "slack": 3 + "slack": 0 }, "reportMissingParameterType": { "baseline": 3933, @@ -121,7 +121,7 @@ }, "reportReturnType": { "baseline": 126, - "slack": 13 + "slack": 100 }, "reportTypedDictNotRequiredAccess": { "baseline": 20, @@ -157,7 +157,7 @@ }, "reportUnnecessaryComparison": { "baseline": 683, - "slack": 10 + "slack": 100 }, "reportUnnecessaryContains": { "baseline": 4, diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e591a4a2adb..50ef55e3261 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index eafbd23fd90..ab02b43d0f9 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,6 +1,6 @@ # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 8486e37384e..af4870bb1a5 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -13,7 +13,9 @@ from litellm import Router, verbose_logger from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_metadata, +) from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, @@ -981,9 +983,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): target_model_names_list: List[str], ) -> OpenAIFileObject: ## GET THE FILE TYPE FROM THE CREATE FILE REQUEST - file_data = extract_file_data(create_file_request["file"]) - - file_type = file_data["content_type"] + _, file_type = extract_file_metadata(create_file_request["file"]) output_file_id = file_objects[0].id model_id = file_objects[0]._hidden_params.get("model_id") diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 19c8a10fdfe..716b2fa09d1 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/litellm-rust/.cargo/config.toml b/litellm-rust/.cargo/config.toml new file mode 100644 index 00000000000..c7fb2592542 --- /dev/null +++ b/litellm-rust/.cargo/config.toml @@ -0,0 +1,10 @@ +# PyO3 cdylib (`litellm-python-bridge`) links against the host interpreter's +# symbols, which are not present at link time when building an extension module. +# On macOS, tell the linker to resolve undefined `_Py*` symbols dynamically at +# load time (the standard pyo3 extension-module flag) so the cdylib links without +# a libpython on the link line. +[target.x86_64-apple-darwin] +rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"] + +[target.aarch64-apple-darwin] +rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"] diff --git a/litellm-rust/.gitignore b/litellm-rust/.gitignore new file mode 100644 index 00000000000..b83d22266ac --- /dev/null +++ b/litellm-rust/.gitignore @@ -0,0 +1 @@ +/target/ diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md new file mode 100644 index 00000000000..2fa81798605 --- /dev/null +++ b/litellm-rust/ADDING_A_PROVIDER.md @@ -0,0 +1,9 @@ +# Adding a provider / route to litellm-rust + +Three layers, same for every route (see `ocr` and `realtime` as references): + +1. **Transform contract (pure)** — `crates/core/src//transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth. +2. **Provider config (pure)** — `crates/providers/src///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. +3. **HTTP / transport (the host)** — `crates/providers/src/.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O. + +**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`. diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md new file mode 100644 index 00000000000..86dd2c92744 --- /dev/null +++ b/litellm-rust/AGENTS.md @@ -0,0 +1,17 @@ +# AGENTS.md + +litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers. + +## Crates + +| Crate | Role | Pure / I/O | +|-------|------|------------| +| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure | +| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding | + +Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. + +Adding a crate: default to a MODULE. New crate ONLY on 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. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md new file mode 100644 index 00000000000..7c723e570ef --- /dev/null +++ b/litellm-rust/CLAUDE.md @@ -0,0 +1,109 @@ +# CLAUDE.md + +This file defines the rules for Rust work in LiteLLM. + +## Crates (exactly three — see AGENTS.md) + +`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge` +exposes it to the Python SDK. A crate is a **layer**, not a route — add modules, not crates. + +## Core Boundary + +`litellm-core` is the pure translation layer; the `litellm-ai-gateway` host executes work. + +Route-level Rust structure mirrors LiteLLM's Python responsibilities: +- `core/src//` owns the route contract, shared types, and provider + template traits. For OCR, this means `core/src/ocr`. +- `core/src/providers///transformation.rs` owns the + provider-specific transform. For Mistral OCR, this means + `core/src/providers/mistral/ocr/transformation.rs`. +- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`), + never inside `core`. + +Allowed in `core`: +- Pure request transforms +- Pure response transforms +- Pure stream chunk normalization +- Shared data types and validation errors +- Deterministic token/cost helper logic + +Not allowed in `core`: +- Network calls +- Environment variable or secret reads +- Filesystem access +- Database or cache access +- Provider SDK signing or auth flows +- Logging callbacks, spend writes, or custom callbacks +- Global mutable runtime state + +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. + +## 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. + +## Host I/O Rules + +These rules apply when adding future crates or modules that execute network I/O, +such as `ai-gateway`, router hosts, or standalone servers: + +- 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. + +## 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 +# the ai-gateway binary + server code is behind the `server` feature +cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings +cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings +cargo test --workspace +``` + +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 new file mode 100644 index 00000000000..6fe84f1cfbc --- /dev/null +++ b/litellm-rust/Cargo.lock @@ -0,0 +1,1873 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "axum" +version = "0.7.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" +dependencies = [ + "async-trait", + "axum-core", + "base64", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "rustversion", + "serde", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sha1", + "sync_wrapper", + "tokio", + "tokio-tungstenite", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" +dependencies = [ + "async-trait", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "rustversion", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" + +[[package]] +name = "cc" +version = "1.2.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-io", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "litellm-ai-gateway" +version = "0.1.0" +dependencies = [ + "axum", + "futures-channel", + "futures-util", + "litellm-core", + "pyo3", + "reqwest", + "serde", + "serde_json", + "sha2", + "subtle", + "tokio", + "tokio-tungstenite", +] + +[[package]] +name = "litellm-core" +version = "0.1.0" +dependencies = [ + "rand 0.8.6", + "serde", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "litellm-python-bridge" +version = "0.1.0" +dependencies = [ + "litellm-ai-gateway", + "litellm-core", + "pyo3", + "serde_json", +] + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "matchit" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +dependencies = [ + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.4", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.60.2", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + +[[package]] +name = "rustls" +version = "0.23.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a" +dependencies = [ + "byteorder", + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.8.6", + "rustls", + "rustls-pki-types", + "sha1", + "thiserror 1.0.69", + "utf-8", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.75" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml new file mode 100644 index 00000000000..25ee2213040 --- /dev/null +++ b/litellm-rust/Cargo.toml @@ -0,0 +1,28 @@ +[workspace] +members = [ + "crates/core", + "crates/ai-gateway", + "crates/python-bridge", +] +resolver = "2" + +[workspace.package] +edition = "2021" +license = "MIT" +repository = "https://github.com/BerriAI/litellm" + +[workspace.dependencies] +litellm-core = { path = "crates/core" } +litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } +axum = "0.7" +pyo3 = "0.23.5" +rand = "0.8" +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sha2 = "0.10" +subtle = "2" +thiserror = "2.0" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } +tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } +futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } diff --git a/litellm-rust/README.md b/litellm-rust/README.md new file mode 100644 index 00000000000..1646c90ad76 --- /dev/null +++ b/litellm-rust/README.md @@ -0,0 +1,44 @@ +# LiteLLM Rust + +This workspace contains the staged Rust implementation for LiteLLM. + +Rust starts as a pure transform core used by the existing Python host. Python +continues to own auth, configuration, network I/O, retries, routing, logging, +callbacks, spend tracking, and customer plugins until each Rust path has parity +coverage and production evidence. + +## Crates + +| Crate | Role | Pure / I/O | +|-------|------|------------| +| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure | +| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding | + +Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. + +## Layout + +```text +crates/ + core/ Route contracts, shared pure types, errors, and templates. + src/ocr/ + providers/ Provider-specific pure transforms. + src/mistral/ocr/transformation.rs + python-bridge/ PyO3 bridge for Python LiteLLM. +``` + +The folder shape should follow the Python provider tree: +`providers/src///transformation.rs`. The bridge should expose +one function per top-level route, starting with `ocr(payload)`. + +## Checks + +Run these before pushing Rust changes. GitHub Actions runs the same checks for +changes under `litellm-rust/`. + +```bash +cargo fmt --check +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace +``` diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md new file mode 100644 index 00000000000..d9e6e1adde5 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/AGENTS.md @@ -0,0 +1,50 @@ +# ai-gateway — folder architecture + +The Axum server that fronts the Rust gateway. It owns transport + config + auth +only; deployment selection lives in `core::router`, transforms in `core`/`providers`. + +``` +src/ + main.rs # entrypoint: build AppState (router + master key), bind, serve + state.rs # AppState — shared Arc + master_key + gil.rs # GIL-activity tracker (records Python acquisitions) + 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 + gil.rs # simple route (one file): router() + GET /health/gil + 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 + python/ # Python interop (feature: python-config) — load-time only + mod.rs, config.rs, AGENTS.md +``` + +## 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. +- **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 + +Anything that calls into Python lives in `python/` and is **load-time only** — see +`python/AGENTS.md`. 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 new file mode 100644 index 00000000000..733953bbdb3 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md @@ -0,0 +1,12 @@ +# 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] +``` diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml new file mode 100644 index 00000000000..b08fb89d5e8 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -0,0 +1,42 @@ +[package] +name = "litellm-ai-gateway" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[lib] +name = "litellm_ai_gateway" + +[[bin]] +name = "litellm-ai-gateway" +path = "src/main.rs" +required-features = ["server"] + +[dependencies] +litellm-core.workspace = true +# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the +# Python proxy callbacks API. +reqwest.workspace = true +# `sync` powers the bounded mpsc channel the realtime logger drains. +tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } +tokio-tungstenite.workspace = true +futures-util.workspace = true +serde_json.workspace = true +axum = { workspace = true, features = ["ws"], optional = true } +serde = { workspace = true, optional = true } +subtle = { workspace = true, optional = true } +# sha2 hashes the master key into user_api_key_hash (matches the proxy's +# SHA-256 hash_token) so the plaintext credential never enters a log payload. +sha2 = { workspace = true, optional = true } +pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } + +[features] +default = [] +server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"] +# Build the gateway's config from the proxy YAML via an embedded Python +# interpreter (links libpython; requires `litellm` importable at runtime). +python-config = ["dep:pyo3"] + +[dev-dependencies] +futures-channel = "0.3" diff --git a/litellm-rust/crates/ai-gateway/Dockerfile b/litellm-rust/crates/ai-gateway/Dockerfile new file mode 100644 index 00000000000..adf6fca0741 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/Dockerfile @@ -0,0 +1,86 @@ +# Multi-stage build for the LiteLLM Rust AI Gateway (realtime WebSocket proxy). +# +# Build context is the **repo root** so we can install `litellm` from this repo's +# source (the gateway loads its model_list via litellm.proxy.read_model_list, +# which is not in any PyPI release yet) AND build the rust workspace under +# litellm-rust/. +# +# docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway . +# +# No secrets live in this file. Runtime config (LITELLM_MASTER_KEY, +# OPENAI_API_KEY referenced by config.yaml, etc.) is injected as environment +# variables at deploy time. + +# ---- Chef ------------------------------------------------------------------- +# cargo-chef caches the dependency build so only the gateway crate recompiles on +# a source-only change. python3-dev is present in every rust stage because the +# `python-config` feature links libpython via pyo3 (even in the cook step). +FROM rust:1.90-slim-bookworm AS chef +ENV PYO3_PYTHON=python3.11 +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + python3 python3-dev pkg-config libssl-dev clang \ + && rm -rf /var/lib/apt/lists/* \ + && cargo install cargo-chef --locked --version 0.1.77 +WORKDIR /build/litellm-rust + +# ---- Planner ---------------------------------------------------------------- +# Produce the dependency recipe from the rust workspace manifests + Cargo.lock. +FROM chef AS planner +COPY litellm-rust/ . +RUN cargo chef prepare --recipe-path recipe.json + +# ---- Builder ---------------------------------------------------------------- +FROM chef AS builder +# Cook (compile) just the dependencies first — this layer is cached and reused +# whenever only gateway source changes. +COPY --from=planner /build/litellm-rust/recipe.json recipe.json +RUN cargo chef cook --locked --release \ + -p litellm-ai-gateway --features python-config \ + --recipe-path recipe.json +# Now copy the real sources and build the gateway binary. Deps are already cooked +# above, so this step only recompiles the gateway crate. +COPY litellm-rust/ . +RUN cargo build --locked --release -p litellm-ai-gateway --features python-config + +# ---- Runtime ---------------------------------------------------------------- +# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3 +# 3.11 ABI so the embedded interpreter links and imports cleanly. +FROM python:3.11-slim-bookworm AS runtime + +# CA certificates for outbound TLS to the OpenAI realtime endpoint. +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so +# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the +# package + packaging metadata, then pip install the proxy extra. +COPY pyproject.toml README.md LICENSE ./ +COPY litellm/ ./litellm/ +RUN pip install --no-cache-dir ".[proxy]" + +# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time +# only). +COPY --from=builder /build/litellm-rust/target/release/litellm-ai-gateway /usr/local/bin/litellm-ai-gateway + +# Default config.yaml. A real deploy can override this (e.g. mount a Render +# secret file at the same path) — never bake secrets into the image. +COPY litellm-rust/crates/ai-gateway/config.yaml /app/config.yaml + +# Bind to all interfaces (Render routes to 0.0.0.0:$PORT) and load the model_list +# from config.yaml via the embedded python config reader. +ENV HOST=0.0.0.0 \ + LITELLM_CONFIG_PATH=/app/config.yaml + +# Drop to a non-root user. The realtime hot path needs no root privileges, so +# running unprivileged limits blast radius if the process is ever compromised. +# The binary in /usr/local/bin is world-executable (COPY default mode 755); we +# only need /app (and the config.yaml it reads) owned by the unprivileged user. +RUN useradd --system --no-create-home --uid 10001 appuser \ + && chown -R appuser:appuser /app +USER appuser + +ENTRYPOINT ["/usr/local/bin/litellm-ai-gateway"] diff --git a/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore new file mode 100644 index 00000000000..030ee6a37c5 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/Dockerfile.dockerignore @@ -0,0 +1,45 @@ +# Dockerfile-specific ignore-file for the Rust AI Gateway build. +# +# The build context is the repo root (so the image can pip install litellm from +# source AND build the rust workspace). BuildKit honors `.dockerignore` +# next to the Dockerfile and it takes precedence over the repo-root `.dockerignore`, +# so this file shrinks the (large) repo-root context for THIS build only without +# touching the root `.dockerignore` used by the main litellm images. +# +# Strategy: ignore everything, then re-include only what the build needs: +# - litellm/ (pip install . needs the full package + proxy reader) +# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources) +# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install) +* + +# --- re-include the build inputs --- +!litellm/ +!litellm-rust/ +!pyproject.toml +!README.md +!LICENSE + +# --- prune heavy / irrelevant subpaths back out of the re-included trees --- +# Rust build artifacts (huge; regenerated in the builder). +**/target/ +# Python caches and compiled bytecode. +**/__pycache__/ +**/*.pyc +**/*.pyo +**/.pytest_cache/ +**/.ruff_cache/ +**/.mypy_cache/ +# Node / UI build output bundled under the python package (not needed to import +# litellm.proxy.read_model_list). +**/node_modules/ +litellm/proxy/_experimental/out/ +# Tests, logs, and local scratch. +**/tests/ +**/test/ +*.log +log.txt +*.tgz +# VCS / editor / CI metadata that may live under re-included trees. +**/.git/ +.git/ +**/.DS_Store diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md new file mode 100644 index 00000000000..f913beff6d5 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/README.md @@ -0,0 +1,198 @@ +# LiteLLM Rust AI Gateway + +A minimal Axum service that fronts OpenAI's realtime API. Clients open a +WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a deployment, +dials OpenAI upstream, and splices the two sockets frame-by-frame. + +## Crates + +`litellm-rust` is exactly three crates (a crate is a **layer**, not a route): + +| Crate | Role | Pure / I/O | +|-------|------|------------| +| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. Builds requests/responses; no network. | Pure | +| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under `io/`) plus the Axum server binary (behind the `server` feature). | I/O | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding | + +Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. + +- **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) +- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) +- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil` +- **Request logs:** POSTed to a LiteLLM proxy at `/v1/rust_control_plane/logs` (see [Request logging](#request-logging)) + +> **Realtime serving is pure Rust.** Python is used at **load time only** — to +> read the config once at boot. The realtime hot path never touches Python. + +## Configuration (config.yaml) + +The gateway loads its `model_list` from a **config.yaml**, the same as the +LiteLLM proxy. Point `LITELLM_CONFIG_PATH` at the file: + +```yaml +# config.yaml +model_list: + - model_name: gpt-realtime + litellm_params: + model: openai/gpt-realtime + api_key: os.environ/OPENAI_API_KEY +``` + +```bash +LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway +``` + +At boot the gateway calls into `litellm.proxy.read_model_list`, which reuses the +**real proxy config reader** (`ProxyConfig.get_config`). That means everything +the proxy supports in config.yaml works here too: + +- `include:` to merge in other config files, +- `os.environ/VAR` secret references (resolved via the secret manager, never + inlined), +- DB-stored models (when a database is configured). + +Secrets stay out of the config — reference them with `os.environ/...` and set +the env var at deploy time. The shipped Docker image is built with the +`python-config` feature and **bundles litellm**, so config loading works out of +the box; the default baked config lives at `/app/config.yaml` and can be +overridden at deploy time (e.g. a Render secret file mounted at the same path). + +### Environment variables + +| Var | Required | Default | Purpose | +|---|---|---|---| +| `LITELLM_CONFIG_PATH` | yes (config mode) | — | Path to the config.yaml the gateway loads its `model_list` from. The Docker image defaults this to `/app/config.yaml`. | +| `LITELLM_MASTER_KEY` | yes | — | Bearer token clients must send. Unset ⇒ all `/v1/realtime` requests are rejected (fail closed). | +| `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. | +| `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. | +| `PORT` | no | `4001` | Listen port. Render and most PaaS inject this automatically. | +| `LITELLM_PROXY_BASE_URL` | no | `http://localhost:4000` | LiteLLM proxy that request logs are POSTed to. See [Request logging](#request-logging). | + +> Secrets (`LITELLM_MASTER_KEY`, `OPENAI_API_KEY`) are never baked into the image +> or `render.yaml` — inject them at deploy time only. + +### Lean env stand-in (fallback) + +If the binary is built **without** `python-config` (default features), or +`LITELLM_CONFIG_PATH` is unset, the gateway falls back to a single-deployment +stand-in built from the environment: + +| Var | Default | Purpose | +|---|---|---| +| `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). | + +This mode links no libpython and needs no config file, but it only supports one +hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the +stand-in only for the leanest possible build. + +## Request logging + +The gateway runs no spend logic. When a session ends it builds one +`StandardLoggingPayload` and POSTs it to `{LITELLM_PROXY_BASE_URL}/v1/rust_control_plane/logs` +(admin-only, bearer = `LITELLM_MASTER_KEY`), and the proxy replays it through its +normal callbacks (spend logs, Langfuse, etc.). The POST is non-blocking: a bounded +channel drained by a background worker, dropping with a counter if the proxy is +down. It sends one payload per session. Both env vars are in the table above. + +Worker tuning, rarely needed: `LITELLM_LOG_CHANNEL_CAPACITY` (4096), +`LITELLM_LOG_BATCH_SIZE` (256), `LITELLM_LOG_FLUSH_INTERVAL_MS` (500). + +## Build & run with Docker + +The image is built `--features python-config` and installs litellm **from this +repo's source** (the config reader is newer than any PyPI release), so the build +**context is the repo root**: + +```bash +# from the repo root +docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway . + +docker run --rm -p 4001:4001 \ + -e HOST=0.0.0.0 -e PORT=4001 \ + -e LITELLM_MASTER_KEY=sk-local \ + -e OPENAI_API_KEY=$OPENAI_API_KEY \ + litellm-ai-gateway # LITELLM_CONFIG_PATH defaults to /app/config.yaml + +# smoke test +curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/health/readiness # -> 200 +curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/v1/realtime # -> 401 (auth fails closed) +``` + +On boot you should see `loaded model_list from /app/config.yaml via python +config reader` — that confirms the config path (not the env stand-in fallback). +To use your own config, mount it over the default: + +```bash +docker run --rm -p 4001:4001 \ + -e HOST=0.0.0.0 -e LITELLM_MASTER_KEY=sk-local -e OPENAI_API_KEY=$OPENAI_API_KEY \ + -v $(pwd)/my-config.yaml:/app/config.yaml:ro \ + litellm-ai-gateway +``` + +### Cargo-only (no Docker) + +```bash +# config.yaml mode — needs litellm importable in the active python env +LITELLM_CONFIG_PATH=./crates/ai-gateway/config.yaml \ + cargo run --release -p litellm-ai-gateway --features python-config + +# env stand-in mode — no python, no config +cargo run --release -p litellm-ai-gateway +``` + +## Deploy on Render + +The service is a Docker **web service**; Render terminates TLS and supports +WebSockets, so the public endpoint is `wss://.onrender.com/v1/realtime`. + +### Option A — Blueprint (`render.yaml`) + +`crates/ai-gateway/render.yaml` describes the service (Docker runtime, +`healthCheckPath: /health/readiness`, repo-root `dockerContext: .`, +`dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile`, +`LITELLM_CONFIG_PATH: /app/config.yaml`). `LITELLM_MASTER_KEY` and +`OPENAI_API_KEY` are `sync: false` — set them in the dashboard after the first +deploy. To use a non-default model_list, mount a **Render Secret File** at +`/app/config.yaml`. Point a Render Blueprint at this repo/branch and apply. + +### Option B — Render API + +```bash +# create a Docker web service from this repo+branch, then set env vars: +curl -X POST https://api.render.com/v1/services \ + -H "Authorization: Bearer $RENDER_API_KEY" -H "Content-Type: application/json" \ + -d '{ + "type": "web_service", "name": "litellm-rust-ai-gateway", + "ownerId": "", "repo": "https://github.com/BerriAI/litellm", + "branch": "", + "serviceDetails": { + "env": "docker", + "envSpecificDetails": { + "dockerfilePath": "./litellm-rust/crates/ai-gateway/Dockerfile", + "dockerContext": "." + }, + "healthCheckPath": "/health/readiness" + } + }' +# then set env vars LITELLM_MASTER_KEY, OPENAI_API_KEY, HOST=0.0.0.0, +# LITELLM_CONFIG_PATH=/app/config.yaml +``` + +Health check path **must** be `/health/readiness`. `autoDeploy` is off by default +in the blueprint — trigger deploys manually (or flip it on) to pick up new commits. + +## Scaling + +Concurrency is what matters, not total connections: each in-flight session holds +one client socket + one upstream socket. To scale, raise the instance count / +enable autoscaling on the Render service (e.g. baseline 10, max 100). Each +instance needs file descriptors for `2 × peak_concurrent_sessions` — raise +`ulimit -n` if you push very high concurrency. + +## Latency note + +The gateway adds the cost of one extra hop: client→gateway, then a fresh +gateway→OpenAI realtime handshake (TLS + WS upgrade + `session.created`). In +benchmarks this is ~100–150 ms of added session-establishment time; first-audio +and steady-state streaming add no measurable overhead. To minimize it, deploy the +gateway in the Render region with the lowest RTT to OpenAI's realtime endpoint. diff --git a/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md b/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md new file mode 100644 index 00000000000..84e926af243 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md @@ -0,0 +1,55 @@ +# 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/config.yaml b/litellm-rust/crates/ai-gateway/config.yaml new file mode 100644 index 00000000000..ac598c220dd --- /dev/null +++ b/litellm-rust/crates/ai-gateway/config.yaml @@ -0,0 +1,13 @@ +# Sample realtime config for the LiteLLM Rust AI Gateway. +# +# The gateway loads this model_list at boot via the embedded python config +# reader (litellm.proxy.read_model_list), which reuses the proxy's own reader — +# so include:, os.environ/ secrets, and DB-stored models all work here too. +# +# Secrets are referenced (never inlined) via os.environ/. A real deploy can +# override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH). +model_list: + - model_name: gpt-realtime + litellm_params: + model: openai/gpt-realtime + api_key: os.environ/OPENAI_API_KEY diff --git a/litellm-rust/crates/ai-gateway/render.yaml b/litellm-rust/crates/ai-gateway/render.yaml new file mode 100644 index 00000000000..4170849f65d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/render.yaml @@ -0,0 +1,35 @@ +# Render blueprint for the LiteLLM Rust AI Gateway (realtime WebSocket proxy). +# +# Single instance for now (no autoscaling). The public endpoint is a +# WebSocket served over TLS: wss://.onrender.com/v1/realtime +# +# Paths are relative to the **repo root** (Render's convention). The build +# context is the repo root so the image can install litellm from source — the +# gateway loads its model_list via litellm.proxy.read_model_list at boot. +# +# Secrets (LITELLM_MASTER_KEY, OPENAI_API_KEY) are marked sync: false — set +# them in the Render dashboard or via the API, never inline here. +services: + - type: web + name: litellm-rust-ai-gateway + runtime: docker + plan: standard + dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile + dockerContext: . + healthCheckPath: /health/readiness + numInstances: 1 + envVars: + # The gateway loads its model_list from this config.yaml via the embedded + # python config reader. The image bakes a default config at /app/config.yaml; + # a real deploy can override it by mounting a Render secret file at this + # same path (Dashboard → Environment → Secret Files) — never inline secrets. + - key: LITELLM_CONFIG_PATH + value: /app/config.yaml + - key: HOST + value: 0.0.0.0 + # Bearer token clients must send on /v1/realtime (fail closed if unset). + - key: LITELLM_MASTER_KEY + sync: false + # Referenced by config.yaml as os.environ/OPENAI_API_KEY for the upstream dial. + - key: OPENAI_API_KEY + sync: false diff --git a/litellm-rust/crates/ai-gateway/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs new file mode 100644 index 00000000000..438a0513057 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/auth/mod.rs @@ -0,0 +1,93 @@ +//! Gateway authentication, as an axum **extractor** (the idiomatic pattern — +//! keeps handlers clean and auth testable). +//! +//! For now this is a single **master key**: any caller presenting it as +//! `Authorization: Bearer ` may invoke the gateway. Per-key auth, budgets, +//! and rate limits are delegated to the Python proxy in a later phase. +//! +//! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then +//! runs during extraction, before the handler body. Routes never re-implement it. + +use axum::extract::FromRequestParts; +use axum::http::header::AUTHORIZATION; +use axum::http::request::Parts; +use axum::http::StatusCode; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; + +use crate::state::AppState; + +/// SHA-256 hex digest of a token — the exact transform the Python proxy applies +/// (`litellm.proxy.utils.hash_token`). +/// +/// STRICT REQUIREMENT: a raw key (`LITELLM_MASTER_KEY`, a virtual key, …) must +/// **never** leave this gateway in a log payload. Spend logs and every callback +/// integration receive `user_api_key_hash`, so that field must be this hash, not +/// the credential. Hashing here also means the value matches the key's hash in +/// `LiteLLM_SpendLogs.api_key`, so realtime spend joins with the rest of LiteLLM. +pub fn hash_token(token: &str) -> String { + let digest = Sha256::digest(token.as_bytes()); + let mut hex = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write; + let _ = write!(hex, "{byte:02x}"); + } + hex +} + +/// Extractor that requires the configured master key as a bearer token. +/// +/// Rejections: `500` when no master key is configured (permanent +/// misconfiguration, not a transient outage); `401` on a missing/incorrect +/// token. The comparison is constant-time. +pub struct RequireMasterKey; + +#[axum::async_trait] +impl FromRequestParts for RequireMasterKey { + type Rejection = (StatusCode, String); + + async fn from_request_parts( + parts: &mut Parts, + state: &AppState, + ) -> Result { + let Some(expected) = state.master_key.as_deref() else { + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(), + )); + }; + let provided = parts + .headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .map(str::trim); + match provided { + Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self), + _ => Err(( + StatusCode::UNAUTHORIZED, + "missing or invalid bearer token".to_string(), + )), + } + } +} + +#[cfg(test)] +mod tests { + use super::hash_token; + + #[test] + fn hash_token_matches_python_sha256_hexdigest() { + // Must equal hashlib.sha256("sk-1234".encode()).hexdigest() — the value + // the proxy stores in LiteLLM_SpendLogs.api_key. + assert_eq!( + hash_token("sk-1234"), + "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + ); + // 64 lowercase hex chars, and never the raw input. + let h = hash_token("sk-secret"); + assert_eq!(h.len(), 64); + assert!(h.chars().all(|c| c.is_ascii_hexdigit())); + assert_ne!(h, "sk-secret"); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs new file mode 100644 index 00000000000..3116a4c9932 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -0,0 +1,29 @@ +//! Crate-level constants for the ai-gateway. +//! +//! Per `litellm-rust/CLAUDE.md`, magic numbers and fixed strings live here +//! (the Rust mirror of Python's `litellm/constants.py`), not inline in feature +//! modules. Env-overridable tunables keep their `DEFAULT_*` value here; the env +//! read + fallback happens at the host/config layer. + +/// Default LiteLLM control-plane base URL for request-log egress when +/// `LITELLM_PROXY_BASE_URL` is unset. +pub(crate) const DEFAULT_PROXY_BASE_URL: &str = "http://localhost:4000"; + +/// The logs ingest path appended to the proxy base. Not a tunable; it is the +/// proxy's API contract (the rust-control-plane router on the Python proxy). +pub(crate) const RUST_CONTROL_PLANE_LOGS_PATH: &str = "/v1/rust_control_plane/logs"; + +/// Default bounded channel depth for the log-egress worker. +/// Override: `LITELLM_LOG_CHANNEL_CAPACITY`. +pub(crate) const DEFAULT_CHANNEL_CAPACITY: usize = 4096; + +/// Default max records POSTed per request to the control plane. +/// Override: `LITELLM_LOG_BATCH_SIZE`. +pub(crate) const DEFAULT_MAX_BATCH_SIZE: usize = 256; + +/// Default partial-batch flush cadence, in ms. +/// Override: `LITELLM_LOG_FLUSH_INTERVAL_MS`. +pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; + +/// Provider attributed to realtime sessions in the logging payload. +pub(crate) const DEFAULT_PROVIDER: &str = "openai"; diff --git a/litellm-rust/crates/ai-gateway/src/gil.rs b/litellm-rust/crates/ai-gateway/src/gil.rs new file mode 100644 index 00000000000..c749f722c73 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/gil.rs @@ -0,0 +1,58 @@ +//! GIL-activity tracking. +//! +//! Every acquisition of the Python GIL is recorded here so the `/health/gil` +//! endpoint can report whether Python was touched recently. The design goal is +//! that the GIL is acquired **only at load time** (config read) and never on the +//! realtime hot path — polling this endpoint during traffic should show the +//! count holding steady and `acquired_last_30s` falling to `false`. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +/// Window (seconds) for the "recently acquired" signal. +pub const RECENT_WINDOW_SECS: u64 = 30; + +static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0); +/// Unix seconds of the last acquisition; `0` means "never". +static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0); + +fn now_unix_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// Record that the GIL was just acquired. Call immediately before taking the GIL. +/// +/// Only invoked under the `python-config` feature; without it the gateway never +/// touches Python, so the recorder is unused (and the endpoint reports zero). +#[cfg_attr(not(feature = "python-config"), allow(dead_code))] +pub fn record_acquisition() { + GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed); + LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed); +} + +/// Point-in-time view of GIL activity. +pub struct GilSnapshot { + pub total_acquisitions: u64, + pub seconds_since_last: Option, + pub acquired_last_30s: bool, +} + +/// Read the current GIL-activity snapshot. +pub fn snapshot() -> GilSnapshot { + let total = GIL_ACQUISITIONS.load(Ordering::Relaxed); + let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed); + let seconds_since_last = if last == 0 { + None + } else { + Some(now_unix_secs().saturating_sub(last)) + }; + let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS); + GilSnapshot { + total_acquisitions: total, + seconds_since_last, + acquired_last_30s, + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs new file mode 100644 index 00000000000..53b599d8c98 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs @@ -0,0 +1,24 @@ +//! The `CustomLogger` trait — the Rust mirror of Python +//! `litellm/integrations/custom_logger.py::CustomLogger`. +//! +//! Synchronous (no `async_trait`): callbacks are O(1) enqueue-and-return so the +//! realtime splice never blocks on a logger. Default bodies are no-ops so a +//! logger can implement only the events it cares about. + +use crate::integrations::types::{LogError, LoggingError, StandardLoggingPayload}; + +pub trait CustomLogger: Send + Sync { + /// Record a successful call. Default: no-op. + fn log_success_event(&self, _payload: &StandardLoggingPayload) -> Result<(), LogError> { + Ok(()) + } + + /// Record a failed call. Default: no-op. + fn log_failure_event( + &self, + _payload: &StandardLoggingPayload, + _error: &LoggingError, + ) -> Result<(), LogError> { + Ok(()) + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs new file mode 100644 index 00000000000..165a90d8dbe --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs @@ -0,0 +1,209 @@ +//! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's +//! `/v1/rust_control_plane/logs` endpoint. +//! +//! The callback path is non-blocking: `log_success_event` / `log_failure_event` +//! build a `LogRecord` and `try_send` it onto a bounded channel, returning a +//! `LogError` (never panicking, never awaiting) if the channel is full or the +//! worker has gone away. A spawned background worker drains the channel, batches +//! records into `{"records":[...]}`, and POSTs them to the proxy with a pooled +//! `reqwest::Client`. + +use std::sync::Arc; +use std::time::Duration; + +use reqwest::Client; +use tokio::sync::mpsc::{self, Receiver, Sender}; +use tokio::time::interval; + +use crate::constants::{ + DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE, + DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH, +}; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::{ + CallbackLogsRequest, LogError, LogRecord, LoggingError, StandardLoggingPayload, +}; + +/// Egress worker tunables. Each field defaults to the matching `DEFAULT_*` const +/// in `crate::constants` and is overridable via an env var (read once at logger +/// construction). +struct EgressTunables { + channel_capacity: usize, + max_batch_size: usize, + flush_interval: Duration, +} + +impl EgressTunables { + fn from_env() -> Self { + Self { + channel_capacity: env_positive( + "LITELLM_LOG_CHANNEL_CAPACITY", + DEFAULT_CHANNEL_CAPACITY, + ), + max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE), + flush_interval: Duration::from_millis(env_positive( + "LITELLM_LOG_FLUSH_INTERVAL_MS", + DEFAULT_FLUSH_INTERVAL_MS, + )), + } + } +} + +/// Parse a positive integer env var, falling back to `default` on missing, +/// unparseable, or non-positive values. Generic over the integer type so one +/// helper serves both the `usize` capacities and the `u64` interval. +fn env_positive(name: &str, default: T) -> T +where + T: std::str::FromStr + PartialOrd + From, +{ + let zero = T::from(0u8); + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|n| *n > zero) + .unwrap_or(default) +} + +/// Ships realtime logging events to the LiteLLM Python proxy. +pub struct LiteLLMPythonProxyAPILogger { + sink: Sender, +} + +impl LiteLLMPythonProxyAPILogger { + /// Spawn the background worker and return a logger handle. `base` is the + /// proxy base URL (no trailing path); `master_key` is sent as a bearer token. + pub fn start(base: String, master_key: String) -> Arc { + let tunables = EgressTunables::from_env(); + let (sink, receiver) = mpsc::channel::(tunables.channel_capacity); + let url = format!( + "{}{}", + base.trim_end_matches('/'), + RUST_CONTROL_PLANE_LOGS_PATH + ); + let client = Client::new(); + tokio::spawn(worker_loop( + receiver, + client, + url, + master_key, + tunables.max_batch_size, + tunables.flush_interval, + )); + Arc::new(Self { sink }) + } + + /// Build a logger from the environment: `LITELLM_PROXY_BASE_URL` (default + /// `http://localhost:4000`) and `LITELLM_MASTER_KEY`. + /// + /// `LITELLM_PROXY_BASE_URL` is treated as the full base and the route is + /// appended verbatim, so if the proxy runs under a `SERVER_ROOT_PATH` + /// (e.g. served at `https://host/litellm`), include it in the base + /// (`LITELLM_PROXY_BASE_URL=https://host/litellm`) and the POST lands at + /// `https://host/litellm/v1/rust_control_plane/logs`. + pub fn from_env() -> Arc { + let base = std::env::var("LITELLM_PROXY_BASE_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_PROXY_BASE_URL.to_string()); + let key = std::env::var("LITELLM_MASTER_KEY").unwrap_or_default(); + Self::start(base, key) + } + + fn enqueue(&self, record: LogRecord) -> Result<(), LogError> { + self.sink.try_send(record).map_err(|err| match err { + mpsc::error::TrySendError::Full(_) => LogError::channel_full(), + mpsc::error::TrySendError::Closed(_) => LogError::channel_closed(), + }) + } +} + +impl CustomLogger for LiteLLMPythonProxyAPILogger { + fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> { + self.enqueue(LogRecord { + status: "success".to_string(), + payload: payload.clone(), + error: None, + }) + } + + fn log_failure_event( + &self, + payload: &StandardLoggingPayload, + error: &LoggingError, + ) -> Result<(), LogError> { + self.enqueue(LogRecord { + status: "failure".to_string(), + payload: payload.clone(), + error: Some(format!("{}: {}", error.kind, error.message)), + }) + } +} + +/// Drain the channel, batching records and POSTing them to the proxy. Exits when +/// the channel is closed (all senders dropped) and drained. +async fn worker_loop( + mut receiver: Receiver, + client: Client, + url: String, + master_key: String, + max_batch_size: usize, + flush_interval: Duration, +) { + let mut ticker = interval(flush_interval); + let mut batch: Vec = Vec::with_capacity(max_batch_size); + + loop { + tokio::select! { + maybe_record = receiver.recv() => { + match maybe_record { + Some(record) => { + batch.push(record); + if batch.len() >= max_batch_size { + flush(&client, &url, &master_key, &mut batch).await; + } + } + None => { + // Channel closed: flush remaining and exit. + flush(&client, &url, &master_key, &mut batch).await; + break; + } + } + } + _ = ticker.tick() => { + flush(&client, &url, &master_key, &mut batch).await; + } + } + } +} + +/// POST the current batch (if any), clearing it. Errors are logged, not fatal. +async fn flush(client: &Client, url: &str, master_key: &str, batch: &mut Vec) { + if batch.is_empty() { + return; + } + let records = std::mem::take(batch) + .into_iter() + .map(LogRecord::into_callback_record) + .collect(); + let body = CallbackLogsRequest { records }; + + let response = client + .post(url) + .bearer_auth(master_key) + .json(&body) + .send() + .await; + + match response { + Ok(resp) if resp.status().is_success() => {} + Ok(resp) => { + eprintln!( + "litellm-ai-gateway: callback logs POST returned {} to {url}", + resp.status() + ); + } + Err(err) => { + eprintln!("litellm-ai-gateway: callback logs POST failed to {url}: {err}"); + } + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs new file mode 100644 index 00000000000..8799be0c040 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs @@ -0,0 +1,10 @@ +//! Pure-Rust logging integrations. Names map 1:1 to Python +//! `litellm/integrations/`: +//! - [`custom_logger::CustomLogger`] — the callback trait +//! - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events +//! to the Python proxy's `/v1/callbacks/logs` endpoint +//! - [`types`] — the typed `StandardLoggingPayload` wire contract + +pub mod custom_logger; +pub mod litellm_python_proxy_api; +pub mod types; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/types.rs new file mode 100644 index 00000000000..d61a1f816a7 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/types.rs @@ -0,0 +1,164 @@ +//! Typed payloads for the LiteLLM `/v1/callbacks/logs` realtime-logging contract. +//! +//! Field names below are the EXACT JSON keys the Python replay path + spend-logs +//! builder read. Note the deliberate mix: +//! - `startTime` / `endTime` are camelCase (epoch f64 seconds) +//! - `response_cost` / `prompt_tokens` / etc. are snake_case +//! +//! Mirrors Python `litellm/integrations/` + the proxy `CallbackLogsRequest` +//! contract 1:1. + +use serde::Serialize; +use serde_json::Value; +use std::collections::HashMap; + +/// Cumulative token usage for a realtime session. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Usage { + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, +} + +/// Cost-attribution metadata threaded from the authenticated request. +#[derive(Clone, Debug, Default)] +pub struct RequestMetadata { + pub user_api_key_hash: Option, + pub user_api_key_user_id: Option, + pub user_api_key_team_id: Option, +} + +/// A logging-callback failure (e.g. a custom logger raised). Mirrors the Python +/// failure-event shape: a message plus an exception kind/class name. +#[derive(Clone, Debug)] +pub struct LoggingError { + pub message: String, + pub kind: String, +} + +/// A non-fatal error returned by a `CustomLogger` when it cannot enqueue an +/// event (channel full or the background worker has shut down). +#[derive(Clone, Debug)] +pub struct LogError { + pub message: String, + pub kind: String, +} + +impl LogError { + pub fn channel_full() -> Self { + Self { + message: "logging channel is full; dropping record".to_string(), + kind: "ChannelFull".to_string(), + } + } + + pub fn channel_closed() -> Self { + Self { + message: "logging channel is closed; worker has shut down".to_string(), + kind: "ChannelClosed".to_string(), + } + } +} + +impl std::fmt::Display for LogError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.kind, self.message) + } +} + +impl std::error::Error for LogError {} + +/// Batch wrapper — the top-level request body. +/// Matches Python `CallbackLogsRequest { records: list[CallbackLogRecord] }`. +#[derive(Serialize)] +pub struct CallbackLogsRequest { + pub records: Vec, +} + +/// One finished logging event. +/// Matches `CallbackLogRecord { status, standard_logging_payload, error? }`. +#[derive(Serialize)] +pub struct CallbackLogRecord { + /// "success" | "failure". On "failure", `error` (or payload.error_str) + /// becomes the replayed exception string. + pub status: String, + + pub standard_logging_payload: StandardLoggingPayload, + + /// Only meaningful when status == "failure". Omitted on success. + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +/// The self-describing payload. Field names are the EXACT JSON keys the Python +/// replay path + spend-logs builder read. +#[derive(Clone, Debug, Serialize)] +pub struct StandardLoggingPayload { + pub id: String, + pub litellm_call_id: String, + + /// e.g. "realtime", "acompletion". Falls back to "acompletion" if absent. + pub call_type: String, + + pub model: String, + pub custom_llm_provider: String, + + /// Spend ($) written to LiteLLM_SpendLogs.spend. + pub response_cost: f64, + + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, + + /// EPOCH SECONDS as float — camelCase keys, NOT snake_case. + #[serde(rename = "startTime")] + pub start_time: f64, + #[serde(rename = "endTime")] + pub end_time: f64, + + pub stream: bool, + + pub metadata: StandardLoggingMetadata, + + /// Optional; stored as request input on the spend log row. + #[serde(skip_serializing_if = "Option::is_none")] + pub messages: Option, +} + +/// Cost-attribution keys. The replayer maps these into litellm_params.metadata, +/// which the spend-logs builder reads to set user / team_id / organization_id. +#[derive(Clone, Debug, Serialize, Default)] +pub struct StandardLoggingMetadata { + pub user_api_key_hash: Option, // -> SpendLogs.api_key + pub user_api_key_user_id: Option, // -> SpendLogs.user + pub user_api_key_team_id: Option, // -> SpendLogs.team_id + + // Optional but read by the builder; include when known: + #[serde(skip_serializing_if = "Option::is_none")] + pub user_api_key_alias: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub user_api_key_org_id: Option, // -> SpendLogs.organization_id + #[serde(skip_serializing_if = "Option::is_none")] + pub user_api_key_end_user_id: Option, // -> SpendLogs.end_user + #[serde(skip_serializing_if = "Option::is_none")] + pub spend_logs_metadata: Option>, +} + +/// The unit handed to a `CustomLogger` sink: a finished payload plus its status +/// and (on failure) the replayed error string. +#[derive(Clone, Debug)] +pub struct LogRecord { + pub status: String, + pub payload: StandardLoggingPayload, + pub error: Option, +} + +impl LogRecord { + pub fn into_callback_record(self) -> CallbackLogRecord { + CallbackLogRecord { + status: self.status, + standard_logging_payload: self.payload, + error: self.error, + } + } +} diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs new file mode 100644 index 00000000000..3b566027646 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -0,0 +1,3 @@ +pub mod ocr; +pub mod realtime; +pub mod realtime_pool; diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs new file mode 100644 index 00000000000..5c32157bc6f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -0,0 +1,127 @@ +//! End-to-end OCR orchestration. +//! +//! Owns the whole Mistral OCR call so the Python side stays a thin bridge: +//! resolve the API key, build the URL + body via the pure transforms, POST it, +//! and normalize the response. The HTTP client is built once and reused. + +use std::sync::OnceLock; +use std::time::Duration; + +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::CoreResult; +use serde_json::{Map, Value}; + +use litellm_core::providers::mistral::ocr::transformation as mistral; +use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +/// OCR over large documents can take a while; bound it generously rather than +/// hanging forever on an unresponsive upstream. The client-level limit is the +/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``. +const OCR_TIMEOUT_SECS: u64 = 600; + +/// Maximum upstream body characters retained in error messages. OCR responses +/// can echo document contents and prompts; keep enough for debugging without +/// forwarding sensitive payloads across the host boundary. +const ERROR_BODY_MAX_CHARS: usize = 256; + +/// Process-wide blocking HTTP client (connection pool + TLS reused across calls). +fn http_client() -> &'static reqwest::blocking::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::blocking::Client::builder() + .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .build() + .expect("failed to build reqwest client") + }) +} + +fn truncate_error_body(body: &str) -> String { + if body.chars().count() <= ERROR_BODY_MAX_CHARS { + return body.to_string(); + } + let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect(); + format!("{truncated}... (truncated)") +} + +/// Perform a Mistral OCR call end to end and return the normalized response as +/// JSON (the shape the Python `OCRResponse` model expects). +/// +/// Blocking: intended to be called with the GIL released from the Python bridge. +pub fn run_ocr( + model: &str, + document: Value, + api_key: Option<&str>, + api_base: Option<&str>, + optional_params: Map, + timeout: Option, +) -> CoreResult { + let config = &MISTRAL_OCR_CONFIG; + + let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?; + let url = mistral::complete_url(api_base); + let filtered_params = config.map_ocr_params(&optional_params); + let body = config + .transform_ocr_request(model, document, filtered_params)? + .data; + + let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body); + if let Some(duration) = timeout { + request = request.timeout(duration); + } + + let response = request + .send() + .map_err(|err| CoreError::Network(err.to_string()))?; + + let status = response.status(); + let text = response + .text() + .map_err(|err| CoreError::Network(err.to_string()))?; + + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + + let response_json: Value = serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; + + Ok(config + .transform_ocr_response(model, response_json)? + .into_json()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn truncate_error_body_passes_short_strings_through() { + let body = "Unauthorized"; + assert_eq!(truncate_error_body(body), "Unauthorized"); + } + + #[test] + fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50); + let truncated = truncate_error_body(&body); + + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS); + } + + #[test] + fn truncate_error_body_does_not_split_multibyte_chars() { + let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10); + let truncated = truncate_error_body(&body); + assert!(truncated.is_char_boundary(truncated.len())); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs new file mode 100644 index 00000000000..4047de5cb26 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -0,0 +1,390 @@ +//! End-to-end OpenAI realtime invocation. +//! +//! The host-facing entry point, mirroring `crate::io::ocr::run_ocr`: open the +//! WebSocket to OpenAI, then splice a client realtime stream to the upstream, +//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms. +//! Network, auth header, key resolution, and wire (de)serialization live here so +//! the `transformation` module stays pure and typed. +//! +//! The dial and splice steps are factored out ([`dial_upstream`], [`splice`]) so +//! the connection pool ([`crate::io::realtime_pool`]) can pre-establish an upstream, +//! buffer its `session.created`, and later hand the live socket to the same +//! splice loop a fresh dial uses. + +use std::time::Duration; + +use futures_util::stream::{SplitSink, SplitStream}; +use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::error::CoreError; +use litellm_core::realtime::transformation::RealtimeProviderConfig; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::CoreResult; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; + +use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; + +/// Environment variable holding the OpenAI API key (last-resort fallback). +const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; + +const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; + +/// Default **idle** timeout: if neither side sends a frame for this long, the +/// session is reaped. It resets on any activity, so it does not cap a healthy +/// (continuously streaming) session — it only frees a stalled one (e.g. a +/// half-open upstream that keeps the socket open but stops sending). +const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 300; + +/// The concrete upstream WebSocket type (TLS or plain). Shared by the dial path +/// and the pool so warm sockets and fresh sockets are the exact same type. +pub type UpstreamWs = WebSocketStream>; +pub(crate) type UpstreamTx = SplitSink; +pub(crate) type UpstreamRx = SplitStream; + +/// Resolve the OpenAI API key from the explicit param or the environment. +/// +/// Blank/whitespace values are treated as absent (guard at resolution time). +pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { + api_key + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + std::env::var(OPENAI_API_KEY_ENV) + .ok() + .filter(|key| !key.trim().is_empty()) + }) + .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) +} + +/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`. +/// +/// This is the dial half of [`realtime`], factored out so the pool can +/// pre-establish sockets ahead of any client. `api_key` here is already resolved +/// (non-blank) — the pool resolves it once when it is created. +pub(crate) async fn dial_upstream( + model: &str, + api_key: &str, + api_base: Option<&str>, +) -> CoreResult { + let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model); + + let mut request = url + .as_str() + .into_client_request() + .map_err(|err| CoreError::Network(err.to_string()))?; + // GA realtime: only Authorization. The legacy OpenAI-Beta header triggers + // beta_api_shape_disabled, so we do not send it. + request.headers_mut().insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {api_key}")) + .map_err(|err| CoreError::Auth(err.to_string()))?, + ); + + let (upstream, _response) = connect_async(request) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + Ok(upstream) +} + +/// Read the next text frame from the upstream and decode it as a typed event. +/// +/// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an +/// error on a non-text frame, a closed socket, or undecodable JSON so the pool can +/// discard a misbehaving socket rather than warm it. +pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult { + loop { + let message = upstream_rx + .next() + .await + .ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))? + .map_err(|err| CoreError::Network(err.to_string()))?; + match message { + Message::Text(text) => { + return serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(err.to_string())); + } + // Ignore protocol frames (ping/pong) while waiting for the first event. + Message::Ping(_) | Message::Pong(_) => continue, + Message::Close(_) => { + return Err(CoreError::Network( + "upstream closed before first event".to_string(), + )) + } + _ => continue, + } + } +} + +/// Splice an already-connected upstream to the client streams. +/// +/// `prelude` is relayed to the client first (the pool passes the buffered +/// `session.created` here; the fresh-dial path passes `None` and lets the upstream +/// deliver it). Then a single select loop forwards both directions through the +/// transforms until either side closes or the idle timeout fires. +/// `observe` is invoked on **upstream→client** events only (the trusted side that +/// carries `session.created` and `response.done` usage) — never on client events, +/// so a client cannot fabricate usage into its own logs. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn splice( + model: &str, + mut upstream_tx: UpstreamTx, + mut upstream_rx: UpstreamRx, + prelude: Option, + idle_timeout: Option, + mut observe: impl FnMut(&RealtimeEvent) + Send, + mut client_in: In, + mut client_out: Out, +) -> CoreResult<()> +where + In: Stream + Unpin + Send, + Out: Sink + Unpin + Send, + >::Error: std::fmt::Display, +{ + let config = &OPENAI_REALTIME_CONFIG; + + // Relay a buffered backend event (warm handoff's session.created) first, so a + // warm session looks identical to a fresh one from the client's view. + if let Some(event) = prelude { + for outbound in config.transform_realtime_response(&event, model)?.events { + client_out + .send(outbound) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + } + } + + let idle = idle_timeout.unwrap_or(Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS)); + + // One loop forwarding both directions. The `sleep(idle)` arm is rebuilt every + // iteration, so any frame (either way) resets it — it fires only when the + // session has been fully idle for `idle`, reaping a stalled connection + // (task + upstream TCP socket) instead of leaking it. + loop { + tokio::select! { + // client -> upstream + client_event = client_in.next() => { + let Some(event) = client_event else { break }; // client disconnected + // NOTE: do NOT observe client events. session.created / response.done + // (carrying usage) are server→client events; observing the client arm + // would let an authenticated client POST a fabricated response.done and + // inflate its own spend log. Logging observes upstream events only. + for outbound in config.transform_realtime_request(&event, model)?.events { + let payload = serde_json::to_string(&outbound) + .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + upstream_tx + .send(Message::Text(payload)) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + } + } + // upstream -> client + upstream_message = upstream_rx.next() => { + let Some(message) = upstream_message else { break }; // upstream closed + match message.map_err(|err| CoreError::Network(err.to_string()))? { + Message::Text(text) => { + let event: RealtimeEvent = serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + observe(&event); + for outbound in config.transform_realtime_response(&event, model)?.events { + client_out + .send(outbound) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + } + } + Message::Close(_) => break, + _ => {} + } + } + // idle timeout: no activity from either side within `idle` + _ = tokio::time::sleep(idle) => break, + } + } + Ok(()) +} + +/// Splice a client realtime stream to OpenAI: forward client events upstream +/// (via `transform_realtime_request`) and backend events downstream (via +/// `transform_realtime_response`). Returns when either side closes. +/// +/// Generic over the client transport (typed events) so this crate stays +/// framework-agnostic; the gateway adapts its axum socket to these. This is the +/// fresh-dial path: dial, then splice. The pool's warm-handoff path skips the dial +/// and calls [`splice`] directly with a buffered `session.created`. +#[allow(clippy::too_many_arguments)] +pub async fn realtime( + model: &str, + api_key: Option<&str>, + api_base: Option<&str>, + idle_timeout: Option, + observe: impl FnMut(&RealtimeEvent) + Send, + client_in: In, + client_out: Out, +) -> CoreResult<()> +where + In: Stream + Unpin + Send, + Out: Sink + Unpin + Send, + >::Error: std::fmt::Display, +{ + let api_key = resolve_api_key(api_key)?; + let upstream = dial_upstream(model, &api_key, api_base).await?; + let (upstream_tx, upstream_rx) = upstream.split(); + splice( + model, + upstream_tx, + upstream_rx, + None, + idle_timeout, + observe, + client_in, + client_out, + ) + .await +} + +/// Splice a pre-warmed upstream (taken from [`crate::io::realtime_pool`]) to the +/// client. Relays the buffered `session.created` first, then splices exactly like +/// the fresh-dial path — so a warm session is indistinguishable from a fresh one. +#[allow(clippy::too_many_arguments)] +pub async fn realtime_warm( + model: &str, + handoff: crate::io::realtime_pool::WarmHandoff, + idle_timeout: Option, + observe: impl FnMut(&RealtimeEvent) + Send, + client_in: In, + client_out: Out, +) -> CoreResult<()> +where + In: Stream + Unpin + Send, + Out: Sink + Unpin + Send, + >::Error: std::fmt::Display, +{ + splice( + model, + handoff.tx, + handoff.rx, + Some(handoff.session_created), + idle_timeout, + observe, + client_in, + client_out, + ) + .await +} + +#[cfg(test)] +mod tests { + use super::*; + + fn event(raw: &str) -> RealtimeEvent { + serde_json::from_str(raw).expect("valid event json") + } + + #[test] + fn resolve_api_key_prefers_param_then_blank_falls_through() { + assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test"); + // A blank param with no env set should error. + if std::env::var(OPENAI_API_KEY_ENV).is_err() { + assert!(resolve_api_key(Some(" ")).is_err()); + } + } + + /// Live end-to-end check against OpenAI. Ignored by default (CI never runs + /// it); run explicitly with `OPENAI_API_KEY` set: + /// `cargo test -p litellm-ai-gateway --features server realtime_invokes_openai -- --ignored --nocapture` + #[tokio::test] + #[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"] + async fn realtime_invokes_openai_and_responds() { + use futures_channel::mpsc; + + let key = + std::env::var(OPENAI_API_KEY_ENV).expect("set OPENAI_API_KEY to run this ignored test"); + + // client -> provider (we hold `client_tx` to push events upstream) + let (mut client_tx, client_in) = mpsc::unbounded::(); + // provider -> client (we hold `backend_rx` to read backend events) + let (client_out, mut backend_rx) = mpsc::unbounded::(); + + // Clone the key so the spawned task owns its `String` (no borrow across await). + let key_owned = key.clone(); + let call = tokio::spawn(async move { + realtime( + "gpt-realtime", + Some(&key_owned), + None, + None, + |_| {}, + client_in, + client_out, + ) + .await + }); + + // 1. First backend event should be session.created. + let first = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()) + .await + .expect("timed out waiting for session.created") + .expect("backend stream closed before session.created"); + assert_eq!( + first.event_type, "session.created", + "expected session.created, got: {}", + first.event_type + ); + + // 2. Ask for a short audio response. + client_tx + .send(event( + r#"{"type":"conversation.item.create","item":{"type":"message","role":"user","content":[{"type":"input_text","text":"Say hi."}]}}"#, + )) + .await + .expect("send conversation.item.create"); + client_tx + .send(event(r#"{"type":"response.create"}"#)) + .await + .expect("send response.create"); + + // 3. Read backend events; require a non-empty audio delta, then response.done. + let mut saw_audio_delta = false; + let mut saw_done = false; + for _ in 0..500 { + let next = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()).await; + let event = match next { + Ok(Some(event)) => event, + Ok(None) => break, + Err(_) => panic!("timed out waiting for backend events"), + }; + match event.event_type.as_str() { + "response.output_audio.delta" => { + let delta = event + .data + .get("delta") + .and_then(|value| value.as_str()) + .unwrap_or(""); + if !delta.is_empty() { + saw_audio_delta = true; + } + } + "response.done" => { + saw_done = true; + break; + } + _ => {} + } + } + + assert!( + saw_audio_delta, + "expected a response.output_audio.delta with non-empty delta" + ); + assert!(saw_done, "expected a response.done event"); + + // Drop the client sender so the provider's to_upstream side finishes. + drop(client_tx); + let _ = call.await; + } +} diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs new file mode 100644 index 00000000000..bf8041f31d7 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -0,0 +1,712 @@ +//! Pre-warmed upstream realtime connection pool. +//! +//! The gateway's realtime overhead lives entirely in session establishment: on +//! every client connect it dials a fresh upstream WS to OpenAI and waits for +//! `session.created` before it can serve. This pool keeps a small set of upstream +//! sockets **already connected and already past `session.created`** so a connect +//! can be served from a warm socket and the handshake is off the critical path. +//! +//! Layering: this lives in the gateway's `io` module next to the dial/splice it +//! reuses. The gateway holds an `Arc` in its state and asks for a +//! warm socket per connect; on a miss it fresh-dials exactly as before. The pool +//! is a latency optimization, never a correctness dependency — see the gateway's +//! `src/routes/realtime/README.md`. +//! +//! ## Caveats (enforced here) +//! - One warm socket serves exactly one session (realtime isn't multiplexed), so +//! the pool is sized to the connect *rate*, not concurrent connections. +//! - `session.created` is pre-read once and buffered; nothing else is read from a +//! warm socket before handoff, so a warm session starts at OpenAI defaults just +//! like a fresh one (`session.update` semantics unchanged). +//! - Warm sockets are short-lived (`max_idle`) and liveness-checked at handoff to +//! bound idle billing / dodge OpenAI's idle timeout. +//! - On miss or dead socket the caller fresh-dials; the pool never blocks or fails +//! a connect because it is empty. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use futures_util::StreamExt; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::CoreResult; + +use crate::io::realtime::{ + dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs, +}; + +/// Default target warm sockets per key when pooling is enabled. +pub const DEFAULT_POOL_SIZE: usize = 4; + +/// Default max time a warm socket may sit before it is closed and replaced. +pub const DEFAULT_MAX_IDLE: Duration = Duration::from_secs(30); + +/// Env var: target warm sockets per key. `0` disables pooling (fresh-dial only). +pub const POOL_SIZE_ENV: &str = "REALTIME_POOL_SIZE"; + +/// Env var: max warm-socket idle lifetime, in seconds. +pub const MAX_IDLE_ENV: &str = "REALTIME_POOL_MAX_IDLE_SECS"; + +/// How often the background replenisher wakes to top up and reap stale sockets. +const REPLENISH_TICK: Duration = Duration::from_millis(250); + +/// Backoff floor after a key's warm-up dials all fail. The first failed pass +/// waits this long before retrying that key. +const BACKOFF_BASE: Duration = Duration::from_millis(500); + +/// Backoff ceiling. A key that keeps failing (invalid credentials, an +/// unreachable upstream) is retried at most once per this interval — instead of +/// firing `needed` concurrent TLS dials every 250 ms tick, which would hammer +/// the upstream and risk rate-limit exhaustion that degrades valid cold-path +/// traffic. Backoff resets the moment a dial for the key succeeds. +const BACKOFF_MAX: Duration = Duration::from_secs(30); + +/// Identifies an upstream connection: the tuple that fully determines the dial. +/// `api_key` is included so a warm socket is only ever reused for the same key +/// (no cross-tenant reuse). +#[derive(Clone, PartialEq, Eq, Hash)] +pub struct UpstreamKey { + pub model: String, + pub api_key: String, + pub api_base: Option, +} + +impl std::fmt::Debug for UpstreamKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UpstreamKey") + .field("model", &self.model) + .field("api_key", &"[REDACTED]") + .field("api_base", &self.api_base) + .finish() + } +} + +/// A warm upstream: split halves + the buffered `session.created` + when it was +/// warmed (for `max_idle` expiry). +struct WarmConnection { + tx: UpstreamTx, + rx: UpstreamRx, + session_created: RealtimeEvent, + warmed_at: Instant, +} + +/// A live upstream taken from the pool, ready to splice. The caller relays +/// `session_created` to the client first, then splices `(tx, rx)` as usual. +pub struct WarmHandoff { + pub tx: UpstreamTx, + pub rx: UpstreamRx, + pub session_created: RealtimeEvent, +} + +/// Pool configuration, resolved once at startup from the environment. +#[derive(Clone, Copy, Debug)] +pub struct PoolConfig { + /// Target warm sockets per key. `0` disables pooling. + pub target_size: usize, + /// Max time a warm socket may sit before it is closed and replaced. + pub max_idle: Duration, +} + +impl Default for PoolConfig { + fn default() -> Self { + Self { + target_size: DEFAULT_POOL_SIZE, + max_idle: DEFAULT_MAX_IDLE, + } + } +} + +impl PoolConfig { + /// Read config from the environment, falling back to defaults. An invalid + /// value warns and uses the default rather than failing startup. + pub fn from_env() -> Self { + let target_size = match std::env::var(POOL_SIZE_ENV) { + Ok(raw) => raw.trim().parse().unwrap_or_else(|_| { + eprintln!("warning: {POOL_SIZE_ENV}={raw:?} is not a valid size; using {DEFAULT_POOL_SIZE}"); + DEFAULT_POOL_SIZE + }), + Err(_) => DEFAULT_POOL_SIZE, + }; + let max_idle = match std::env::var(MAX_IDLE_ENV) { + Ok(raw) => raw + .trim() + .parse() + .map(Duration::from_secs) + .unwrap_or_else(|_| { + eprintln!( + "warning: {MAX_IDLE_ENV}={raw:?} is not a valid number of seconds; using {}s", + DEFAULT_MAX_IDLE.as_secs() + ); + DEFAULT_MAX_IDLE + }), + Err(_) => DEFAULT_MAX_IDLE, + }; + Self { + target_size, + max_idle, + } + } + + /// Whether pooling is on (`target_size > 0`). + pub fn enabled(&self) -> bool { + self.target_size > 0 + } +} + +/// Per-key warm sockets, behind a single `Mutex`. Realtime warm sockets are few +/// (the pool is small), so a plain mutex over a `VecDeque`-ish `Vec` is simpler +/// and faster than sharding; contention is negligible at this scale. +type Warm = HashMap>; + +/// Per-key replenish backoff. Absent (or `consecutive_failures == 0`) means the +/// key is healthy and replenished every tick. After a pass whose dials all fail, +/// `retry_after` is pushed out with exponential backoff so a broken key (invalid +/// credentials, unreachable upstream) is not re-dialed on every 250 ms tick. +#[derive(Default)] +struct Backoff { + /// Don't attempt warm-up dials for this key until this instant. `None` = + /// eligible now. + retry_after: Option, + consecutive_failures: u32, +} + +type Backoffs = HashMap; + +/// Pre-warmed upstream realtime connection pool. +/// +/// Cheap to clone-via-`Arc`. The background replenisher is spawned by +/// [`RealtimePool::spawn`]; a pool built with [`RealtimePool::disabled`] never +/// warms anything and every `take` misses (callers fresh-dial). +pub struct RealtimePool { + config: PoolConfig, + warm: Mutex, + /// Per-key replenish backoff so a broken key doesn't trigger unbounded + /// concurrent dials every tick. Separate lock from `warm` so the request + /// hot path (`take`) never contends on it. + backoff: Mutex, +} + +impl RealtimePool { + /// A disabled pool: no background task, every `take` returns `None`. + pub fn disabled() -> Arc { + Arc::new(Self { + config: PoolConfig { + target_size: 0, + ..PoolConfig::default() + }, + warm: Mutex::new(HashMap::new()), + backoff: Mutex::new(HashMap::new()), + }) + } + + /// Build a pool from config **without** the background replenisher. The pool + /// only warms when [`RealtimePool::warm_now`] is called. Used by deterministic + /// unit tests; production uses [`RealtimePool::spawn`]. + #[cfg(test)] + fn new_unspawned(config: PoolConfig) -> Arc { + Arc::new(Self { + config, + warm: Mutex::new(HashMap::new()), + backoff: Mutex::new(HashMap::new()), + }) + } + + /// Build a pool from config and, if enabled, spawn the background replenisher. + /// Returns the shared handle the gateway stores in its state. + pub fn spawn(config: PoolConfig) -> Arc { + let pool = Arc::new(Self { + config, + warm: Mutex::new(HashMap::new()), + backoff: Mutex::new(HashMap::new()), + }); + if config.enabled() { + let weak = Arc::downgrade(&pool); + tokio::spawn(async move { + let mut tick = tokio::time::interval(REPLENISH_TICK); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tick.tick().await; + // Stop once the gateway has dropped its handle. + let Some(pool) = weak.upgrade() else { break }; + pool.replenish_all().await; + } + }); + } + pool + } + + /// Resolved config (test/inspection). + pub fn config(&self) -> PoolConfig { + self.config + } + + /// Register a key so the replenisher starts warming it. Idempotent. The + /// gateway calls this once per known deployment at startup; the pool only + /// warms keys it has seen, so it never dials a model nobody asked for. + pub fn register(&self, key: UpstreamKey) { + if !self.config.enabled() { + return; + } + self.warm.lock().unwrap().entry(key).or_default(); + } + + /// Take a warm, live socket for `key`, or `None` on miss / dead socket. + /// + /// Pops the freshest non-expired socket and liveness-checks it; a socket that + /// is too old or already dead is dropped (closing it) and the next candidate + /// tried. Never blocks: if nothing warm is live, returns `None` so the caller + /// fresh-dials. + pub fn take(&self, key: &UpstreamKey) -> Option { + if !self.config.enabled() { + return None; + } + loop { + let mut candidate = { + let mut warm = self.warm.lock().unwrap(); + let bucket = warm.get_mut(key)?; + bucket.pop()? + }; + // Discard sockets past their warm lifetime (idle-billing guard). + if candidate.warmed_at.elapsed() > self.config.max_idle { + continue; // drops `candidate`, closing the socket + } + // Liveness: a non-blocking check that the socket hasn't already + // delivered a Close/Err. A warm socket should be silent after + // session.created, so anything pending means it is unhealthy. + if is_dead(&mut candidate.rx) { + continue; + } + return Some(WarmHandoff { + tx: candidate.tx, + rx: candidate.rx, + session_created: candidate.session_created, + }); + } + } + + /// One replenish pass over every registered key: reap stale sockets, then + /// dial up to `target_size`. Dials run concurrently; failures are swallowed + /// (a key that can't be warmed just keeps fresh-dialing on the request path) + /// and put the key into exponential backoff so a broken key isn't re-dialed + /// on every tick. + async fn replenish_all(&self) { + let keys: Vec = { self.warm.lock().unwrap().keys().cloned().collect() }; + for key in keys { + self.reap_stale(&key); + // Skip keys still in backoff from a prior all-failed pass — this is + // what bounds dials against an invalid/unreachable key to once per + // `BACKOFF_MAX` instead of `needed` dials every 250 ms tick. + if self.in_backoff(&key) { + continue; + } + let needed = { + let warm = self.warm.lock().unwrap(); + let have = warm.get(&key).map(Vec::len).unwrap_or(0); + self.config.target_size.saturating_sub(have) + }; + if needed == 0 { + continue; + } + // Dial the missing sockets CONCURRENTLY. A sequential loop here makes + // a full refill cost `needed × handshake` (~needed × 350 ms), which + // can't keep up with a high connect rate — the pool drains faster + // than it refills and most connects miss. Firing the dials together + // refills in ~one handshake window, keeping warm supply ≈ peak + // concurrent connects so the sub-ms warm handoff becomes the median, + // not the lucky-hit tail. + let dials = (0..needed).map(|_| warm_one(&key)); + let results = futures_util::future::join_all(dials).await; + let mut any_ok = false; + // `.flatten()` keeps only the successful dials; a key that can't be + // warmed just keeps fresh-dialing on the request path. + for conn in results.into_iter().flatten() { + any_ok = true; + self.warm + .lock() + .unwrap() + .entry(key.clone()) + .or_default() + .push(conn); + } + // Reset backoff on any success; otherwise grow it. We only ever enter + // backoff when a pass that *attempted* dials produced none — a `needed + // == 0` pass is handled by the `continue` above and never touches it. + self.record_replenish_outcome(&key, any_ok); + } + } + + /// Whether `key` is currently in a backoff window (a prior pass failed and + /// the retry time hasn't arrived). Eligible keys are pruned from the backoff + /// map so it doesn't grow unbounded for healthy keys. + fn in_backoff(&self, key: &UpstreamKey) -> bool { + let mut backoff = self.backoff.lock().unwrap(); + match backoff.get(key).and_then(|b| b.retry_after) { + Some(retry_after) if Instant::now() < retry_after => true, + Some(_) => { + // Window elapsed — allow the attempt. Keep the failure count so a + // still-broken key backs off further, but clear the gate so this + // tick proceeds. + if let Some(b) = backoff.get_mut(key) { + b.retry_after = None; + } + false + } + None => false, + } + } + + /// Update a key's backoff after a replenish attempt. Success clears it; + /// failure grows the retry delay exponentially up to `BACKOFF_MAX`. + fn record_replenish_outcome(&self, key: &UpstreamKey, any_ok: bool) { + let mut backoff = self.backoff.lock().unwrap(); + if any_ok { + backoff.remove(key); + return; + } + let entry = backoff.entry(key.clone()).or_default(); + entry.consecutive_failures = entry.consecutive_failures.saturating_add(1); + // Exponential: BASE * 2^(failures-1), saturating at MAX. `min` of the + // shift exponent keeps the doubling from overflowing. + let shift = (entry.consecutive_failures - 1).min(16); + let delay = BACKOFF_BASE.saturating_mul(1u32 << shift).min(BACKOFF_MAX); + entry.retry_after = Some(Instant::now() + delay); + } + + /// Drop sockets past `max_idle` or already dead for a key. + fn reap_stale(&self, key: &UpstreamKey) { + let mut warm = self.warm.lock().unwrap(); + if let Some(bucket) = warm.get_mut(key) { + bucket.retain_mut(|conn| { + conn.warmed_at.elapsed() <= self.config.max_idle && !is_dead(&mut conn.rx) + }); + } + } + + /// Test/inspection: number of warm sockets currently held for `key`. + #[cfg(test)] + pub fn warm_len(&self, key: &UpstreamKey) -> usize { + self.warm + .lock() + .unwrap() + .get(key) + .map(Vec::len) + .unwrap_or(0) + } + + /// Test/inspection: consecutive replenish failures recorded for `key` (0 if + /// the key is healthy / has no backoff entry). + #[cfg(test)] + pub fn backoff_failures(&self, key: &UpstreamKey) -> u32 { + self.backoff + .lock() + .unwrap() + .get(key) + .map(|b| b.consecutive_failures) + .unwrap_or(0) + } + + /// Test helper: synchronously warm `target_size` sockets for `key` (no + /// background task). Lets tests assert handoff behavior deterministically. + #[cfg(test)] + pub async fn warm_now(&self, key: &UpstreamKey) { + let needed = { + let warm = self.warm.lock().unwrap(); + let have = warm.get(key).map(Vec::len).unwrap_or(0); + self.config.target_size.saturating_sub(have) + }; + for _ in 0..needed { + if let Ok(conn) = warm_one(key).await { + self.warm + .lock() + .unwrap() + .entry(key.clone()) + .or_default() + .push(conn); + } + } + } + + /// Test helper: insert an already-built warm connection (used to inject a + /// dead socket and assert it is discarded at handoff). + #[cfg(test)] + fn insert_warm(&self, key: UpstreamKey, conn: WarmConnection) { + self.warm.lock().unwrap().entry(key).or_default().push(conn); + } +} + +/// Dial one upstream and pre-read its `session.created` into a [`WarmConnection`]. +/// +/// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends +/// unprompted is `session.created`; we buffer exactly that and read nothing more. +async fn warm_one(key: &UpstreamKey) -> CoreResult { + let upstream: UpstreamWs = + dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?; + let (tx, mut rx) = upstream.split(); + let session_created = read_event(&mut rx).await?; + Ok(WarmConnection { + tx, + rx, + session_created, + warmed_at: Instant::now(), + }) +} + +/// Resolve a deployment's API key into the pool key, returning `None` when no key +/// can be resolved (those deployments simply aren't pooled — the request path +/// still fresh-dials and surfaces the auth error there). +pub fn upstream_key( + model: &str, + api_key: Option<&str>, + api_base: Option<&str>, +) -> Option { + let api_key = resolve_api_key(api_key).ok()?; + Some(UpstreamKey { + model: model.to_string(), + api_key, + api_base: api_base.map(str::to_string), + }) +} + +/// Non-blocking liveness check: poll the upstream once. A warm socket is silent +/// after `session.created`, so a pending `Close`/`Err`/`None` means it is dead. +/// A pending data frame (shouldn't happen pre-handoff) is also treated as +/// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an +/// unexpected state. `Pending` (the healthy case) returns `false`. +fn is_dead(rx: &mut UpstreamRx) -> bool { + use futures_util::task::noop_waker_ref; + use futures_util::Stream; + use std::pin::Pin; + use std::task::{Context, Poll}; + + let mut cx = Context::from_waker(noop_waker_ref()); + match Pin::new(rx).poll_next(&mut cx) { + Poll::Pending => false, + Poll::Ready(None) => true, + Poll::Ready(Some(Err(_))) => true, + // Any frame arriving before handoff is unexpected for a silent warm + // socket; treat it as unhealthy. + Poll::Ready(Some(Ok(_))) => true, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures_util::SinkExt; + use std::net::SocketAddr; + use tokio::net::TcpListener; + use tokio_tungstenite::tungstenite::Message; + + /// An in-process fake OpenAI realtime WS server. On connect it sends + /// `session.created`; on `response.create` it sends `response.created` + + /// `response.output_audio.delta` + `response.done`. Returns its `ws://` base. + async fn spawn_fake_openai() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr: SocketAddr = listener.local_addr().unwrap(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + tokio::spawn(handle_fake_conn(stream)); + } + }); + format!("ws://{addr}") + } + + async fn handle_fake_conn(stream: tokio::net::TcpStream) { + let mut ws = match tokio_tungstenite::accept_async(stream).await { + Ok(ws) => ws, + Err(_) => return, + }; + // Unprompted session.created, exactly like OpenAI. + let _ = ws + .send(Message::Text( + r#"{"type":"session.created","session":{"id":"sess_fake"}}"#.to_string(), + )) + .await; + while let Some(Ok(msg)) = ws.next().await { + if let Message::Text(text) = msg { + if text.contains("response.create") { + for frame in [ + r#"{"type":"response.created"}"#, + r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#, + r#"{"type":"response.done"}"#, + ] { + let _ = ws.send(Message::Text(frame.to_string())).await; + } + } + } + } + } + + fn test_config() -> PoolConfig { + PoolConfig { + target_size: 2, + max_idle: Duration::from_secs(30), + } + } + + fn key_for(base: &str) -> UpstreamKey { + UpstreamKey { + model: "gpt-realtime".to_string(), + api_key: "sk-test".to_string(), + api_base: Some(base.to_string()), + } + } + + #[tokio::test] + async fn warm_handoff_relays_buffered_session_created() { + let base = spawn_fake_openai().await; + let pool = RealtimePool::new_unspawned(test_config()); + let key = key_for(&base); + pool.register(key.clone()); + pool.warm_now(&key).await; + assert_eq!(pool.warm_len(&key), 2); + + let handoff = pool.take(&key).expect("a warm socket should be available"); + assert_eq!(handoff.session_created.event_type, "session.created"); + assert_eq!( + handoff + .session_created + .data + .get("session") + .and_then(|s| s.get("id")) + .and_then(|v| v.as_str()), + Some("sess_fake") + ); + // Taking one leaves one. + assert_eq!(pool.warm_len(&key), 1); + } + + #[tokio::test] + async fn pool_miss_returns_none_for_fresh_dial_fallback() { + let base = spawn_fake_openai().await; + let pool = RealtimePool::new_unspawned(test_config()); + let key = key_for(&base); + // Registered but never warmed → empty bucket → miss. + pool.register(key.clone()); + assert!(pool.take(&key).is_none()); + + // Unknown key → miss. + let other = key_for("ws://127.0.0.1:1"); + assert!(pool.take(&other).is_none()); + } + + #[tokio::test] + async fn disabled_pool_never_hands_off() { + let pool = RealtimePool::disabled(); + let key = key_for("ws://127.0.0.1:1"); + pool.register(key.clone()); + assert_eq!(pool.warm_len(&key), 0); + assert!(pool.take(&key).is_none()); + } + + #[tokio::test] + async fn dead_warm_socket_is_discarded() { + let base = spawn_fake_openai().await; + let pool = RealtimePool::new_unspawned(test_config()); + let key = key_for(&base); + pool.register(key.clone()); + + // Build one real warm connection, then kill the upstream by dropping the + // server side: easiest is to dial, read session.created, then close our + // own rx's peer. Instead we forge "dead" via an already-closed socket: + // dial a connection and immediately send a Close from the client side so + // the server closes back, then warm it. Simpler: warm normally, then + // mark it stale by backdating warmed_at past max_idle and confirm it's + // dropped — that exercises the same discard path. + let mut conn = warm_one(&key).await.expect("warm one"); + conn.warmed_at = Instant::now() - Duration::from_secs(3600); // past max_idle + pool.insert_warm(key.clone(), conn); + assert_eq!(pool.warm_len(&key), 1); + + // take() must discard the stale socket and report a miss. + assert!(pool.take(&key).is_none()); + assert_eq!(pool.warm_len(&key), 0); + } + + #[tokio::test] + async fn background_replenisher_tops_up_registered_key() { + let base = spawn_fake_openai().await; + let pool = RealtimePool::spawn(test_config()); + let key = key_for(&base); + pool.register(key.clone()); + + // Wait (bounded) for the background task to reach the target size. + let mut warmed = 0; + for _ in 0..40 { + tokio::time::sleep(Duration::from_millis(50)).await; + warmed = pool.warm_len(&key); + if warmed >= test_config().target_size { + break; + } + } + assert_eq!( + warmed, + test_config().target_size, + "background replenisher should warm up to target_size" + ); + let handoff = pool.take(&key).expect("a warm socket should be available"); + assert_eq!(handoff.session_created.event_type, "session.created"); + } + + #[tokio::test] + async fn closed_upstream_socket_is_detected_dead() { + // A genuinely dead socket: dial the fake, read session.created, then drop + // the server by closing from our side and waiting for the close to land. + let base = spawn_fake_openai().await; + let pool = RealtimePool::new_unspawned(test_config()); + let key = key_for(&base); + pool.register(key.clone()); + + let mut conn = warm_one(&key).await.expect("warm one"); + // Close the upstream from the client side; the server echoes a close. + let _ = conn.tx.send(Message::Close(None)).await; + // Give the close a moment to arrive on rx. + tokio::time::sleep(Duration::from_millis(50)).await; + pool.insert_warm(key.clone(), conn); + + // Liveness check at take() should detect the close and discard it. + assert!(pool.take(&key).is_none()); + assert_eq!(pool.warm_len(&key), 0); + } + + #[tokio::test] + async fn broken_key_backs_off_instead_of_dialing_every_tick() { + // A key whose upstream is unreachable: every warm-up dial fails. + let pool = RealtimePool::new_unspawned(test_config()); + let key = key_for("ws://127.0.0.1:1"); // nothing listens here + pool.register(key.clone()); + + // First pass attempts dials, they all fail → key enters backoff, no warm + // sockets, one recorded failure. + pool.replenish_all().await; + assert_eq!(pool.warm_len(&key), 0); + assert_eq!(pool.backoff_failures(&key), 1); + assert!( + pool.in_backoff(&key), + "a key whose dials all failed must be in backoff" + ); + + // An immediate next pass must be SKIPPED (still in the backoff window), so + // it does NOT fire another round of dials — the failure count is unchanged. + pool.replenish_all().await; + assert_eq!( + pool.backoff_failures(&key), + 1, + "replenish during the backoff window must not re-dial the broken key" + ); + } + + #[tokio::test] + async fn healthy_key_never_enters_backoff_and_clears_after_recovery() { + let base = spawn_fake_openai().await; + let pool = RealtimePool::new_unspawned(test_config()); + let key = key_for(&base); + pool.register(key.clone()); + + // A reachable upstream: the pass succeeds, so the key is never backed off. + pool.replenish_all().await; + assert_eq!(pool.warm_len(&key), test_config().target_size); + assert_eq!(pool.backoff_failures(&key), 0); + assert!(!pool.in_backoff(&key)); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs new file mode 100644 index 00000000000..6c04fbb7626 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -0,0 +1,38 @@ +//! LiteLLM AI Gateway library. +//! +//! Two layers, split by feature so the Python `cdylib` can depend on the I/O +//! without pulling in the HTTP server: +//! +//! - [`io`]: all network I/O (OCR HTTP call, realtime WebSocket splice, the +//! pre-warmed realtime pool). Always available — no feature required. The +//! Python bridge links this for `run_ocr`. +//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling +//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` +//! binary turns on. The `python-config` feature additionally pulls in [`python`] +//! for the load-time config reader. + +pub mod io; + +/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and +/// the `python-config` reader, so it is available without either feature. +pub mod gil; + +#[cfg(feature = "server")] +pub mod auth; +#[cfg(feature = "server")] +pub mod routes; +#[cfg(feature = "server")] +pub mod state; + +// Realtime request logging. Only the server serves realtime, so these are +// `server`-gated; `io::realtime` exposes the generic `observe` hook while the +// collector and callback fan-out live here. +#[cfg(feature = "server")] +mod constants; +#[cfg(feature = "server")] +pub mod integrations; +#[cfg(feature = "server")] +mod realtime; + +#[cfg(feature = "python-config")] +pub mod python; diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs new file mode 100644 index 00000000000..f9ce97801d3 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -0,0 +1,162 @@ +//! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router. +//! +//! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment +//! (simple-shuffle) → `io::realtime::realtime()` invokes OpenAI. The +//! server owns transport + config; routing lives in the `router` crate. +//! +//! The binary requires the `server` feature (declared in `Cargo.toml` via +//! `required-features`), so cargo skips it unless that feature is on. Everything +//! the binary needs lives in the library (`litellm_ai_gateway`); `main` just +//! wires startup. + +use std::sync::Arc; + +use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool}; +use litellm_ai_gateway::routes; +use litellm_ai_gateway::state::AppState; +use litellm_core::router::{Deployment, LiteLLMParams, Router}; + +use litellm_ai_gateway::integrations::custom_logger::CustomLogger; +use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger; +#[cfg(feature = "python-config")] +use litellm_ai_gateway::python; + +/// Bind to localhost by default so the gateway is not a public, unauthenticated +/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`). +const DEFAULT_HOST: &str = "127.0.0.1"; +const DEFAULT_PORT: u16 = 4001; + +#[tokio::main] +async fn main() { + // Trim before storing so it matches the trimmed bearer token in `auth` + // (avoids a silent auth failure when the env var has surrounding whitespace). + let master_key: Option> = std::env::var("LITELLM_MASTER_KEY") + .ok() + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty()) + .map(Arc::from); + if master_key.is_none() { + eprintln!( + "warning: LITELLM_MASTER_KEY is not set; /v1/realtime will reject all requests (fail closed)" + ); + } + + // Spawn the realtime-logging worker (drains a channel → POSTs batches to the + // Python proxy's /v1/callbacks/logs). Built here so the spawn lands on the + // tokio runtime. `from_env` reads LITELLM_PROXY_BASE_URL + LITELLM_MASTER_KEY. + let proxy_logger = LiteLLMPythonProxyAPILogger::from_env(); + let loggers: Vec> = vec![proxy_logger]; + + let router = Arc::new(build_router()); + + // Build the pre-warmed realtime pool and register each deployment's upstream + // so the background replenisher starts warming it. `REALTIME_POOL_SIZE=0` + // yields a disabled pool → every connect fresh-dials (original behavior). + let pool_config = PoolConfig::from_env(); + let realtime_pool = RealtimePool::spawn(pool_config); + if pool_config.enabled() { + register_deployments(&router, &realtime_pool); + eprintln!( + "realtime connection pool enabled: target {} warm sockets/key, max idle {}s", + pool_config.target_size, + pool_config.max_idle.as_secs() + ); + } else { + eprintln!( + "realtime connection pool disabled (REALTIME_POOL_SIZE=0); fresh-dialing each connect" + ); + } + + let state = AppState { + router, + master_key, + loggers: Arc::new(loggers), + realtime_pool, + }; + + let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string()); + let port = resolve_port(); + + let listener = tokio::net::TcpListener::bind((host.as_str(), port)) + .await + .expect("failed to bind listener"); + eprintln!("litellm-ai-gateway listening on {host}:{port}"); + axum::serve(listener, routes::app(state)) + .await + .expect("server error"); +} + +/// Register every deployment's upstream key with the pool so the replenisher +/// pre-warms it. Mirrors `service::run`'s key derivation (strip `openai/`, resolve +/// api_key); deployments whose key can't be resolved are skipped (they fresh-dial +/// and surface the auth error on the request path, as before). +fn register_deployments(router: &Router, pool: &RealtimePool) { + for deployment in router.deployments() { + let params = &deployment.litellm_params; + let provider_model = params + .model + .strip_prefix("openai/") + .unwrap_or(¶ms.model); + if let Some(key) = upstream_key( + provider_model, + params.api_key.as_deref(), + params.api_base.as_deref(), + ) { + pool.register(key); + } + } +} + +/// Resolve `PORT`, warning (rather than silently defaulting) on an invalid value. +fn resolve_port() -> u16 { + match std::env::var("PORT") { + Ok(raw) => raw.parse().unwrap_or_else(|_| { + eprintln!("warning: PORT={raw:?} is not a valid port; using {DEFAULT_PORT}"); + DEFAULT_PORT + }), + Err(_) => DEFAULT_PORT, + } +} + +/// Build the router. With the `python-config` feature and `LITELLM_CONFIG_PATH` +/// set, load the resolved `model_list` from the proxy config via the embedded +/// Python reader (load time only). Otherwise fall back to the env stand-in. +fn build_router() -> Router { + #[cfg(feature = "python-config")] + if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") { + match python::config::load_router_from_config(&config_path) { + Ok(router) => { + eprintln!("loaded model_list from {config_path} via python config reader"); + return router; + } + Err(err) => { + eprintln!("config load failed ({err}); falling back to env deployment"); + } + } + } + build_router_from_env() +} + +/// Build a minimal single-deployment `model_list` from the environment. +/// +/// A real deployment loads `model_list` from config; this is the minimal stand-in +/// so the gateway has one OpenAI deployment to route to. +fn build_router_from_env() -> Router { + let model = + std::env::var("OPENAI_REALTIME_MODEL").unwrap_or_else(|_| "gpt-realtime".to_string()); + let api_key = std::env::var("OPENAI_API_KEY").ok(); + if api_key.is_none() { + eprintln!( + "warning: OPENAI_API_KEY is not set; realtime requests will fail with auth errors" + ); + } + let deployment = Deployment { + model_name: model.clone(), + litellm_params: LiteLLMParams { + model, + api_key, + api_base: None, + }, + }; + Router::new(vec![deployment]) +} diff --git a/litellm-rust/crates/ai-gateway/src/python/AGENTS.md b/litellm-rust/crates/ai-gateway/src/python/AGENTS.md new file mode 100644 index 00000000000..47aa117e0b9 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/python/AGENTS.md @@ -0,0 +1,27 @@ +# ai-gateway/src/python — Python interop (load-time only) + +Functions here embed the Python interpreter (pyo3) and take the GIL to call into +`litellm` (e.g. read the proxy `model_list`). Compiled only under the +`python-config` feature. + +## Hard rule: non-hot-path functions only + +Everything in this folder MUST run **at most once per process lifetime — at +startup / load time** (config read, warm-up). NEVER call into Python on the +request path: + +- No GIL acquisition per request, per connection, or per realtime event. +- No Python call inside a route handler, the router's hot path, or any loop that + scales with traffic. + +**Why:** the GIL serializes execution and would cap throughput; the realtime data +path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll +`GET /health/gil`, and `total_acquisitions` MUST stay flat under load. + +## How to add one + +Resolve whatever Python-derived data you need **once at boot** and hand the rest +of the gateway an owned, plain-Rust value (e.g. build a `Router` from the +resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()` +immediately before taking the GIL. If a function would need to run per request, +it does not belong here — move the work to Rust, or pre-resolve it at startup. diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs new file mode 100644 index 00000000000..6ec9595469d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -0,0 +1,39 @@ +//! Build the router by calling the Python proxy config reader (load time only). +//! +//! Embeds the interpreter via pyo3 and calls +//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's +//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot** +//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python. +//! +//! Compiled only under the `python-config` feature. + +use litellm_core::error::CoreError; +use litellm_core::router::{Deployment, Router}; +use litellm_core::CoreResult; +use pyo3::prelude::*; + +use crate::gil; + +/// Load the router's `model_list` from `config_path` via the Python reader. +pub fn load_router_from_config(config_path: &str) -> CoreResult { + gil::record_acquisition(); + Python::with_gil(|py| { + let model_list = py + .import("litellm.proxy.read_model_list") + .and_then(|module| module.getattr("read_model_list")) + .and_then(|reader| reader.call1((config_path,))) + .map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?; + + let model_list_json: String = py + .import("json") + .and_then(|json| json.getattr("dumps")) + .and_then(|dumps| dumps.call1((model_list,))) + .and_then(|encoded| encoded.extract()) + .map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?; + + let deployments: Vec = serde_json::from_str(&model_list_json) + .map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?; + + Ok(Router::new(deployments)) + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/python/mod.rs b/litellm-rust/crates/ai-gateway/src/python/mod.rs new file mode 100644 index 00000000000..a677bade676 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/python/mod.rs @@ -0,0 +1,4 @@ +//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path +//! only.** Compiled only under the `python-config` feature. + +pub mod config; diff --git a/litellm-rust/crates/ai-gateway/src/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/realtime/mod.rs new file mode 100644 index 00000000000..82be596ba86 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/realtime/mod.rs @@ -0,0 +1,4 @@ +//! Realtime logging collector. Observes the realtime event stream and emits a +//! `StandardLoggingPayload` to the registered callbacks on session close. + +pub mod streaming; diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs new file mode 100644 index 00000000000..34c82897808 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -0,0 +1,352 @@ +//! `RealTimeStreaming` — the realtime logging collector. +//! +//! Mirrors Python `litellm.realtime_api.main.RealTimeStreaming`: it observes the +//! event stream in O(1) (never buffering frames), accumulating just the fields +//! the spend log needs (model, id, cumulative usage), then on session close +//! builds a `StandardLoggingPayload` and fans it out to every registered +//! `CustomLogger`. + +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::realtime::types::RealtimeEvent; +use serde_json::Value; + +use crate::constants::DEFAULT_PROVIDER; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::{ + RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, Usage, +}; + +/// Current wall-clock time as epoch seconds (float), matching the Python +/// `startTime`/`endTime` contract. +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) +} + +/// Status of a finished realtime session, mapped to the callback record status. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SessionStatus { + Success, + Failure, +} + +/// Accumulates realtime session state and emits a logging payload on close. +pub struct RealTimeStreaming { + callbacks: Vec>, + /// REQUEST-ID RULE: the SpendLogs `request_id` == the OpenAI realtime session + /// id (`sess_…`), captured from `session.created`. Both `id` and + /// `litellm_call_id` are set to that value so the Python writer logs the same + /// id regardless of which field it reads. The gateway-generated `rt-…` id + /// (the constructor seed) is only a fallback for sessions that fail before + /// `session.created` arrives. + litellm_call_id: String, + /// See the request-id rule above — mirrors `litellm_call_id`. + id: String, + model: String, + custom_llm_provider: String, + usage: Usage, + response_cost: f64, + start_time: f64, + end_time: f64, + metadata: RequestMetadata, + /// Count of logging callbacks that failed to enqueue (non-fatal). + dropped: u64, +} + +impl RealTimeStreaming { + /// Create a collector for one session. `litellm_call_id` is the gateway's + /// per-connection id; `model` is the requested model (a sane default until + /// `session.created` reports the upstream model). + pub fn new( + callbacks: Vec>, + litellm_call_id: String, + model: String, + metadata: RequestMetadata, + ) -> Self { + let now = epoch_seconds(); + Self { + callbacks, + id: litellm_call_id.clone(), + litellm_call_id, + model, + custom_llm_provider: DEFAULT_PROVIDER.to_string(), + usage: Usage::default(), + response_cost: 0.0, + start_time: now, + end_time: now, + metadata, + dropped: 0, + } + } + + /// Number of logging callbacks that failed to enqueue so far (test/observ.). + #[allow(dead_code)] + pub fn dropped(&self) -> u64 { + self.dropped + } + + /// Observe one realtime event. O(1): updates accumulated state only; never + /// buffers frames. Safe to call on every event in either direction. + pub fn observe(&mut self, event: &RealtimeEvent) { + match event.event_type.as_str() { + "session.created" | "session.updated" => self.on_session(event), + "response.done" => self.on_response_done(event), + _ => {} + } + } + + /// `session.created` / `session.updated` → capture upstream id + model. + /// Per the request-id rule, the OpenAI session id becomes BOTH `id` and + /// `litellm_call_id`, replacing the gateway-generated fallback. + fn on_session(&mut self, event: &RealtimeEvent) { + let session = event.data.get("session").and_then(Value::as_object); + if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) { + if !id.is_empty() { + self.id = id.to_string(); + self.litellm_call_id = id.to_string(); + } + } + if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) { + if !model.is_empty() { + self.model = model.to_string(); + } + } + } + + /// `response.done` → add this response's usage to the cumulative totals. + fn on_response_done(&mut self, event: &RealtimeEvent) { + let usage = event + .data + .get("response") + .and_then(Value::as_object) + .and_then(|r| r.get("usage")) + .and_then(Value::as_object); + let Some(usage) = usage else { return }; + + let input = usage.get("input_tokens").and_then(Value::as_u64); + let output = usage.get("output_tokens").and_then(Value::as_u64); + let total = usage.get("total_tokens").and_then(Value::as_u64); + + if let Some(input) = input { + self.usage.prompt_tokens += input; + } + if let Some(output) = output { + self.usage.completion_tokens += output; + } + // Prefer the upstream-reported total; otherwise derive it. + match total { + Some(total) => self.usage.total_tokens += total, + None => { + self.usage.total_tokens += input.unwrap_or(0) + output.unwrap_or(0); + } + } + } + + /// Set the per-session response cost ($). Cost computation is Python-side in + /// the proxy; the gateway forwards 0.0 by default and lets the proxy price. + /// Public API (exercised in tests) for the future path where the gateway + /// prices realtime sessions itself. + #[allow(dead_code)] + pub fn set_response_cost(&mut self, cost: f64) { + self.response_cost = cost; + } + + /// Build the `StandardLoggingPayload` from accumulated state. + pub fn build_payload(&self) -> StandardLoggingPayload { + StandardLoggingPayload { + id: self.id.clone(), + litellm_call_id: self.litellm_call_id.clone(), + call_type: "realtime".to_string(), + model: self.model.clone(), + custom_llm_provider: self.custom_llm_provider.clone(), + response_cost: self.response_cost, + prompt_tokens: self.usage.prompt_tokens, + completion_tokens: self.usage.completion_tokens, + total_tokens: self.usage.total_tokens, + start_time: self.start_time, + end_time: self.end_time, + stream: true, + metadata: StandardLoggingMetadata { + user_api_key_hash: self.metadata.user_api_key_hash.clone(), + user_api_key_user_id: self.metadata.user_api_key_user_id.clone(), + user_api_key_team_id: self.metadata.user_api_key_team_id.clone(), + ..Default::default() + }, + messages: None, + } + } + + /// Finish the session: stamp the end time and fan the payload out to every + /// callback. On a logger enqueue error we bump a non-fatal counter (the + /// realtime session has already ended; a dropped log must never propagate). + pub fn log_messages(&mut self, status: SessionStatus) { + self.end_time = epoch_seconds(); + let payload = self.build_payload(); + + match status { + SessionStatus::Success => { + for callback in &self.callbacks { + if let Err(err) = callback.log_success_event(&payload) { + self.dropped += 1; + eprintln!("litellm-ai-gateway: log_success_event dropped: {err}"); + } + } + } + SessionStatus::Failure => { + let error = crate::integrations::types::LoggingError { + message: "realtime session ended in failure".to_string(), + kind: "RealtimeSessionError".to_string(), + }; + for callback in &self.callbacks { + if let Err(err) = callback.log_failure_event(&payload, &error) { + self.dropped += 1; + eprintln!("litellm-ai-gateway: log_failure_event dropped: {err}"); + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::types::{LogError, LoggingError}; + use std::sync::atomic::{AtomicU64, Ordering}; + + fn event(raw: &str) -> RealtimeEvent { + serde_json::from_str(raw).expect("valid event json") + } + + /// A test logger that records the last payload it saw. + #[derive(Default)] + struct CapturingLogger { + calls: AtomicU64, + last_model: std::sync::Mutex>, + last_total_tokens: AtomicU64, + } + + impl CustomLogger for CapturingLogger { + fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> { + self.calls.fetch_add(1, Ordering::SeqCst); + *self.last_model.lock().unwrap() = Some(payload.model.clone()); + self.last_total_tokens + .store(payload.total_tokens, Ordering::SeqCst); + Ok(()) + } + } + + #[test] + fn observe_accumulates_model_and_tokens_then_logs() { + let logger = Arc::new(CapturingLogger::default()); + let callbacks: Vec> = vec![logger.clone()]; + let mut streaming = RealTimeStreaming::new( + callbacks, + "call_abc".to_string(), + "gpt-realtime".to_string(), + RequestMetadata { + user_api_key_hash: Some("hash123".to_string()), + user_api_key_user_id: Some("user-1".to_string()), + user_api_key_team_id: Some("team-1".to_string()), + }, + ); + + streaming.observe(&event( + r#"{"type":"session.created","session":{"id":"sess_001","model":"gpt-realtime-2025"}}"#, + )); + streaming.observe(&event( + r#"{"type":"response.done","response":{"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}"#, + )); + // A second response.done accumulates. + streaming.observe(&event( + r#"{"type":"response.done","response":{"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}"#, + )); + + let payload = streaming.build_payload(); + assert_eq!(payload.model, "gpt-realtime-2025"); + // Request-id rule: session.created's id becomes BOTH id and + // litellm_call_id (replacing the "call_abc" gateway fallback), so the + // SpendLogs request_id is always the OpenAI session id. + assert_eq!(payload.id, "sess_001"); + assert_eq!(payload.litellm_call_id, "sess_001"); + assert_eq!(payload.prompt_tokens, 13); + assert_eq!(payload.completion_tokens, 7); + assert_eq!(payload.total_tokens, 20); + assert_eq!(payload.response_cost, 0.0); + assert_eq!(payload.call_type, "realtime"); + assert_eq!(payload.custom_llm_provider, "openai"); + assert_eq!( + payload.metadata.user_api_key_hash.as_deref(), + Some("hash123") + ); + + streaming.log_messages(SessionStatus::Success); + assert_eq!(logger.calls.load(Ordering::SeqCst), 1); + assert_eq!( + logger.last_model.lock().unwrap().as_deref(), + Some("gpt-realtime-2025") + ); + assert_eq!(logger.last_total_tokens.load(Ordering::SeqCst), 20); + assert_eq!(streaming.dropped(), 0); + } + + #[test] + fn payload_serializes_with_camelcase_times_and_realtime_call_type() { + let mut streaming = RealTimeStreaming::new( + Vec::new(), + "call_xyz".to_string(), + "gpt-realtime".to_string(), + RequestMetadata::default(), + ); + streaming.observe(&event( + r#"{"type":"response.done","response":{"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}"#, + )); + streaming.set_response_cost(0.0042); + let payload = streaming.build_payload(); + let json = serde_json::to_string(&payload).expect("serialize payload"); + + assert!(json.contains("\"startTime\""), "missing startTime: {json}"); + assert!(json.contains("\"endTime\""), "missing endTime: {json}"); + assert!( + json.contains("\"call_type\":\"realtime\""), + "missing call_type realtime: {json}" + ); + assert!( + json.contains("\"response_cost\""), + "missing response_cost: {json}" + ); + assert_eq!(payload.response_cost, 0.0042); + } + + /// A logger whose enqueue always fails should bump the dropped counter, not + /// panic or propagate. + #[test] + fn failing_logger_bumps_dropped_counter() { + struct FailingLogger; + impl CustomLogger for FailingLogger { + fn log_success_event(&self, _p: &StandardLoggingPayload) -> Result<(), LogError> { + Err(LogError::channel_full()) + } + fn log_failure_event( + &self, + _p: &StandardLoggingPayload, + _e: &LoggingError, + ) -> Result<(), LogError> { + Err(LogError::channel_closed()) + } + } + let callbacks: Vec> = vec![Arc::new(FailingLogger)]; + let mut streaming = RealTimeStreaming::new( + callbacks, + "call_1".to_string(), + "gpt-realtime".to_string(), + RequestMetadata::default(), + ); + streaming.log_messages(SessionStatus::Success); + assert_eq!(streaming.dropped(), 1); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md new file mode 100644 index 00000000000..02c5f18c4f3 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md @@ -0,0 +1,38 @@ +# routes/ — the route template + +Every route follows the **same shape** so the layout is predictable. The rule: + +> **Each route module exposes `pub fn router() -> Router`.** +> `routes/mod.rs::app` merges them all and applies state once. Adding a route is: +> create the module, then add one `.merge(::router())` line. + +## Default: one file +A route is a single file containing `router()` + its handler(s) (handlers stay +private). This is the norm — don't split until it hurts. +``` +pub fn router() -> Router { Router::new().route(PATH, get(handle)) } +async fn handle(...) -> impl IntoResponse { ... } +``` +`health.rs` and `gil.rs` are examples. + +## Split out `service` when there's real logic +When a route has business logic worth testing without axum, put it in a sibling +`service` (a file, or a folder if the route grows). The route file stays the +**axum surface** (router + handler + any socket/SSE adapter); `service` is plain +Rust with **no axum types**. `realtime/` is the example: +``` +realtime/ + mod.rs # axum surface: router() + handler + the WS<->events adapter + service.rs # pure logic: select deployment + call provider (no axum) — testable +``` +Split `service` further (or add `transport`, `repo`, …) only once a single file +genuinely gets hard to read. + +## Invariants +- **Auth is an extractor, not a manual call.** A handler requires auth by adding + `crate::auth::RequireMasterKey` to its arguments; it runs during extraction. + Never re-implement the check per route. +- **Handlers contain no business logic; `service` contains no axum types.** +- A route owns its paths in its own `router()`; `mod.rs` only merges. +- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`, + not duplicated in handlers. diff --git a/litellm-rust/crates/ai-gateway/src/routes/gil.rs b/litellm-rust/crates/ai-gateway/src/routes/gil.rs new file mode 100644 index 00000000000..0db0c6f0b14 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/gil.rs @@ -0,0 +1,30 @@ +//! `GET /health/gil` — poll to confirm Python is only touched at load time. +//! Simple-route template: a `router()` plus its handler, in one file. + +use axum::routing::get; +use axum::{Json, Router}; +use serde::Serialize; + +use crate::gil; +use crate::state::AppState; + +/// This route's contribution to the app router. +pub fn router() -> Router { + Router::new().route("/health/gil", get(status)) +} + +#[derive(Debug, Serialize)] +struct GilStatusResponse { + gil_acquired_last_30s: bool, + total_acquisitions: u64, + seconds_since_last: Option, +} + +async fn status() -> Json { + let snapshot = gil::snapshot(); + Json(GilStatusResponse { + gil_acquired_last_30s: snapshot.acquired_last_30s, + total_acquisitions: snapshot.total_acquisitions, + seconds_since_last: snapshot.seconds_since_last, + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/health.rs b/litellm-rust/crates/ai-gateway/src/routes/health.rs new file mode 100644 index 00000000000..15c67fea325 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/health.rs @@ -0,0 +1,24 @@ +//! Health probes. Simple-route template: a `router()` plus its handlers, in one file. + +use axum::http::StatusCode; +use axum::routing::get; +use axum::Router; + +use crate::state::AppState; + +/// This route's contribution to the app router. +pub fn router() -> Router { + Router::new() + .route("/health/liveness", get(liveness)) + .route("/health/readiness", get(readiness)) +} + +/// The process is up. +async fn liveness() -> StatusCode { + StatusCode::OK +} + +/// The server is ready to accept traffic. +async fn readiness() -> StatusCode { + StatusCode::OK +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/mod.rs new file mode 100644 index 00000000000..c6b9573781a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/mod.rs @@ -0,0 +1,23 @@ +//! HTTP routes. +//! +//! **Template:** every route module exposes `pub fn router() -> Router` +//! that mounts its own paths; [`app`] merges them. A trivial route is a single +//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with +//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md. + +pub mod gil; +pub mod health; +pub mod realtime; + +use axum::Router; + +use crate::state::AppState; + +/// Assemble the application router by merging every route module's `router()`. +pub fn app(state: AppState) -> Router { + Router::new() + .merge(health::router()) + .merge(gil::router()) + .merge(realtime::router()) + .with_state(state) +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md b/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md new file mode 100644 index 00000000000..3301576bb85 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md @@ -0,0 +1,87 @@ +# Realtime route (`GET /v1/realtime`) + +Proxies OpenAI's realtime WebSocket. `mod.rs` is the axum surface (handler + +socket↔events adapter); `service.rs` is the pure logic (select a deployment, then +splice client ↔ upstream). The pool itself lives in +`crates/providers/src/realtime_pool.rs`. + +## Connection pooling + +### The problem + +The gateway's realtime overhead lives **entirely in session establishment**. On each +client connect it dials a *fresh* upstream WS to OpenAI and waits for +`session.created` before it can serve. Measured at 5000 calls / 500 concurrency, the +fresh-dial session phase is **~360 ms** vs **~7 ms** direct; dial, first-audio, and +streaming add ~0. So the one lever is removing that per-connect handshake from the +critical path. + +### The idea + +Keep a few upstream OpenAI sockets **already connected and already past +`session.created`** (buffered). On a client connect, hand off a warm socket — relay +its buffered `session.created` instantly (a local `Vec::pop`, sub-millisecond) and +splice exactly as a fresh dial would. A background task keeps the pool topped up. On +a miss or dead socket we fall back to fresh-dial: the pool is a latency optimization, +never a correctness dependency. + +``` + ┌───────────────────────────────────────┐ + client connect ──────► │ routes/realtime → service::run │ + │ pool.take(key) │ + │ hit → relay buffered │ + │ session.created, then splice │ + │ miss → fresh dial (original path) │ + └───────────────┬───────────────────────┘ + │ replenish (async, concurrent) + ┌───────────────▼───────────────────────┐ + background task ─────► │ RealtimePool: per-key warm sockets │ + │ each = { ws, buffered session.created}│ + │ liveness-checked before handoff │ + └─────────────────────────────────────────┘ +``` + +A warm session is indistinguishable from a fresh one: OpenAI sends `session.created` +unprompted on connect, we pre-read exactly that one frame and relay it on handoff, +and we send nothing else on the socket before a client exists — so the client's first +`session.update` behaves identically either way. + +### Sizing + +Each warm socket serves **exactly one** session (realtime isn't multiplexed), so the +pool is sized to the **peak concurrent connects per instance**, not total live +connections: + +``` +REALTIME_POOL_SIZE ≈ peak_concurrency / instance_count +``` + +e.g. 500 concurrency over 10 instances → ~50–64 per instance. The replenisher dials +the missing sockets **concurrently**, so a drained pool refills in ~one handshake +window and keeps supply close to the connect rate. Over-provisioning just burns idle +upstream sockets, which is why warm sockets are short-lived +(`REALTIME_POOL_MAX_IDLE_SECS`). + +### Config + +| env | default | meaning | +| ----------------------------- | ------- | --------------------------------------------------------------- | +| `REALTIME_POOL_SIZE` | `4` | target warm sockets per key. `0` disables pooling (fresh-dial). | +| `REALTIME_POOL_MAX_IDLE_SECS` | `30` | max time a warm socket sits before it's closed and replaced. | + +### Notes + +- **Miss / dead socket → fresh dial.** Burst beyond warm supply, or a socket that + died, never blocks or fails — it falls back to the original path. The pool can only + make a connect faster, never slower or more fragile. +- **Auth scope.** The pool key includes `api_key`, so a warm socket is only handed to + a request resolving to the same key — no cross-tenant reuse. +- **Idle billing.** Warm sockets are liveness-checked at handoff and capped at + `REALTIME_POOL_MAX_IDLE_SECS` to bound idle billing and dodge OpenAI's idle timeout. +- **Replenish backoff.** If a key's warm-up dials all fail (invalid credentials, an + unreachable upstream), the replenisher puts that key into exponential backoff + (500 ms → 30 s cap) instead of re-dialing it every tick. This bounds connection + attempts against a broken key so it can't exhaust upstream rate limits and degrade + valid cold-path traffic; the backoff resets the moment a dial succeeds. + +Benchmarks and repro: `../../benchmarks/realtime/README.md`. diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs new file mode 100644 index 00000000000..899ad73829f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -0,0 +1,166 @@ +//! `GET /v1/realtime` (WebSocket). +//! +//! This file is the **axum surface**: `router()`, the handler, and the small +//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is +//! the `RequireMasterKey` extractor, so the handler stays thin. + +mod service; + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::io::realtime_pool::RealtimePool; +use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::response::Response; +use axum::routing::get; +use axum::Router; +use futures_util::{SinkExt, StreamExt}; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::router::Router as ModelRouter; +use serde::Deserialize; + +use crate::auth::RequireMasterKey; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::RequestMetadata; +use crate::realtime::streaming::{RealTimeStreaming, SessionStatus}; +use crate::state::AppState; + +/// Process-local monotonic counter, mixed into the per-session call id so two +/// sessions opened in the same nanosecond still get distinct ids. +static CALL_SEQ: AtomicU64 = AtomicU64::new(0); + +/// Generate a per-connection `litellm_call_id`. No external uuid dep: epoch +/// nanos + a process-local sequence is unique enough for log correlation. +fn new_call_id() -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + let seq = CALL_SEQ.fetch_add(1, Ordering::Relaxed); + format!("rt-{nanos:x}-{seq:x}") +} + +/// This route's contribution to the app router. +pub fn router() -> Router { + Router::new().route("/v1/realtime", get(handle)) +} + +#[derive(Debug, Deserialize)] +struct RealtimeQuery { + model: String, +} + +/// Auth runs via the `RequireMasterKey` extractor. We validate the model BEFORE +/// the upgrade so failures are clean HTTP (400/404), not a socket that opens then +/// closes, then hand the socket to `bridge`. +async fn handle( + _auth: RequireMasterKey, + ws: WebSocketUpgrade, + State(state): State, + Query(query): Query, +) -> Result { + if query.model.trim().is_empty() { + return Err(( + StatusCode::BAD_REQUEST, + "missing 'model' query param".to_string(), + )); + } + if !state.router.has_deployment(&query.model) { + return Err(( + StatusCode::NOT_FOUND, + format!("no deployment for model '{}'", query.model), + )); + } + + let router = state.router.clone(); + let pool = state.realtime_pool.clone(); + let loggers = state.loggers.clone(); + let master_key = state.master_key.clone(); + let model = query.model; + Ok(ws.on_upgrade(move |socket| bridge(socket, router, pool, loggers, master_key, model))) +} + +/// Adapt the axum socket (text frames) to the typed-event `Stream`/`Sink` the +/// service wants, keeping axum types out of `service`. +/// +/// This is also the realtime-logging seam: every upstream→client event (the +/// direction carrying `session.created` and `response.done` with usage) is fed +/// to a [`RealTimeStreaming`] collector via the splice's `observe` callback. The +/// observe is O(1) and never buffers frames. When the splice returns (any of the +/// three break paths — client disconnect, upstream close, idle timeout), we flush +/// one logging payload to the registered callbacks. +async fn bridge( + socket: WebSocket, + router: Arc, + pool: Arc, + loggers: Arc>>, + master_key: Option>, + model: String, +) { + let (ws_sink, ws_stream) = socket.split(); + + // Attribute the spend log to the key that authenticated this session (the + // master key — the gateway is master-key auth). A non-null user_api_key_hash + // is required for the Python spend logger to write a SpendLogs row. + // + // SECURITY: hash the key — never send the raw credential. This field fans out + // to spend logs and every callback integration; the SHA-256 (matching the + // proxy's hash_token) keeps the plaintext master key out of all of them while + // still matching the key's hash in LiteLLM_SpendLogs. + let metadata = RequestMetadata { + user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token), + ..RequestMetadata::default() + }; + + // Owned by THIS task only. The splice observes it via a synchronous `&mut` + // callback (below), so there is no Arc/Mutex/atomic on the per-frame hot + // path — just a monomorphized FnMut mutating stack-local fields. This is + // what lets observe scale: 10K concurrent sessions = 10K independent + // collectors, zero cross-task synchronization. + let mut collector = RealTimeStreaming::new( + loggers.as_ref().clone(), + new_call_id(), + model.clone(), + metadata, + ); + + let client_in = ws_stream.filter_map(|message| async move { + match message { + Ok(Message::Text(text)) => serde_json::from_str::(&text).ok(), + _ => None, + } + }); + // Plain forwarding sink — no observe here anymore. + let client_out = ws_sink.with(|event: RealtimeEvent| async move { + Ok::(Message::Text( + serde_json::to_string(&event).unwrap_or_default(), + )) + }); + + futures_util::pin_mut!(client_in, client_out); + + // The observe closure borrows `&mut collector` for the duration of the + // splice; the borrow ends when `run` returns, freeing the collector for the + // single post-session `log_messages` flush. `run` picks a pooled (warm) or + // fresh upstream — observe fires on the upstream arm either way. + let result = service::run( + &router, + &pool, + &model, + None, + |event: &RealtimeEvent| collector.observe(event), + client_in, + client_out, + ) + .await; + + let status = if result.is_ok() { + SessionStatus::Success + } else { + SessionStatus::Failure + }; + collector.log_messages(status); +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs new file mode 100644 index 00000000000..d6c31edd454 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -0,0 +1,79 @@ +//! Business logic: select a deployment with the (pure) core router, then call the +//! provider splice. The seam between `core::router` (selection only) and +//! `io` (the actual WebSocket I/O). +//! +//! On connect we try a pre-warmed upstream from the pool (handshake already paid, +//! `session.created` buffered) and relay it instantly. On a pool miss or dead warm +//! socket we fresh-dial exactly as before — the pool is never on the critical path +//! for correctness, only latency. + +use std::time::Duration; + +use crate::io::realtime_pool::{upstream_key, RealtimePool}; +use futures_util::{Sink, Stream}; +use litellm_core::error::CoreError; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::router::Router; +use litellm_core::CoreResult; + +/// Select a deployment for `model` and splice the client stream to the provider. +/// +/// `pool` supplies a pre-warmed upstream when one is available; otherwise we +/// fresh-dial. A disabled pool always misses, so this collapses to the original +/// fresh-dial behavior. +pub async fn run( + router: &Router, + pool: &RealtimePool, + model: &str, + idle_timeout: Option, + observe: impl FnMut(&RealtimeEvent) + Send, + client_in: In, + client_out: Out, +) -> CoreResult<()> +where + In: Stream + Unpin + Send, + Out: Sink + Unpin + Send, + >::Error: std::fmt::Display, +{ + let deployment = router.get_available_deployment(model).ok_or_else(|| { + CoreError::Routing(format!("no deployment available for model '{model}'")) + })?; + let params = &deployment.litellm_params; + // Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model. + let provider_model = params + .model + .strip_prefix("openai/") + .unwrap_or(¶ms.model); + + // Warm path: take a pooled upstream (handshake already paid) and relay its + // buffered session.created immediately. On miss/dead socket fall through. + if let Some(key) = upstream_key( + provider_model, + params.api_key.as_deref(), + params.api_base.as_deref(), + ) { + if let Some(handoff) = pool.take(&key) { + return crate::io::realtime::realtime_warm( + provider_model, + handoff, + idle_timeout, + observe, + client_in, + client_out, + ) + .await; + } + } + + // Cold path: fresh dial (the original behavior). + crate::io::realtime::realtime( + provider_model, + params.api_key.as_deref(), + params.api_base.as_deref(), + idle_timeout, + observe, + client_in, + client_out, + ) + .await +} diff --git a/litellm-rust/crates/ai-gateway/src/state.rs b/litellm-rust/crates/ai-gateway/src/state.rs new file mode 100644 index 00000000000..3b61d8309ea --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/state.rs @@ -0,0 +1,21 @@ +use std::sync::Arc; + +use crate::io::realtime_pool::RealtimePool; +use litellm_core::router::Router; + +use crate::integrations::custom_logger::CustomLogger; + +/// Shared application state handed to every route handler. +#[derive(Clone)] +pub struct AppState { + pub router: Arc, + /// The gateway master key. Any caller presenting it as a bearer token may + /// invoke the gateway. `None` → auth not configured (routes fail closed). + pub master_key: Option>, + /// Logging callbacks fanned out at the end of each realtime session. + pub loggers: Arc>>, + /// Pre-warmed upstream realtime connection pool. Disabled + /// (`RealtimePool::disabled()`) when `REALTIME_POOL_SIZE=0`, in which case + /// every realtime connect fresh-dials exactly as before. + pub realtime_pool: Arc, +} diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md new file mode 100644 index 00000000000..8740dccaf01 --- /dev/null +++ b/litellm-rust/crates/core/AGENTS.md @@ -0,0 +1,3 @@ +litellm-core is the PURE translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. No network, no I/O, no env reads. + +Routes (ocr, realtime) and providers (mistral, openai) are modules, not crates. diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md new file mode 100644 index 00000000000..20873878967 --- /dev/null +++ b/litellm-rust/crates/core/CLAUDE.md @@ -0,0 +1,47 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/core`. + +## Responsibility + +`core` owns shared data types, typed errors, and deterministic helper contracts. +It must stay pure and host-independent. + +Allowed: +- 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 + `ocr::transformation::OcrProviderConfig`. + +Not allowed: +- Network, filesystem, database, cache, or environment access. +- Secret reads or auth/header construction. +- 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/`: `ocr`, future `messages`, +`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not +invent broad names like `engine` for route contracts. + +## 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 new file mode 100644 index 00000000000..1881bcfa602 --- /dev/null +++ b/litellm-rust/crates/core/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "litellm-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +rand.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs new file mode 100644 index 00000000000..9b29260cca4 --- /dev/null +++ b/litellm-rust/crates/core/src/error.rs @@ -0,0 +1,35 @@ +use thiserror::Error; + +pub type CoreResult = Result; + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum CoreError { + #[error("expected {expected}, got {actual}")] + InvalidType { + expected: &'static str, + actual: &'static str, + }, + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error("invalid response: {0}")] + InvalidResponse(String), + #[error("{0}")] + Auth(String), + #[error("OCR request failed with status {status}: {body}")] + Http { status: u16, body: String }, + #[error("OCR network error: {0}")] + Network(String), + #[error("routing error: {0}")] + Routing(String), +} + +pub fn json_type_name(value: &serde_json::Value) -> &'static str { + match value { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "bool", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs new file mode 100644 index 00000000000..2ac479cc725 --- /dev/null +++ b/litellm-rust/crates/core/src/lib.rs @@ -0,0 +1,7 @@ +pub mod error; +pub mod ocr; +pub mod providers; +pub mod realtime; +pub mod router; + +pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs new file mode 100644 index 00000000000..7353d9d22c4 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -0,0 +1,32 @@ +use serde_json::{Map, Value}; + +use crate::CoreResult; + +use super::types::{OcrRequestData, OcrResponseData}; + +pub trait OcrProviderConfig { + fn supported_ocr_params(&self) -> &'static [&'static str]; + + fn map_ocr_params(&self, non_default_params: &Map) -> Map { + let mut mapped_params = Map::new(); + for (param, value) in non_default_params { + if self.supported_ocr_params().contains(¶m.as_str()) { + mapped_params.insert(param.clone(), value.clone()); + } + } + mapped_params + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult; + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs new file mode 100644 index 00000000000..1a72b8f1d66 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -0,0 +1,29 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrRequestData { + pub data: Value, + pub files: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrResponseData { + pub pages: Vec, + pub model: String, + pub document_annotation: Option, + pub usage_info: Option, + pub object: String, +} + +impl OcrResponseData { + pub fn into_json(self) -> Value { + serde_json::json!({ + "pages": self.pages, + "model": self.model, + "document_annotation": self.document_annotation, + "usage_info": self.usage_info, + "object": self.object, + }) + } +} diff --git a/litellm-rust/crates/core/src/providers/mistral/mod.rs b/litellm-rust/crates/core/src/providers/mistral/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mistral/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs new file mode 100644 index 00000000000..d5155991448 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -0,0 +1,292 @@ +use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{Map, Value}; + +const SUPPORTED_OCR_PARAMS: &[&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", + "id", +]; + +/// Default Mistral API base, used when the caller does not override `api_base`. +pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1"; + +/// Environment variable holding the Mistral API key. +pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; + +/// Error message raised when no Mistral API key can be resolved. +pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params"; + +/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`. +/// +/// Blank/whitespace `api_base` is treated as absent (guard at resolution time). +pub fn complete_url(api_base: Option<&str>) -> String { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(MISTRAL_DEFAULT_API_BASE) + .trim_end_matches('/'); + + if base.ends_with("/v1") { + format!("{base}/ocr") + } else { + format!("{base}/v1/ocr") + } +} + +/// Resolve the Mistral API key from the explicit param or the environment. +/// +/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth` +/// when no usable key is available. +/// +/// Note: the env fallback only reads the process environment. Secret-manager +/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in +/// via `api_key`; this fallback is a last resort for direct/standalone use. +pub fn resolve_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + api_key + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) +} + +pub struct MistralOcrConfig; + +pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig; + +impl OcrProviderConfig for MistralOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + if !document.is_object() { + return Err(CoreError::InvalidType { + expected: "object", + actual: json_type_name(&document), + }); + } + + let mut data = Map::new(); + data.insert("model".to_string(), Value::String(model.to_string())); + data.insert("document".to_string(), document); + for (param, value) in optional_params { + data.insert(param, value); + } + + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response_object = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + + let pages = response_object + .get("pages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let model = response_object + .get("model") + .and_then(Value::as_str) + .unwrap_or(model) + .to_string(); + let document_annotation = response_object.get("document_annotation").cloned(); + let usage_info = response_object.get("usage_info").cloned(); + + Ok(OcrResponseData { + pages, + model, + document_annotation, + usage_info, + object: "ocr".to_string(), + }) + } +} + +pub fn supported_ocr_params() -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() +} + +pub fn map_ocr_params(non_default_params: &Map) -> Map { + MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params) +} + +pub fn transform_ocr_request( + model: &str, + document: Value, + optional_params: Map, +) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) +} + +pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn supported_params_match_python_mistral_ocr_config() { + assert_eq!( + supported_ocr_params(), + &[ + "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", + "id", + ] + ); + } + + #[test] + fn map_ocr_params_drops_unknown_params() { + let params = json!({ + "extract_header": true, + "unsupported_param": "value", + "pages": [0, 1] + }); + let mapped = map_ocr_params(params.as_object().unwrap()); + + assert_eq!(mapped.get("extract_header"), Some(&json!(true))); + assert_eq!(mapped.get("pages"), Some(&json!([0, 1]))); + assert!(!mapped.contains_key("unsupported_param")); + } + + #[test] + fn transform_ocr_request_builds_mistral_body() { + let document = json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }); + let optional_params = json!({ + "include_image_base64": true, + "table_format": "html" + }) + .as_object() + .unwrap() + .clone(); + + let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params) + .expect("request should transform"); + + assert_eq!( + result.data, + json!({ + "model": "mistral-ocr-latest", + "document": document, + "include_image_base64": true, + "table_format": "html" + }) + ); + assert_eq!(result.files, None); + } + + #[test] + fn transform_ocr_request_rejects_non_object_document() { + let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new()) + .expect_err("string document should be rejected"); + + assert_eq!( + err, + CoreError::InvalidType { + expected: "object", + actual: "string", + } + ); + } + + #[test] + fn transform_ocr_response_normalizes_mistral_json() { + let response = json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": null, + "usage_info": {"pages_processed": 1} + }); + + let result = transform_ocr_response("mistral-ocr-latest", response) + .expect("response should transform"); + + assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]); + assert_eq!(result.model, "mistral-ocr-2505-completion"); + assert_eq!(result.document_annotation, Some(Value::Null)); + assert_eq!(result.usage_info, Some(json!({"pages_processed": 1}))); + assert_eq!(result.object, "ocr"); + } + + #[test] + fn complete_url_defaults_and_dedupes_v1() { + assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr"); + assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr"); + assert_eq!( + complete_url(Some("https://proxy.internal")), + "https://proxy.internal/v1/ocr" + ); + assert_eq!( + complete_url(Some("https://proxy.internal/v1/")), + "https://proxy.internal/v1/ocr" + ); + } + + #[test] + fn resolve_api_key_prefers_param_then_env() { + let no_env = |_: &str| None; + assert_eq!( + resolve_api_key(Some("sk-param"), &no_env).unwrap(), + "sk-param" + ); + + let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string()); + assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env"); + // Blank param falls through to the environment. + assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env"); + } + + #[test] + fn resolve_api_key_errors_when_absent() { + let err = resolve_api_key(None, &|_| None).expect_err("missing key should error"); + assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string())); + } +} diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs new file mode 100644 index 00000000000..42207f0de0a --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -0,0 +1,2 @@ +pub mod mistral; +pub mod openai; diff --git a/litellm-rust/crates/core/src/providers/openai/mod.rs b/litellm-rust/crates/core/src/providers/openai/mod.rs new file mode 100644 index 00000000000..403e32975cf --- /dev/null +++ b/litellm-rust/crates/core/src/providers/openai/mod.rs @@ -0,0 +1 @@ +pub mod realtime; diff --git a/litellm-rust/crates/core/src/providers/openai/realtime/mod.rs b/litellm-rust/crates/core/src/providers/openai/realtime/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/openai/realtime/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs new file mode 100644 index 00000000000..626e4014ff9 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -0,0 +1,189 @@ +use crate::realtime::transformation::RealtimeProviderConfig; +use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; +use crate::CoreResult; + +/// Default OpenAI API base, used when the caller does not override `api_base`. +pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com"; + +/// Path appended to the resolved host base to reach the realtime endpoint. +pub const OPENAI_REALTIME_PATH: &str = "/v1/realtime"; + +/// Percent-encode a query value, escaping any char outside the RFC 3986 +/// unreserved set (`A-Za-z0-9-._~`). Keeps us dependency-free; common realtime +/// model slugs have no special chars, but this stays correct for the rest. +fn percent_encode(value: &str) -> String { + let mut encoded = String::with_capacity(value.len()); + for byte in value.bytes() { + let unreserved = byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~'); + if unreserved { + encoded.push(byte as char); + } else { + encoded.push('%'); + encoded.push_str(&format!("{byte:02X}")); + } + } + encoded +} + +/// Build the realtime WebSocket URL, porting Python's `OpenAIRealtime._construct_url`. +/// +/// Blank/whitespace `api_base` is treated as absent (guard at resolution time), +/// falling back to the default. The scheme is swapped to its WebSocket +/// equivalent (`https://`→`wss://`, `http://`→`ws://`); bases already using +/// `ws`/`wss` are left untouched. A bare host or unrecognized scheme defaults to +/// secure `wss://` so we never hand a scheme-less URL to the connector (this is +/// a deliberate hardening over Python's `_construct_url`, which would emit a +/// scheme-less URL here). A trailing `/` is trimmed before the path and +/// `?model=` are appended. +pub fn complete_url(api_base: Option<&str>, model: &str) -> String { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(OPENAI_REALTIME_DEFAULT_API_BASE); + + let base = if let Some(rest) = base.strip_prefix("https://") { + format!("wss://{rest}") + } else if let Some(rest) = base.strip_prefix("http://") { + format!("ws://{rest}") + } else if base.starts_with("wss://") || base.starts_with("ws://") { + base.to_string() + } else { + format!("wss://{base}") + }; + + let base = base.trim_end_matches('/'); + + format!( + "{base}{OPENAI_REALTIME_PATH}?model={}", + percent_encode(model) + ) +} + +pub struct OpenAiRealtimeConfig; + +pub const OPENAI_REALTIME_CONFIG: OpenAiRealtimeConfig = OpenAiRealtimeConfig; + +impl RealtimeProviderConfig for OpenAiRealtimeConfig { + fn complete_url(&self, api_base: Option<&str>, model: &str) -> String { + complete_url(api_base, model) + } + + fn transform_realtime_request( + &self, + event: &RealtimeEvent, + _model: &str, + ) -> CoreResult { + Ok(RealtimeTransformResult::passthrough(event.clone())) + } + + fn transform_realtime_response( + &self, + event: &RealtimeEvent, + _model: &str, + ) -> CoreResult { + Ok(RealtimeTransformResult::passthrough(event.clone())) + } +} + +pub fn transform_realtime_request( + event: &RealtimeEvent, + model: &str, +) -> CoreResult { + OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model) +} + +pub fn transform_realtime_response( + event: &RealtimeEvent, + model: &str, +) -> CoreResult { + OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn complete_url_defaults_to_openai_wss() { + assert_eq!( + complete_url(None, "gpt-4o-realtime-preview"), + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_blank_base_uses_default() { + assert_eq!( + complete_url(Some(" "), "gpt-4o-realtime-preview"), + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_swaps_http_to_ws() { + assert_eq!( + complete_url(Some("http://localhost:8080"), "gpt-4o-realtime-preview"), + "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_dedupes_trailing_slash() { + assert_eq!( + complete_url(Some("https://api.openai.com/"), "gpt-4o-realtime-preview"), + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_custom_base() { + assert_eq!( + complete_url(Some("https://oai.azure.example"), "gpt-4o-realtime-preview"), + "wss://oai.azure.example/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_preserves_existing_wss_scheme() { + assert_eq!( + complete_url(Some("wss://api.openai.com"), "gpt-realtime"), + "wss://api.openai.com/v1/realtime?model=gpt-realtime" + ); + } + + #[test] + fn complete_url_bare_host_defaults_to_wss() { + assert_eq!( + complete_url(Some("api.openai.com"), "gpt-realtime"), + "wss://api.openai.com/v1/realtime?model=gpt-realtime" + ); + } + + #[test] + fn complete_url_percent_encodes_model_space() { + assert_eq!( + complete_url(None, "gpt 4o"), + "wss://api.openai.com/v1/realtime?model=gpt%204o" + ); + } + + #[test] + fn transform_realtime_request_passthrough_preserves_event() { + let event: RealtimeEvent = + serde_json::from_str(r#"{"type":"session.update","session":{"voice":"alloy"}}"#) + .expect("valid event"); + let result = + transform_realtime_request(&event, "gpt-realtime").expect("passthrough is infallible"); + assert_eq!(result.events, vec![event]); + } + + #[test] + fn transform_realtime_response_passthrough_preserves_event() { + let event: RealtimeEvent = + serde_json::from_str(r#"{"type":"response.output_audio.delta","delta":"abc=="}"#) + .expect("valid event"); + let result = + transform_realtime_response(&event, "gpt-realtime").expect("passthrough is infallible"); + assert_eq!(result.events, vec![event]); + } +} diff --git a/litellm-rust/crates/core/src/realtime/mod.rs b/litellm-rust/crates/core/src/realtime/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/realtime/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs new file mode 100644 index 00000000000..a4baa27a6c2 --- /dev/null +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -0,0 +1,22 @@ +use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; +use crate::CoreResult; + +pub trait RealtimeProviderConfig { + /// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`). + /// Pure string construction only — no network, no env. + fn complete_url(&self, api_base: Option<&str>, model: &str) -> String; + + /// Transform a client → backend event before it is forwarded upstream. + fn transform_realtime_request( + &self, + event: &RealtimeEvent, + model: &str, + ) -> CoreResult; + + /// Transform a backend → client event before it is forwarded downstream. + fn transform_realtime_response( + &self, + event: &RealtimeEvent, + model: &str, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/realtime/types.rs b/litellm-rust/crates/core/src/realtime/types.rs new file mode 100644 index 00000000000..3b59224b6e9 --- /dev/null +++ b/litellm-rust/crates/core/src/realtime/types.rs @@ -0,0 +1,60 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// A single realtime event exchanged over the WebSocket. +/// +/// The `type` discriminator is a typed field; the remaining fields are +/// preserved losslessly in `data` so a transform can pass an event through, or +/// inspect/modify specific fields, without enumerating every event variant. +/// Wire (de)serialization happens at the host edge — `core`/`providers` operate +/// only on this typed form. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct RealtimeEvent { + #[serde(rename = "type")] + pub event_type: String, + #[serde(flatten)] + pub data: Map, +} + +/// One or more typed events produced by a realtime transform. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct RealtimeTransformResult { + pub events: Vec, +} + +impl RealtimeTransformResult { + /// Forward a single event unchanged (the OpenAI baseline). + pub fn passthrough(event: RealtimeEvent) -> Self { + Self { + events: vec![event], + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn event(raw: &str) -> RealtimeEvent { + serde_json::from_str(raw).expect("valid event json") + } + + #[test] + fn realtime_event_round_trips_type_and_extra_fields() { + let raw = r#"{"type":"response.output_text.delta","delta":"hi","response_id":"r1"}"#; + let parsed = event(raw); + assert_eq!(parsed.event_type, "response.output_text.delta"); + assert_eq!(parsed.data.get("delta"), Some(&Value::String("hi".into()))); + // Re-serializing yields a semantically-equal event (key order may differ). + let reparsed: RealtimeEvent = + serde_json::from_str(&serde_json::to_string(&parsed).unwrap()).unwrap(); + assert_eq!(parsed, reparsed); + } + + #[test] + fn passthrough_produces_single_element_vec() { + let parsed = event(r#"{"type":"session.update"}"#); + let result = RealtimeTransformResult::passthrough(parsed.clone()); + assert_eq!(result.events, vec![parsed]); + } +} diff --git a/litellm-rust/crates/core/src/router/deployment.rs b/litellm-rust/crates/core/src/router/deployment.rs new file mode 100644 index 00000000000..1ee88e682a3 --- /dev/null +++ b/litellm-rust/crates/core/src/router/deployment.rs @@ -0,0 +1,44 @@ +//! `model_list` data types, mirroring Python's deployment dict. Deserialize-ready +//! so a deployment can be loaded straight from the proxy config's `model_list`. + +use serde::Deserialize; + +/// Per-deployment call parameters, mirroring Python's `litellm_params`. +#[derive(Clone, Debug, Deserialize)] +pub struct LiteLLMParams { + /// Provider model, e.g. `gpt-realtime` or `openai/gpt-realtime`. + pub model: String, + #[serde(default)] + pub api_key: Option, + #[serde(default)] + pub api_base: Option, +} + +/// One entry of the `model_list`, mirroring Python's deployment dict. +#[derive(Clone, Debug, Deserialize)] +pub struct Deployment { + /// Public alias clients request, e.g. `gpt-realtime`. + pub model_name: String, + pub litellm_params: LiteLLMParams, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn deserializes_from_model_list_entry() { + let entry = r#"{ + "model_name": "gpt-realtime", + "litellm_params": {"model": "openai/gpt-realtime", "api_base": "https://x"} + }"#; + let deployment: Deployment = serde_json::from_str(entry).expect("valid entry"); + assert_eq!(deployment.model_name, "gpt-realtime"); + assert_eq!(deployment.litellm_params.model, "openai/gpt-realtime"); + assert_eq!(deployment.litellm_params.api_key, None); + assert_eq!( + deployment.litellm_params.api_base.as_deref(), + Some("https://x") + ); + } +} diff --git a/litellm-rust/crates/core/src/router/mod.rs b/litellm-rust/crates/core/src/router/mod.rs new file mode 100644 index 00000000000..96bc91bc6b5 --- /dev/null +++ b/litellm-rust/crates/core/src/router/mod.rs @@ -0,0 +1,93 @@ +//! Minimal Rust port of LiteLLM's `router.py` deployment selection. +//! +//! A [`Router`] is built from a `model_list` of [`Deployment`]s +//! (`{ model_name, litellm_params: { model, api_key, api_base } }`) and selects +//! one per request via a [`RoutingStrategy`]. For now the only strategy is +//! `simple-shuffle` — a uniform random pick within a `model_name` group. +//! +//! This stays pure (no I/O): it only *chooses* a deployment. The host (the +//! gateway) takes the chosen deployment and performs the actual provider call. +//! +//! - [`deployment`] — the `model_list` data types. +//! - [`strategy`] — how a deployment is chosen. + +mod deployment; +mod strategy; + +pub use deployment::{Deployment, LiteLLMParams}; +pub use strategy::RoutingStrategy; + +/// Load-balancing router over a `model_list`. +#[derive(Clone, Debug, Default)] +pub struct Router { + model_list: Vec, + routing_strategy: RoutingStrategy, +} + +impl Router { + /// Build a router from a `model_list` using the default `simple-shuffle` strategy. + pub fn new(model_list: Vec) -> Self { + Self { + model_list, + routing_strategy: RoutingStrategy::SimpleShuffle, + } + } + + /// All deployments in the `model_list`. Read-only; used by the host to + /// enumerate upstreams (e.g. to pre-warm a connection pool per deployment). + pub fn deployments(&self) -> &[Deployment] { + &self.model_list + } + + /// Whether any deployment is registered under `model`. + pub fn has_deployment(&self, model: &str) -> bool { + self.model_list + .iter() + .any(|deployment| deployment.model_name == model) + } + + /// Pick a deployment for `model` per the routing strategy. Returns `None` + /// when no deployment is registered under that `model_name`. + pub fn get_available_deployment(&self, model: &str) -> Option<&Deployment> { + let candidates: Vec<&Deployment> = self + .model_list + .iter() + .filter(|deployment| deployment.model_name == model) + .collect(); + self.routing_strategy.select(&candidates) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn deployment(name: &str, model: &str) -> Deployment { + Deployment { + model_name: name.to_string(), + litellm_params: LiteLLMParams { + model: model.to_string(), + api_key: None, + api_base: None, + }, + } + } + + #[test] + fn selects_a_matching_deployment() { + let router = Router::new(vec![ + deployment("gpt-realtime", "gpt-realtime"), + deployment("other", "other-model"), + ]); + let chosen = router + .get_available_deployment("gpt-realtime") + .expect("a deployment should match"); + assert_eq!(chosen.model_name, "gpt-realtime"); + } + + #[test] + fn unknown_model_returns_none() { + let router = Router::new(vec![deployment("gpt-realtime", "gpt-realtime")]); + assert!(router.get_available_deployment("missing").is_none()); + } +} diff --git a/litellm-rust/crates/core/src/router/strategy/mod.rs b/litellm-rust/crates/core/src/router/strategy/mod.rs new file mode 100644 index 00000000000..7e8ac217db3 --- /dev/null +++ b/litellm-rust/crates/core/src/router/strategy/mod.rs @@ -0,0 +1,26 @@ +//! Routing policy: how the router picks one deployment from a model group. +//! +//! One module per strategy; [`RoutingStrategy::select`] dispatches to it. New +//! strategies (least-busy, latency-based, …) get their own file here. + +mod simple_shuffle; + +use super::Deployment; + +/// How the router chooses among the deployments sharing a `model_name`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RoutingStrategy { + /// Uniform random pick among the matching deployments. + #[default] + SimpleShuffle, +} + +impl RoutingStrategy { + /// Choose one deployment from `candidates` (all sharing the requested + /// `model_name`). Returns `None` when there are no candidates. + pub fn select<'a>(&self, candidates: &[&'a Deployment]) -> Option<&'a Deployment> { + match self { + RoutingStrategy::SimpleShuffle => simple_shuffle::select(candidates), + } + } +} diff --git a/litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs b/litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs new file mode 100644 index 00000000000..74ce0c21e80 --- /dev/null +++ b/litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs @@ -0,0 +1,47 @@ +//! `simple-shuffle`: a uniform random pick among the candidate deployments. + +use rand::seq::SliceRandom; + +use crate::router::Deployment; + +/// Uniform random choice among `candidates` (all sharing the requested +/// `model_name`). Returns `None` when there are no candidates. +pub fn select<'a>(candidates: &[&'a Deployment]) -> Option<&'a Deployment> { + candidates.choose(&mut rand::thread_rng()).copied() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::router::{Deployment, LiteLLMParams}; + + fn deployment(model: &str) -> Deployment { + Deployment { + model_name: "gpt-realtime".to_string(), + litellm_params: LiteLLMParams { + model: model.to_string(), + api_key: None, + api_base: None, + }, + } + } + + #[test] + fn picks_from_candidates() { + let a = deployment("key-a"); + let b = deployment("key-b"); + let candidates = vec![&a, &b]; + for _ in 0..20 { + let chosen = select(&candidates).expect("non-empty"); + assert!(matches!( + chosen.litellm_params.model.as_str(), + "key-a" | "key-b" + )); + } + } + + #[test] + fn empty_candidates_select_none() { + assert!(select(&[]).is_none()); + } +} diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs new file mode 100644 index 00000000000..a56d19b8242 --- /dev/null +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -0,0 +1,93 @@ +//! Enforcement: the litellm-rust workspace has exactly three crates. +//! +//! `core` (pure translation), `ai-gateway` (routes + all network I/O), 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/ai-gateway", "crates/python-bridge"]; + +/// The crate subdirectory names that must exist under `crates/`. +const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "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 immediate subdirectory names under `crates/`. +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)) + .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 new file mode 100644 index 00000000000..d6d3d90e6ab --- /dev/null +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -0,0 +1,3 @@ +litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over litellm-ai-gateway. + +Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call into litellm-ai-gateway. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md new file mode 100644 index 00000000000..efa1a554c9c --- /dev/null +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -0,0 +1,36 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/python-bridge`. + +## Responsibility + +`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. +Keep this crate thin. It adapts Python objects to Rust payloads and returns +Python-compatible dictionaries. + +## Bridge Shape + +- Prefer one stable method per top-level LiteLLM route, for example + `ocr(payload)`. +- Do not add one exported PyO3 function per provider helper unless there is a + measured reason. +- Provider dispatch belongs in Rust route modules such as + `litellm_providers::ocr`, not in this PyO3 crate. +- Python owns rollout state and fallback. Rust should return errors; Python + decides whether to raise or fall back. + +## Data Handling + +- OCR payloads can contain personal data and large base64 images. Do not log + payloads or provider responses. +- Avoid copying large payloads more than needed. The current JSON round-trip is + acceptable for the first scaffold, but future performance work should evaluate + direct PyO3 conversion before expanding Rust coverage to image-heavy paths. +- Do not expose raw Rust errors that include document contents or upstream + bodies. + +## Tests + +- `cargo test --workspace` must compile this crate. +- Python tests must cover bridge disabled, bridge enabled, and module-missing + fallback behavior for every exposed route. diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml new file mode 100644 index 00000000000..f5b29f49cfd --- /dev/null +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "litellm-python-bridge" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[lib] +name = "litellm_python_bridge" +crate-type = ["cdylib"] + +[dependencies] +litellm-core.workspace = true +litellm-ai-gateway = { workspace = true, default-features = false } +pyo3 = { workspace = true, features = ["extension-module"] } +serde_json.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/gil.rs b/litellm-rust/crates/python-bridge/src/gil.rs new file mode 100644 index 00000000000..dc1b591735c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/gil.rs @@ -0,0 +1,32 @@ +//! GIL accounting. +//! +//! A single chokepoint for releasing the GIL around blocking work. Every +//! blocking call in the bridge goes through [`release_gil`] instead of calling +//! `Python::allow_threads` directly, so the release count stays accurate and we +//! have one place to extend later (timing histograms, per-call labels, etc.). + +use std::sync::atomic::{AtomicU64, Ordering}; + +use pyo3::prelude::*; + +/// Number of times the bridge has released the GIL since process start. +static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); + +/// Release the GIL around `f`, recording the release. +/// +/// `f` must not touch any Python state — that is what makes releasing the GIL +/// safe. Returning the value back to Python re-acquires the GIL at the call +/// site, after `f` has finished. +pub fn release_gil(py: Python<'_>, f: F) -> T +where + F: FnOnce() -> T + Send, + T: Send, +{ + GIL_RELEASES.fetch_add(1, Ordering::Relaxed); + py.allow_threads(f) +} + +/// Total GIL releases performed by the bridge so far. +pub fn release_count() -> u64 { + GIL_RELEASES.load(Ordering::Relaxed) +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs new file mode 100644 index 00000000000..50ec7fceeac --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -0,0 +1,100 @@ +use std::time::Duration; + +use litellm_ai_gateway::io::ocr::run_ocr; +use litellm_core::error::CoreError; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyDict}; +use serde_json::{Map, Value}; + +mod gil; + +fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { + let json = py.import("json")?; + let encoded: String = json.call_method1("dumps", (value,))?.extract()?; + serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) +} + +fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { + let json = py.import("json")?; + let encoded = + serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; + Ok(json.call_method1("loads", (encoded,))?.unbind()) +} + +/// Map a core error to the closest Python exception. Caller-input problems +/// (auth, bad types, missing fields) -> `ValueError`; everything else +/// (network, upstream status, parse failures) -> `RuntimeError`. +fn core_error_to_pyerr(err: CoreError) -> PyErr { + match err { + CoreError::Auth(message) => PyValueError::new_err(message), + CoreError::InvalidType { .. } | CoreError::MissingField(_) => { + PyValueError::new_err(err.to_string()) + } + other => PyRuntimeError::new_err(other.to_string()), + } +} + +/// Perform a Mistral OCR call end to end and return the response as a dict. +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))] +fn ocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let document = py_to_json(py, document.bind(py))?; + + let optional_params = match optional_params { + Some(params) => match py_to_json(py, params.bind(py))? { + Value::Object(map) => map, + _ => return Err(PyValueError::new_err("optional_params must be a dict")), + }, + None => Map::new(), + }; + + let timeout = timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }); + + // Release the GIL during the blocking HTTP call (counted for observability). + let result = gil::release_gil(py, || { + run_ocr( + &model, + document, + api_key.as_deref(), + api_base.as_deref(), + optional_params, + timeout, + ) + }); + + match result { + Ok(value) => json_to_py(py, value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe +/// how often the bridge has dropped the GIL for blocking work. +#[pyfunction] +fn gil_stats(py: Python<'_>) -> PyResult> { + let stats = PyDict::new(py); + stats.set_item("releases", gil::release_count())?; + Ok(stats.into_any().unbind()) +} + +#[pymodule] +fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + Ok(()) +} diff --git a/litellm/__init__.py b/litellm/__init__.py index d21234d2a81..9650dc12c97 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -80,6 +80,7 @@ from litellm.constants import ( WANDB_MODELS, REPEATED_STREAMING_CHUNK_LIMIT, request_timeout, + request_timeout_explicitly_set as request_timeout_explicitly_set, open_ai_embedding_models, cohere_embedding_models, bedrock_embedding_models, @@ -389,12 +390,8 @@ require_managed_files: bool = ( enable_caching_on_provider_specific_optional_params: bool = ( False # feature-flag for caching on optional params - e.g. 'top_k' ) -caching: bool = ( - False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -caching_with_models: bool = ( - False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) +caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 cache: Optional["Cache"] = ( None # cache object <- use this - https://docs.litellm.ai/docs/caching ) @@ -415,9 +412,7 @@ forward_traceparent_to_llm_provider: bool = False _current_cost = 0.0 # private variable, used if max budget is set error_logs: Dict = {} -add_function_to_prompt: bool = ( - False # if function calling not supported by api, append function call details to system prompt -) +add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt client_session: Optional[httpx.Client] = None aclient_session: Optional[httpx.AsyncClient] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' @@ -484,9 +479,7 @@ prometheus_end_user_metrics_cleanup_interval_seconds: Optional[float] = 60.0 disable_add_prefix_to_prompt: bool = ( False # used by anthropic, to disable adding prefix to prompt ) -disable_copilot_system_to_assistant: bool = ( - False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. -) +disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. public_mcp_servers: Optional[List[str]] = None public_mcp_hub_strict_whitelist: bool = True public_model_groups: Optional[List[str]] = None @@ -506,17 +499,13 @@ if TYPE_CHECKING: ######## Networking Settings ######## -use_aiohttp_transport: bool = ( - True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. -) +use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead disable_aiohttp_trust_env: bool = ( False # When False, aiohttp will respect HTTP(S)_PROXY env vars ) -force_ipv4: bool = ( - False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. -) +force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. network_mock: bool = False # When True, use mock transport — no real network calls ####### STOP SEQUENCE LIMIT ####### @@ -550,12 +539,12 @@ output_parse_pii: bool = False from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map model_cost = get_model_cost_map(url=model_cost_map_url) -cost_discount_config: Dict[str, float] = ( - {} -) # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount -cost_margin_config: Dict[str, Union[float, Dict[str, float]]] = ( - {} -) # Provider-specific or global cost margins. Examples: +cost_discount_config: Dict[ + str, float +] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount +cost_margin_config: Dict[ + str, Union[float, Dict[str, float]] +] = {} # Provider-specific or global cost margins. Examples: # Percentage: {"openai": 0.10} = 10% margin # Fixed: {"openai": {"fixed_amount": 0.001}} = $0.001 per request # Global: {"global": 0.05} = 5% global margin on all providers @@ -673,6 +662,7 @@ elevenlabs_models: Set = set() dashscope_models: Set = set() moonshot_models: Set = set() publicai_models: Set = set() +darkbloom_models: Set = set() v0_models: Set = set() morph_models: Set = set() lambda_ai_models: Set = set() @@ -927,6 +917,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): moonshot_models.add(key) elif value.get("litellm_provider") == "publicai": publicai_models.add(key) + elif value.get("litellm_provider") == "darkbloom": + darkbloom_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -1075,6 +1067,7 @@ model_list = list( | dashscope_models | moonshot_models | publicai_models + | darkbloom_models | v0_models | morph_models | lambda_ai_models @@ -1179,6 +1172,7 @@ models_by_provider: dict = { "modelscope": modelscope_models, "moonshot": moonshot_models, "publicai": publicai_models, + "darkbloom": darkbloom_models, "v0": v0_models, "morph": morph_models, "lambda_ai": lambda_ai_models, @@ -1400,6 +1394,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * +from .ocr.rust_bridge import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * @@ -1450,9 +1445,9 @@ from . import rag from .types.llms.custom_llm import CustomLLMItem custom_provider_map: List[CustomLLMItem] = [] -_custom_providers: List[str] = ( - [] -) # internal helper util, used to track names of custom providers +_custom_providers: List[ + str +] = [] # internal helper util, used to track names of custom providers disable_hf_tokenizer_download: Optional[bool] = ( None # disable huggingface tokenizer download. Defaults to openai clk100 ) @@ -1922,9 +1917,6 @@ if TYPE_CHECKING: from .llms.fireworks_ai.completion.transformation import ( FireworksAITextCompletionConfig as FireworksAITextCompletionConfig, ) - from .llms.fireworks_ai.audio_transcription.transformation import ( - FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig, - ) from .llms.fireworks_ai.embed.fireworks_ai_transformation import ( FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e653b40fd04..4f131354d2e 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -260,7 +260,6 @@ LLM_CONFIG_NAMES = ( "SambaNovaEmbeddingConfig", "FireworksAIConfig", "FireworksAITextCompletionConfig", - "FireworksAIAudioTranscriptionConfig", "FireworksAIEmbeddingConfig", "FriendliaiChatConfig", "JinaAIEmbeddingConfig", @@ -1027,10 +1026,6 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.fireworks_ai.completion.transformation", "FireworksAITextCompletionConfig", ), - "FireworksAIAudioTranscriptionConfig": ( - ".llms.fireworks_ai.audio_transcription.transformation", - "FireworksAIAudioTranscriptionConfig", - ), "FireworksAIEmbeddingConfig": ( ".llms.fireworks_ai.embed.fireworks_ai_transformation", "FireworksAIEmbeddingConfig", diff --git a/litellm/_redis.py b/litellm/_redis.py index 1b6e1a5e4b0..fdae674d55d 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -327,7 +327,9 @@ def _get_redis_client_logic(**env_overrides): **env_overrides, } - _startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore + _startup_nodes: Optional[Union[str, list]] = redis_kwargs.get( + "startup_nodes", None + ) or get_secret( # type: ignore "REDIS_CLUSTER_NODES" ) @@ -338,7 +340,9 @@ def _get_redis_client_logic(**env_overrides): elif _startup_nodes is None: redis_kwargs.pop("startup_nodes", None) - _sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore + _sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get( + "sentinel_nodes", None + ) or get_secret( # type: ignore "REDIS_SENTINEL_NODES" ) @@ -609,7 +613,8 @@ def get_redis_async_client( # Create async RedisCluster with IAM token as password if available cluster_client = async_redis.RedisCluster( - startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore + startup_nodes=new_startup_nodes, + **cluster_kwargs, # type: ignore ) return cluster_client diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index cb9375e6b84..b7dfed6b169 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -184,7 +184,9 @@ def get_assistants( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) @@ -198,7 +200,9 @@ def get_assistants( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) @@ -394,7 +398,9 @@ def create_assistants( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) if response is None: @@ -761,7 +767,9 @@ def create_thread( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response # type: ignore @@ -916,7 +924,9 @@ def get_thread( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response # type: ignore @@ -1103,7 +1113,9 @@ def add_message( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) @@ -1263,7 +1275,9 @@ def get_messages( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) @@ -1478,7 +1492,9 @@ def run_thread( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response # type: ignore diff --git a/litellm/assistants/utils.py b/litellm/assistants/utils.py index f8fc6ee0af7..bde279602ad 100644 --- a/litellm/assistants/utils.py +++ b/litellm/assistants/utils.py @@ -71,9 +71,7 @@ def get_optional_params_add_message( if custom_llm_provider == "openai": optional_params = non_default_params elif custom_llm_provider == "azure": - supported_params = ( - litellm.AzureOpenAIAssistantsAPIConfig().get_supported_openai_create_message_params() - ) + supported_params = litellm.AzureOpenAIAssistantsAPIConfig().get_supported_openai_create_message_params() _check_valid_arg(supported_params=supported_params) optional_params = litellm.AzureOpenAIAssistantsAPIConfig().map_openai_params_create_message_params( non_default_params=non_default_params, optional_params=optional_params diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 74e753b09ea..aeec58f1dfc 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,5 +1,5 @@ import json -from typing import Any, List, Literal, Optional, Tuple +from typing import Any, Iterator, List, Literal, Optional, Tuple import litellm from litellm._logging import verbose_logger @@ -314,6 +314,70 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: raise e +def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]: + """ + Yield non-empty JSONL lines (unparsed) one at a time, so a caller can parse + each row in its own try/except and a single malformed line cannot abort the + whole pass. Peak memory stays bounded for large batch files. + """ + start, length, newline = 0, len(file_content), ord("\n") + while start < length: + idx = file_content.find(newline, start) + if idx == -1: + chunk, start = file_content[start:], length + else: + chunk, start = file_content[start:idx], idx + 1 + line = chunk.strip() + if line: + yield line + + +def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]: + """ + Yield parsed batch input JSONL entries one at a time without materializing the + whole file as a list, so peak memory stays bounded. Raises on a malformed line; + callers that must survive bad rows should iterate ``_iter_batch_input_lines`` + and parse per-row instead. + """ + for line in _iter_batch_input_lines(file_content): + yield json.loads(line) + + +# A batch request's input tokens scale roughly with its serialized size, so this +# is a conservative per-row fallback when the token counter cannot measure a row. +_BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN = 4 + + +def _estimate_batch_entry_tokens(raw_line: bytes) -> int: + """Conservative token estimate for a batch row the token counter cannot measure + (or that cannot be parsed). Keeps the batch token total non-zero so a crafted + row cannot evade the TPM limit, without hard-rejecting a legitimate batch.""" + return max(1, len(raw_line) // _BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN) + + +def _count_entry_tokens( + entry: dict, + model_name: Optional[str] = None, +) -> int: + """Token-count a single batch input entry's body (chat / text / embedding).""" + body = entry.get("body", {}) or {} + model = body.get("model", model_name or "") + + messages = body.get("messages") + if messages: + return token_counter(model=model, messages=messages) + + prompt = body.get("prompt") + if prompt: + return _count_prompt_or_input_tokens(model=model, value=prompt) + + input_data = body.get("input") + if input_data: + return _count_prompt_or_input_tokens(model=model, value=input_data) + + return 0 + + def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], custom_llm_provider: Literal[ @@ -396,70 +460,6 @@ def _get_batch_job_total_usage_from_file_content( ) -def _get_models_from_batch_input_file_content( - file_content_dictionary: List[dict], -) -> List[str]: - """Extract the distinct ``body.model`` values from a batch *input* file. - - Used by the proxy's batch pre-call hook to enforce that the caller is - authorized for every model named inside the JSONL — not just the one - on the outer request — so the proxy's per-key model allowlist isn't - bypassed by smuggling expensive models into the batch file. - """ - models: List[str] = [] - seen: set = set() - for _item in file_content_dictionary: - body = _item.get("body") or {} - model = body.get("model") - if model and model not in seen: - seen.add(model) - models.append(model) - return models - - -def _get_batch_job_input_file_usage( - file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - model_name: Optional[str] = None, -) -> Usage: - """ - Count the number of tokens in the input file - - Used for batch rate limiting to count the number of tokens in the input file - """ - prompt_tokens: int = 0 - completion_tokens: int = 0 - - for _item in file_content_dictionary: - body = _item.get("body", {}) - model = body.get("model", model_name or "") - - # Chat completion payloads. - messages = body.get("messages") - if messages: - prompt_tokens += token_counter(model=model, messages=messages) - continue - - # Text completion payloads (`prompt`). - prompt = body.get("prompt") - if prompt: - prompt_tokens += _count_prompt_or_input_tokens(model=model, value=prompt) - continue - - # Embedding payloads (`input`). - input_data = body.get("input") - if input_data: - prompt_tokens += _count_prompt_or_input_tokens( - model=model, value=input_data - ) - - return Usage( - total_tokens=prompt_tokens + completion_tokens, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - ) - - def _count_prompt_or_input_tokens(model: str, value: Any) -> int: """Token-count a ``prompt`` / ``input`` field that the OpenAI batch schema allows in four shapes: diff --git a/litellm/batches/main.py b/litellm/batches/main.py index f124882b5a4..c5f1c86a4e1 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -359,7 +359,9 @@ def create_batch( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_batch", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -553,7 +555,9 @@ def _handle_retrieve_batch_providers_without_provider_config( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="retrieve_batch", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -819,7 +823,11 @@ def list_batches( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore + api_base = ( + optional_params.api_base + or litellm.api_base + or get_secret_str("AZURE_API_BASE") + ) # type: ignore api_version = ( optional_params.api_version or litellm.api_version @@ -887,7 +895,9 @@ def list_batches( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -1097,7 +1107,9 @@ def cancel_batch( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="cancel_batch", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="cancel_batch", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index bbebb6042cb..915c22b90ec 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -67,9 +67,7 @@ class BudgetManager: ) response = response.json() if response["status"] == "error": - self.user_dict = ( - {} - ) # assume this means the user dict hasn't been stored yet + self.user_dict = {} # assume this means the user dict hasn't been stored yet else: self.user_dict = response["data"] diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 2a8bd856040..1ff4ee04080 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -79,9 +79,7 @@ class CachingHandlerResponse(BaseModel): cached_result: Optional[Any] = None final_embedding_cached_response: Optional[EmbeddingResponse] = None - embedding_all_elements_cache_hit: bool = ( - False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call - ) + embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call in_memory_cache_obj = InMemoryCache() @@ -165,10 +163,11 @@ class LLMCachingHandler: """ # Check if caching should be performed BEFORE doing expensive operations if ( - (kwargs.get("caching", None) is None and litellm.cache is not None) - or kwargs.get("caching", False) is True - ) and ( - kwargs.get("cache", {}).get("no-cache", False) is not True + ( + (kwargs.get("caching", None) is None and litellm.cache is not None) + or kwargs.get("caching", False) is True + ) + and (kwargs.get("cache", {}).get("no-cache", False) is not True) ): # allow users to control returning cached responses from the completion function args = args or () final_embedding_cached_response: Optional[EmbeddingResponse] = None diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index b0f5754f58e..2dc9224e715 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -79,7 +79,8 @@ class RedisClusterCache(RedisCache): # Create a fresh Redis Cluster client with current settings redis_client = redis_async.RedisCluster( - startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore + startup_nodes=new_startup_nodes, + **cluster_kwargs, # type: ignore ) # Test the connection diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index d27cfefda73..87ac5d132d7 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -151,7 +151,9 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider=custom_llm_provider, ) - def completion(self, *args, **kwargs) -> Union[ + def completion( + self, *args, **kwargs + ) -> Union[ Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], "ModelResponse", "CustomStreamWrapper", diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 3fa6b983e5f..f7ac67927b6 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -205,9 +205,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if provider_specific_fields: tool_call_dict["provider_specific_fields"] = provider_specific_fields # Also add to function's provider_specific_fields for consistency - tool_call_dict["function"][ - "provider_specific_fields" - ] = provider_specific_fields + tool_call_dict["function"]["provider_specific_fields"] = ( + provider_specific_fields + ) msg = Message( content=None, @@ -301,7 +301,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): { "type": "message", "role": role, - "content": self._convert_content_to_responses_format(content, cast(str, role)), # type: ignore[arg-type] + "content": self._convert_content_to_responses_format( + content, cast(str, role) + ), # type: ignore[arg-type] } ) @@ -1021,7 +1023,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # If string is passed, map with optional summary based on flag/env var if reasoning_effort == "none": - return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") # type: ignore + return ( + Reasoning(effort="none", summary="detailed") + if auto_summary_enabled + else Reasoning(effort="none") + ) # type: ignore elif reasoning_effort == "high": return ( Reasoning(effort="high", summary="detailed") @@ -1029,7 +1035,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else Reasoning(effort="high") ) elif reasoning_effort == "xhigh": - return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") # type: ignore[typeddict-item] + return ( + Reasoning(effort="xhigh", summary="detailed") + if auto_summary_enabled + else Reasoning(effort="xhigh") + ) # type: ignore[typeddict-item] elif reasoning_effort == "medium": return ( Reasoning(effort="medium", summary="detailed") diff --git a/litellm/constants.py b/litellm/constants.py index c0e265c0e4a..d2f2e89eca3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -30,6 +30,9 @@ DEFAULT_SQS_BATCH_SIZE = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) SQS_SEND_MESSAGE_ACTION = "SendMessage" SQS_API_VERSION = "2012-11-05" DEFAULT_MAX_RETRIES = int(os.getenv("DEFAULT_MAX_RETRIES", 2)) +# Max records accepted in one POST /v1/callbacks/logs batch. Bounds the blast +# radius: each record fans out to spend logs + every callback integration. +MAX_CALLBACK_LOG_RECORDS = 1000 DEFAULT_MAX_RECURSE_DEPTH = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100)) DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int( os.getenv("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10) @@ -201,6 +204,18 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int( # Provider-specific API base URLs XAI_API_BASE = "https://api.x.ai/v1" +OPEN_SANDBOX_API_BASE_ENV_VAR = "OPEN_SANDBOX_API_BASE" +OPEN_SANDBOX_API_KEY_ENV_VAR = "OPEN_SANDBOX_API_KEY" +OPEN_SANDBOX_DEFAULT_TEMPLATE = "opensandbox/code-interpreter:v1.1.0" +_OPEN_SANDBOX_FALLBACK_ENTRYPOINT = "/opt/code-interpreter/code-interpreter.sh" +OPEN_SANDBOX_DEFAULT_ENTRYPOINT = (_OPEN_SANDBOX_FALLBACK_ENTRYPOINT,) +OPEN_SANDBOX_DEFAULT_LANGUAGE = "python" +OPEN_SANDBOX_DEFAULT_CPU_LIMIT = "1" +OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT = "2Gi" +OPEN_SANDBOX_EXECD_PORT = 44772 +OPEN_SANDBOX_DEFAULT_TIMEOUT = 300 +OPEN_SANDBOX_READY_TIMEOUT = 30.0 +OPEN_SANDBOX_POLL_INTERVAL = 0.2 DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024) @@ -281,7 +296,8 @@ DEFAULT_SSL_CIPHERS = os.getenv( "ECDHE-ECDSA-AES256-GCM-SHA384:" "ECDHE-ECDSA-AES128-GCM-SHA256:" # Priority 3: Additional modern ciphers (good balance) - "ECDHE-RSA-CHACHA20-POLY1305:" "ECDHE-ECDSA-CHACHA20-POLY1305:" + "ECDHE-RSA-CHACHA20-POLY1305:" + "ECDHE-ECDSA-CHACHA20-POLY1305:" # Priority 4: Widely compatible fallbacks (slower but universally supported) "ECDHE-RSA-AES256-SHA384:" # Common fallback "ECDHE-RSA-AES128-SHA256:" # Very widely supported @@ -456,6 +472,7 @@ HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: float = 5.0 request_timeout: float = float( os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))) ) +request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ DEFAULT_A2A_AGENT_TIMEOUT: float = float( os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000) ) # 10 minutes @@ -867,32 +884,31 @@ openai_compatible_providers: List = [ "docker_model_runner", "ragflow", "pinstripes", # Pinstripes - JSON-configured provider + "darkbloom", +] +openai_text_completion_compatible_providers: List = [ # providers that support `/v1/completions` + "together_ai", + "fireworks_ai", + "hosted_vllm", + "meta_llama", + "llamafile", + "featherless_ai", + "nebius", + "dashscope", + "modelscope", + "moonshot", + "publicai", + "synthetic", + "tensormesh", + "apertis", + "nano-gpt", + "poe", + "chutes", + "v0", + "lambda_ai", + "hyperbolic", + "wandb", ] -openai_text_completion_compatible_providers: List = ( - [ # providers that support `/v1/completions` - "together_ai", - "fireworks_ai", - "hosted_vllm", - "meta_llama", - "llamafile", - "featherless_ai", - "nebius", - "dashscope", - "modelscope", - "moonshot", - "publicai", - "synthetic", - "tensormesh", - "apertis", - "nano-gpt", - "poe", - "chutes", - "v0", - "lambda_ai", - "hyperbolic", - "wandb", - ] -) _openai_like_providers: List = [ "predibase", "databricks", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 27a146df7bf..f95deda0f81 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1020,7 +1020,7 @@ def _apply_cost_discount( if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug( - f"Applied {discount_percent*100}% discount to {custom_llm_provider}: " + f"Applied {discount_percent * 100}% discount to {custom_llm_provider}: " f"${original_cost:.6f} -> ${final_cost:.6f} (saved ${discount_amount:.6f})" ) @@ -1088,7 +1088,7 @@ def _apply_cost_margin( verbose_logger.debug( f"Applied margin to {custom_llm_provider or 'global'}: " f"${original_cost:.6f} -> ${final_cost:.6f} " - f"(margin: {margin_percent*100 if margin_percent > 0 else 0}% + ${margin_fixed_amount:.6f} = ${margin_total_amount:.6f})" + f"(margin: {margin_percent * 100 if margin_percent > 0 else 0}% + ${margin_fixed_amount:.6f} = ${margin_total_amount:.6f})" ) return final_cost, margin_percent, margin_fixed_amount, margin_total_amount @@ -1621,7 +1621,9 @@ def completion_cost( model in litellm.replicate_models or "replicate" in model ) and model not in litellm.model_cost: # for unmapped replicate model, default to replicate's time tracking logic - return get_replicate_completion_pricing(completion_response, total_time) # type: ignore + return get_replicate_completion_pricing( + completion_response, total_time + ) # type: ignore if model is None: raise ValueError( diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index c6d427e7f09..5baa7cbc9c5 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -224,6 +224,7 @@ class MCPClient: extra_headers: Optional[Dict[str, str]] = None, ssl_verify: Optional[VerifyTypes] = None, aws_auth: Optional[httpx.Auth] = None, + resolved_auth: Optional[httpx.Auth] = None, sampling_callback: Optional[Callable] = None, elicitation_callback: Optional[Callable] = None, logging_callback: Optional[Callable] = None, @@ -237,6 +238,9 @@ class MCPClient: self.extra_headers: Optional[Dict[str, str]] = extra_headers self.ssl_verify: Optional[VerifyTypes] = ssl_verify self._aws_auth: Optional[httpx.Auth] = aws_auth + # A pre-resolved httpx.Auth (e.g. from the v2 credential resolver) attached to the + # upstream client's auth= slot, taking precedence over the SigV4 aws_auth. + self._resolved_auth: Optional[httpx.Auth] = resolved_auth self._last_initialize_instructions: Optional[str] = None self._sampling_callback: Optional[Callable] = sampling_callback self._elicitation_callback: Optional[Callable] = elicitation_callback @@ -482,11 +486,15 @@ class MCPClient: verbose_logger.debug( f"MCP client using SSL configuration: {type(ssl_config).__name__}" ) - # Use SigV4 auth if configured and no explicit auth provided. - # The MCP SDK's sse_client and streamable_http_client call this - # factory without passing auth=, so self._aws_auth is used. - # For non-SigV4 clients, self._aws_auth is None — no behavior change. - effective_auth = auth if auth is not None else self._aws_auth + # The MCP SDK's sse_client and streamable_http_client call this factory without + # passing auth=, so the fallback is used: a v2-resolved auth if present, else the + # SigV4 aws_auth. Both are None for the common case — no behavior change. + fallback_auth = ( + self._resolved_auth + if self._resolved_auth is not None + else self._aws_auth + ) + effective_auth = auth if auth is not None else fallback_auth return httpx.AsyncClient( headers=headers, timeout=timeout, diff --git a/litellm/files/main.py b/litellm/files/main.py index 669d50dde41..582d9d5cdd0 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -264,7 +264,9 @@ def create_file( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_file", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_file", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -435,7 +437,10 @@ def file_retrieve( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), ) @@ -618,7 +623,10 @@ def file_delete( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), ) return cast(FileDeleted, response) @@ -786,7 +794,9 @@ def file_list( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="file_list", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="file_list", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -1037,7 +1047,9 @@ def file_content( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -1118,7 +1130,9 @@ def file_content_streaming( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) diff --git a/litellm/files/utils.py b/litellm/files/utils.py index a2b9a42c154..a0df7a89b0f 100644 --- a/litellm/files/utils.py +++ b/litellm/files/utils.py @@ -3,6 +3,22 @@ from typing import Optional from litellm.types.llms.openai import CreateFileRequest from litellm.types.utils import ExtractedFileData +# MIME types a .jsonl batch upload is plausibly labeled with. Clients are +# inconsistent (text/plain, application/json, octet-stream, ndjson, ...), so a +# batch file must not silently bypass the streaming path just because of its +# declared type. ``purpose == "batch"`` is the authoritative signal; non-JSONL +# content still fails loudly when the rows are parsed. +_BATCH_JSONL_CONTENT_TYPES = frozenset( + { + "application/jsonl", + "application/json", + "application/octet-stream", + "application/x-ndjson", + "application/x-jsonlines", + "text/plain", + } +) + class FilesAPIUtils: """ @@ -24,9 +40,24 @@ class FilesAPIUtils: and extracted_file_data.get("content") is not None ) + @staticmethod + def is_batch_jsonl_request( + create_file_data: CreateFileRequest, content_type: Optional[str] + ) -> bool: + """ + Batch-jsonl check from metadata only, so the body can stay a streamable + Path/handle instead of being read into memory. + """ + return ( + create_file_data.get("purpose") == "batch" + and FilesAPIUtils.valid_content_type(content_type) + and create_file_data.get("file") is not None + ) + @staticmethod def valid_content_type(content_type: Optional[str]) -> bool: """ - Check if the content type is valid + Whether the upload's MIME type is one a batch JSONL file is plausibly + sent as (see ``_BATCH_JSONL_CONTENT_TYPES``). """ - return content_type in set(["application/jsonl", "application/octet-stream"]) + return content_type in _BATCH_JSONL_CONTENT_TYPES diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index 08373cda782..846a8a504a8 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -245,7 +245,11 @@ def create_fine_tuning_job( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore + api_base = ( + optional_params.api_base + or litellm.api_base + or get_secret_str("AZURE_API_BASE") + ) # type: ignore api_version = ( optional_params.api_version @@ -340,7 +344,9 @@ def create_fine_tuning_job( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -458,7 +464,11 @@ def cancel_fine_tuning_job( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_base = ( + optional_params.api_base + or litellm.api_base + or get_secret("AZURE_API_BASE") + ) # type: ignore api_version = ( optional_params.api_version @@ -500,7 +510,9 @@ def cancel_fine_tuning_job( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -621,7 +633,11 @@ def list_fine_tuning_jobs( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore + api_base = ( + optional_params.api_base + or litellm.api_base + or get_secret_str("AZURE_API_BASE") + ) # type: ignore api_version = ( optional_params.api_version @@ -664,7 +680,9 @@ def list_fine_tuning_jobs( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="create_thread", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) return response @@ -776,7 +794,11 @@ def retrieve_fine_tuning_job( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore + api_base = ( + optional_params.api_base + or litellm.api_base + or get_secret_str("AZURE_API_BASE") + ) # type: ignore api_version = ( optional_params.api_version @@ -818,7 +840,10 @@ def retrieve_fine_tuning_job( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="retrieve_fine_tuning_job", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="retrieve_fine_tuning_job", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), ) return response diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index c5d9fd124fa..d8f9f1feb0b 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -726,9 +726,9 @@ class GoogleGenAIAdapter: wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name if args_chunk: - wrapper.accumulated_tool_calls[tool_call_index][ - "arguments" - ] += args_chunk + wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += ( + args_chunk + ) # Attempt to parse and emit a complete tool call accumulated_data = wrapper.accumulated_tool_calls[tool_call_index] diff --git a/litellm/images/main.py b/litellm/images/main.py index 8b108ded4c9..34e77ffe8ab 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -880,20 +880,20 @@ def image_edit( local_vars.update(kwargs) # Get ImageEditOptionalRequestParams with only valid parameters - image_edit_optional_params: ( - ImageEditOptionalRequestParams - ) = _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( - local_vars + image_edit_optional_params: ImageEditOptionalRequestParams = ( + _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( + local_vars + ) ) # Get optional parameters for the responses API - image_edit_request_params: ( - Dict - ) = _get_ImageEditRequestUtils().get_optional_params_image_edit( - model=model, - image_edit_provider_config=image_edit_provider_config, - image_edit_optional_params=image_edit_optional_params, - drop_params=kwargs.get("drop_params"), - additional_drop_params=kwargs.get("additional_drop_params"), + image_edit_request_params: Dict = ( + _get_ImageEditRequestUtils().get_optional_params_image_edit( + model=model, + image_edit_provider_config=image_edit_provider_config, + image_edit_optional_params=image_edit_optional_params, + drop_params=kwargs.get("drop_params"), + additional_drop_params=kwargs.get("additional_drop_params"), + ) ) # Pre Call logging diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index 98f1eb2d551..8c0da9bb4fb 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -104,10 +104,10 @@ class AlertingHangingRequestCheck: ) for request_id in hanging_requests: - hanging_request_data: Optional[HangingRequestData] = ( - await self.hanging_request_cache.async_get_cache( - key=request_id, - ) + hanging_request_data: Optional[ + HangingRequestData + ] = await self.hanging_request_cache.async_get_cache( + key=request_id, ) if hanging_request_data is None: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 2108ebae312..35731306b93 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -245,7 +245,7 @@ class SlackAlerting(CustomBatchLogger): return for api_base, latency in _deployment_latency_map.items(): - _message_to_send += f"\n{api_base}: {round(latency,2)}s" + _message_to_send += f"\n{api_base}: {round(latency, 2)}s" _message_to_send = "```" + _message_to_send + "```" return _message_to_send @@ -272,7 +272,7 @@ class SlackAlerting(CustomBatchLogger): if litellm.turn_off_message_logging or litellm.redact_messages_in_exceptions: messages = "Message not logged. litellm.redact_messages_in_exceptions=True" request_info = f"\nRequest Model: `{model}`\nAPI Base: `{api_base}`\nMessages: `{messages}`" - slow_message = f"`Responses are slow - {round(time_difference_float,2)}s response time > Alerting threshold: {self.alerting_threshold}s`" + slow_message = f"`Responses are slow - {round(time_difference_float, 2)}s response time > Alerting threshold: {self.alerting_threshold}s`" alerting_metadata: dict = {} if time_difference_float > self.alerting_threshold: # add deployment latencies to alert @@ -460,7 +460,7 @@ class SlackAlerting(CustomBatchLogger): if api_base is None: api_base = "" value = replaced_failed_values[top_5_failed[i]] - message += f"\t{i+1}. Deployment: `{deployment_name}`, Failed Requests: `{value}`, API Base: `{api_base}`\n" + message += f"\t{i + 1}. Deployment: `{deployment_name}`, Failed Requests: `{value}`, API Base: `{api_base}`\n" message += "\n\n*😅 Top Slowest Deployments:*\n\n" if not top_5_slowest: @@ -479,7 +479,7 @@ class SlackAlerting(CustomBatchLogger): ), ) value = round(replaced_slowest_values[top_5_slowest[i]], 3) - message += f"\t{i+1}. Deployment: `{deployment_name}`, Latency per output token: `{value}s/token`, API Base: `{api_base}`\n\n" + message += f"\t{i + 1}. Deployment: `{deployment_name}`, Latency per output token: `{value}s/token`, API Base: `{api_base}`\n\n" # cache cleanup -> reset values to 0 latency_cache_keys = [(key, 0) for key in latency_keys] @@ -595,9 +595,7 @@ class SlackAlerting(CustomBatchLogger): "projected_limit_exceeded", "soft_budget_crossed", ] - ] = ( - "projected_limit_exceeded" if type == "projected_limit_exceeded" else None - ) + ] = "projected_limit_exceeded" if type == "projected_limit_exceeded" else None webhook_event: Optional[WebhookEvent] = None @@ -854,9 +852,9 @@ class SlackAlerting(CustomBatchLogger): ### UNIQUE CACHE KEY ### cache_key = provider + region_name - outage_value: Optional[ProviderRegionOutageModel] = ( - await self.internal_usage_cache.async_get_cache(key=cache_key) - ) + outage_value: Optional[ + ProviderRegionOutageModel + ] = await self.internal_usage_cache.async_get_cache(key=cache_key) # Convert deployment_ids back to set if it was stored as a list if outage_value is not None: @@ -981,7 +979,9 @@ class SlackAlerting(CustomBatchLogger): max_alerts_size = 10 """ try: - outage_value: Optional[OutageModel] = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore + outage_value: Optional[ + OutageModel + ] = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore if ( getattr(exception, "status_code", None) is None or ( diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 6a6313f72e1..a69785ba2e3 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -52,9 +52,9 @@ class BraintrustLogger(CustomLogger): "Authorization": "Bearer " + self.api_key, "Content-Type": "application/json", } - self._project_id_cache: Dict[str, str] = ( - {} - ) # Cache mapping project names to IDs + self._project_id_cache: Dict[ + str, str + ] = {} # Cache mapping project names to IDs self.global_braintrust_http_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py index 59e0988a10a..1af14deeab6 100644 --- a/litellm/integrations/braintrust_mock_client.py +++ b/litellm/integrations/braintrust_mock_client.py @@ -157,7 +157,7 @@ def create_mock_braintrust_client(): create_mock_braintrust_factory_client() verbose_logger.debug( - f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms" + f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms" ) verbose_logger.debug( "[BRAINTRUST MOCK] Braintrust mock client initialization complete" diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index da8149eab9b..362581937d7 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -9,9 +9,11 @@ captured stdout back through the typed agentic loop plan. import json import time import uuid -from typing import Any, cast +from typing import Any, Literal, TypedDict, cast import litellm +from pydantic import ValidationError + from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.integrations.code_interpreter_interception import ( @@ -20,15 +22,93 @@ from litellm.types.integrations.code_interpreter_interception import ( from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + is_interception_internal_key, +) +from litellm.types.llms.openai import ( + ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, + ChatCompletionToolMessage, +) +from litellm.types.utils import ( + CallTypes, + ChatCompletionMessageToolCall, + ModelResponse, ) -from litellm.types.utils import CallTypes LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" _INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +_CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream" +_LITELLM_METADATA_KEY = "litellm_metadata" _CACHE_TTL_SECONDS = 15 * 60 +class CodeExecutionToolCall(TypedDict, total=False): + id: str | None + call_id: str | None + type: Literal["function"] + name: str + arguments: str + + +class CodeInterpreterLogOutput(TypedDict): + type: Literal["logs"] + logs: str + + +class CodeInterpreterCall(TypedDict): + id: str + type: Literal["code_interpreter_call"] + status: Literal["completed"] + code: str + container_id: str | None + outputs: list[CodeInterpreterLogOutput] + + +class CodeExecutionFunctionParameters(TypedDict): + type: Literal["object"] + properties: dict[str, dict[str, str]] + required: list[str] + + +class ResponsesFunctionTool(TypedDict): + type: Literal["function"] + name: str + description: str + parameters: CodeExecutionFunctionParameters + + +class ChatCompletionFunctionDefinition(TypedDict): + name: str + description: str + parameters: CodeExecutionFunctionParameters + + +class ChatCompletionFunctionTool(TypedDict): + type: Literal["function"] + function: ChatCompletionFunctionDefinition + + +CodeExecutionFunctionTool = ResponsesFunctionTool | ChatCompletionFunctionTool + + +class ResponsesFunctionToolChoice(TypedDict): + type: Literal["function"] + name: str + + +class ChatCompletionFunctionToolChoice(TypedDict): + type: Literal["function"] + function: dict[str, str] + + +CodeExecutionFunctionToolChoice = ( + ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice +) + + def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None: try: from litellm.sandbox.sandbox_tools import resolve_sandbox_tool @@ -97,9 +177,15 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if not kwargs.get("_agentic_loop_depth"): kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None) kwargs.pop(_SANDBOX_KEY, None) + self._strip_interception_metadata(kwargs) if not self.enabled: return None - if call_type not in (CallTypes.responses, CallTypes.aresponses): + if call_type not in ( + CallTypes.responses, + CallTypes.aresponses, + CallTypes.completion, + CallTypes.acompletion, + ): return None if ( self.enabled_providers is not None @@ -120,18 +206,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): kwargs[_SANDBOX_KEY] = uuid.uuid4().hex if kwargs.get("stream"): kwargs["stream"] = False - kwargs["_code_interpreter_interception_converted_stream"] = True + kwargs[_CONVERTED_STREAM_KEY] = True + self._write_interception_metadata(kwargs) - function_tool = { - "type": "function", - "name": LITELLM_CODE_EXECUTION_TOOL_NAME, - "description": "Execute python code in a sandbox and return stdout.", - "parameters": { - "type": "object", - "properties": {"code": {"type": "string"}}, - "required": ["code"], - }, - } + function_tool = self._get_function_tool(call_type=call_type) kwargs["tools"] = [ ( function_tool @@ -141,19 +219,90 @@ class CodeInterpreterInterceptionLogger(CustomLogger): for tool in tools ] if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): - kwargs["tool_choice"] = { - "type": "function", - "name": LITELLM_CODE_EXECUTION_TOOL_NAME, - } + kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type) return kwargs + @staticmethod + def _strip_interception_metadata(kwargs: dict[str, Any]) -> None: + metadata = kwargs.get(_LITELLM_METADATA_KEY) + if not isinstance(metadata, dict): + return + filtered_metadata = { + key: value + for key, value in metadata.items() + if not is_interception_internal_key(key) + and not key.startswith("_agentic_loop") + and key != "max_agentic_loops" + } + if filtered_metadata: + kwargs[_LITELLM_METADATA_KEY] = filtered_metadata + else: + kwargs.pop(_LITELLM_METADATA_KEY, None) + + @staticmethod + def _write_interception_metadata(kwargs: dict[str, Any]) -> None: + metadata = kwargs.get(_LITELLM_METADATA_KEY) + metadata = dict(metadata) if isinstance(metadata, dict) else {} + for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _CONVERTED_STREAM_KEY): + if key in kwargs: + metadata[key] = kwargs[key] + kwargs[_LITELLM_METADATA_KEY] = metadata + + @staticmethod + def _get_function_parameters() -> CodeExecutionFunctionParameters: + return { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + } + + def _get_function_tool( + self, call_type: CallTypes | None + ) -> CodeExecutionFunctionTool: + description = "Execute python code in a sandbox and return stdout." + if call_type in (CallTypes.completion, CallTypes.acompletion): + return { + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": description, + "parameters": self._get_function_parameters(), + }, + } + return { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": description, + "parameters": self._get_function_parameters(), + } + + @staticmethod + def _get_function_tool_choice( + call_type: CallTypes | None, + ) -> CodeExecutionFunctionToolChoice: + if call_type in (CallTypes.completion, CallTypes.acompletion): + return { + "type": "function", + "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}, + } + return { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + } + @staticmethod def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool: if not isinstance(tool_choice, dict): return False + function = tool_choice.get("function") return ( tool_choice.get("type") == "code_interpreter" or tool_choice.get("name") == "code_interpreter" + or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + or ( + isinstance(function, dict) + and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + ) ) def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None: @@ -188,7 +337,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ): return False, {} - tool_calls = self._extract_code_execution_tool_calls(response=response) + tool_calls = ( + self._extract_chat_completion_code_execution_tool_calls(response=response) + if kwargs.get("_agentic_loop_api_surface") + == CHAT_COMPLETION_AGENTIC_SURFACE + else self._extract_code_execution_tool_calls(response=response) + ) if not tool_calls: return False, {} @@ -206,15 +360,24 @@ class CodeInterpreterInterceptionLogger(CustomLogger): stream: bool, kwargs: dict, ) -> AgenticLoopPlan: + if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE: + return await self._build_chat_completion_agentic_loop_plan( + tools=tools, + model=model, + messages=messages, + optional_params=anthropic_messages_optional_request_params, + kwargs=kwargs, + ) + await self._prune_expired_cache() - tool_calls = cast(list[dict[str, Any]], tools.get("tool_calls", [])) + tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) sandbox_key = kwargs.get(_SANDBOX_KEY) container, params = await self._get_or_create_container(cache_key=sandbox_key) try: - container_id = getattr(container, "id", None) + container_id = cast(str | None, getattr(container, "id", None)) input_list = self._normalize_messages(messages) - code_interpreter_calls = [] + code_interpreter_calls: list[CodeInterpreterCall] = [] for tool_call in tool_calls: arguments = tool_call.get("arguments", "") code = self._parse_code(arguments) @@ -256,9 +419,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger): request_patch = AgenticLoopRequestPatch( model=model, messages=input_list, - tools=optional_params.get("tools"), - optional_params={k: v for k, v in optional_params.items() if k != "tools"}, - kwargs={k: v for k, v in kwargs.items() if k != "litellm_logging_obj"}, + tools=self._get_followup_tools( + tools=optional_params.get("tools"), + call_type=CallTypes.responses, + ), + optional_params=self._get_followup_optional_params(optional_params), + kwargs=self._filter_agentic_loop_kwargs(kwargs), ) return AgenticLoopPlan( @@ -271,12 +437,134 @@ class CodeInterpreterInterceptionLogger(CustomLogger): }, ) + async def _build_chat_completion_agentic_loop_plan( + self, + tools: dict[str, object], + model: str, + messages: list[dict], + optional_params: dict[str, object], + kwargs: dict[str, object], + ) -> AgenticLoopPlan: + await self._prune_expired_cache() + tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) + sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY)) + container, params = await self._get_or_create_container(cache_key=sandbox_key) + + try: + container_id = cast(str | None, getattr(container, "id", None)) + tool_results = [ + await self._build_chat_completion_tool_result( + container=container, + params=params, + tool_call=tool_call, + container_id=container_id, + ) + for tool_call in tool_calls + ] + except Exception: + await self._delete_container_for_cache_key(sandbox_key) + raise + tool_messages = [result[0] for result in tool_results] + code_interpreter_calls = [result[1] for result in tool_results] + + request_patch = AgenticLoopRequestPatch( + model=model, + messages=list(messages) + + [self._build_chat_completion_assistant_message(tool_calls)] + + tool_messages, + tools=self._get_followup_tools( + tools=optional_params.get("tools"), + call_type=CallTypes.completion, + ), + optional_params=self._get_followup_optional_params(optional_params), + kwargs=self._filter_agentic_loop_kwargs(kwargs), + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={ + "tool_type": "code_interpreter", + "sandbox_key": sandbox_key or "", + "code_interpreter_calls": code_interpreter_calls, + "response_format": "openai", + }, + ) + + async def _build_chat_completion_tool_result( + self, + container: object, + params: dict[str, Any] | None, + tool_call: CodeExecutionToolCall, + container_id: str | None, + ) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]: + arguments = tool_call.get("arguments", "") + code = self._parse_code(arguments) + stdout = await self._run_tool_call( + container=container, params=params, arguments=arguments + ) + tool_call_id = ( + tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex + ) + return ( + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": stdout, + }, + { + "id": f"ci_{uuid.uuid4().hex}", + "type": "code_interpreter_call", + "status": "completed", + "code": code, + "container_id": container_id, + "outputs": [{"type": "logs", "logs": stdout}] if stdout else [], + }, + ) + async def async_agentic_loop_cleanup_hook( self, plan: AgenticLoopPlan, kwargs: dict ) -> None: metadata = plan.metadata or {} if plan else {} await self._delete_container_for_cache_key(metadata.get("sandbox_key")) + @staticmethod + def _filter_agentic_loop_kwargs(kwargs: dict[str, object]) -> dict[str, object]: + return { + k: v + for k, v in kwargs.items() + if k not in {"litellm_logging_obj", "acompletion"} + and not is_interception_internal_key( + k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES + ) + } + + def _get_followup_tools( + self, tools: object, call_type: CallTypes | None + ) -> list[dict[str, Any]] | None: + if not isinstance(tools, list): + return None + return [ + ( + self._get_function_tool(call_type=call_type) + if isinstance(tool, dict) and tool.get("type") == "code_interpreter" + else tool + ) + for tool in tools + ] + + def _get_followup_optional_params( + self, optional_params: dict[str, object] + ) -> dict[str, object]: + drop_tool_choice = self._tool_choice_targets_code_interpreter( + optional_params.get("tool_choice") + ) + return { + k: v + for k, v in optional_params.items() + if k != "tools" and not (k == "tool_choice" and drop_tool_choice) + } + async def async_post_agentic_loop_response_hook( self, response: Any, plan: AgenticLoopPlan, kwargs: dict ) -> Any: @@ -420,7 +708,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return list(messages) return [] - def _extract_code_execution_tool_calls(self, response: Any) -> list[dict[str, Any]]: + def _extract_code_execution_tool_calls( + self, response: object + ) -> list[CodeExecutionToolCall]: if isinstance(response, dict): output = response.get("output", []) else: @@ -446,6 +736,82 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if self._is_code_execution_call(item) ] + def _extract_chat_completion_code_execution_tool_calls( + self, response: ModelResponse | dict[str, Any] + ) -> list[CodeExecutionToolCall]: + model_response = self._to_model_response(response) + if model_response is None: + return [] + choices = model_response.choices or [] + if not choices: + return [] + message = choices[0].message + tool_calls = message.tool_calls or [] + + return [ + normalized + for tool_call in tool_calls + if (normalized := self._normalize_chat_completion_tool_call(tool_call)) + is not None + ] + + @staticmethod + def _normalize_chat_completion_tool_call( + tool_call: ChatCompletionMessageToolCall, + ) -> CodeExecutionToolCall | None: + if ( + tool_call.type != "function" + or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME + ): + return None + + arguments = tool_call.function.arguments + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + elif not isinstance(arguments, str): + arguments = "" if arguments is None else str(arguments) + + return { + "id": tool_call.id, + "call_id": tool_call.id, + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": arguments, + } + + @staticmethod + def _build_chat_completion_assistant_message( + tool_calls: list[CodeExecutionToolCall], + ) -> ChatCompletionAssistantMessage: + return { + "role": "assistant", + "tool_calls": [ + cast( + ChatCompletionAssistantToolCall, + { + "id": tool_call.get("id"), + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": tool_call.get("arguments", ""), + }, + }, + ) + for tool_call in tool_calls + ], + } + + @staticmethod + def _to_model_response( + response: ModelResponse | dict[str, Any], + ) -> ModelResponse | None: + if isinstance(response, ModelResponse): + return response + try: + return ModelResponse(**response) + except (TypeError, ValidationError): + return None + def _is_code_execution_call(self, item: Any) -> bool: if isinstance(item, dict): return ( diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 94fb97dff53..6d65b4ec0d2 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -981,7 +981,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac verbose_logger.debug( f"Incrementing callback failure metric for {callback_name}" ) - callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore + callback_obj.increment_callback_logging_failure( + callback_name=callback_name + ) # type: ignore return verbose_logger.debug( diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index d7847027d7e..bd5c165cba2 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -256,7 +256,9 @@ class DatadogMetricsLogger(CustomBatchLogger): headers["Content-Encoding"] = "gzip" response = await self.async_client.post( - self.upload_url, content=compressed_data, headers=headers # type: ignore + self.upload_url, + content=compressed_data, + headers=headers, # type: ignore ) response.raise_for_status() diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index c58e955984c..4e3dd2b6d8b 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -157,7 +157,7 @@ class FocusVantageDestination(FocusDestination): async def _upload_csv( self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str ) -> None: - url = f"{self.base_url}/v2/integrations/" f"{self.integration_token}/costs.csv" + url = f"{self.base_url}/v2/integrations/{self.integration_token}/costs.csv" headers = { "Authorization": f"Bearer {self.api_key}", } diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 1761fe010c9..9161455a246 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -154,7 +154,10 @@ def create_mock_gcs_client(): This function is idempotent - it only initializes mocks once, even if called multiple times. """ - global _original_async_handler_get, _original_async_handler_delete, _mocks_initialized + global \ + _original_async_handler_get, \ + _original_async_handler_delete, \ + _mocks_initialized # Use factory for POST handler _create_mock_gcs_post() @@ -179,7 +182,7 @@ def create_mock_gcs_client(): verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.delete") verbose_logger.debug( - f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms" + f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms" ) verbose_logger.debug("[GCS MOCK] GCS mock client initialization complete") diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index a468741aead..99d9d9b285b 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -372,7 +372,9 @@ class GitLabPromptManager(CustomPromptManagement): if parsed_messages: final_messages: List[AllMessageValues] = parsed_messages else: - final_messages = [{"role": "user", "content": rendered_prompt}] + messages # type: ignore + final_messages = [ + {"role": "user", "content": rendered_prompt} + ] + messages # type: ignore if litellm_params is None: litellm_params = {} @@ -412,24 +414,41 @@ class GitLabPromptManager(CustomPromptManagement): low = line.lower() if low.startswith("system:"): if current_role and current_content: - messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + messages.append( + { + "role": current_role, + "content": "\n".join(current_content).strip(), + } + ) # type: ignore current_role = "system" current_content = [line[7:].strip()] elif low.startswith("user:"): if current_role and current_content: - messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + messages.append( + { + "role": current_role, + "content": "\n".join(current_content).strip(), + } + ) # type: ignore current_role = "user" current_content = [line[5:].strip()] elif low.startswith("assistant:"): if current_role and current_content: - messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + messages.append( + { + "role": current_role, + "content": "\n".join(current_content).strip(), + } + ) # type: ignore current_role = "assistant" current_content = [line[10:].strip()] else: current_content.append(line) if current_role and current_content: - messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + messages.append( + {"role": current_role, "content": "\n".join(current_content).strip()} + ) # type: ignore if not messages and prompt_content.strip(): messages = [{"role": "user", "content": prompt_content.strip()}] # type: ignore return messages diff --git a/litellm/integrations/lago.py b/litellm/integrations/lago.py index b881193e869..c7c010f9976 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -131,9 +131,9 @@ class LagoLogger(CustomLogger): def log_success_event(self, kwargs, response_obj, start_time, end_time): _url = os.getenv("LAGO_API_BASE") - assert _url is not None and isinstance( - _url, str - ), "LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={}".format(_url) + assert _url is not None and isinstance(_url, str), ( + "LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={}".format(_url) + ) if _url.endswith("/"): _url += "api/v1/events" else: @@ -165,10 +165,10 @@ class LagoLogger(CustomLogger): try: verbose_logger.debug("ENTERS LAGO CALLBACK") _url = os.getenv("LAGO_API_BASE") - assert _url is not None and isinstance( - _url, str - ), "LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={}".format( - _url + assert _url is not None and isinstance(_url, str), ( + "LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={}".format( + _url + ) ) if _url.endswith("/"): _url += "api/v1/events" diff --git a/litellm/integrations/langfuse/langfuse_handler.py b/litellm/integrations/langfuse/langfuse_handler.py index 4a809726424..797c1609f80 100644 --- a/litellm/integrations/langfuse/langfuse_handler.py +++ b/litellm/integrations/langfuse/langfuse_handler.py @@ -86,9 +86,9 @@ class LangFuseHandler: if globalLangfuseLogger is not None: return globalLangfuseLogger - credentials_dict: Dict[str, Any] = ( - {} - ) # the global langfuse logger uses Environment Variables, there are no dynamic credentials + credentials_dict: Dict[ + str, Any + ] = {} # the global langfuse logger uses Environment Variables, there are no dynamic credentials globalLangfuseLogger = in_memory_dynamic_logger_cache.get_cache( credentials=credentials_dict, service_name="langfuse", diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 81570e462c4..15f92e8b322 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -65,8 +65,7 @@ class LangsmithLogger(CustomBatchLogger): langsmith_tenant_id=langsmith_tenant_id, ) self.sampling_rate: float = ( - langsmith_sampling_rate - or float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore + langsmith_sampling_rate or float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore if os.getenv("LANGSMITH_SAMPLING_RATE") is not None and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore else 1.0 diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py index 9b912ce70c8..3b013f96a79 100644 --- a/litellm/integrations/mock_client_factory.py +++ b/litellm/integrations/mock_client_factory.py @@ -232,7 +232,11 @@ def create_mock_client_factory(config: MockClientConfig): # Create mock client initialization function def create_mock_client(): """Initialize the mock client by patching HTTP handlers.""" - nonlocal _original_async_handler_post, _original_sync_client_post, _original_http_handler_post, _mocks_initialized + nonlocal \ + _original_async_handler_post, \ + _original_sync_client_post, \ + _original_http_handler_post, \ + _mocks_initialized if _mocks_initialized: return @@ -261,7 +265,7 @@ def create_mock_client_factory(config: MockClientConfig): verbose_logger.debug(f"[{config.name} MOCK] Patched HTTPHandler.post") verbose_logger.debug( - f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms" + f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms" ) verbose_logger.debug( f"[{config.name} MOCK] {config.name} mock client initialization complete" diff --git a/litellm/integrations/opik/opik.py b/litellm/integrations/opik/opik.py index 7b687d34d1c..b15b024ff09 100644 --- a/litellm/integrations/opik/opik.py +++ b/litellm/integrations/opik/opik.py @@ -174,7 +174,9 @@ class OpikLogger(CustomBatchLogger): ) -> None: try: response = self.sync_httpx_client.post( - url=url, headers=headers, json=batch # type: ignore + url=url, + headers=headers, + json=batch, # type: ignore ) response.raise_for_status() if response.status_code != 204: @@ -264,7 +266,9 @@ class OpikLogger(CustomBatchLogger): ) -> None: try: response = await self.async_httpx_client.post( - url=url, headers=headers, json=batch # type: ignore + url=url, + headers=headers, + json=batch, # type: ignore ) response.raise_for_status() diff --git a/litellm/integrations/otel/mappers/__init__.py b/litellm/integrations/otel/mappers/__init__.py index 012e63f1bee..9c728678f47 100644 --- a/litellm/integrations/otel/mappers/__init__.py +++ b/litellm/integrations/otel/mappers/__init__.py @@ -38,7 +38,7 @@ def resolve_mappers(names: Iterable[str]) -> list[AttributeMapper]: factory = _MAPPER_BY_NAME.get(name) if factory is None: raise ValueError( - f"unknown mapper name {name!r}; known: " f"{sorted(_MAPPER_BY_NAME)}" + f"unknown mapper name {name!r}; known: {sorted(_MAPPER_BY_NAME)}" ) out.append(factory()) return out diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index d9be68a06c2..75330ebd1bd 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -35,7 +35,6 @@ from litellm.integrations.otel.model.spans import db_system class GenAIMapper: - _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { GenAI.OPERATION_NAME: lambda d: d.operation.value, GenAI.PROVIDER_NAME: lambda d: d.provider or None, @@ -81,9 +80,13 @@ class GenAIMapper: f"{LiteLLM.COST_PREFIX}original": lambda d: d.cost.original, f"{LiteLLM.COST_PREFIX}discount_amount": lambda d: d.cost.discount_amount, f"{LiteLLM.COST_PREFIX}discount_percent": lambda d: d.cost.discount_percent, - f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount, + f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: ( + d.cost.margin_fixed_amount + ), f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent, - f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount, + f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: ( + d.cost.margin_total_amount + ), LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming, } diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index 14c9fd01d05..d0af460f752 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -27,7 +27,6 @@ from litellm.integrations.otel.model.payloads import ( class LangfuseMapper: - _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { "langfuse.observation.type": lambda d: "generation", "langfuse.observation.model.name": lambda d: d.request_model or None, diff --git a/litellm/integrations/otel/mappers/langtrace.py b/litellm/integrations/otel/mappers/langtrace.py index 7c0f30e57dd..ec595439fe5 100644 --- a/litellm/integrations/otel/mappers/langtrace.py +++ b/litellm/integrations/otel/mappers/langtrace.py @@ -20,7 +20,6 @@ from litellm.integrations.otel.model.payloads import LLMCallSpanData class LangtraceMapper: - _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { "gen_ai.operation.name": lambda d: "chat", "langtrace.service.name": lambda d: d.provider or None, diff --git a/litellm/integrations/otel/model/baggage.py b/litellm/integrations/otel/model/baggage.py index ecab643a26b..77ace736c2f 100644 --- a/litellm/integrations/otel/model/baggage.py +++ b/litellm/integrations/otel/model/baggage.py @@ -29,13 +29,15 @@ _PROMOTABLE: Final[ ] = { LiteLLM.TEAM_ID: lambda identity, model, team_metadata_keys: identity.team_id, LiteLLM.TEAM_ALIAS: lambda identity, model, team_metadata_keys: identity.team_alias, - LiteLLM.TEAM_METADATA: lambda identity, model, team_metadata_keys: _filtered_team_metadata_json( - identity.team_metadata, team_metadata_keys + LiteLLM.TEAM_METADATA: lambda identity, model, team_metadata_keys: ( + _filtered_team_metadata_json(identity.team_metadata, team_metadata_keys) ), LiteLLM.KEY_HASH: lambda identity, model, team_metadata_keys: identity.key_hash, LiteLLM.END_USER: lambda identity, model, team_metadata_keys: identity.end_user, GenAI.REQUEST_MODEL: lambda identity, model, team_metadata_keys: model, - LiteLLM.PROVIDER_MODEL: lambda identity, model, team_metadata_keys: identity.provider_model, + LiteLLM.PROVIDER_MODEL: lambda identity, model, team_metadata_keys: ( + identity.provider_model + ), } # Keys promoted by default (a subset of ``_PROMOTABLE``). ``END_USER`` is diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index c63f114514a..18fb5d9491e 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3237,7 +3237,9 @@ class PrometheusLogger(CustomLogger): ) return - async def fetch_keys(page_size: int, page: int) -> Tuple[ + async def fetch_keys( + page_size: int, page: int + ) -> Tuple[ List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]], Optional[int], ]: diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index af396ecdc73..922d8f71cf0 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -550,8 +550,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "response": response_data, } verbose_logger.debug( - f"Sending request to tool blocking service: " - f"{self.tool_blocking_endpoint}" + f"Sending request to tool blocking service: {self.tool_blocking_endpoint}" ) http_response = await self.tool_blocking_client.post( self.tool_blocking_endpoint, diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 482a19c5d72..c54b6e4cced 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -88,12 +88,12 @@ class VectorStorePreCallHook(CustomLogger): pass # Use database fallback to ensure synchronization across instances - vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = ( - await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( - non_default_params=non_default_params, - tools=tools, - prisma_client=prisma_client, - ) + vector_stores_to_run: List[ + LiteLLM_ManagedVectorStore + ] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=tools, + prisma_client=prisma_client, ) if not vector_stores_to_run: diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 173d4ca8764..0ff1a97cd0b 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -300,9 +300,6 @@ class LiteLLMResponsesInteractionsConfig: "total_output_tokens": getattr(usage, "output_tokens", 0), } - # Add role - interactions_response_dict["role"] = "model" - # Add updated (same as created for now) interactions_response_dict["updated"] = created diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py new file mode 100644 index 00000000000..938e892bd50 --- /dev/null +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -0,0 +1,332 @@ +# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api + +import json +from typing import cast + +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + AgenticLoopPlan, + AgenticLoopRequestPatch, + is_interception_internal_key, +) +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper + +_FOLLOWUP_INTERNAL_PARAMS = frozenset( + ( + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + "_agentic_loop_api_surface", + ) +) + + +def _gate_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_should_run_agentic_loop + func = type(callback).async_should_run_agentic_loop + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _build_plan_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_build_agentic_loop_plan + func = type(callback).async_build_agentic_loop_plan + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _post_hook_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_post_agentic_loop_response_hook + func = type(callback).async_post_agentic_loop_response_hook + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _coerce_int(value: object, default: int) -> int: + return int(value) if isinstance(value, (int, str)) else default + + +def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]: + depth = _coerce_int(kwargs.get("_agentic_loop_depth"), 0) + max_loops = max(_coerce_int(kwargs.get("max_agentic_loops"), 3), 1) + raw_fingerprints = kwargs.get("_agentic_loop_fingerprints") + fingerprints = ( + [str(fp) for fp in raw_fingerprints] + if isinstance(raw_fingerprints, list) + else [] + ) + return depth, max_loops, fingerprints + + +def _fingerprint_tools(tool_calls: object) -> str: + try: + return json.dumps(tool_calls, sort_keys=True, default=str) + except Exception: + return str(tool_calls) + + +def _check_agentic_loop_safety( + tool_calls: object, + fingerprints: list[str], + depth: int, + max_loops: int, + model: str, +) -> str: + fingerprint = _fingerprint_tools(tool_calls) + if fingerprint in fingerprints: + raise ValueError( + "Agentic loop detected repeated tool-call fingerprint; aborting rerun" + ) + if depth >= max_loops: + raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}") + return fingerprint + + +def _wrap_response_as_fake_stream(response: object) -> object: + if getattr(response, "object", None) == "chat.completion.chunk": + return response + if not hasattr(response, "choices"): + return response + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) + + return convert_model_response_to_streaming(cast(ModelResponse, response)) + + +def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None: + metadata = kwargs_for_followup.get("litellm_metadata") + metadata = dict(metadata) if isinstance(metadata, dict) else {} + for key, value in kwargs_for_followup.items(): + if ( + key.startswith("_agentic_loop") + or key == "max_agentic_loops" + or is_interception_internal_key(key) + ): + metadata[key] = value + kwargs_for_followup["litellm_metadata"] = metadata + + +def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: + return { + k: v + for k, v in source.items() + if not is_interception_internal_key( + k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES + ) + and k not in _FOLLOWUP_INTERNAL_PARAMS + } + + +async def _execute_chat_completion_agentic_plan( + *, + plan: AgenticLoopPlan, + callback: CustomLogger, + model: str, + optional_params: dict[str, object], + kwargs: dict[str, object], + logging_obj: object, + custom_llm_provider: str, + depth: int, + max_loops: int, + fingerprints: list[str], + fingerprint: str, +) -> object: + import litellm + + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = patch.model or model + if "/" not in full_model_name: + full_model_name = f"{custom_llm_provider}/{full_model_name}" + + optional_params_for_followup = {**optional_params, **patch.optional_params} + if patch.tools is not None: + optional_params_for_followup["tools"] = patch.tools + if "tool_choice" not in patch.optional_params: + optional_params_for_followup.pop("tool_choice", None) + + kwargs_for_followup = _filter_followup_kwargs(kwargs) + kwargs_for_followup.update( + { + k: v + for k, v in _filter_followup_kwargs(patch.kwargs).items() + if k not in optional_params_for_followup + } + ) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + _add_agentic_loop_metadata(kwargs_for_followup) + + try: + response_followup = await litellm.acompletion( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) + if _post_hook_overridden(callback): + try: + response_followup = ( + await callback.async_post_agentic_loop_response_hook( + response=response_followup, plan=plan, kwargs=kwargs + ) + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + if kwargs.get("_code_interpreter_interception_converted_stream") and not depth: + return _wrap_response_as_fake_stream(response_followup) + return response_followup + finally: + try: + await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + +async def maybe_run_chat_completion_agentic_loop( + *, + response: ModelResponse, + model: str, + messages: list, + optional_params: dict, + kwargs: dict, + logging_obj: object, + custom_llm_provider: str, + stream: bool, +) -> ModelResponse | CustomStreamWrapper | None: + import litellm + + callbacks = litellm.callbacks + ( + getattr(logging_obj, "dynamic_success_callbacks", None) or [] + ) + depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) + tools = optional_params.get("tools", []) + + for callback in callbacks: + if not isinstance(callback, CustomLogger): + continue + if not _gate_overridden(callback): + continue + + gate_kwargs = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + try: + should_run, tool_calls = await callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=gate_kwargs, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in chat completion agentic gate: %s", + str(e), + ) + continue + + if not should_run: + continue + + fingerprint = _check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + plan_kwargs = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + if not _build_plan_overridden(callback): + return await callback.async_run_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=plan_kwargs, + ) + + plan = await callback.async_build_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=plan_kwargs, + ) + + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + return response + if not plan.run_agentic_loop: + continue + + return await _execute_chat_completion_agentic_plan( + plan=plan, + callback=callback, + model=model, + optional_params=optional_params, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s", + str(e), + ) + + if ( + kwargs.get("_code_interpreter_interception_converted_stream") + and not depth + and hasattr(response, "choices") + ): + return cast( + "ModelResponse | CustomStreamWrapper", + _wrap_response_as_fake_stream(response), + ) + return None diff --git a/litellm/litellm_core_utils/completion_timeout.py b/litellm/litellm_core_utils/completion_timeout.py index 5350d88e593..70c6896a323 100644 --- a/litellm/litellm_core_utils/completion_timeout.py +++ b/litellm/litellm_core_utils/completion_timeout.py @@ -6,10 +6,7 @@ from typing import Callable, Optional, Union import httpx -from litellm.constants import ( - COMPLETION_HTTP_FALLBACK_SECONDS, - DEFAULT_REQUEST_TIMEOUT_SECONDS, -) +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS class CompletionTimeout: @@ -22,17 +19,13 @@ class CompletionTimeout: """ Used when ``model_timeout`` and kwargs timeouts are all unset. - ``global_timeout`` is :attr:`litellm.request_timeout` (numeric / string), not - :class:`httpx.Timeout`. - - If it equals :data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS` (6000), - return :data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`. Same if - ``None``. Otherwise return ``float(global_timeout)``. + ``global_timeout`` is the explicitly-configured ``litellm.request_timeout`` + (numeric / string) or ``None`` when it was never set. ``None`` falls back to + :data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`; any explicit value + (including ``6000``) is honored. """ if global_timeout is None: return COMPLETION_HTTP_FALLBACK_SECONDS - if float(global_timeout) == float(DEFAULT_REQUEST_TIMEOUT_SECONDS): - return COMPLETION_HTTP_FALLBACK_SECONDS return float(global_timeout) @staticmethod @@ -50,11 +43,10 @@ class CompletionTimeout: 1. ``model_timeout`` (call argument / merged ``litellm_params``) 2. ``kwargs["timeout"]`` 3. ``kwargs["request_timeout"]`` - 4. Fallback from ``global_timeout`` (:attr:`litellm.request_timeout`) — if it is - the package default (6000), use 600 instead. + 4. ``global_timeout`` (the explicitly-configured ``litellm.request_timeout``), + or 600 when nothing was configured. Coerce :class:`httpx.Timeout` when the provider does not support it. - Explicit ``6000`` on the model or in kwargs is kept as ``6000``. """ resolved: Union[float, str, httpx.Timeout] if model_timeout is not None: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 98b792efa59..0cdd721598c 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -361,9 +361,9 @@ def safe_deep_copy(data): "litellm_metadata" in data and "litellm_parent_otel_span" in data["litellm_metadata"] ): - data["litellm_metadata"][ - "litellm_parent_otel_span" - ] = litellm_parent_otel_span + data["litellm_metadata"]["litellm_parent_otel_span"] = ( + litellm_parent_otel_span + ) return new_data diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 9b2a9af4126..a0f9bb00dd2 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -272,8 +272,8 @@ def _map_openai_exception( else: message = str(original_exception) - if message is not None and isinstance( - message, str + if ( + message is not None and isinstance(message, str) ): # done to prevent user-confusion. Relevant issue - https://github.com/BerriAI/litellm/issues/1414 message = message.replace("OPENAI", custom_llm_provider.upper()) message = message.replace( @@ -726,7 +726,6 @@ def _map_openai_like_exception( extra_information: str, ) -> None: if "authorization denied for" in error_str: - # Predibase returns the raw API Key in the response - this block ensures it's not returned in the exception if ( error_str is not None @@ -1161,7 +1160,9 @@ def _map_vertex_exception( response=httpx.Response( status_code=500, content=str(original_exception), - request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="completion", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), litellm_debug_info=extra_information, ) @@ -1327,7 +1328,9 @@ def _map_vertex_exception( response=httpx.Response( status_code=500, content=str(original_exception), - request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="completion", url="https://github.com/BerriAI/litellm" + ), # type: ignore ), ) if original_exception.status_code == 502: @@ -1965,13 +1968,9 @@ def _map_azure_exception( # content policy violation even when the top-level # code is generic (e.g. "invalid_request_error"). if azure_error_code != "content_policy_violation": - _inner = body_dict["error"].get( - "inner_error" - ) or body_dict[ # type: ignore[index] + _inner = body_dict["error"].get("inner_error") or body_dict[ # type: ignore[index] "error" - ].get( - "innererror" - ) # type: ignore[index] + ].get("innererror") # type: ignore[index] if ( isinstance(_inner, dict) and _inner.get("code") == "ResponsibleAIPolicyViolation" diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index fc3c25e0d95..c88f8b77dc2 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -14,6 +14,7 @@ OPTIONAL_KWARGS_KEYS = frozenset( "azure_password", "azure_scope", "timeout", + "gcs_bucket_name", "bucket_name", "vertex_credentials", "vertex_project", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index bb8b1a82996..061316a69d6 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -616,7 +616,11 @@ def _get_openai_compatible_provider_info( return model, "aiohttp_openai", api_key, api_base elif custom_llm_provider == "anyscale": # anyscale is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = api_base or get_secret_str("ANYSCALE_API_BASE") or "https://api.endpoints.anyscale.com/v1" # type: ignore + api_base = ( + api_base + or get_secret_str("ANYSCALE_API_BASE") + or "https://api.endpoints.anyscale.com/v1" + ) # type: ignore dynamic_api_key = api_key or get_secret_str("ANYSCALE_API_KEY") elif custom_llm_provider == "deepinfra": ( @@ -709,9 +713,7 @@ def _get_openai_compatible_provider_info( ) # type: ignore dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif custom_llm_provider == "ollama": - api_base = ( - api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" - ) # type: ignore + api_base = api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" # type: ignore dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") elif (custom_llm_provider == "ai21_chat") or ( custom_llm_provider == "ai21" and model in litellm.ai21_chat_models diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index e87042b9101..c22d3b99705 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -86,9 +86,7 @@ def get_supported_openai_params( model=model ) elif request_type == "transcription": - return litellm.FireworksAIAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return None else: return litellm.FireworksAIConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "nvidia_nim": @@ -191,7 +189,9 @@ def get_supported_openai_params( ) elif custom_llm_provider == "sambanova": if request_type == "embeddings": - litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(model=model) + return litellm.SambaNovaEmbeddingConfig().get_supported_openai_params( + model=model + ) else: return litellm.SambanovaConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "nebius": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d750a509054..e860f48a7bd 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -293,7 +293,17 @@ def _get_cached_prometheus_logger(): class Logging(LiteLLMLoggingBaseClass): - global supabaseClient, promptLayerLogger, weightsBiasesLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app + global \ + supabaseClient, \ + promptLayerLogger, \ + weightsBiasesLogger, \ + logfireLogger, \ + capture_exception, \ + add_breadcrumb, \ + lunaryLogger, \ + logfireLogger, \ + prometheusLogger, \ + slack_app custom_pricing: bool = False stream_options = None litellm_request_debug: bool = False @@ -359,9 +369,9 @@ class Logging(LiteLLMLoggingBaseClass): ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_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 # Initialize dynamic callbacks @@ -903,8 +913,11 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["prompt_integration"] = logger.__class__.__name__ return logger - if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( - non_default_params + if ( + anthropic_cache_control_logger + := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( + non_default_params + ) ): self.model_call_details["prompt_integration"] = ( anthropic_cache_control_logger.__class__.__name__ @@ -978,9 +991,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["api_key"] = api_key self.model_call_details["additional_args"] = additional_args self.model_call_details["log_event_type"] = "pre_api_call" - if ( - model - ): # if model name was changes pre-call, overwrite the initial model call name with the new one + if model: # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model self.model_call_details["litellm_params"]["api_base"] = ( self._get_masked_api_base(additional_args.get("api_base", "")) @@ -1359,13 +1370,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[MCPPostCallResponseObject] = ( - await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, - ) + response: Optional[ + MCPPostCallResponseObject + ] = await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1865,7 +1876,9 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["litellm_params"].setdefault("metadata", {}) if self.model_call_details["litellm_params"]["metadata"] is None: self.model_call_details["litellm_params"]["metadata"] = {} - self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = getattr(logging_result, "_hidden_params", {}) # type: ignore + self.model_call_details["litellm_params"]["metadata"][ + "hidden_params" + ] = getattr(logging_result, "_hidden_params", {}) # type: ignore if self.model_call_details.get("cache_hit") is True: self.model_call_details["response_cost"] = 0.0 @@ -1944,7 +1957,9 @@ class Logging(LiteLLMLoggingBaseClass): ) result = result.model_copy() - transformed_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object(result.usage) # type: ignore + transformed_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object( + result.usage + ) # type: ignore setattr(result, "usage", transformed_usage) return result @@ -2948,7 +2963,9 @@ class Logging(LiteLLMLoggingBaseClass): for callback_obj in all_callbacks: if hasattr(callback_obj, "increment_callback_logging_failure"): - callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore + callback_obj.increment_callback_logging_failure( + callback_name=callback_name + ) # type: ignore break # Only increment once except Exception as e: @@ -3292,9 +3309,7 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.exception( "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure \ - logging {}\nCallback={}".format( - str(e), callback - ) + logging {}\nCallback={}".format(str(e), callback) ) # Track callback logging failures in Prometheus self._handle_callback_failure(callback=callback) @@ -3762,7 +3777,29 @@ def set_callbacks(callback_list, function_id=None): """ Globally sets the callback client """ - global sentry_sdk_instance, capture_exception, add_breadcrumb, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, supabaseClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, logfireLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger, greenscaleLogger, openMeterLogger, deepevalLogger + global \ + sentry_sdk_instance, \ + capture_exception, \ + add_breadcrumb, \ + slack_app, \ + alerts_channel, \ + traceloopLogger, \ + athinaLogger, \ + heliconeLogger, \ + supabaseClient, \ + lunaryLogger, \ + promptLayerLogger, \ + langFuseLogger, \ + customLogger, \ + weightsBiasesLogger, \ + logfireLogger, \ + dynamoLogger, \ + s3Logger, \ + dataDogLogger, \ + prometheusLogger, \ + greenscaleLogger, \ + openMeterLogger, \ + deepevalLogger try: for callback in callback_list: @@ -4607,7 +4644,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: litellm.logging_callback_manager.add_litellm_callback(phoenix_logger) verbose_logger.info( - "Auto-initialized Arize Phoenix logger alongside otel " "(endpoint=%s)", + "Auto-initialized Arize Phoenix logger alongside otel (endpoint=%s)", arize_phoenix_config.endpoint, ) except Exception as e: @@ -5781,7 +5818,8 @@ def get_standard_logging_object_payload( id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id saved_cache_cost = ( logging_obj._response_cost_calculator( - result=init_response_obj, cache_hit=False # type: ignore + result=init_response_obj, + cache_hit=False, # type: ignore ) or 0.0 ) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 7a7fde3087e..19b711f6748 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -131,8 +131,10 @@ def _generic_cost_per_character( assert ( "input_cost_per_character" in model_info and model_info["input_cost_per_character"] is not None - ), "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( - model, model_info + ), ( + "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( + model, model_info + ) ) custom_prompt_cost = model_info["input_cost_per_character"] @@ -152,8 +154,10 @@ def _generic_cost_per_character( assert ( "output_cost_per_character" in model_info and model_info["output_cost_per_character"] is not None - ), "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( - model, model_info + ), ( + "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( + model, model_info + ) ) custom_completion_cost = model_info["output_cost_per_character"] completion_cost = completion_characters * custom_completion_cost diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 016bb6b1e22..79c6c665684 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -1,5 +1,6 @@ import asyncio import json +import re import time import traceback from typing import Dict, Iterable, List, Literal, Optional, Tuple, Union, cast @@ -118,7 +119,44 @@ def convert_tool_call_to_json_mode( return None, None -async def convert_to_streaming_response_async(response_object: Optional[dict] = None): +# Whitespace-preserving word splitter used by the cache-hit replay generators. +# Each match is any leading whitespace plus a non-whitespace run plus any +# trailing whitespace, so concatenating the matches losslessly reconstructs +# the original string (including content that starts with whitespace). +_REPLAY_CONTENT_SLICE_RE = re.compile(r"\s*\S+\s*", re.UNICODE) + + +def _split_assembled_content_for_replay(content: Optional[str]) -> list[str]: + """ + Slice an assembled cached completion's ``content`` into word-shaped pieces + for cadence-preserving streaming replay. The split is lossless: + ``"".join(_split_assembled_content_for_replay(s)) == s`` for every + non-empty ``s``. Returns ``[]`` for ``None`` / empty / all-whitespace + content. + """ + if not content or content.isspace(): + # isspace() guard: on all-whitespace content the regex backtracks + # quadratically before returning no matches. + return [] + return _REPLAY_CONTENT_SLICE_RE.findall(content) + + +def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: + # Rebuild the delta as content-only so every accumulate-able field (role, + # tool_calls, reasoning_content, thinking_blocks, audio, images, + # annotations, ...) is dropped on later slices instead of an enumerated + # subset; repeating any of them makes downstream handlers that accumulate + # streamed deltas collect it once per slice, and a field added to Delta + # later can't silently re-introduce the duplication. + choice.delta = Delta(content=choice.delta.content) + choice.logprobs = None # type: ignore[assignment] + if hasattr(choice, "enhancements"): + del choice.enhancements + + +async def convert_to_streaming_response_async( + response_object: Optional[dict] = None, +): """ Asynchronously converts a response object to a streaming response. @@ -215,11 +253,45 @@ async def convert_to_streaming_response_async(response_object: Optional[dict] = if "model" in response_object: model_response_object.model = response_object["model"] - yield model_response_object - await asyncio.sleep(0) + # Replay cached content with per-word cadence so stream=true cache hits + # don't arrive as a single SSE frame. Multi-choice (n>1) responses and + # unsplittable content (None/empty/whitespace-free) keep the original + # single-yield behavior. + slices: list[str] = [] + if len(model_response_object.choices) == 1: + slices = _split_assembled_content_for_replay( + model_response_object.choices[0].delta.content + ) + if len(slices) <= 1: + yield model_response_object + await asyncio.sleep(0) + return + + # Detach usage from the base object so we can re-attach it only to the + # final slice chunk. A non-None usage always lives in __pydantic_extra__ + # here (set via setattr above), so delattr cannot fail. + original_usage = getattr(model_response_object, "usage", None) + if original_usage is not None: + delattr(model_response_object, "usage") + original_finish_reason = model_response_object.choices[0].finish_reason + last_idx = len(slices) - 1 + for i, piece in enumerate(slices): + slice_chunk = model_response_object.model_copy(deep=True) + slice_chunk.choices[0].delta.content = piece + if i > 0: + _clear_later_replay_slice_metadata(slice_chunk.choices[0]) + slice_chunk.choices[0].finish_reason = ( + original_finish_reason if i == last_idx else None # type: ignore[assignment] + ) + if i == last_idx and original_usage is not None: + setattr(slice_chunk, "usage", original_usage) + yield slice_chunk + await asyncio.sleep(0) -def convert_to_streaming_response(response_object: Optional[dict] = None): +def convert_to_streaming_response( + response_object: Optional[dict] = None, +): # used for yielding Cache hits when stream == True if response_object is None: raise Exception("Error in response object format") @@ -261,9 +333,15 @@ def convert_to_streaming_response(response_object: Optional[dict] = None): if "usage" in response_object and response_object["usage"] is not None: setattr(model_response_object, "usage", Usage()) - model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) # type: ignore - model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) # type: ignore - model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) # type: ignore + model_response_object.usage.completion_tokens = response_object["usage"].get( + "completion_tokens", 0 + ) # type: ignore + model_response_object.usage.prompt_tokens = response_object["usage"].get( + "prompt_tokens", 0 + ) # type: ignore + model_response_object.usage.total_tokens = response_object["usage"].get( + "total_tokens", 0 + ) # type: ignore if "id" in response_object: model_response_object.id = response_object["id"] @@ -278,7 +356,35 @@ def convert_to_streaming_response(response_object: Optional[dict] = None): if "model" in response_object: model_response_object.model = response_object["model"] - yield model_response_object + + # Replay cached content with per-word cadence on sync cache-hit paths + # (S3Cache, sync completion()). See convert_to_streaming_response_async + # for the full rationale — this mirrors its tail. + slices: list[str] = [] + if len(model_response_object.choices) == 1: + slices = _split_assembled_content_for_replay( + model_response_object.choices[0].delta.content + ) + if len(slices) <= 1: + yield model_response_object + return + + original_usage = getattr(model_response_object, "usage", None) + if original_usage is not None: + delattr(model_response_object, "usage") + original_finish_reason = model_response_object.choices[0].finish_reason + last_idx = len(slices) - 1 + for i, piece in enumerate(slices): + slice_chunk = model_response_object.model_copy(deep=True) + slice_chunk.choices[0].delta.content = piece + if i > 0: + _clear_later_replay_slice_metadata(slice_chunk.choices[0]) + slice_chunk.choices[0].finish_reason = ( + original_finish_reason if i == last_idx else None # type: ignore[assignment] + ) + if i == last_idx and original_usage is not None: + setattr(slice_chunk, "usage", original_usage) + yield slice_chunk from collections import defaultdict @@ -748,9 +854,15 @@ def convert_to_model_response_object( model_response_object.data = response_object["data"] if "usage" in response_object and response_object["usage"] is not None: - model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) # type: ignore - model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) # type: ignore - model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) # type: ignore + model_response_object.usage.completion_tokens = response_object[ + "usage" + ].get("completion_tokens", 0) # type: ignore + model_response_object.usage.prompt_tokens = response_object[ + "usage" + ].get("prompt_tokens", 0) # type: ignore + model_response_object.usage.total_tokens = response_object["usage"].get( + "total_tokens", 0 + ) # type: ignore if start_time is not None and end_time is not None: model_response_object._response_ms = ( # type: ignore diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index b7adda3a9a4..f58126ec901 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -68,7 +68,8 @@ class LoggingCallbackManager: Ensures no duplicates are added. """ self._safe_add_callback_to_list( - callback=callback, parent_list=litellm.callbacks # type: ignore + callback=callback, + parent_list=litellm.callbacks, # type: ignore ) def add_litellm_success_callback( diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index fe34731759f..bf9ce3b0acb 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -757,6 +757,46 @@ def update_responses_tools_with_model_file_ids( return updated_tools +def extract_file_metadata(file_data: FileTypes) -> Tuple[Optional[str], Optional[str]]: + """ + Resolve (filename, content_type) without reading the file body. + + Mirrors extract_file_data's metadata resolution but never calls .read(), so + it stays O(1) on large uploads. Use this when only metadata is needed (batch + detection, GCS object naming) and the body must remain a streamable Path/handle. + """ + filename: Optional[str] = None + content_type: Optional[str] = None + file_content: Any = None + + if isinstance(file_data, tuple): + if len(file_data) == 2: + filename, file_content = file_data + elif len(file_data) == 3: + filename, file_content, content_type = file_data + elif len(file_data) == 4: + filename, file_content, content_type, _ = file_data + elif isinstance(file_data, InMemoryFile): + filename = file_data.name + content_type = file_data.content_type + else: + file_content = file_data + + if filename is None: + if isinstance(file_content, PathLike): + filename = Path(file_content).name + elif isinstance(file_content, io.IOBase): + name_attr = getattr(file_content, "name", None) + if isinstance(name_attr, str): + filename = Path(name_attr).name + + if not content_type: + guessed = mimetypes.guess_type(filename)[0] if filename else None + content_type = guessed or "application/octet-stream" + + return filename, content_type + + def extract_file_data(file_data: FileTypes) -> ExtractedFileData: """ Extracts and processes file data from various input formats. diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b95b73398ac..b4f492bc26f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1586,7 +1586,9 @@ def convert_to_gemini_tool_call_result( file_data = ( file_content.get("file_data", "") if isinstance(file_content, dict) - else file_content if isinstance(file_content, str) else "" + else file_content + if isinstance(file_content, str) + else "" ) if file_data: @@ -2556,9 +2558,7 @@ def anthropic_messages_pt( ChatCompletionToolMessage, ChatCompletionUserMessage, ChatCompletionFunctionMessage, - ] = messages[ - msg_i - ] # type: ignore + ] = messages[msg_i] # type: ignore if user_message_types_block["role"] == "user": if isinstance(user_message_types_block["content"], list): for m in user_message_types_block["content"]: @@ -4926,8 +4926,10 @@ class BedrockConverseMessagesProcessor: image_url = element["image_url"]["url"] else: image_url = element["image_url"] - assistants_part = await BedrockImageProcessor.process_image_async( # type: ignore - image_url=image_url + assistants_part = ( + await BedrockImageProcessor.process_image_async( # type: ignore + image_url=image_url + ) ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c56a70177bf..ca7208458d8 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -169,7 +169,9 @@ class RealTimeStreaming: try: event_type = message_obj.get("type", "") if event_type in self._SESSION_EVENT_TYPES: - typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore + typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents( + **message_obj + ) # type: ignore else: # Catch-all base object so unknown/new event names never raise. typed_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore diff --git a/litellm/litellm_core_utils/request_timeout_resolver.py b/litellm/litellm_core_utils/request_timeout_resolver.py new file mode 100644 index 00000000000..146c39ce9f3 --- /dev/null +++ b/litellm/litellm_core_utils/request_timeout_resolver.py @@ -0,0 +1,29 @@ +"""Single source of truth for whether ``litellm.request_timeout`` was configured. + +``litellm.request_timeout`` always holds a value (the package default, +:data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS`), so a bare read can't +tell "user asked for this" from "nobody set it". This resolver answers that: + +* ``request_timeout_explicitly_set`` is the authoritative signal, set when the + value comes from the ``REQUEST_TIMEOUT`` env var or ``litellm_settings``. +* A runtime value that differs from the package default (e.g. ``litellm.request_timeout + = 300`` in SDK code) is also treated as explicit, for backwards compatibility. +""" + +from __future__ import annotations + +from typing import Optional + +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS + + +def get_configured_request_timeout() -> Optional[float]: + """Return the explicitly-configured ``litellm.request_timeout``, else ``None``.""" + import litellm + + timeout = float(litellm.request_timeout) + if litellm.request_timeout_explicitly_set: + return timeout + if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS): + return timeout + return None diff --git a/litellm/litellm_core_utils/rules.py b/litellm/litellm_core_utils/rules.py index 717ff55ab22..75c177c7249 100644 --- a/litellm/litellm_core_utils/rules.py +++ b/litellm/litellm_core_utils/rules.py @@ -33,7 +33,11 @@ class Rules: if callable(rule): decision = rule(input) if decision is False: - raise litellm.APIResponseValidationError(message="LLM Response failed post-call-rule check", llm_provider="", model=model) # type: ignore + raise litellm.APIResponseValidationError( + message="LLM Response failed post-call-rule check", + llm_provider="", + model=model, + ) # type: ignore return True def post_call_rules(self, input: Optional[str], model: str) -> bool: @@ -44,12 +48,18 @@ class Rules: decision = rule(input) if isinstance(decision, bool): if decision is False: - raise litellm.APIResponseValidationError(message="LLM Response failed post-call-rule check", llm_provider="", model=model) # type: ignore + raise litellm.APIResponseValidationError( + message="LLM Response failed post-call-rule check", + llm_provider="", + model=model, + ) # type: ignore elif isinstance(decision, dict): decision_val = decision.get("decision", True) decision_message = decision.get( "message", "LLM Response failed post-call-rule check" ) if decision_val is False: - raise litellm.APIResponseValidationError(message=decision_message, llm_provider="", model=model) # type: ignore + raise litellm.APIResponseValidationError( + message=decision_message, llm_provider="", model=model + ) # type: ignore return True diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 4928dd08386..3135b5f831a 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -12,6 +12,7 @@ class SensitiveDataMasker: visible_prefix: int = 4, visible_suffix: int = 4, mask_char: str = "*", + mask_short_values: bool = True, ): self.sensitive_patterns = sensitive_patterns or { "password", @@ -38,19 +39,24 @@ class SensitiveDataMasker: self.visible_prefix = visible_prefix self.visible_suffix = visible_suffix self.mask_char = mask_char + self.mask_short_values = mask_short_values def _mask_value(self, value: str) -> str: - if not value or len(str(value)) < (self.visible_prefix + self.visible_suffix): - return value - value_str = str(value) + if not value_str: + return value + if len(value_str) <= (self.visible_prefix + self.visible_suffix): + return ( + self.mask_char * len(value_str) if self.mask_short_values else value_str + ) + masked_length = len(value_str) - (self.visible_prefix + self.visible_suffix) # Handle the case where visible_suffix is 0 to avoid showing the entire string if self.visible_suffix == 0: - return f"{value_str[:self.visible_prefix]}{self.mask_char * masked_length}" + return f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}" else: - return f"{value_str[:self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix:]}" + return f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix :]}" def is_sensitive_key( self, key: str, excluded_keys: Optional[Set[str]] = None diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 04f6b1241c3..66abb824610 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -213,9 +213,9 @@ class ChunkProcessor: self, tool_call_chunks: List[Dict[str, Any]] ) -> List[ChatCompletionMessageToolCall]: tool_calls_list: List[ChatCompletionMessageToolCall] = [] - tool_call_map: Dict[int, Dict[str, Any]] = ( - {} - ) # Map to store tool calls by index + tool_call_map: Dict[ + int, Dict[str, Any] + ] = {} # Map to store tool calls by index for chunk in tool_call_chunks: choices = chunk["choices"] @@ -720,15 +720,16 @@ class ChunkProcessor: returned_usage.prompt_tokens = prompt_tokens or token_counter( model=model, messages=messages ) - except ( - Exception - ): # don't allow this failing to block a complete streaming response from being returned + except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 - returned_usage.completion_tokens = completion_tokens or token_counter( - model=model, - text=completion_output, - count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + returned_usage.completion_tokens = ( + completion_tokens + or token_counter( + model=model, + text=completion_output, + count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + ) ) returned_usage.total_tokens = ( returned_usage.prompt_tokens + returned_usage.completion_tokens diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index e278483d689..7d7a8e562cb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -189,9 +189,7 @@ class CustomStreamWrapper: True if self.check_send_stream_usage(self.stream_options) else False ) self.tool_call = False - self.chunks: List = ( - [] - ) # keep track of the returned chunks - used for calculating the input/output tokens for stream options + self.chunks: List = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options self._repeated_messages_count = 1 self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: Optional[int] = None @@ -1438,10 +1436,11 @@ class CustomStreamWrapper: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "cached_response": chunk = cast(ModelResponseStream, chunk) + chunk_finish_reason = chunk.choices[0].finish_reason response_obj = { "text": chunk.choices[0].delta.content, - "is_finished": True, - "finish_reason": chunk.choices[0].finish_reason, + "is_finished": chunk_finish_reason is not None, + "finish_reason": chunk_finish_reason, "original_chunk": chunk, "tool_calls": ( chunk.choices[0].delta.tool_calls @@ -1860,8 +1859,10 @@ class CustomStreamWrapper: Caches the streaming response """ if not cache_hit and self.logging_obj._llm_caching_handler is not None: - await self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( - processed_chunk + await ( + self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( + processed_chunk + ) ) def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool): @@ -2216,7 +2217,9 @@ class CustomStreamWrapper: ) ) # Add MCP metadata to final chunk if present (after hooks) - processed_chunk = self._add_mcp_metadata_to_final_chunk(processed_chunk) # type: ignore[reportArgumentType] + processed_chunk = self._add_mcp_metadata_to_final_chunk( + processed_chunk + ) # type: ignore[reportArgumentType] return processed_chunk raise StopAsyncIteration @@ -2228,7 +2231,9 @@ class CustomStreamWrapper: ): chunk = self.completion_stream else: - chunk = await asyncio.to_thread(_next_sync_or_exhausted, self.completion_stream) # type: ignore[arg-type] + chunk = await asyncio.to_thread( + _next_sync_or_exhausted, self.completion_stream + ) # type: ignore[arg-type] if chunk is _SYNC_ITER_EXHAUSTED: raise StopAsyncIteration if chunk is not None and chunk != b"": diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 710342bbc78..fead8a79bc3 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -63,9 +63,9 @@ def get_cost_for_web_search_request( return None -def discover_guardrail_translation_mappings() -> ( - Dict[CallTypes, Type["BaseTranslation"]] -): +def discover_guardrail_translation_mappings() -> Dict[ + CallTypes, Type["BaseTranslation"] +]: """ Discover guardrail translation mappings by scanning the llms directory structure. diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 74dadee5ecb..7d5944b4cc2 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -288,9 +288,9 @@ class AnthropicMessagesHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - messages[msg_idx]["content"][content_idx_optional][ - "text" - ] = guardrail_response + messages[msg_idx]["content"][content_idx_optional]["text"] = ( + guardrail_response + ) async def process_output_response( self, diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 5d14f3cc4ae..e7e0f12c455 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -624,7 +624,9 @@ class ModelResponseIterator: speed=self.speed, ) - def _content_block_delta_helper(self, chunk: dict) -> Tuple[ + def _content_block_delta_helper( + self, chunk: dict + ) -> Tuple[ str, Optional[ChatCompletionToolCallChunk], List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]], diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index c24c990f356..2b2264e6da9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -229,6 +229,11 @@ DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING = ( "Sonnet 4.6+, and Mythos Preview." ) +DROP_UNSUPPORTED_SPEED_WARNING = ( + "Dropping unsupported `speed` for model=%s " + "(drop_params=True). Fast mode is only supported on select Opus models." +) + class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ @@ -374,6 +379,51 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): for level in ("low", "minimal", "medium", "high", "xhigh", "max") ) + @staticmethod + def _model_supports_speed_param( + model: str, custom_llm_provider: Optional[str] = None + ) -> bool: + """Whether the model accepts Anthropic's ``speed`` parameter (fast mode). + + Fast mode is direct Anthropic API-only (not Bedrock, Vertex, or Azure). + Those providers strip their prefix before this shared transform runs, so a + bare ``claude-opus-4-8`` would otherwise resolve to the direct-API entry; + the routed provider is checked explicitly to keep them out. + """ + if custom_llm_provider is not None and custom_llm_provider != "anthropic": + return False + return ( + AnthropicModelInfo._get_exact_model_capability(model, "supports_speed") + is True + ) + + @staticmethod + def _maybe_drop_speed_param( + model: str, + optional_params: dict, + drop_params: bool, + custom_llm_provider: Optional[str] = None, + ) -> None: + if "speed" not in optional_params: + return + if AnthropicConfig._model_supports_speed_param(model, custom_llm_provider): + return + if not (litellm.drop_params or drop_params): + speed_value = optional_params.get("speed") + raise litellm.utils.UnsupportedParamsError( + message=( + f"{model} does not support speed={speed_value!r}. " + "To drop unsupported params, set " + "`litellm.drop_params = True`." + ), + status_code=400, + ) + litellm.verbose_logger.warning( + DROP_UNSUPPORTED_SPEED_WARNING, + model, + ) + optional_params.pop("speed", None) + @staticmethod def _raise_invalid_reasoning_effort( model: str, value: Any, llm_provider: str @@ -700,7 +750,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): additional_tool_params[k] = v returned_tool = AnthropicHostedTools( - type=tool["type"], name=function_name, **additional_tool_params # type: ignore + type=tool["type"], + name=function_name, + **additional_tool_params, # type: ignore ) elif tool["type"] == "url": # mcp server tool mcp_server = AnthropicMcpServerTool(**tool) # type: ignore @@ -1569,8 +1621,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_context_management ) elif param == "speed" and isinstance(value, str): - # Pass through Anthropic-specific speed parameter for fast mode optional_params["speed"] = value + AnthropicConfig._maybe_drop_speed_param( + model=model, + optional_params=optional_params, + drop_params=drop_params, + custom_llm_provider=self.custom_llm_provider, + ) elif param == "cache_control" and isinstance(value, dict): # Pass through top-level cache_control for automatic prompt caching optional_params["cache_control"] = value @@ -1875,6 +1932,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "has no thinking_blocks. The model won't use extended thinking for this turn." ) + AnthropicConfig._maybe_drop_speed_param( + model=model, + optional_params=optional_params, + drop_params=litellm.drop_params + or litellm_params.get("drop_params") is True, + custom_llm_provider=self.custom_llm_provider, + ) + headers = self.update_headers_with_optional_anthropic_beta( headers=headers, optional_params=optional_params ) @@ -2081,7 +2146,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices] return None, filtered_tools, extra_content - def extract_response_content(self, completion_response: dict) -> Tuple[ + def extract_response_content( + self, completion_response: dict + ) -> Tuple[ str, Optional[List[Any]], Optional[ diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 5741513903c..eaa762204f2 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -3,6 +3,7 @@ This file contains common utils for anthropic calls. """ import copy +import re from typing import Any, Dict, List, Optional, Union import httpx @@ -11,6 +12,9 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, ) +from litellm.litellm_core_utils.prompt_templates.factory import ( + THOUGHT_SIGNATURE_SEPARATOR, +) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import ( @@ -367,6 +371,16 @@ class AnthropicModelInfo(BaseLLMModelInfo): pass return None + @staticmethod + def _get_exact_model_capability(model: str, key: str) -> Optional[bool]: + """Read boolean capability ``key`` from the exact model-map entry only. + + Unlike ``_get_model_capability``, does not walk stripped provider aliases. + Use when a feature is tied to a specific host (e.g. Anthropic API fast mode). + """ + value = litellm.model_cost.get(model, {}).get(key) + return value if isinstance(value, bool) else None + @staticmethod def _supports_model_capability(model: str, key: str) -> bool: """Check a boolean capability ``key`` in the model map. @@ -488,7 +502,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): "computer_20241022": "computer-use-2024-10-22", } return computer_tool_beta_mapping.get( - computer_tool_version, "computer-use-2024-10-22" # Default fallback + computer_tool_version, + "computer-use-2024-10-22", # Default fallback ) def get_anthropic_beta_list( @@ -989,6 +1004,67 @@ def _is_empty_text_block(block: Any) -> bool: return not isinstance(text, str) or not text.strip() +def normalize_anthropic_tool_use_id(raw_id: str) -> str: + """ + Normalize a tool_use / tool_result id for Anthropic's ``^[a-zA-Z0-9_-]+$`` + pattern. + + Strips Gemini thought-signature suffixes (``__thought__``) first, then + replaces any remaining invalid characters with underscores. + """ + base_id = ( + raw_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] + if THOUGHT_SIGNATURE_SEPARATOR in raw_id + else raw_id + ) + sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", base_id) + return sanitized or "tool_use_id" + + +def _sanitize_tool_use_id_content_block(block: Any) -> Any: + if not isinstance(block, dict): + return block + block_type = block.get("type") + if block_type in ("tool_use", "server_tool_use"): + raw_id = block.get("id") + if isinstance(raw_id, str): + normalized = normalize_anthropic_tool_use_id(raw_id) + if normalized != raw_id: + return {**block, "id": normalized} + elif block_type == "tool_result": + raw_id = block.get("tool_use_id") + if isinstance(raw_id, str): + normalized = normalize_anthropic_tool_use_id(raw_id) + if normalized != raw_id: + return {**block, "tool_use_id": normalized} + return block + + +def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any]: + """ + Return a new message list with ``tool_use`` / ``server_tool_use`` ``id`` and + ``tool_result`` ``tool_use_id`` values rewritten to satisfy Anthropic's + ``^[a-zA-Z0-9_-]+$`` requirement. + + Cross-provider clients (e.g. Claude Code routed through kimi) may replay + conversation history containing ids like ``functions.Bash:0`` with ``.`` + and ``:`` — valid on the upstream provider but rejected by Anthropic when + the session is switched to a native Anthropic deployment. + """ + out: list[Any] = [] + for m in messages: + if not isinstance(m, dict) or not isinstance(m.get("content"), list): + out.append(m) + continue + content = m["content"] + new_content = [_sanitize_tool_use_id_content_block(b) for b in content] + if new_content == content: + out.append(m) + else: + out.append({**m, "content": new_content}) + return out + + def process_anthropic_headers(headers: Union[httpx.Headers, dict]) -> dict: openai_headers = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index a8e2fceb4ee..9860997b1cd 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -184,9 +184,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # class level) so concurrent streams don't share the same mutable dict # — `_should_start_new_content_block` mutates `tool_block["name"]` in # place, which would otherwise leak across streams. - self.current_content_block_start: ( - "AnthropicStreamWrapper.ContentBlockContentBlockDict" - ) = self.TextBlock( + self.current_content_block_start: "AnthropicStreamWrapper.ContentBlockContentBlockDict" = self.TextBlock( type="text", text="", ) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 75a8acdfcc3..46564a565ce 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -76,6 +76,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, ) +from litellm.llms.anthropic.common_utils import normalize_anthropic_tool_use_id from litellm.llms.anthropic.experimental_pass_through.context_management import ( PolyfillResult, ) @@ -594,9 +595,9 @@ class LiteLLMAnthropicMessagesAdapter: ## ASSISTANT MESSAGE ## assistant_message_str: Optional[str] = None - assistant_content_list: List[Dict[str, Any]] = ( - [] - ) # For content blocks with cache_control + assistant_content_list: List[ + Dict[str, Any] + ] = [] # For content blocks with cache_control has_cache_control_in_text = False tool_calls: List[ChatCompletionAssistantToolCall] = [] thinking_blocks: List[ @@ -1024,7 +1025,9 @@ class LiteLLMAnthropicMessagesAdapter: if openai_system_content: new_messages.insert( 0, - ChatCompletionSystemMessage(role="system", content=openai_system_content), # type: ignore + ChatCompletionSystemMessage( + role="system", content=openai_system_content + ), # type: ignore ) def _translate_metadata_to_openai( @@ -1363,18 +1366,12 @@ class LiteLLMAnthropicMessagesAdapter: else truncated_name ) - # Strip Gemini thought-signature suffix from id (mirrors streaming - # path below); base64 chars (+ / =) violate Anthropic's - # `^[a-zA-Z0-9_-]+$` tool_use.id pattern when replayed. + # Strip Gemini thought-signature suffix and normalize id chars + # (e.g. ``functions.Bash:0`` from cross-provider clients). raw_id = tool_call.id or "" - base_id = ( - raw_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] - if THOUGHT_SIGNATURE_SEPARATOR in raw_id - else raw_id - ) tool_use_block = AnthropicResponseContentBlockToolUse( type="tool_use", - id=base_id, + id=normalize_anthropic_tool_use_id(raw_id), name=original_name, input=parse_tool_call_arguments( tool_call.function.arguments, @@ -1461,7 +1458,9 @@ class LiteLLMAnthropicMessagesAdapter: "input_tokens": uncached_input_tokens, "output_tokens": usage.completion_tokens or 0, } - anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [message_iteration] # type: ignore[typeddict-unknown-key] + anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [ + message_iteration + ] # type: ignore[typeddict-unknown-key] translated_obj = AnthropicMessagesResponse( id=response.id, @@ -1501,15 +1500,13 @@ class LiteLLMAnthropicMessagesAdapter: ): raw_id = choice.delta.tool_calls[0].id or str(uuid.uuid4()) tool_name = choice.delta.tool_calls[0].function.name or "" - base_id = raw_id thought_sig: Optional[str] = None if THOUGHT_SIGNATURE_SEPARATOR in raw_id: parts = raw_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1) - base_id = parts[0] thought_sig = parts[1] if len(parts) > 1 else None tool_block: Dict[str, Any] = { "type": "tool_use", - "id": base_id, + "id": normalize_anthropic_tool_use_id(raw_id), "name": tool_name, "input": {}, } @@ -1679,7 +1676,9 @@ class LiteLLMAnthropicMessagesAdapter: else: usage_delta = UsageDelta(input_tokens=0, output_tokens=0) message_block = MessageBlockDelta( - type="message_delta", delta=delta, usage=usage_delta # type: ignore + type="message_delta", + delta=delta, + usage=usage_delta, # type: ignore ) if applied_edits: message_block["context_management"] = ContextManagementResponse( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index a3ac465c463..cb61c196fd8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -23,6 +23,7 @@ from typing import ( import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( + sanitize_tool_use_ids_in_anthropic_messages, strip_empty_text_blocks_from_anthropic_messages, ) from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -214,6 +215,9 @@ async def anthropic_messages( # already handles this in anthropic_messages_pt; sanitize the native # Anthropic Messages path here for the same guarantee. See #22930. messages = strip_empty_text_blocks_from_anthropic_messages(messages) + # Replay of cross-provider tool history (e.g. kimi -> Anthropic) may carry + # ids like ``functions.Bash:0`` that violate Anthropic's id pattern. + messages = sanitize_tool_use_ids_in_anthropic_messages(messages) original_stream = stream or kwargs.get( "_websearch_interception_converted_stream", False @@ -397,6 +401,7 @@ def anthropic_messages_handler( # full-messages scan. Pop it so it never leaks into provider params. if not kwargs.pop("_litellm_messages_presanitized", False): messages = strip_empty_text_blocks_from_anthropic_messages(messages) + messages = sanitize_tool_use_ids_in_anthropic_messages(messages) metadata = validate_anthropic_api_metadata(metadata) @@ -507,7 +512,10 @@ def anthropic_messages_handler( local_vars.update(kwargs) anthropic_messages_optional_request_params = ( AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( - params=local_vars + params=local_vars, + model=model, + drop_params=litellm_params.get("drop_params") is True, + custom_llm_provider=custom_llm_provider, ) ) if is_reasoning_auto_summary_enabled(): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 978eaab65d8..25109765772 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -87,7 +87,7 @@ class BaseAnthropicMessagesStreamingIterator: """ if isinstance(chunk, dict): event_type: str = str(chunk.get("type", "message")) - payload = f"event: {event_type}\n" f"data: {json.dumps(chunk)}\n\n" + payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" return payload.encode() else: # For non-dict chunks, return as is diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 88832fb3f63..42167e0fdaa 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -23,12 +23,19 @@ class AnthropicMessagesRequestUtils: @staticmethod def get_requested_anthropic_messages_optional_param( params: Dict[str, Any], + *, + model: str | None = None, + drop_params: bool = False, + custom_llm_provider: str | None = None, ) -> AnthropicMessagesRequestOptionalParams: """ Filter parameters to only include those defined in AnthropicMessagesRequestOptionalParams. Args: params: Dictionary of parameters to filter + model: Resolved model id; when set, unsupported params may be dropped + drop_params: Per-request drop_params flag (also respects litellm.drop_params) + custom_llm_provider: Routed provider; fast mode is gated to direct Anthropic Returns: AnthropicMessagesRequestOptionalParams instance with only the valid parameters @@ -37,6 +44,15 @@ class AnthropicMessagesRequestUtils: filtered_params = { k: v for k, v in params.items() if k in valid_keys and v is not None } + if model is not None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + AnthropicConfig._maybe_drop_speed_param( + model=model, + optional_params=filtered_params, + drop_params=drop_params, + custom_llm_provider=custom_llm_provider, + ) return cast(AnthropicMessagesRequestOptionalParams, filtered_params) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 04819a416a2..f400dc7804e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -35,9 +35,9 @@ class AnthropicResponsesStreamWrapper: # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: Dict[str, int] = {} # Track open function_call items by item_id so we can emit tool_use start - self._pending_tool_ids: Dict[str, str] = ( - {} - ) # item_id -> call_id / name accumulator + self._pending_tool_ids: Dict[ + str, str + ] = {} # item_id -> call_id / name accumulator self._sent_message_start = False self._sent_message_stop = False self._chunk_queue: deque = deque() diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index 1eb7d34c875..bc11875ba12 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -53,7 +53,13 @@ class APISerpentSearchConfig(BaseSearchConfig): api_base: Optional[str] = None, **kwargs, ) -> Dict: - api_key = api_key or get_secret_str("APISERPENT_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("APISERPENT_API_KEY",), + base_env_var="APISERPENT_API_BASE", + default_api_base=APISERPENT_BASE, + ) if not api_key: raise ValueError( "APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable." diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 271cd698e7b..750087b722e 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -203,8 +203,11 @@ class AzureAssistantsAPI(BaseAzureLLM): litellm_params=litellm_params, ) - thread_message: OpenAIMessage = await openai_client.beta.threads.messages.create( # type: ignore - thread_id, **message_data # type: ignore + thread_message: OpenAIMessage = ( + await openai_client.beta.threads.messages.create( # type: ignore + thread_id, + **message_data, # type: ignore + ) ) response_obj: Optional[OpenAIMessage] = None @@ -292,7 +295,8 @@ class AzureAssistantsAPI(BaseAzureLLM): ) thread_message: OpenAIMessage = openai_client.beta.threads.messages.create( # type: ignore - thread_id, **message_data # type: ignore + thread_id, + **message_data, # type: ignore ) response_obj: Optional[OpenAIMessage] = None diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 70b2f1ccc08..591fe9d03a2 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -79,7 +79,8 @@ class AzureAudioTranscription(AzureChatCompletion): ) response = azure_client.audio.transcriptions.create( - **data, timeout=timeout # type: ignore + **data, + timeout=timeout, # type: ignore ) if isinstance(response, BaseModel): @@ -95,7 +96,12 @@ class AzureAudioTranscription(AzureChatCompletion): original_response=stringified_response, ) hidden_params = {"model": model, "custom_llm_provider": "azure"} - final_response: TranscriptionResponse = convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, hidden_params=hidden_params, response_type="audio_transcription") # type: ignore + final_response: TranscriptionResponse = convert_to_model_response_object( + response_object=stringified_response, + model_response_object=model_response, + hidden_params=hidden_params, + response_type="audio_transcription", + ) # type: ignore return final_response async def async_audio_transcriptions( diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 5be3ce22832..8edab95ef6c 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -817,7 +817,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) ## COMPLETION CALL - raw_response = azure_client.embeddings.with_raw_response.create(**data, timeout=timeout) # type: ignore + raw_response = azure_client.embeddings.with_raw_response.create( + **data, timeout=timeout + ) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() if isinstance(response, str): @@ -833,7 +835,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): original_response=response, ) - return convert_to_model_response_object(response_object=response.model_dump(), model_response_object=model_response, response_type="embedding", _response_headers=process_azure_headers(headers)) # type: ignore + return convert_to_model_response_object( + response_object=response.model_dump(), + model_response_object=model_response, + response_type="embedding", + _response_headers=process_azure_headers(headers), + ) # type: ignore except AzureOpenAIError as e: raise e except Exception as e: @@ -1296,7 +1303,18 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): is_async=False, ) if aimg_generation is True: - return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers, model=model) # type: ignore + return self.aimage_generation( + data=data, + input=input, + logging_obj=logging_obj, + model_response=model_response, + api_key=api_key, + client=client, + azure_client_params=azure_client_params, + timeout=timeout, + headers=headers, + model=model, + ) # type: ignore img_gen_api_base = self.create_azure_base_url( azure_client_params=azure_client_params, @@ -1348,7 +1366,11 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): original_response=response, ) # return response - return convert_to_model_response_object(response_object=response, model_response_object=model_response, response_type="image_generation") # type: ignore + return convert_to_model_response_object( + response_object=response, + model_response_object=model_response, + response_type="image_generation", + ) # type: ignore except AzureOpenAIError as e: raise e except Exception as e: diff --git a/litellm/llms/azure/batches/handler.py b/litellm/llms/azure/batches/handler.py index 6da3670b34a..ea6722839e6 100644 --- a/litellm/llms/azure/batches/handler.py +++ b/litellm/llms/azure/batches/handler.py @@ -75,7 +75,9 @@ class AzureBatchesAPI(BaseAzureLLM): return self.acreate_batch( # type: ignore create_batch_data=create_batch_data, azure_client=azure_client ) - response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.create(**create_batch_data) # type: ignore[arg-type] + response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.create( + **create_batch_data + ) # type: ignore[arg-type] return LiteLLMBatch(**response.model_dump()) async def aretrieve_batch( diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index 72cbcba8a9a..cca83b8e6fd 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -45,7 +45,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> OpenAIFileObject: verbose_logger.debug("create_file_data=%s", create_file_data) - response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) # type: ignore[arg-type] + response = await openai_client.files.create( + **self._prepare_create_file_data(create_file_data) + ) # type: ignore[arg-type] verbose_logger.debug("create_file_response=%s", response) return OpenAIFileObject(**response.model_dump()) @@ -86,7 +88,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): return self.acreate_file( create_file_data=create_file_data, openai_client=openai_client ) - response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.create(**self._prepare_create_file_data(create_file_data)) # type: ignore[arg-type] + response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.create( + **self._prepare_create_file_data(create_file_data) + ) # type: ignore[arg-type] return OpenAIFileObject(**response.model_dump()) async def afile_content( diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index c3abfafc552..85016c7a5c4 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union import httpx from openai.types.file_deleted import FileDeleted @@ -32,6 +32,22 @@ else: Router = Any +class BaseFileUploadStream(ABC): + """Re-iterable request body that yields an upload's bytes lazily. + + A provider returns one of these (inside the upload config from + ``transform_create_file_request``) when the upload body can be produced + incrementally; the HTTP handler then sends it in bounded chunks instead of + buffering the whole payload, which is what exhausts memory on large uploads. + + ``iter_bytes`` must return a fresh iterator each call so the body can be + replayed if the upload is retried. + """ + + @abstractmethod + def iter_bytes(self) -> Iterator[bytes]: ... + + class BaseFilesConfig(BaseConfig): @property @abstractmethod diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index 6ad945f47a3..1c012a15fdb 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -8,10 +8,14 @@ run code -> delete container; `code_interpreter_tool` combines all three. from typing import Any, Union +import httpx + from pydantic import Field, PrivateAttr from litellm.types.llms.base import LiteLLMPydanticObjectBase +SANDBOX_MAX_OUTPUT_BYTES = 10 * 1024 * 1024 + class ContainerHandle(LiteLLMPydanticObjectBase): """A live sandbox container. Carries everything needed to reach it again.""" @@ -53,7 +57,7 @@ class BaseSandboxConfig: *, template: str | None = None, timeout: int | None = None, - allow_internet_access: bool = True, + allow_internet_access: bool | None = None, api_key: str | None = None, **kwargs, ) -> ContainerHandle: @@ -77,3 +81,16 @@ class BaseSandboxConfig: **kwargs, ) -> bool: raise NotImplementedError("adelete_sandbox must be implemented by provider") + + async def _read_capped_lines(self, response: httpx.Response) -> list[str]: + lines: list[str] = [] + total = 0 + async for line in response.aiter_lines(): + total += len(line.encode("utf-8")) + if total > SANDBOX_MAX_OUTPUT_BYTES: + raise ValueError( + f"Sandbox output exceeded {SANDBOX_MAX_OUTPUT_BYTES} bytes; aborting " + "to avoid unbounded memory use." + ) + lines.append(line) + return lines diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 4dfe86685fb..1581d8bb064 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -3,11 +3,13 @@ Base Search transformation configuration. """ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from urllib.parse import urlsplit import httpx from pydantic import PrivateAttr from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -16,6 +18,29 @@ else: LiteLLMLoggingObj = Any +def _search_host(url: str) -> str: + return urlsplit(url).netloc.lower() + + +def _is_trusted_search_api_base( + caller_api_base: str, + default_api_base: str | None, + base_env_var: str | None, +) -> bool: + candidate = _search_host(caller_api_base) + if not candidate: + return False + trusted = { + _search_host(base) + for base in ( + default_api_base, + get_secret_str(base_env_var) if base_env_var else None, + ) + if base + } + return candidate in trusted + + class SearchResult(LiteLLMPydanticObjectBase): """Single search result.""" @@ -86,6 +111,60 @@ class BaseSearchConfig: "max_tokens_per_page", } + def _assert_trusted_api_base_for_server_credential( + self, + caller_api_base: str | None, + default_api_base: str | None, + base_env_var: str | None, + credential_name: str, + ) -> None: + """ + Block sending a server-managed credential to a caller-chosen host. + + A caller-supplied api_base is honored when constructing the request URL, so + falling back to a server-configured secret while the caller controls the host + leaks that secret. The provider default and the operator's own api_base + override are the only trusted destinations for a server-managed credential. + """ + if not caller_api_base: + return + if _is_trusted_search_api_base(caller_api_base, default_api_base, base_env_var): + return + raise ValueError( + f"Refusing to send the server-configured {credential_name} to the " + f"caller-supplied api_base '{caller_api_base}'. Pass an explicit api_key " + f"when overriding api_base for this search provider." + ) + + def resolve_server_api_key( + self, + *, + caller_api_key: str | None, + caller_api_base: str | None, + key_env_vars: tuple[str, ...], + base_env_var: str | None, + default_api_base: str | None, + ) -> str | None: + """ + Resolve a single-secret search API key, falling back to a server-managed + secret only when the request targets a trusted host. + + Returns the caller's key when provided, otherwise the first set + server-managed secret (or None when none is set, for keyless providers). + """ + if caller_api_key: + return caller_api_key + server_key = next( + (key for key in (get_secret_str(var) for var in key_env_vars) if key), + None, + ) + if server_key is None: + return None + self._assert_trusted_api_base_for_server_credential( + caller_api_base, default_api_base, base_env_var, key_env_vars[0] + ) + return server_key + def validate_environment( self, headers: Dict, diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 3cd3a249c33..542af61dd3f 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -403,8 +403,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # skip strategy matching and fall back to raw JSON string if not isinstance(response_json, dict): verbose_logger.warning( - "AgentCore: JSON response is not a dict. " - "Returning raw JSON as content." + "AgentCore: JSON response is not a dict. Returning raw JSON as content." ) return AgentCoreParsedResponse( content=json.dumps(response_json), @@ -940,9 +939,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) parsed = self._parse_json_response(response_json) - async def _json_as_async_stream() -> ( - AsyncGenerator[ModelResponseStream, None] - ): + async def _json_as_async_stream() -> AsyncGenerator[ + ModelResponseStream, None + ]: # Content chunk content_chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7b1064ccef9..040193f6bcc 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -51,19 +51,19 @@ def make_sync_call( ) if fake_stream: - model_response: ( - ModelResponse - ) = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, + model_response: ModelResponse = ( + litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + encoding=litellm.encoding, + ) ) # type: ignore completion_stream: Any = MockResponseIterator( model_response=model_response, json_mode=json_mode diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index bb261ec85b2..a700c07d87a 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1963,7 +1963,9 @@ class AmazonConverseConfig(BaseConfig): return message, returned_finish_reason - def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tuple[ + def _translate_message_content( + self, content_blocks: List[ContentBlock] + ) -> Tuple[ str, List[ChatCompletionToolCallChunk], Optional[List[BedrockConverseReasoningContentBlock]], diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 75b560b4d6d..29e97068100 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -70,6 +70,7 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( BedrockError, ModelResponseIterator, + build_bedrock_stream_error, get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -224,19 +225,19 @@ async def make_call( raise BedrockError(status_code=response.status_code, message=response.text) if fake_stream: - model_response: ( - ModelResponse - ) = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, + model_response: ModelResponse = ( + litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + encoding=litellm.encoding, + ) ) # type: ignore completion_stream: Any = MockResponseIterator( model_response=model_response, json_mode=json_mode @@ -320,19 +321,19 @@ def make_sync_call( raise BedrockError(status_code=response.status_code, message=response.text) if fake_stream: - model_response: ( - ModelResponse - ) = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, + model_response: ModelResponse = ( + litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + encoding=litellm.encoding, + ) ) # type: ignore completion_stream: Any = MockResponseIterator( model_response=model_response, json_mode=json_mode @@ -1299,7 +1300,9 @@ class BedrockLLM(BaseAWSLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = get_async_httpx_client(params=_params, llm_provider=litellm.LlmProviders.BEDROCK) # type: ignore + client = get_async_httpx_client( + params=_params, llm_provider=litellm.LlmProviders.BEDROCK + ) # type: ignore else: client = client # type: ignore @@ -1841,23 +1844,7 @@ class AWSEventStreamDecoder: parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - decoded_body = response_dict["body"].decode() - if isinstance(decoded_body, dict): - error_message = decoded_body.get("message") - elif isinstance(decoded_body, str): - error_message = decoded_body - else: - error_message = "" - exception_status = response_dict["headers"].get(":exception-type") - error_message = exception_status + " " + error_message - raise BedrockError( - status_code=response_dict["status_code"], - message=( - json.dumps(error_message) - if isinstance(error_message, dict) - else error_message - ), - ) + raise build_bedrock_stream_error(response_dict, response_stream_shape) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index cbed2232be5..93306025b02 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, List, Optional from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.bedrock.common_utils import build_mantle_messages_url from litellm.types.llms.openai import AllMessageValues if TYPE_CHECKING: @@ -21,10 +22,6 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -MANTLE_ENDPOINT_TEMPLATE = ( - "https://bedrock-mantle.{region}.api.aws/anthropic/v1/messages" -) - class AmazonMantleConfig(AmazonAnthropicClaudeConfig): """ @@ -46,7 +43,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): stream: Optional[bool] = None, ) -> str: region = self._get_aws_region_name(optional_params=optional_params, model=model) - return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + return build_mantle_messages_url( + api_base=api_base, + aws_bedrock_runtime_endpoint=optional_params.get( + "aws_bedrock_runtime_endpoint" + ), + region=region, + ) def validate_environment( self, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index bdc5da321c6..5e97394f459 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -7,9 +7,21 @@ Common utilities used across bedrock chat/embedding/image generation import functools import json import os -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Mapping, + Optional, + TypedDict, + Union, +) if TYPE_CHECKING: + from botocore.model import Shape + from litellm.types.llms.bedrock import BedrockCreateBatchRequest import httpx @@ -610,6 +622,31 @@ def strip_bedrock_throughput_suffix(model: str) -> str: return model +MANTLE_MESSAGES_PATH = "/anthropic/v1/messages" + + +def build_mantle_messages_url( + api_base: Optional[str], + aws_bedrock_runtime_endpoint: Optional[str], + region: str, +) -> str: + """Build the bedrock-mantle Anthropic /messages URL. + + Honors an explicit endpoint override (``api_base``, then + ``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle + endpoints are reachable; otherwise falls back to the public regional host. + The mantle messages path is appended unless the override already carries it, + so callers can pass either the host or the full messages URL. + """ + override = api_base or aws_bedrock_runtime_endpoint + if override: + base = override.rstrip("/") + if base.endswith(MANTLE_MESSAGES_PATH): + return base + return f"{base}{MANTLE_MESSAGES_PATH}" + return f"https://bedrock-mantle.{region}.api.aws{MANTLE_MESSAGES_PATH}" + + def get_bedrock_base_model(model: str) -> str: """ Get the base model from the given model name. @@ -1132,6 +1169,39 @@ def get_bedrock_response_stream_shape(): return _load_bedrock_response_stream_shape() +class BedrockEventStreamResponseDict(TypedDict): + status_code: int + headers: Mapping[str, str] + body: bytes + + +def build_bedrock_stream_error( + response_dict: BedrockEventStreamResponseDict, + response_stream_shape: Shape | None, +) -> BedrockError: + """Build a BedrockError for a non-200 event-stream error event. + + botocore hard-codes HTTP 400 on every mid-stream error event, so the modeled + ResponseStream member's httpStatusCode is the real status. Resolve it from the + shape and fall back to the raw status when the type is not modeled. + """ + exception_type = response_dict["headers"].get(":exception-type") + decoded_body = response_dict["body"].decode() + message = f"{exception_type} {decoded_body}" if exception_type else decoded_body + + status_code = response_dict["status_code"] + if exception_type is not None and response_stream_shape is not None: + member = response_stream_shape.members.get(exception_type) + if member is not None: + modeled_status = ( + (member.metadata or {}).get("error", {}).get("httpStatusCode") + ) + if modeled_status is not None: + status_code = int(modeled_status) + + return BedrockError(status_code=status_code, message=message) + + class BedrockEventStreamDecoderBase: """ Base class for event stream decoding for Bedrock @@ -1156,23 +1226,7 @@ class BedrockEventStreamDecoderBase: parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - decoded_body = response_dict["body"].decode() - if isinstance(decoded_body, dict): - error_message = decoded_body.get("message") - elif isinstance(decoded_body, str): - error_message = decoded_body - else: - error_message = "" - exception_status = response_dict["headers"].get(":exception-type") - error_message = exception_status + " " + error_message - raise BedrockError( - status_code=response_dict["status_code"], - message=( - json.dumps(error_message) - if isinstance(error_message, dict) - else error_message - ), - ) + raise build_bedrock_stream_error(response_dict, response_stream_shape) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index b6aa99842d7..bc72f04deac 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -142,7 +142,9 @@ class BedrockEmbedding(BaseAWSLLM): client = client try: - response = await client.post(url=api_base, headers=headers, data=json.dumps(data)) # type: ignore + response = await client.post( + url=api_base, headers=headers, data=json.dumps(data) + ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -451,10 +453,10 @@ class BedrockEmbedding(BaseAWSLLM): batch_data = [] for i in input: if model == "amazon.titan-embed-image-v1": - transformed_request: ( - AmazonEmbeddingRequest - ) = AmazonTitanMultimodalEmbeddingG1Config()._transform_request( - input=i, inference_params=inference_params + transformed_request: AmazonEmbeddingRequest = ( + AmazonTitanMultimodalEmbeddingG1Config()._transform_request( + input=i, inference_params=inference_params + ) ) elif model == "amazon.titan-embed-text-v1": transformed_request = AmazonTitanG1Config()._transform_request( diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 90344310746..04fa5f803bb 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -112,7 +112,11 @@ class BedrockImageEdit(BaseAWSLLM): if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client() try: - response = client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body) # type: ignore + response = client.post( + url=prepared_request.endpoint_url, + headers=prepared_request.prepped.headers, + data=prepared_request.body, + ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -150,7 +154,11 @@ class BedrockImageEdit(BaseAWSLLM): ) try: - response = await async_client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body) # type: ignore + response = await async_client.post( + url=prepared_request.endpoint_url, + headers=prepared_request.prepped.headers, + data=prepared_request.body, + ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index d00d62a8530..e76b2885a88 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -122,7 +122,9 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): if k in param_mapping: # Map param if mapping exists and value is valid if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: - mapped_params[param_mapping[k]] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] # type: ignore + mapped_params[param_mapping[k]] = ( + OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] + ) # type: ignore # Don't copy "size" itself to final dict elif k == "n": # Store for logic but do not add to outgoing params diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index 87ef469beb5..9e06a8e747d 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -111,8 +111,10 @@ class AmazonNovaCanvasConfig: **color_guided_generation_params, } try: - color_guided_generation_params_typed = AmazonNovaCanvasColorGuidedGenerationParams( - **color_guided_generation_params # type: ignore + color_guided_generation_params_typed = ( + AmazonNovaCanvasColorGuidedGenerationParams( + **color_guided_generation_params # type: ignore + ) ) except Exception as e: raise ValueError( @@ -171,8 +173,9 @@ class AmazonNovaCanvasConfig: _size = non_default_params.get("size") if _size is not None: width, height = _size.split("x") - optional_params["width"], optional_params["height"] = int(width), int( - height + optional_params["width"], optional_params["height"] = ( + int(width), + int(height), ) if non_default_params.get("n") is not None: optional_params["numberOfImages"] = non_default_params.get("n") diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index d6053278cbd..0c594d0b142 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -115,7 +115,11 @@ class BedrockImageGeneration(BaseAWSLLM): if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client() try: - response = client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body) # type: ignore + response = client.post( + url=prepared_request.endpoint_url, + headers=prepared_request.prepped.headers, + data=prepared_request.body, + ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -154,7 +158,11 @@ class BedrockImageGeneration(BaseAWSLLM): ) try: - response = await async_client.post(url=prepared_request.endpoint_url, headers=prepared_request.prepped.headers, data=prepared_request.body) # type: ignore + response = await async_client.post( + url=prepared_request.endpoint_url, + headers=prepared_request.prepped.headers, + data=prepared_request.body, + ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index 900d9aa97d8..94e7f90b719 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -8,6 +8,7 @@ stripping that are specific to the bedrock-mantle endpoint. from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from litellm.llms.bedrock.common_utils import build_mantle_messages_url from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) @@ -20,10 +21,6 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -MANTLE_ENDPOINT_TEMPLATE = ( - "https://bedrock-mantle.{region}.api.aws/anthropic/v1/messages" -) - class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): """ @@ -43,7 +40,13 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): stream: Optional[bool] = None, ) -> str: region = self._get_aws_region_name(optional_params=optional_params, model=model) - return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + return build_mantle_messages_url( + api_base=api_base, + aws_bedrock_runtime_endpoint=optional_params.get( + "aws_bedrock_runtime_endpoint" + ), + region=region, + ) def validate_anthropic_messages_environment( self, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 0e2e06cf62c..5d22f4b3cd9 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -237,9 +237,7 @@ class BedrockRealtime(BaseAWSLLM): # Transform Bedrock format to OpenAI format from litellm.types.realtime import RealtimeResponseTransformInput - realtime_response_transform_input: ( - RealtimeResponseTransformInput - ) = { + realtime_response_transform_input: RealtimeResponseTransformInput = { "current_output_item_id": session_state.get( "current_output_item_id" ), diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 812ca116c27..276c6f23c33 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -96,7 +96,13 @@ class BedrockRerankHandler(BaseAWSLLM): ) if _is_async: - return self.arerank(prepared_request, timeout=timeout, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None) # type: ignore + return self.arerank( + prepared_request, + timeout=timeout, + client=client + if client is not None and isinstance(client, AsyncHTTPHandler) + else None, + ) # type: ignore if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client() diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index ec20d76102b..a479d148064 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -253,9 +253,9 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" ) - retrieval_config.setdefault("vectorSearchConfiguration", {})[ - "filter" - ] = filters + retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = ( + filters + ) if retrieval_config: request_body["retrievalConfiguration"] = cast( BedrockKBRetrievalConfiguration, retrieval_config diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py index 9dfcd6bc75a..8ffe7dcb126 100644 --- a/litellm/llms/brave/search/transformation.py +++ b/litellm/llms/brave/search/transformation.py @@ -115,7 +115,13 @@ class BraveSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("BRAVE_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("BRAVE_API_KEY",), + base_env_var="BRAVE_API_BASE", + default_api_base=self.BRAVE_API_BASE, + ) if not api_key: raise ValueError( diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 66e253f304d..68f08741cc5 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -1,26 +1,15 @@ -import json -import time -from typing import AsyncIterator, Iterator, List, Optional, Union +from typing import List, Optional, Union import httpx -import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segments -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator -from litellm.llms.base_llm.chat.transformation import ( - BaseConfig, - BaseLLMException, - LiteLLMLoggingObj, +from litellm._logging import verbose_logger +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.secret_managers.main import ( + get_secret_str, + normalize_nonempty_secret_str, ) -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ( - ChatCompletionToolCallChunk, - ChatCompletionUsageBlock, - GenericStreamingChunk, - ModelResponse, - Usage, -) class CloudflareError(BaseLLMException): @@ -34,26 +23,46 @@ class CloudflareError(BaseLLMException): message=message, request=self.request, response=self.response, - ) # Call the base class constructor with the parameters it needs + ) -class CloudflareChatConfig(BaseConfig): - max_tokens: Optional[int] = None - stream: Optional[bool] = None - - def __init__( +class CloudflareChatConfig(OpenAIGPTConfig): + def get_complete_url( self, - max_tokens: Optional[int] = None, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, stream: Optional[bool] = None, - ) -> None: - locals_ = locals().copy() - for key, value in locals_.items(): - if key != "self" and value is not None: - setattr(self.__class__, key, value) + ) -> str: + return super().get_complete_url( + api_base=self._resolve_api_base(api_base), + api_key=api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + stream=stream, + ) - @classmethod - def get_config(cls): - return super().get_config() + @staticmethod + def _resolve_api_base(api_base: Optional[str]) -> str: + if not api_base: + account_id = normalize_nonempty_secret_str( + get_secret_str("CLOUDFLARE_ACCOUNT_ID") + ) + if account_id is None: + raise ValueError( + "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID in the environment or pass api_base explicitly" + ) + return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1" + trimmed = api_base.rstrip("/") + if trimmed.endswith("/ai/run"): + verbose_logger.warning( + "Cloudflare api_base ending in '/ai/run' is the legacy Workers AI path and no longer serves OpenAI-compatible requests; rewriting to the '/ai/v1' endpoint" + ) + return f"{trimmed[: -len('/ai/run')]}/ai/v1" + return api_base def validate_environment( self, @@ -67,107 +76,18 @@ class CloudflareChatConfig(BaseConfig): ) -> dict: if api_key is None: raise ValueError( - "Missing CloudflareError API Key - A call is being made to cloudflare but no key is set either in the environment variables or via params" + "Missing Cloudflare API Key - A call is being made to cloudflare but no key is set either in the environment variables or via params" ) - headers = { - "accept": "application/json", - "content-type": "apbplication/json", - "Authorization": "Bearer " + api_key, - } - return headers - - def get_complete_url( - self, - api_base: Optional[str], - api_key: Optional[str], - model: str, - optional_params: dict, - litellm_params: dict, - stream: Optional[bool] = None, - ) -> str: - if api_base is None: - account_id = get_secret_str("CLOUDFLARE_ACCOUNT_ID") - api_base = ( - f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/" - ) - encoded_model = encode_url_path_segments(model, field_name="model") - return api_base + encoded_model - - def get_supported_openai_params(self, model: str) -> List[str]: - return [ - "stream", - "max_tokens", - ] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_openai_params = self.get_supported_openai_params(model=model) - for param, value in non_default_params.items(): - if param == "max_completion_tokens": - optional_params["max_tokens"] = value - elif param in supported_openai_params: - optional_params[param] = value - return optional_params - - def transform_request( - self, - model: str, - messages: List[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - config = litellm.CloudflareChatConfig.get_config() - for k, v in config.items(): - if k not in optional_params: - optional_params[k] = v - - data = { - "messages": messages, - **optional_params, - } - return data - - def transform_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ModelResponse, - logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: List[AllMessageValues], - optional_params: dict, - litellm_params: dict, - encoding: str, - api_key: Optional[str] = None, - json_mode: Optional[bool] = None, - ) -> ModelResponse: - completion_response = raw_response.json() - - # Support both "response" and "response_text" keys (newer models like Nemotron use "response_text") - result = completion_response["result"] - model_response.choices[0].message.content = result.get("response") if result.get("response") is not None else result.get("response_text", "") # type: ignore - - prompt_tokens = litellm.utils.get_token_count(messages=messages, model=model) - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) + return super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, ) - model_response.created = int(time.time()) - model_response.model = "cloudflare/" + model - usage = Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) - setattr(model_response, "usage", usage) - return model_response - def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: @@ -175,48 +95,3 @@ class CloudflareChatConfig(BaseConfig): status_code=status_code, message=error_message, ) - - def get_model_response_iterator( - self, - streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], - sync_stream: bool, - json_mode: Optional[bool] = False, - ): - return CloudflareChatResponseIterator( - streaming_response=streaming_response, - sync_stream=sync_stream, - json_mode=json_mode, - ) - - -class CloudflareChatResponseIterator(BaseModelResponseIterator): - def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: - try: - text = "" - tool_use: Optional[ChatCompletionToolCallChunk] = None - is_finished = False - finish_reason = "" - usage: Optional[ChatCompletionUsageBlock] = None - provider_specific_fields = None - - index = int(chunk.get("index", 0)) - - if "response" in chunk and chunk["response"] is not None: - text = chunk["response"] - elif "response_text" in chunk and chunk["response_text"] is not None: - text = chunk["response_text"] - - returned_chunk = GenericStreamingChunk( - text=text, - tool_use=tool_use, - is_finished=is_finished, - finish_reason=finish_reason, - usage=usage, - index=index, - provider_specific_fields=provider_specific_fields, - ) - - return returned_chunk - - except json.JSONDecodeError: - raise ValueError(f"Failed to decode JSON from chunk: {chunk}") diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 01c94476431..33601b47609 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -42,6 +42,9 @@ from litellm.constants import ( HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, ) from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -69,9 +72,9 @@ _AIOHTTP_SUPPORTS_SOCKET_FACTORY = ( ) -def _build_aiohttp_keepalive_socket_factory() -> ( - Optional[Callable[[Tuple[Any, ...]], socket.socket]] -): +def _build_aiohttp_keepalive_socket_factory() -> Optional[ + Callable[[Tuple[Any, ...]], socket.socket] +]: """ Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets. @@ -134,6 +137,18 @@ _DEFAULT_TIMEOUT = httpx.Timeout( timeout=COMPLETION_HTTP_FALLBACK_SECONDS, connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, ) + + +def _default_cached_client_timeout() -> httpx.Timeout: + """Timeout for cached default httpx clients; honors an explicit litellm.request_timeout.""" + configured = get_configured_request_timeout() + if configured is None: + return _DEFAULT_TIMEOUT + return httpx.Timeout( + timeout=configured, connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS + ) + + _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS = 5.0 _STREAMING_ERROR_BODY_READ_EXECUTOR = concurrent.futures.ThreadPoolExecutor( max_workers=50, @@ -704,7 +719,14 @@ class AsyncHTTPHandler: ) req = self.client.build_request( - "PUT", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore + "PUT", + url, + data=request_data, + json=json, + params=params, + headers=headers, + timeout=timeout, + content=request_content, # type: ignore ) response = await self.client.send(req) response.raise_for_status() @@ -765,7 +787,14 @@ class AsyncHTTPHandler: ) req = self.client.build_request( - "PATCH", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore + "PATCH", + url, + data=request_data, + json=json, + params=params, + headers=headers, + timeout=timeout, + content=request_content, # type: ignore ) response = await self.client.send(req) response.raise_for_status() @@ -826,7 +855,14 @@ class AsyncHTTPHandler: ) req = self.client.build_request( - "DELETE", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore + "DELETE", + url, + data=request_data, + json=json, + params=params, + headers=headers, + timeout=timeout, + content=request_content, # type: ignore ) response = await self.client.send(req, stream=stream) response.raise_for_status() @@ -873,7 +909,13 @@ class AsyncHTTPHandler: request_data, request_content = _prepare_request_data_and_content(data, content) req = client.build_request( - "POST", url, data=request_data, json=json, params=params, headers=headers, content=request_content # type: ignore + "POST", + url, + data=request_data, + json=json, + params=params, + headers=headers, + content=request_content, # type: ignore ) response = await client.send(req, stream=stream) response.raise_for_status() @@ -1185,7 +1227,14 @@ class HTTPHandler: ) else: req = self.client.build_request( - "POST", url, data=request_data, json=json, params=params, headers=headers, files=files, content=request_content # type: ignore + "POST", + url, + data=request_data, + json=json, + params=params, + headers=headers, + files=files, + content=request_content, # type: ignore ) response = self.client.send(req, stream=stream) response.raise_for_status() @@ -1220,11 +1269,24 @@ class HTTPHandler: if timeout is not None: req = self.client.build_request( - "PATCH", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore + "PATCH", + url, + data=request_data, + json=json, + params=params, + headers=headers, + timeout=timeout, + content=request_content, # type: ignore ) else: req = self.client.build_request( - "PATCH", url, data=request_data, json=json, params=params, headers=headers, content=request_content # type: ignore + "PATCH", + url, + data=request_data, + json=json, + params=params, + headers=headers, + content=request_content, # type: ignore ) response = self.client.send(req, stream=stream) response.raise_for_status() @@ -1259,11 +1321,24 @@ class HTTPHandler: if timeout is not None: req = self.client.build_request( - "PUT", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore + "PUT", + url, + data=request_data, + json=json, + params=params, + headers=headers, + timeout=timeout, + content=request_content, # type: ignore ) else: req = self.client.build_request( - "PUT", url, data=request_data, json=json, params=params, headers=headers, content=request_content # type: ignore + "PUT", + url, + data=request_data, + json=json, + params=params, + headers=headers, + content=request_content, # type: ignore ) response = self.client.send(req, stream=stream) return response @@ -1297,11 +1372,24 @@ class HTTPHandler: if timeout is not None: req = self.client.build_request( - "DELETE", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore + "DELETE", + url, + data=request_data, + json=json, + params=params, + headers=headers, + timeout=timeout, + content=request_content, # type: ignore ) else: req = self.client.build_request( - "DELETE", url, data=request_data, json=json, params=params, headers=headers, content=request_content # type: ignore + "DELETE", + url, + data=request_data, + json=json, + params=params, + headers=headers, + content=request_content, # type: ignore ) response = self.client.send(req, stream=stream) response.raise_for_status() @@ -1379,7 +1467,7 @@ def get_async_httpx_client( _new_client = AsyncHTTPHandler(**handler_params) else: _new_client = AsyncHTTPHandler( - timeout=_DEFAULT_TIMEOUT, + timeout=_default_cached_client_timeout(), shared_session=shared_session, ) @@ -1428,7 +1516,7 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: } _new_client = HTTPHandler(**handler_params) else: - _new_client = HTTPHandler(timeout=_DEFAULT_TIMEOUT) + _new_client = HTTPHandler(timeout=_default_cached_client_timeout()) cache.set_cache( key=_cache_key_name, diff --git a/litellm/llms/custom_httpx/httpx_handler.py b/litellm/llms/custom_httpx/httpx_handler.py index ce587946710..6e9f29151cd 100644 --- a/litellm/llms/custom_httpx/httpx_handler.py +++ b/litellm/llms/custom_httpx/httpx_handler.py @@ -54,7 +54,10 @@ class HTTPHandler: ): try: response = await self.client.post( - url, data=data, params=params, headers=headers # type: ignore + url, + data=data, + params=params, + headers=headers, # type: ignore ) return response except Exception as e: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 790bd0519d7..d33ec295e94 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,19 +1,23 @@ +import asyncio import json import ssl -from urllib.parse import parse_qs, urlencode, urlparse, urlunparse +from functools import lru_cache from typing import ( TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, + Iterator, List, Literal, Optional, Tuple, Union, cast, + get_type_hints, ) +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx # type: ignore from openai.types.file_deleted import FileDeleted @@ -25,6 +29,7 @@ import litellm.types.utils from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -101,8 +106,10 @@ from litellm.types.llms.openai import ( HttpxBinaryResponseContent, OpenAIFileObject, ResponseInputParam, + ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ) +from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult from litellm.types.router import GenericLiteLLMParams @@ -128,13 +135,13 @@ from litellm.types.vector_stores import ( VectorStoreSearchOptionalRequestParams, VectorStoreSearchResponse, ) -from litellm.types.realtime import RealtimeQueryParams from litellm.types.videos.main import VideoObject from litellm.utils import ( CustomStreamWrapper, ImageResponse, ModelResponse, ProviderConfigManager, + async_pre_call_deployment_hook, ) from .http_handler import get_shared_realtime_ssl_context @@ -184,6 +191,47 @@ def _google_genai_streaming_hidden_params( } +@lru_cache(maxsize=None) +def _responses_api_optional_request_param_names() -> frozenset[str]: + return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) + + +def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import ( + get_custom_logger_compatible_class, + ) + + dynamic_success_callbacks = getattr(logging_obj, "dynamic_success_callbacks", None) + callbacks = list(litellm.callbacks) + if isinstance(dynamic_success_callbacks, (list, tuple)): + callbacks.extend(dynamic_success_callbacks) + + custom_loggers: list[Any] = [] + for cb in callbacks: + if isinstance(cb, str): + resolved = get_custom_logger_compatible_class(cb) # type: ignore[arg-type] + if resolved is None: + continue + cb = resolved + if isinstance(cb, CustomLogger): + custom_loggers.append(cb) + return custom_loggers + + +def _has_pre_call_deployment_hook(logging_obj: Any) -> bool: + from litellm.integrations.custom_logger import CustomLogger + + base_func = CustomLogger.async_pre_call_deployment_hook + for cb in _custom_logger_callbacks(logging_obj): + cb_func = getattr(type(cb), "async_pre_call_deployment_hook", base_func) + if getattr(cb_func, "__func__", cb_func) is not getattr( + base_func, "__func__", base_func + ): + return True + return False + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -1833,6 +1881,9 @@ class BaseLLMHTTPHandler: data = provider_config.transform_search_request( query=query, optional_params=optional_params, + api_key=api_key, + api_base=api_base, + headers=headers or {}, ) # Get complete URL (pass data for providers that need request body for URL construction) @@ -2224,12 +2275,92 @@ class BaseLLMHTTPHandler: ) raise ValueError("anthropic_messages_handler is not implemented for sync calls") + def _run_sync_responses_pre_call_deployment_hook( + self, + *, + model: str, + input: Union[str, ResponseInputParam], + custom_llm_provider: str, + response_api_optional_request_params: dict[str, Any], + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + ) -> tuple[ + str, + Union[str, ResponseInputParam], + str, + dict[str, Any], + GenericLiteLLMParams, + ]: + if not _has_pre_call_deployment_hook(logging_obj): + return ( + model, + input, + custom_llm_provider, + response_api_optional_request_params, + litellm_params, + ) + + modified_kwargs = run_async_function( + async_pre_call_deployment_hook, + { + **dict(litellm_params), + **response_api_optional_request_params, + "model": model, + "input": input, + "custom_llm_provider": custom_llm_provider, + }, + CallTypes.responses.value, + ) + if modified_kwargs is None: + return ( + model, + input, + custom_llm_provider, + response_api_optional_request_params, + litellm_params, + ) + + optional_param_names = _responses_api_optional_request_param_names() + updated_response_params = { + **response_api_optional_request_params, + **{ + key: value + for key, value in modified_kwargs.items() + if key in optional_param_names + }, + } + updated_litellm_params = GenericLiteLLMParams( + **{ + **dict(litellm_params), + **{ + key: value + for key, value in modified_kwargs.items() + if key not in optional_param_names + and key not in {"model", "input", "custom_llm_provider"} + }, + } + ) + return ( + str(modified_kwargs["model"]) if "model" in modified_kwargs else model, + cast( + Union[str, ResponseInputParam], + modified_kwargs["input"] if "input" in modified_kwargs else input, + ), + ( + str(modified_kwargs["custom_llm_provider"]) + if "custom_llm_provider" in modified_kwargs + else custom_llm_provider + ), + updated_response_params, + updated_litellm_params, + ) + def response_api_handler( self, model: str, input: Union[str, ResponseInputParam], responses_api_provider_config: BaseResponsesAPIConfig, - response_api_optional_request_params: Dict, + response_api_optional_request_params: dict[str, Any], custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, @@ -2276,6 +2407,21 @@ class BaseLLMHTTPHandler: shared_session=shared_session, ) + ( + model, + input, + custom_llm_provider, + response_api_optional_request_params, + litellm_params, + ) = self._run_sync_responses_pre_call_deployment_hook( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + logging_obj=logging_obj, + ) + if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client( params={"ssl_verify": litellm_params.get("ssl_verify", None)} @@ -2414,9 +2560,27 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) ) - # Responses agentic interception (e.g. code interpreter) runs the follow-up - # loop via the async hook, so it is async-only for now; the sync path returns - # the initial response unchanged. + + if self._has_agentic_completion_hook(logging_obj): + final_response = run_async_function( + self._call_agentic_completion_hooks, + response=initial_response, + model=model, + messages=( + input + if isinstance(input, list) + else [{"role": "user", "content": input}] + ), + anthropic_messages_provider_config=responses_api_provider_config, + anthropic_messages_optional_request_params=response_api_optional_request_params, + logging_obj=logging_obj, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=dict(litellm_params), + api_surface="responses", + ) + return final_response if final_response is not None else initial_response + return initial_response async def async_response_api_handler( @@ -3276,6 +3440,23 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) + elif ( + isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + ): + try: + upload_response = self._resumable_chunked_upload( + client=sync_httpx_client, + initiate_url=api_base, + base_headers=headers, + config=cast(Dict[str, Any], transformed_request)[ + "resumable_chunked_upload" + ], + timeout=timeout, + ) + except Exception as e: + verbose_logger.exception(f"Error creating file: {e}") + raise self._handle_error(e=e, provider_config=provider_config) elif isinstance(transformed_request, str) or isinstance( transformed_request, bytes ): @@ -3357,7 +3538,15 @@ class BaseLLMHTTPHandler: input="", api_key="", additional_args={ - "complete_input_dict": transformed_request, + # A resumable upload config holds a reference to the (potentially + # huge) upload payload; logging deep-copies additional_args, so log + # a placeholder instead of re-materializing the payload. + "complete_input_dict": ( + "" + if isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + else transformed_request + ), "api_base": api_base, "headers": headers, }, @@ -3434,6 +3623,23 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) + elif ( + isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + ): + try: + upload_response = await self._aresumable_chunked_upload( + client=async_httpx_client, + initiate_url=api_base, + base_headers=headers, + config=cast(Dict[str, Any], transformed_request)[ + "resumable_chunked_upload" + ], + timeout=timeout, + ) + except Exception as e: + verbose_logger.exception(f"Error creating file: {e}") + raise self._handle_error(e=e, provider_config=provider_config) elif isinstance(transformed_request, str) or isinstance( transformed_request, bytes ): @@ -3477,6 +3683,224 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) + # 8 MiB; a 256 KiB multiple, which GCS requires for every non-final chunk. + _RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024 + + @staticmethod + def _iter_resumable_chunks( + byte_iter: Iterator[bytes], chunk_size: int + ) -> Iterator[bytes]: + """Regroup a byte stream into ``chunk_size`` pieces, yielding a final + partial piece only when it is non-empty. Every full piece is exactly + ``chunk_size`` bytes (kept a 256 KiB multiple for GCS) and never more than + one chunk is buffered. An exactly chunk-aligned stream yields only full + chunks, so the upload finalizes on its last data chunk instead of making + an extra empty request; a 0-byte stream yields nothing and the caller + finalizes with a single empty request. + """ + buf = bytearray() + for piece in byte_iter: + buf.extend(piece) + while len(buf) >= chunk_size: + yield bytes(buf[:chunk_size]) + del buf[:chunk_size] + if buf: + yield bytes(buf) + + @staticmethod + def _resumable_content_range(offset: int, data_len: int, is_final: bool) -> str: + if not is_final: + return f"bytes {offset}-{offset + data_len - 1}/*" + total = offset + data_len + if data_len == 0: + return f"bytes */{total}" + return f"bytes {offset}-{total - 1}/{total}" + + @staticmethod + def _resumable_request_kwargs( + headers: dict, + content: bytes, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> dict: + kwargs: Dict[str, Any] = {"headers": headers, "content": content} + if timeout is not None: + kwargs["timeout"] = timeout + return kwargs + + def _resumable_chunked_upload( + self, + *, + client: HTTPHandler, + initiate_url: str, + base_headers: dict, + config: dict, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + """Open a GCS resumable session, then PUT the body in bounded chunks so a + large upload is never held in memory in full.""" + stream = config["body_stream"] + chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) + session_url_header = config.get("session_url_header", "location") + httpx_client = client.client + + init_headers = {**base_headers, **config.get("initiate_headers", {})} + init_req = httpx_client.build_request( + "POST", + initiate_url, + **self._resumable_request_kwargs(init_headers, b"", timeout), + ) + init_resp = httpx_client.send(init_req, follow_redirects=False) + init_resp.read() + if init_resp.status_code not in (200, 201): + init_resp.raise_for_status() + session_url = init_resp.headers.get(session_url_header) + if not session_url: + raise ValueError( + f"resumable upload: no session URL in '{session_url_header}' header" + ) + + offset = 0 + pending: Optional[bytes] = None + for chunk in self._iter_resumable_chunks(stream.iter_bytes(), chunk_size): + if pending is not None: + self._send_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending, + offset, + is_final=False, + timeout=timeout, + ) + offset += len(pending) + pending = chunk + return self._send_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending or b"", + offset, + is_final=True, + timeout=timeout, + ) + + def _send_resumable_chunk( + self, + httpx_client: httpx.Client, + url: str, + base_headers: dict, + data: bytes, + offset: int, + *, + is_final: bool, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + headers = { + **base_headers, + "Content-Range": self._resumable_content_range(offset, len(data), is_final), + } + req = httpx_client.build_request( + "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) + ) + resp = httpx_client.send(req, follow_redirects=False) + resp.read() + if resp.status_code not in ((200, 201) if is_final else (308,)): + # 4xx/5xx raise here; the ValueError catches an unexpected success + # status (e.g. a 200 where the protocol expects a 308 between chunks). + resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected status {resp.status_code}") + return resp + + async def _aresumable_chunked_upload( + self, + *, + client: AsyncHTTPHandler, + initiate_url: str, + base_headers: dict, + config: dict, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + stream = config["body_stream"] + chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) + session_url_header = config.get("session_url_header", "location") + httpx_client = client.client + + init_headers = {**base_headers, **config.get("initiate_headers", {})} + init_req = httpx_client.build_request( + "POST", + initiate_url, + **self._resumable_request_kwargs(init_headers, b"", timeout), + ) + init_resp = await httpx_client.send(init_req, follow_redirects=False) + await init_resp.aread() + if init_resp.status_code not in (200, 201): + init_resp.raise_for_status() + session_url = init_resp.headers.get(session_url_header) + if not session_url: + raise ValueError( + f"resumable upload: no session URL in '{session_url_header}' header" + ) + + offset = 0 + pending: Optional[bytes] = None + # Producing each chunk runs the synchronous per-row transform for that + # chunk's worth of rows. Pull it off the event loop thread so a large + # upload does not block other concurrent requests between PUTs. + chunk_iter = self._iter_resumable_chunks(stream.iter_bytes(), chunk_size) + done = object() + while True: + chunk = await asyncio.to_thread(next, chunk_iter, done) + if chunk is done: + break + if pending is not None: + await self._asend_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending, + offset, + is_final=False, + timeout=timeout, + ) + offset += len(pending) + pending = chunk + return await self._asend_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending or b"", + offset, + is_final=True, + timeout=timeout, + ) + + async def _asend_resumable_chunk( + self, + httpx_client: httpx.AsyncClient, + url: str, + base_headers: dict, + data: bytes, + offset: int, + *, + is_final: bool, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + headers = { + **base_headers, + "Content-Range": self._resumable_content_range(offset, len(data), is_final), + } + req = httpx_client.build_request( + "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) + ) + resp = await httpx_client.send(req, follow_redirects=False) + await resp.aread() + if resp.status_code not in ((200, 201) if is_final else (308,)): + # 4xx/5xx raise here; the ValueError catches an unexpected success + # status (e.g. a 200 where the protocol expects a 308 between chunks). + resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected status {resp.status_code}") + return resp + def create_batch( self, create_batch_data: "CreateBatchRequest", @@ -4772,22 +5196,9 @@ class BaseLLMHTTPHandler: agentic callback is detected too. """ from litellm.integrations.custom_logger import CustomLogger - from litellm.litellm_core_utils.litellm_logging import ( - get_custom_logger_compatible_class, - ) base_func = CustomLogger.async_should_run_agentic_loop - callbacks = litellm.callbacks + ( - getattr(logging_obj, "dynamic_success_callbacks", None) or [] - ) - for cb in callbacks: - if isinstance(cb, str): - resolved = get_custom_logger_compatible_class(cb) # type: ignore[arg-type] - if resolved is None: - continue - cb = resolved - if not isinstance(cb, CustomLogger): - continue + for cb in _custom_logger_callbacks(logging_obj): cb_func = getattr(type(cb), "async_should_run_agentic_loop", base_func) if getattr(cb_func, "__func__", cb_func) is not getattr( base_func, "__func__", base_func @@ -5537,9 +5948,7 @@ class BaseLLMHTTPHandler: import websockets from websockets.asyncio.client import ClientConnection - url = self._append_query_params( - provider_config.get_complete_url(api_base, model, api_key), query_params - ) + url = provider_config.get_complete_url(api_base, model, api_key) headers = provider_config.validate_environment( headers=headers, model=model, diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index 27c10d740b5..701db586b72 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -61,9 +61,18 @@ class DataForSEOSearchConfig(BaseSearchConfig): password = get_secret_str("DATAFORSEO_PASSWORD") # If api_key is provided in "login:password" format, use it + caller_supplied_credentials = bool(api_key and ":" in api_key) if api_key and ":" in api_key: login, password = api_key.split(":", 1) + if not caller_supplied_credentials and login and password: + self._assert_trusted_api_base_for_server_credential( + api_base, + self.DATAFORSEO_API_BASE, + "DATAFORSEO_API_BASE", + "DATAFORSEO_LOGIN", + ) + if not login: raise ValueError( "DATAFORSEO_LOGIN is not set. Set `DATAFORSEO_LOGIN` environment variable or pass credentials in api_key parameter." diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index c279fab22ab..ecfc1642c97 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -16,6 +16,7 @@ from litellm.llms.base_llm.sandbox.transformation import ( BaseSandboxConfig, CodeExecutionResult, ContainerHandle, + SANDBOX_MAX_OUTPUT_BYTES, ) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -29,7 +30,7 @@ E2B_DEFAULT_TEMPLATE = "code-interpreter-v1" E2B_DEFAULT_DOMAIN = "e2b.app" JUPYTER_PORT = 49999 DEFAULT_SANDBOX_TIMEOUT = 300 -MAX_OUTPUT_BYTES = 10 * 1024 * 1024 +MAX_OUTPUT_BYTES = SANDBOX_MAX_OUTPUT_BYTES class E2BSandboxConfig(BaseSandboxConfig): @@ -49,7 +50,7 @@ class E2BSandboxConfig(BaseSandboxConfig): *, template: str | None = None, timeout: int | None = None, - allow_internet_access: bool = True, + allow_internet_access: bool | None = None, api_key: str | None = None, api_base: str | None = None, metadata: dict | None = None, @@ -62,7 +63,9 @@ class E2BSandboxConfig(BaseSandboxConfig): "templateID": template or E2B_DEFAULT_TEMPLATE, "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, "secure": True, - "allow_internet_access": allow_internet_access, + "allow_internet_access": ( + True if allow_internet_access is None else allow_internet_access + ), } if metadata: body["metadata"] = metadata @@ -168,20 +171,6 @@ class E2BSandboxConfig(BaseSandboxConfig): handle._hidden_params = {} return handle - @staticmethod - async def _read_capped_lines(response: httpx.Response) -> list[str]: - lines: list[str] = [] - total = 0 - async for line in response.aiter_lines(): - total += len(line.encode("utf-8")) - if total > MAX_OUTPUT_BYTES: - raise ValueError( - f"Sandbox output exceeded {MAX_OUTPUT_BYTES} bytes; aborting to " - "avoid unbounded memory use." - ) - lines.append(line) - return lines - @staticmethod def _parse_lines(lines: list[str]) -> CodeExecutionResult: def _try_parse(stripped: str): @@ -192,10 +181,9 @@ class E2BSandboxConfig(BaseSandboxConfig): messages = tuple( parsed - for stripped in (line.strip() for line in lines) - if stripped - for parsed in (_try_parse(stripped),) - if parsed is not None + for line in lines + if (stripped := line.strip()) + if (parsed := _try_parse(stripped)) is not None ) def of_type(message_type: str): diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 7a34ededa6b..5cfd14aeaa9 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -65,7 +65,13 @@ class ExaAISearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("EXA_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("EXA_API_KEY",), + base_env_var="EXA_API_BASE", + default_api_base=self.EXA_AI_API_BASE, + ) if not api_key: raise ValueError( "EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable." diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py index ce702266e7b..b571a659cac 100644 --- a/litellm/llms/fastcrw/search/transformation.py +++ b/litellm/llms/fastcrw/search/transformation.py @@ -57,7 +57,13 @@ class FastCRWSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("CRW_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("CRW_API_KEY",), + base_env_var="CRW_API_BASE", + default_api_base=self.FASTCRW_API_BASE, + ) if not api_key: raise ValueError( "CRW_API_KEY is not set. Set `CRW_API_KEY` environment variable." diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index 18cf1d28c4d..7e01ba58706 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -61,7 +61,13 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("FIRECRAWL_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("FIRECRAWL_API_KEY",), + base_env_var="FIRECRAWL_API_BASE", + default_api_base=self.FIRECRAWL_API_BASE, + ) if not api_key: raise ValueError( "FIRECRAWL_API_KEY is not set. Set `FIRECRAWL_API_KEY` environment variable." diff --git a/litellm/llms/fireworks_ai/audio_transcription/transformation.py b/litellm/llms/fireworks_ai/audio_transcription/transformation.py deleted file mode 100644 index 00bb5f26797..00000000000 --- a/litellm/llms/fireworks_ai/audio_transcription/transformation.py +++ /dev/null @@ -1,17 +0,0 @@ -from typing import List - -from litellm.types.llms.openai import OpenAIAudioTranscriptionOptionalParams - -from ...openai.transcriptions.whisper_transformation import ( - OpenAIWhisperAudioTranscriptionConfig, -) -from ..common_utils import FireworksAIMixin - - -class FireworksAIAudioTranscriptionConfig( - FireworksAIMixin, OpenAIWhisperAudioTranscriptionConfig -): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: - return ["language", "prompt", "response_format", "timestamp_granularities"] diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 7e4395959b9..3c93677d6d4 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -578,11 +578,11 @@ class FireworksAIConfig(OpenAIGPTConfig): ## FIREWORKS AI sends tool calls in the content field instead of tool_calls for choice in response.choices: - cast(Choices, choice).message = ( - self._handle_message_content_with_tool_calls( - message=cast(Choices, choice).message, - tool_calls=optional_params.get("tools", None), - ) + cast( + Choices, choice + ).message = self._handle_message_content_with_tool_calls( + message=cast(Choices, choice).message, + tool_calls=optional_params.get("tools", None), ) response._hidden_params = { diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 74f6cd4d831..a8cb40ac6db 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -103,6 +103,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # bypassing spend and budget accounting. self._pending_usage_metadata: Optional[dict] = None + def _include_function_response_id(self) -> bool: + """Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it.""" + return True + @staticmethod def _usage_detail_alias(details: Any, defaults: Dict[str, int]) -> Dict[str, Any]: if not isinstance(details, dict): @@ -604,10 +608,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) # Build Gemini toolResponse format - function_response = { - "id": call_id, - "response": output_dict, - } + function_response: dict[str, Any] = {"response": output_dict} + if self._include_function_response_id() and call_id: + function_response["id"] = call_id if function_name: function_response["name"] = function_name @@ -1752,9 +1755,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) tool_call_temperature = tool_call_generation_config.get("temperature") if tool_call_temperature is not None: - tool_call_done_event["response"][ - "temperature" - ] = tool_call_temperature + tool_call_done_event["response"]["temperature"] = ( + tool_call_temperature + ) tool_call_max_output_tokens = tool_call_generation_config.get( "maxOutputTokens" ) diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index 9de2987b9f6..3785f9c4657 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -182,7 +182,7 @@ class Authenticator: ) except httpx.HTTPStatusError as e: verbose_logger.error( - f"HTTP error refreshing API key (attempt {attempt+1}/{max_retries}): {str(e)}" + f"HTTP error refreshing API key (attempt {attempt + 1}/{max_retries}): {str(e)}" ) except Exception as e: verbose_logger.error(f"Unexpected error refreshing API key: {str(e)}") @@ -318,7 +318,7 @@ class Authenticator: and resp_json.get("error") == "authorization_pending" ): verbose_logger.debug( - f"Authorization pending (attempt {attempt+1}/{max_attempts})" + f"Authorization pending (attempt {attempt + 1}/{max_attempts})" ) else: verbose_logger.warning(f"Unexpected response: {resp_json}") diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index 299f346a7eb..d9f759e0a05 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -58,8 +58,7 @@ def github_copilot_supports_responses_api(model: str) -> bool: ) except Exception as e: verbose_logger.debug( - "github_copilot_supports_responses_api: get_model_info failed " - "for %s: %s", + "github_copilot_supports_responses_api: get_model_info failed for %s: %s", model, e, ) diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py index a8aa109cbf0..5cd3f2085a8 100644 --- a/litellm/llms/google_pse/search/transformation.py +++ b/litellm/llms/google_pse/search/transformation.py @@ -85,7 +85,13 @@ class GooglePSESearchConfig(BaseSearchConfig): Google PSE uses API key as a query parameter, not in headers. This method is called but headers are not used for authentication. """ - api_key = api_key or get_secret_str("GOOGLE_PSE_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("GOOGLE_PSE_API_KEY",), + base_env_var="GOOGLE_PSE_API_BASE", + default_api_base=self.GOOGLE_PSE_API_BASE, + ) if not api_key: raise ValueError( "GOOGLE_PSE_API_KEY is not set. Set `GOOGLE_PSE_API_KEY` environment variable." @@ -137,6 +143,7 @@ class GooglePSESearchConfig(BaseSearchConfig): query: Union[str, List[str]], optional_params: dict, api_key: Optional[str] = None, + api_base: str | None = None, search_engine_id: Optional[str] = None, **kwargs, ) -> Dict: @@ -165,8 +172,16 @@ class GooglePSESearchConfig(BaseSearchConfig): # Google PSE only supports single string queries query = " ".join(query) - # Get API credentials - api_key = api_key or get_secret_str("GOOGLE_PSE_API_KEY") + # Get API credentials. The key is sent as a query param to api_base, so + # resolve it host-aware to avoid leaking a server-managed key to a + # caller-supplied host. + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("GOOGLE_PSE_API_KEY",), + base_env_var="GOOGLE_PSE_API_BASE", + default_api_base=self.GOOGLE_PSE_API_BASE, + ) search_engine_id = search_engine_id or get_secret_str("GOOGLE_PSE_ENGINE_ID") if not api_key: diff --git a/litellm/llms/inception/chat/transformation.py b/litellm/llms/inception/chat/transformation.py index d591f783a99..ff87060449b 100644 --- a/litellm/llms/inception/chat/transformation.py +++ b/litellm/llms/inception/chat/transformation.py @@ -45,7 +45,11 @@ class InceptionChatConfig(OpenAILikeChatConfig): self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: passed_api_base = api_base - api_base = api_base or get_secret_str("INCEPTION_API_BASE") or "https://api.inceptionlabs.ai/v1" # type: ignore + api_base = ( + api_base + or get_secret_str("INCEPTION_API_BASE") + or "https://api.inceptionlabs.ai/v1" + ) # type: ignore dynamic_api_key = api_key if passed_api_base is None or api_key: dynamic_api_key = ( diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index 2b17d5642ac..d27ae038f9e 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -61,7 +61,13 @@ class LinkupSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("LINKUP_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("LINKUP_API_KEY",), + base_env_var="LINKUP_API_BASE", + default_api_base=self.LINKUP_API_BASE, + ) if not api_key: raise ValueError( "LINKUP_API_KEY is not set. Set `LINKUP_API_KEY` environment variable." diff --git a/litellm/llms/llamafile/chat/transformation.py b/litellm/llms/llamafile/chat/transformation.py index 3387a0eb6aa..223f5503c9c 100644 --- a/litellm/llms/llamafile/chat/transformation.py +++ b/litellm/llms/llamafile/chat/transformation.py @@ -27,7 +27,11 @@ class LlamafileChatConfig(OpenAIGPTConfig): If both are None, a default Llamafile server URL is returned. See: https://github.com/Mozilla-Ocho/llamafile/blob/bd1bbe9aabb1ee12dbdcafa8936db443c571eb9d/README.md#L61 """ - return api_base or get_secret_str("LLAMAFILE_API_BASE") or "http://127.0.0.1:8080/v1" # type: ignore + return ( + api_base + or get_secret_str("LLAMAFILE_API_BASE") + or "http://127.0.0.1:8080/v1" + ) # type: ignore def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index f1ad3708236..8d0cf993814 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -247,6 +247,8 @@ class MistralConfig(OpenAIGPTConfig): The above statement is not valid now. Need to plan to remove all the #1,2,3 Mistral API supports content as a list. """ + messages = [self._strip_output_only_fields(m) for m in messages] + ## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling for m in messages: _content_block = m.get("content") @@ -409,6 +411,25 @@ class MistralConfig(OpenAIGPTConfig): return cleaned_tools + @classmethod + def _strip_output_only_fields(cls, message: AllMessageValues) -> AllMessageValues: + """ + ``reasoning_content`` and ``thinking_blocks`` are output-only fields that + LiteLLM attaches to assistant responses. Mistral's input schema forbids + unknown fields, so replaying them verbatim in a follow-up turn triggers a + 422 ``extra_forbidden``. Drop them before the request is sent. + """ + if message["role"] != "assistant": + return message + return cast( + AllMessageValues, + { + k: v + for k, v in message.items() + if k not in ("reasoning_content", "thinking_blocks") + }, + ) + @classmethod def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues: """ diff --git a/litellm/llms/moonshot/chat/transformation.py b/litellm/llms/moonshot/chat/transformation.py index da8687bce72..587fa0ed8d6 100644 --- a/litellm/llms/moonshot/chat/transformation.py +++ b/litellm/llms/moonshot/chat/transformation.py @@ -64,7 +64,11 @@ class MoonshotChatConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" # type: ignore + api_base = ( + api_base + or get_secret_str("MOONSHOT_API_BASE") + or "https://api.moonshot.ai/v1" + ) # type: ignore dynamic_api_key = api_key or get_secret_str("MOONSHOT_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index fc317293acc..05f545c2944 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -234,12 +234,18 @@ class NvidiaNimRerankConfig(BaseRerankConfig): } # Add optional top_k parameter if provided (already mapped from top_n in map_cohere_rerank_params) - if "top_k" in optional_rerank_params and optional_rerank_params.get("top_k") is not None: # type: ignore + if ( + "top_k" in optional_rerank_params + and optional_rerank_params.get("top_k") is not None + ): # type: ignore request_data["top_k"] = optional_rerank_params.get("top_k") # type: ignore # Add Nvidia-specific truncate parameter if provided # This is passed through from non_default_params, not in base OptionalRerankParams - if "truncate" in optional_rerank_params and optional_rerank_params.get("truncate") is not None: # type: ignore + if ( + "truncate" in optional_rerank_params + and optional_rerank_params.get("truncate") is not None + ): # type: ignore truncate_value = optional_rerank_params.get("truncate") # type: ignore if truncate_value in ["NONE", "END"]: request_data["truncate"] = truncate_value # type: ignore diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index ac92fd22aa8..de8a09b7a3b 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -274,7 +274,9 @@ def handle_cohere_response( total_tokens=usage_info.totalTokens, ) else: - model_response.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) # type: ignore[attr-defined] + model_response.usage = Usage( + prompt_tokens=0, completion_tokens=0, total_tokens=0 + ) # type: ignore[attr-defined] return model_response diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index d1248b6e518..35d5aefeacb 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -529,7 +529,8 @@ class OCIChatConfig(BaseConfig): ) else: selected_params["tools"] = adapt_tool_definition_to_oci_standard( # type: ignore[assignment] - selected_params["tools"], vendor # type: ignore[arg-type] + selected_params["tools"], + vendor, # type: ignore[arg-type] ) # Normalise tool_choice to OCI's flat uppercase dict form diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index 8785b1548a5..29c88cbd50f 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -339,7 +339,9 @@ def sign_with_manual_credentials( private_key = ( load_private_key_from_str(oci_key_content) if oci_key_content - else load_private_key_from_file(oci_key_file) if oci_key_file else None + else load_private_key_from_file(oci_key_file) + if oci_key_file + else None ) if private_key is None: diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index e36150a4954..5c0624fe5c9 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -411,7 +411,9 @@ class OllamaChatConfig(BaseConfig): model_response.choices[0].finish_reason = "tool_calls" model_response.created = int(time.time()) model_response.model = "ollama_chat/" + model - prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=messages)) # type: ignore + prompt_tokens = response_json.get( + "prompt_eval_count", litellm.token_counter(messages=messages) + ) # type: ignore completion_tokens = response_json.get( "eval_count", litellm.token_counter(text=response_json["message"]["content"]), diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 7e34af43d43..a8cdbff87d0 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -337,7 +337,8 @@ class OllamaConfig(BaseConfig): model_response.model = "ollama/" + model _prompt = request_data.get("prompt", "") prompt_tokens = response_json.get( - "prompt_eval_count", len(encoding.encode(_prompt, disallowed_special=())) # type: ignore + "prompt_eval_count", + len(encoding.encode(_prompt, disallowed_special=())), # type: ignore ) completion_tokens = response_json.get( "eval_count", len(response_json.get("message", dict()).get("content", "")) diff --git a/litellm/llms/oobabooga/chat/transformation.py b/litellm/llms/oobabooga/chat/transformation.py index e87b70130ce..0c118efbc35 100644 --- a/litellm/llms/oobabooga/chat/transformation.py +++ b/litellm/llms/oobabooga/chat/transformation.py @@ -65,7 +65,9 @@ class OobaboogaConfig(OpenAIGPTConfig): ) else: try: - model_response.choices[0].message.content = completion_response["choices"][0]["message"]["content"] # type: ignore + model_response.choices[0].message.content = completion_response[ + "choices" + ][0]["message"]["content"] # type: ignore except Exception as e: raise OobaboogaError( message=str(e), diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index b8b750b8c12..e4d55404743 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -379,10 +379,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): List[OpenAIMessageContentListBlock], message_content ) for i, content_item in enumerate(message_content_types): - message_content_types[i] = ( - await self._async_transform_content_item( - cast(OpenAIMessageContentListBlock, content_item), - ) + message_content_types[ + i + ] = await self._async_transform_content_item( + cast(OpenAIMessageContentListBlock, content_item), ) return messages @@ -419,7 +419,8 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): for i, message in enumerate(messages): messages[i] = cast( - AllMessageValues, filter_value_from_dict(message, "cache_control") # type: ignore + AllMessageValues, + filter_value_from_dict(message, "cache_control"), # type: ignore ) if tools is not None: for i, tool in enumerate(tools): diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 8c9a8228daf..3b5980023c7 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -259,9 +259,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - messages[msg_idx]["content"][content_idx_optional][ - "text" - ] = guardrail_response + messages[msg_idx]["content"][content_idx_optional]["text"] = ( + guardrail_response + ) async def _apply_guardrail_responses_to_input_tool_calls( self, @@ -746,7 +746,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - choice.message.content[content_idx_optional]["text"] = guardrail_response # type: ignore + choice.message.content[content_idx_optional]["text"] = ( + guardrail_response # type: ignore + ) async def _apply_guardrail_responses_to_output_tool_calls( self, diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 63d39151254..f08ef844bb3 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -96,7 +96,19 @@ class OpenAITextCompletion(BaseLLM): organization=organization, ) else: - return self.acompletion(api_base=api_base, data=data, headers=headers, model_response=model_response, api_key=api_key, logging_obj=logging_obj, model=model, timeout=timeout, max_retries=max_retries, organization=organization, client=client) # type: ignore + return self.acompletion( + api_base=api_base, + data=data, + headers=headers, + model_response=model_response, + api_key=api_key, + logging_obj=logging_obj, + model=model, + timeout=timeout, + max_retries=max_retries, + organization=organization, + client=client, + ) # type: ignore elif optional_params.get("stream", False): return self.streaming( logging_obj=logging_obj, @@ -124,7 +136,9 @@ class OpenAITextCompletion(BaseLLM): else: openai_client = client - raw_response = openai_client.completions.with_raw_response.create(**data) # type: ignore + raw_response = openai_client.completions.with_raw_response.create( + **data + ) # type: ignore response = raw_response.parse() response_json = response.model_dump() diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 8b96fb6ef7a..dae3fa9d457 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -73,7 +73,9 @@ class OpenAIImageVariationsHandler: client=client, init_client_params=init_client_params ) - raw_response = await client.images.with_raw_response.create_variation(**data) # type: ignore + raw_response = await client.images.with_raw_response.create_variation( + **data + ) # type: ignore response = raw_response.parse() response_json = response.model_dump() diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index ea905d8ebca..59acbac6e15 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1132,9 +1132,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data = drop_params_from_unprocessable_entity_error(e, data) else: raise e - except ( - Exception - ) as e: # need to exception handle here. async exceptions don't get caught in sync functions. + except Exception as e: # need to exception handle here. async exceptions don't get caught in sync functions. if isinstance(e, OpenAIError): raise e @@ -1433,7 +1431,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): additional_args={"complete_input_dict": data}, original_response=stringified_response, ) - return convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, response_type="image_generation") # type: ignore + return convert_to_model_response_object( + response_object=stringified_response, + model_response_object=model_response, + response_type="image_generation", + ) # type: ignore except Exception as e: ## LOGGING logging_obj.post_call( @@ -1466,7 +1468,19 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError(status_code=422, message="max retries must be an int") if aimg_generation is True: - return self.aimage_generation(data=data, prompt=prompt, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries, organization=organization, headers=headers) # type: ignore + return self.aimage_generation( + data=data, + prompt=prompt, + logging_obj=logging_obj, + model_response=model_response, + api_base=api_base, + api_key=api_key, + timeout=timeout, + client=client, + max_retries=max_retries, + organization=organization, + headers=headers, + ) # type: ignore openai_client: OpenAI = self._get_openai_client( # type: ignore is_async=False, @@ -1503,7 +1517,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): additional_args={"complete_input_dict": data}, original_response=response, ) - return convert_to_model_response_object(response_object=response, model_response_object=model_response, response_type="image_generation") # type: ignore + return convert_to_model_response_object( + response_object=response, + model_response_object=model_response, + response_type="image_generation", + ) # type: ignore except OpenAIError as e: ## LOGGING logging_obj.post_call( @@ -2517,8 +2535,11 @@ class OpenAIAssistantsAPI(BaseLLM): client=client, ) - thread_message: OpenAIMessage = await openai_client.beta.threads.messages.create( # type: ignore - thread_id, **message_data # type: ignore + thread_message: OpenAIMessage = ( + await openai_client.beta.threads.messages.create( # type: ignore + thread_id, + **message_data, # type: ignore + ) ) response_obj: Optional[OpenAIMessage] = None @@ -2596,7 +2617,8 @@ class OpenAIAssistantsAPI(BaseLLM): ) thread_message: OpenAIMessage = openai_client.beta.threads.messages.create( # type: ignore - thread_id, **message_data # type: ignore + thread_id, + **message_data, # type: ignore ) response_obj: Optional[OpenAIMessage] = None diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b5319797cc6..83ec69a9d7a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -350,9 +350,9 @@ class OpenAIResponsesHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content if isinstance(messages[msg_idx]["content"][content_idx_optional], dict): - messages[msg_idx]["content"][content_idx_optional][ - "text" - ] = guardrail_response + messages[msg_idx]["content"][content_idx_optional]["text"] = ( + guardrail_response + ) async def process_output_response( self, diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index e079a170874..44fb5da8590 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -71,7 +71,9 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): response = raw_response.parse() return headers, response else: - response = openai_client.audio.transcriptions.create(**data, timeout=timeout) # type: ignore + response = openai_client.audio.transcriptions.create( + **data, timeout=timeout + ) # type: ignore return None, response except Exception as e: raise e @@ -160,7 +162,12 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): original_response=stringified_response, ) hidden_params = {"model": model, "custom_llm_provider": "openai"} - final_response: TranscriptionResponse = convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, hidden_params=hidden_params, response_type="audio_transcription") # type: ignore + final_response: TranscriptionResponse = convert_to_model_response_object( + response_object=stringified_response, + model_response_object=model_response, + hidden_params=hidden_params, + response_type="audio_transcription", + ) # type: ignore return final_response async def async_audio_transcriptions( @@ -220,7 +227,12 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): actual_model = data.get("model", "whisper-1") hidden_params = {"model": actual_model, "custom_llm_provider": "openai"} - return convert_to_model_response_object(response_object=stringified_response, model_response_object=model_response, hidden_params=hidden_params, response_type="audio_transcription") # type: ignore + return convert_to_model_response_object( + response_object=stringified_response, + model_response_object=model_response, + hidden_params=hidden_params, + response_type="audio_transcription", + ) # type: ignore except Exception as e: ## LOGGING logging_obj.post_call( diff --git a/litellm/llms/openai_like/embedding/handler.py b/litellm/llms/openai_like/embedding/handler.py index e3884fa56d7..285595d2791 100644 --- a/litellm/llms/openai_like/embedding/handler.py +++ b/litellm/llms/openai_like/embedding/handler.py @@ -118,7 +118,17 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): ) if aembedding is True: - return self.aembedding(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, headers=headers) # type: ignore + return self.aembedding( + data=data, + input=input, + logging_obj=logging_obj, + model_response=model_response, + api_base=api_base, + api_key=api_key, + timeout=timeout, + client=client, + headers=headers, + ) # type: ignore if client is None or isinstance(client, AsyncHTTPHandler): self.client = HTTPHandler(timeout=timeout) # type: ignore else: diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index dba8d6eb4b6..615d693fb69 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -115,6 +115,14 @@ "max_completion_tokens": "max_tokens" } }, + "darkbloom": { + "base_url": "https://api.darkbloom.dev/v1", + "api_key_env": "DARKBLOOM_API_KEY", + "api_base_env": "DARKBLOOM_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + } + }, "neosantara": { "base_url": "https://api.neosantara.xyz/v1", "api_key_env": "NEOSANTARA_API_KEY", diff --git a/litellm/llms/opensandbox/__init__.py b/litellm/llms/opensandbox/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/opensandbox/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/opensandbox/sandbox/__init__.py b/litellm/llms/opensandbox/sandbox/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/opensandbox/sandbox/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py new file mode 100644 index 00000000000..dc9f8440d30 --- /dev/null +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -0,0 +1,598 @@ +import asyncio +import json +import time +from typing import Union, cast + +import httpx + +from litellm.constants import ( + OPEN_SANDBOX_API_BASE_ENV_VAR, + OPEN_SANDBOX_API_KEY_ENV_VAR, + OPEN_SANDBOX_DEFAULT_CPU_LIMIT, + OPEN_SANDBOX_DEFAULT_ENTRYPOINT, + OPEN_SANDBOX_DEFAULT_LANGUAGE, + OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OPEN_SANDBOX_DEFAULT_TIMEOUT, + OPEN_SANDBOX_EXECD_PORT, + OPEN_SANDBOX_POLL_INTERVAL, + OPEN_SANDBOX_READY_TIMEOUT, +) +from litellm.llms.base_llm.sandbox.transformation import ( + BaseSandboxConfig, + CodeExecutionResult, + ContainerHandle, + SANDBOX_MAX_OUTPUT_BYTES, +) +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.custom_http import httpxSpecialProvider + +DEFAULT_SANDBOX_TIMEOUT = OPEN_SANDBOX_DEFAULT_TIMEOUT +DEFAULT_READY_TIMEOUT = OPEN_SANDBOX_READY_TIMEOUT +DEFAULT_POLL_INTERVAL = OPEN_SANDBOX_POLL_INTERVAL +MAX_OUTPUT_BYTES = SANDBOX_MAX_OUTPUT_BYTES + + +class OpenSandboxSandboxConfig(BaseSandboxConfig): + def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler: + if client is not None: + return client + return get_async_httpx_client(llm_provider=httpxSpecialProvider.Sandbox) + + def validate_environment(self, api_key: str | None = None, **kwargs) -> str: + if api_key is not None: + return api_key + return get_secret_str(OPEN_SANDBOX_API_KEY_ENV_VAR) or "" + + async def acreate_sandbox( + self, + *, + template: str | None = None, + timeout: int | None = None, + allow_internet_access: bool | None = None, + api_key: str | None = None, + api_base: str | None = None, + metadata: dict[str, str] | None = None, + env_vars: dict[str, str] | None = None, + resource_limits: dict[str, str] | None = None, + resource_requests: dict[str, str] | None = None, + entrypoint: list[str] | tuple[str, ...] | None = None, + network_policy: dict[str, object] | None = None, + secure_access: bool = False, + use_server_proxy: bool = False, + ready_timeout: float | None = None, + poll_interval: float | None = None, + client: AsyncHTTPHandler | None = None, + **kwargs, + ) -> ContainerHandle: + key = self.validate_environment(api_key=api_key) + base = self._api_base(api_base) + ready_timeout_seconds = ( + float(ready_timeout) if ready_timeout is not None else DEFAULT_READY_TIMEOUT + ) + poll_interval_seconds = ( + float(poll_interval) if poll_interval is not None else DEFAULT_POLL_INTERVAL + ) + body = self._create_body( + template=template, + timeout=timeout, + allow_internet_access=allow_internet_access, + metadata=metadata, + env_vars=env_vars, + resource_limits=resource_limits, + resource_requests=resource_requests, + entrypoint=entrypoint, + network_policy=network_policy, + secure_access=secure_access, + ) + + response = cast( + httpx.Response, + await self._http(client).post( + url=f"{base}/sandboxes", + headers=self._lifecycle_headers(key), + json=body, + ), + ) + data = response.json() + sandbox_id = str(data["id"]) + + if self._sandbox_state(data) != "Running": + await self._wait_until_running( + sandbox_id=sandbox_id, + api_base=base, + headers=self._lifecycle_headers(key), + client=client, + ready_timeout=ready_timeout_seconds, + poll_interval=poll_interval_seconds, + ) + + endpoint, endpoint_headers = await self._wait_for_execd_endpoint( + sandbox_id=sandbox_id, + api_base=base, + headers=self._lifecycle_headers(key), + use_server_proxy=use_server_proxy, + client=client, + ready_timeout=ready_timeout_seconds, + poll_interval=poll_interval_seconds, + ) + + handle = ContainerHandle(id=sandbox_id, provider="opensandbox", domain=base) + handle._hidden_params = { + "api_base": base, + "api_key": key, + "execd_endpoint": endpoint, + "execd_headers": endpoint_headers, + "use_server_proxy": use_server_proxy, + } + return handle + + async def arun_code( + self, + *, + container: Union[ContainerHandle, str], + code: str, + api_key: str | None = None, + api_base: str | None = None, + language: str = OPEN_SANDBOX_DEFAULT_LANGUAGE, + use_server_proxy: bool = False, + ready_timeout: float | None = None, + poll_interval: float | None = None, + client: AsyncHTTPHandler | None = None, + **kwargs, + ) -> CodeExecutionResult: + handle = await self._ensure_handle( + container=container, + api_key=api_key, + api_base=api_base, + use_server_proxy=use_server_proxy, + ready_timeout=( + float(ready_timeout) + if ready_timeout is not None + else DEFAULT_READY_TIMEOUT + ), + poll_interval=( + float(poll_interval) + if poll_interval is not None + else DEFAULT_POLL_INTERVAL + ), + client=client, + ) + endpoint = str(handle._hidden_params["execd_endpoint"]) + endpoint_headers = self._as_str_dict(handle._hidden_params.get("execd_headers")) + base = str( + handle._hidden_params.get("api_base") + or handle.domain + or self._api_base(api_base) + ) + lines = await self._post_code( + url=f"{self._endpoint_base_url(endpoint, base)}/code", + headers={ + "Content-Type": "application/json", + "Accept": "text/event-stream", + "Cache-Control": "no-cache", + **endpoint_headers, + }, + body={ + "code": code, + "context": {"language": language}, + }, + client=client, + ) + return self._parse_lines(lines) + + async def adelete_sandbox( + self, + *, + container: Union[ContainerHandle, str], + api_key: str | None = None, + api_base: str | None = None, + client: AsyncHTTPHandler | None = None, + **kwargs, + ) -> bool: + handle = self._as_handle(container, api_base=api_base) + base = str(handle._hidden_params.get("api_base") or self._api_base(api_base)) + key = self._api_key(api_key=api_key, handle=handle) + try: + response = cast( + httpx.Response, + await self._http(client).delete( + url=f"{base}/sandboxes/{handle.id}", + headers=self._lifecycle_headers(key), + ), + ) + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + return False + raise + return 200 <= response.status_code < 300 + + async def _ensure_handle( + self, + *, + container: Union[ContainerHandle, str], + api_key: str | None, + api_base: str | None, + use_server_proxy: bool, + ready_timeout: float, + poll_interval: float, + client: AsyncHTTPHandler | None, + ) -> ContainerHandle: + handle = self._as_handle(container, api_base=api_base) + if handle._hidden_params.get("execd_endpoint"): + return handle + + base = str(handle._hidden_params.get("api_base") or self._api_base(api_base)) + key = self._api_key(api_key=api_key, handle=handle) + resolved_use_server_proxy = bool( + handle._hidden_params.get("use_server_proxy", use_server_proxy) + ) + endpoint, endpoint_headers = await self._wait_for_execd_endpoint( + sandbox_id=handle.id, + api_base=base, + headers=self._lifecycle_headers(key), + use_server_proxy=resolved_use_server_proxy, + client=client, + ready_timeout=ready_timeout, + poll_interval=poll_interval, + ) + handle.domain = base + handle._hidden_params = { + **handle._hidden_params, + "api_base": base, + "api_key": key, + "execd_endpoint": endpoint, + "execd_headers": endpoint_headers, + "use_server_proxy": resolved_use_server_proxy, + } + return handle + + async def _wait_until_running( + self, + *, + sandbox_id: str, + api_base: str, + headers: dict[str, str], + client: AsyncHTTPHandler | None, + ready_timeout: float, + poll_interval: float, + ) -> None: + deadline = time.monotonic() + ready_timeout + while True: + response = cast( + httpx.Response, + await self._http(client).get( + url=f"{api_base}/sandboxes/{sandbox_id}", + headers=headers, + ), + ) + data = response.json() + state = self._sandbox_state(data) + if state == "Running": + return + if state in {"Failed", "Stopping", "Terminated"}: + raise ValueError(f"OpenSandbox sandbox {sandbox_id} entered {state}") + if time.monotonic() >= deadline: + raise TimeoutError( + f"OpenSandbox sandbox {sandbox_id} was not Running within " + f"{ready_timeout} seconds" + ) + await asyncio.sleep(poll_interval) + + async def _wait_for_execd_endpoint( + self, + *, + sandbox_id: str, + api_base: str, + headers: dict[str, str], + use_server_proxy: bool, + client: AsyncHTTPHandler | None, + ready_timeout: float, + poll_interval: float, + ) -> tuple[str, dict[str, str]]: + deadline = time.monotonic() + ready_timeout + last_error: Exception | None = None + while True: + try: + return await self._get_execd_endpoint( + sandbox_id=sandbox_id, + api_base=api_base, + headers=headers, + use_server_proxy=use_server_proxy, + client=client, + ) + except httpx.HTTPStatusError as e: + if e.response.status_code != 404: + raise + last_error = e + except ValueError as e: + last_error = e + + if time.monotonic() >= deadline: + raise TimeoutError( + f"OpenSandbox execd endpoint for {sandbox_id} was not ready within " + f"{ready_timeout} seconds" + ) from last_error + await asyncio.sleep(poll_interval) + + async def _get_execd_endpoint( + self, + *, + sandbox_id: str, + api_base: str, + headers: dict[str, str], + use_server_proxy: bool, + client: AsyncHTTPHandler | None, + ) -> tuple[str, dict[str, str]]: + response = cast( + httpx.Response, + await self._http(client).get( + url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}", + headers=headers, + params={"use_server_proxy": use_server_proxy}, + ), + ) + data = response.json() + endpoint = data.get("endpoint") + if not endpoint: + raise ValueError( + f"OpenSandbox did not return an execd endpoint for {sandbox_id}" + ) + return str(endpoint), self._as_str_dict(data.get("headers")) + + async def _post_code( + self, + *, + url: str, + headers: dict[str, str], + body: dict[str, object], + client: AsyncHTTPHandler | None, + ) -> list[str]: + timeout = httpx.Timeout(connect=30.0, read=None, write=30.0, pool=None) + response = cast( + httpx.Response, + await self._http(client).post( + url=url, + headers=headers, + timeout=timeout, + json=body, + stream=True, + ), + ) + return await self._read_capped_lines(response) + + def _api_key(self, *, api_key: str | None, handle: ContainerHandle) -> str: + if api_key is not None: + return api_key + if "api_key" in handle._hidden_params: + return str(handle._hidden_params["api_key"]) + return self.validate_environment() + + @staticmethod + def _create_body( + *, + template: str | None, + timeout: int | None, + allow_internet_access: bool | None, + metadata: dict[str, str] | None, + env_vars: dict[str, str] | None, + resource_limits: dict[str, str] | None, + resource_requests: dict[str, str] | None, + entrypoint: list[str] | tuple[str, ...] | None, + network_policy: dict[str, object] | None, + secure_access: bool, + ) -> dict[str, object]: + body: dict[str, object] = { + "image": {"uri": template or OPEN_SANDBOX_DEFAULT_TEMPLATE}, + "entrypoint": list(entrypoint or OPEN_SANDBOX_DEFAULT_ENTRYPOINT), + "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, + "resourceLimits": resource_limits + or OpenSandboxSandboxConfig._default_resource_limits(), + } + if metadata: + body["metadata"] = metadata + if env_vars: + body["env"] = env_vars + if resource_requests: + body["resourceRequests"] = resource_requests + if network_policy is not None: + body["networkPolicy"] = network_policy + elif allow_internet_access is not True: + body["networkPolicy"] = {"defaultAction": "deny", "egress": []} + if secure_access: + body["secureAccess"] = True + return body + + @staticmethod + def _default_resource_limits() -> dict[str, str]: + return { + "cpu": OPEN_SANDBOX_DEFAULT_CPU_LIMIT, + "memory": OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT, + } + + @staticmethod + def _sandbox_state(data: object) -> str | None: + if not isinstance(data, dict): + return None + status = data.get("status") + if not isinstance(status, dict): + return None + state = status.get("state") + return str(state) if state is not None else None + + @staticmethod + def _as_str_dict(value: object) -> dict[str, str]: + if not isinstance(value, dict): + return {} + return {str(k): str(v) for k, v in value.items()} + + @staticmethod + def _api_base(api_base: str | None) -> str: + base = api_base or get_secret_str(OPEN_SANDBOX_API_BASE_ENV_VAR) + if not base: + raise ValueError( + "OpenSandbox api_base is required. Pass api_base or set " + f"{OPEN_SANDBOX_API_BASE_ENV_VAR}." + ) + return str(base).rstrip("/") + + @staticmethod + def _lifecycle_headers(api_key: str) -> dict[str, str]: + headers = {"Content-Type": "application/json"} + if api_key: + headers["OPEN-SANDBOX-API-KEY"] = api_key + return headers + + @staticmethod + def _endpoint_base_url(endpoint: str, api_base: str) -> str: + normalized_endpoint = endpoint.rstrip("/") + if normalized_endpoint.startswith(("http://", "https://")): + return normalized_endpoint + protocol = api_base.split("://", 1)[0] if "://" in api_base else "http" + return f"{protocol}://{normalized_endpoint}" + + @staticmethod + def _as_handle( + container: Union[ContainerHandle, str], *, api_base: str | None + ) -> ContainerHandle: + if isinstance(container, ContainerHandle): + return container + handle = ContainerHandle( + id=str(container), + provider="opensandbox", + domain=OpenSandboxSandboxConfig._api_base(api_base), + ) + handle._hidden_params = {} + return handle + + @staticmethod + def _parse_lines(lines: list[str]) -> CodeExecutionResult: + messages = tuple( + event + for line in lines + if (event := OpenSandboxSandboxConfig._parse_sse_line(line)) is not None + ) + + def of_type(message_type: str): + return (m for m in messages if m.get("type") == message_type) + + error = next( + (OpenSandboxSandboxConfig._normalize_error(m) for m in of_type("error")), + None, + ) + execution_count = next( + ( + OpenSandboxSandboxConfig._as_int(m.get("execution_count")) + for m in of_type("execution_count") + if OpenSandboxSandboxConfig._as_int(m.get("execution_count")) + is not None + ), + None, + ) + + return CodeExecutionResult( + stdout="".join(str(m.get("text", "")) for m in of_type("stdout")), + stderr="".join(str(m.get("text", "")) for m in of_type("stderr")), + results=[ + OpenSandboxSandboxConfig._normalize_result(m) for m in of_type("result") + ], + error=error, + execution_count=execution_count, + ) + + @staticmethod + def _parse_sse_line(line: str) -> dict[str, object] | None: + stripped = line.strip() + if not stripped or stripped.startswith( + ( + ":", + "event:", + "id:", + "retry:", + ) + ): + return None + data = stripped[5:].strip() if stripped.startswith("data:") else stripped + if not data: + return None + try: + parsed = json.loads(data) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + if "type" not in parsed and "code" in parsed and "message" in parsed: + return { + "type": "error", + "error": { + "ename": str(parsed["code"]), + "evalue": str(parsed["message"]), + "traceback": [], + }, + } + return parsed + + @staticmethod + def _normalize_result(message: dict[str, object]) -> dict[str, object]: + results = message.get("results") + if isinstance(results, dict): + return {str(k): v for k, v in results.items()} + return { + str(k): v + for k, v in message.items() + if k not in {"type", "timestamp", "execution_count"} + } + + @staticmethod + def _normalize_error(message: dict[str, object]) -> dict[str, object]: + raw_error = message.get("error") + if isinstance(raw_error, dict): + name = OpenSandboxSandboxConfig._first_non_none_value( + raw_error, "ename", "name", default="" + ) + value = OpenSandboxSandboxConfig._first_non_none_value( + raw_error, "evalue", "value", default="" + ) + traceback = OpenSandboxSandboxConfig._first_non_none_value( + raw_error, "traceback", default=[] + ) + return { + "name": name, + "value": value, + "traceback": traceback, + } + return { + "name": OpenSandboxSandboxConfig._first_non_none_value( + message, "name", default="" + ), + "value": OpenSandboxSandboxConfig._first_non_none_value( + message, "value", "text", default="" + ), + "traceback": OpenSandboxSandboxConfig._first_non_none_value( + message, "traceback", default=[] + ), + } + + @staticmethod + def _as_int(value: object) -> int | None: + if isinstance(value, int): + return value + if isinstance(value, str): + try: + return int(value) + except ValueError: + return None + return None + + @staticmethod + def _first_non_none_value( + values: dict[str, object], *keys: str, default: object + ) -> object: + return next( + (values[key] for key in keys if key in values and values[key] is not None), + default, + ) diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index 85602bf1d86..35a0d84df40 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -67,10 +67,12 @@ class ParallelAISearchConfig(BaseSearchConfig): api_base: Optional[str] = None, **kwargs, ) -> Dict: - api_key = ( - api_key - or get_secret_str("PARALLEL_AI_API_KEY") - or get_secret_str("PARALLEL_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"), + base_env_var="PARALLEL_AI_API_BASE", + default_api_base=self.PARALLEL_AI_API_BASE, ) if not api_key: raise ValueError( diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 48299529ff4..cc0738697d1 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -25,7 +25,11 @@ class PerplexityChatConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" # type: ignore + api_base = ( + api_base + or get_secret_str("PERPLEXITY_API_BASE") + or "https://api.perplexity.ai" + ) # type: ignore dynamic_api_key = ( api_key or get_secret_str("PERPLEXITYAI_API_KEY") diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index bf055f91aa0..ec7ec397ea6 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -98,10 +98,11 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats if isinstance(search_cost_value, dict): - # Use the "low" size as default - tests expect 0.005 / 1000 - search_cost_per_query = ( - _safe_float_cast(search_cost_value.get("search_context_size_low", 0)) - / 1000 + # search_context_cost_per_query stores the per-request price in USD + # (e.g. sonar low = $0.005/request). Use it directly, matching the + # gemini cost calculator which reads the same field per request. + search_cost_per_query = _safe_float_cast( + search_cost_value.get("search_context_size_low", 0) ) else: search_cost_per_query = _safe_float_cast(search_cost_value) diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index ea96f87957c..55de52c5384 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -50,7 +50,13 @@ class PerplexitySearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("PERPLEXITYAI_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("PERPLEXITYAI_API_KEY",), + base_env_var="PERPLEXITY_API_BASE", + default_api_base=self.PERPLEXITY_API_BASE, + ) if not api_key: raise ValueError( "PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_API_KEY` environment variable." diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index cc4c61e397b..a2eddb65a54 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -214,7 +214,9 @@ def completion( headers=headers, http_client=httpx_client, ) - return CustomStreamWrapper(_response, model, logging_obj=logging_obj, custom_llm_provider="replicate") # type: ignore + return CustomStreamWrapper( + _response, model, logging_obj=logging_obj, custom_llm_provider="replicate" + ) # type: ignore else: for retry in range(litellm.DEFAULT_REPLICATE_POLLING_RETRIES): time.sleep( @@ -285,7 +287,9 @@ async def async_completion( headers=headers, http_client=async_handler, ) - return CustomStreamWrapper(_response, model, logging_obj=logging_obj, custom_llm_provider="replicate") # type: ignore + return CustomStreamWrapper( + _response, model, logging_obj=logging_obj, custom_llm_provider="replicate" + ) # type: ignore for retry in range(litellm.DEFAULT_REPLICATE_POLLING_RETRIES): await asyncio.sleep( diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index db8c26b7d96..9899e3be9ad 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -149,7 +149,9 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): def run_env_setup(self, service_key: Optional[str] = None) -> None: try: - self.token_creator, self._base_url, self._resource_group = get_token_creator(service_key) # type: ignore + self.token_creator, self._base_url, self._resource_group = ( + get_token_creator(service_key) + ) # type: ignore except ValueError as err: raise GenAIHubOrchestrationError(status_code=400, message=err.args[0]) @@ -189,7 +191,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): for dep in deployments.get("resources", []): if dep.get("scenarioId") == "orchestration": cfg = client.get( - f'{self.base_url}/lm/configurations/{dep["configurationId"]}', + f"{self.base_url}/lm/configurations/{dep['configurationId']}", headers=self.headers, ).json() if cfg.get("executableId") == "orchestration": diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index dd307ddf496..8cb19f195f2 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -102,20 +102,25 @@ CREDENTIAL_VALUES: Final[List[CredentialsValue]] = [ CredentialsValue( "auth_url", ("url",), - transform_fn=lambda url: url.rstrip("/") - + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX), + transform_fn=lambda url: ( + url.rstrip("/") + + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX) + ), ), CredentialsValue( "base_url", ("serviceurls", "AI_API_URL"), - transform_fn=lambda url: url.rstrip("/") - + ("" if url.endswith("/v2") else "/v2"), + transform_fn=lambda url: ( + url.rstrip("/") + ("" if url.endswith("/v2") else "/v2") + ), ), CredentialsValue( "cert_url", ("certurl",), - transform_fn=lambda url: url.rstrip("/") - + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX), + transform_fn=lambda url: ( + url.rstrip("/") + + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX) + ), ), # file paths (kept for config compatibility) CredentialsValue("cert_file_path"), diff --git a/litellm/llms/sap/embed/transformation.py b/litellm/llms/sap/embed/transformation.py index c74f21c3685..4344c2cc545 100644 --- a/litellm/llms/sap/embed/transformation.py +++ b/litellm/llms/sap/embed/transformation.py @@ -109,7 +109,7 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): if deployment["scenarioId"] == "orchestration": config_details = client.get( self.base_url - + f'/lm/configurations/{deployment["configurationId"]}', + + f"/lm/configurations/{deployment['configurationId']}", headers=self.headers, ).json() if config_details["executableId"] == "orchestration": diff --git a/litellm/llms/searchapi/search/transformation.py b/litellm/llms/searchapi/search/transformation.py index c04e1377f9c..ae8413684cc 100644 --- a/litellm/llms/searchapi/search/transformation.py +++ b/litellm/llms/searchapi/search/transformation.py @@ -74,7 +74,13 @@ class SearchAPIConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("SEARCHAPI_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SEARCHAPI_API_KEY",), + base_env_var="SEARCHAPI_API_BASE", + default_api_base=self.SEARCHAPI_API_BASE, + ) if not api_key: raise ValueError( @@ -114,6 +120,7 @@ class SearchAPIConfig(BaseSearchConfig): query: Union[str, List[str]], optional_params: dict, api_key: Optional[str] = None, + api_base: str | None = None, search_engine_id: Optional[str] = None, **kwargs, ) -> Dict: @@ -137,8 +144,16 @@ class SearchAPIConfig(BaseSearchConfig): if isinstance(query, list): query = " ".join(query) - # Get API key from parameter or environment - api_key = api_key or get_secret_str("SEARCHAPI_API_KEY") + # Get API key from parameter or environment. The key is sent as a query + # param to api_base, so resolve it host-aware to avoid leaking a + # server-managed key to a caller-supplied host. + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SEARCHAPI_API_KEY",), + base_env_var="SEARCHAPI_API_BASE", + default_api_base=self.SEARCHAPI_API_BASE, + ) if not api_key: raise ValueError( "SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable." diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index ee6f3895721..ff68be5709e 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -61,7 +61,13 @@ class SearXNGSearchConfig(BaseSearchConfig): Some instances may require authentication via headers. """ # SearXNG typically doesn't require API keys, but support optional auth - api_key = api_key or get_secret_str("SEARXNG_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SEARXNG_API_KEY",), + base_env_var="SEARXNG_API_BASE", + default_api_base=None, + ) if api_key: headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index 0daccbe652b..dd43f2d2dc9 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -55,7 +55,13 @@ class SerperSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("SERPER_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SERPER_API_KEY",), + base_env_var="SERPER_API_BASE", + default_api_base=self.SERPER_API_BASE, + ) if not api_key: raise ValueError( "SERPER_API_KEY is not set. Set `SERPER_API_KEY` environment variable." diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 522858b8c2a..9c325c3cdda 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -80,7 +80,9 @@ class StabilityImageEditConfig(BaseImageEditConfig): if k in param_mapping: # Map param if mapping exists and value is valid if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: - mapped_params[param_mapping[k]] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] # type: ignore + mapped_params[param_mapping[k]] = ( + OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] + ) # type: ignore # Don't copy "size" itself to final dict elif k == "n": # Store for logic but do not add to outgoing params diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index ec96db96f36..647cfb5fa84 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -64,7 +64,13 @@ class TavilySearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("TAVILY_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("TAVILY_API_KEY",), + base_env_var="TAVILY_API_BASE", + default_api_base=self.TAVILY_API_BASE, + ) if not api_key: raise ValueError( "TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable." diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index c4949380e3a..4a95f519646 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -67,7 +67,13 @@ class TinyfishSearchConfig(BaseSearchConfig): api_base: str | None = None, **kwargs: object, ) -> dict[str, str]: - resolved_key = api_key or get_secret_str("TINYFISH_API_KEY") + resolved_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("TINYFISH_API_KEY",), + base_env_var="TINYFISH_API_BASE", + default_api_base=self.TINYFISH_API_BASE, + ) if not resolved_key: raise ValueError( "TINYFISH_API_KEY is not set. Set `TINYFISH_API_KEY` environment variable." @@ -143,10 +149,8 @@ class TinyfishSearchConfig(BaseSearchConfig): max_results_str: str = "20" if raw_response.request: - raw_param: object = ( - raw_response.request.url.params.get( # any-ok: httpx QueryParams.get() -> Any - "max_results", "20" - ) + raw_param: object = raw_response.request.url.params.get( # any-ok: httpx QueryParams.get() -> Any + "max_results", "20" ) max_results_str = str(raw_param) max_results: int = min(int(max_results_str), 20) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index c627599da8d..00222faf274 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -78,10 +78,8 @@ class VertexAIBatchPrediction(VertexLLM): "Authorization": f"Bearer {access_token}", } - vertex_batch_request: VertexAIBatchPredictionJob = ( - VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( - request=create_batch_data - ) + vertex_batch_request: VertexAIBatchPredictionJob = VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + request=create_batch_data ) if _is_async is True: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5028c0cf5c8..96f016da94d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -779,7 +779,8 @@ def filter_schema_fields( result[key] = filter_schema_fields(value, valid_fields, processed) elif key == "anyOf" and isinstance(value, list): result[key] = [ - filter_schema_fields(item, valid_fields, processed) for item in value # type: ignore + filter_schema_fields(item, valid_fields, processed) + for item in value # type: ignore ] else: result[key] = value diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 103801a1e8d..d29734c0294 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -415,7 +415,9 @@ class ContextCachingEndpoints(VertexBase): try: response = client.post( - url=url, headers=headers, json=cached_content_request_body # type: ignore + url=url, + headers=headers, + json=cached_content_request_body, # type: ignore ) response.raise_for_status() except httpx.HTTPStatusError as err: @@ -566,7 +568,9 @@ class ContextCachingEndpoints(VertexBase): try: response = await client.post( - url=url, headers=headers, json=cached_content_request_body # type: ignore + url=url, + headers=headers, + json=cached_content_request_body, # type: ignore ) response.raise_for_status() except httpx.HTTPStatusError as err: diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py index 9fa57f6bf96..758fc0bae87 100644 --- a/litellm/llms/vertex_ai/cost_calculator.py +++ b/litellm/llms/vertex_ai/cost_calculator.py @@ -105,8 +105,10 @@ def cost_per_character( "input_cost_per_character_above_128k_tokens" in model_info and model_info["input_cost_per_character_above_128k_tokens"] is not None - ), "model info for model={} does not have 'input_cost_per_character_above_128k_tokens'-pricing for > 128k tokens\nmodel_info={}".format( - model, model_info + ), ( + "model info for model={} does not have 'input_cost_per_character_above_128k_tokens'-pricing for > 128k tokens\nmodel_info={}".format( + model, model_info + ) ) prompt_cost = ( prompt_characters @@ -116,8 +118,10 @@ def cost_per_character( assert ( "input_cost_per_character" in model_info and model_info["input_cost_per_character"] is not None - ), "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( - model, model_info + ), ( + "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( + model, model_info + ) ) prompt_cost = prompt_characters * model_info["input_cost_per_character"] except Exception as e: @@ -150,8 +154,10 @@ def cost_per_character( "output_cost_per_character_above_128k_tokens" in model_info and model_info["output_cost_per_character_above_128k_tokens"] is not None - ), "model info for model={} does not have 'output_cost_per_character_above_128k_tokens' pricing\nmodel_info={}".format( - model, model_info + ), ( + "model info for model={} does not have 'output_cost_per_character_above_128k_tokens' pricing\nmodel_info={}".format( + model, model_info + ) ) completion_cost = ( completion_tokens @@ -161,8 +167,10 @@ def cost_per_character( assert ( "output_cost_per_character" in model_info and model_info["output_cost_per_character"] is not None - ), "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( - model, model_info + ), ( + "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( + model, model_info + ) ) completion_cost = ( completion_characters * model_info["output_cost_per_character"] diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index c31bfde69e7..176cfe98411 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -17,17 +17,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( - CreateFileRequest, FileContentRequest, HttpxBinaryResponseContent, - OpenAIFileObject, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES -from .transformation import VertexAIFilesConfig, VertexAIJsonlFilesTransformation - -vertex_ai_files_transformation = VertexAIJsonlFilesTransformation() +from .transformation import VertexAIFilesConfig class VertexAIFilesHandler(GCSBucketBase): @@ -43,82 +39,6 @@ class VertexAIFilesHandler(GCSBucketBase): llm_provider=LlmProviders.VERTEX_AI, ) - async def async_create_file( - self, - create_file_data: CreateFileRequest, - api_base: Optional[str], - vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES], - vertex_project: Optional[str], - vertex_location: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - ) -> OpenAIFileObject: - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - kwargs={} - ) - headers = await self.construct_request_headers( - vertex_instance=gcs_logging_config["vertex_instance"], - service_account_json=gcs_logging_config["path_service_account"], - ) - bucket_name = gcs_logging_config["bucket_name"] - ( - logging_payload, - object_name, - ) = vertex_ai_files_transformation.transform_openai_file_content_to_vertex_ai_file_content( - openai_file_content=create_file_data.get("file") - ) - gcs_upload_response = await self._log_json_data_on_gcs( - headers=headers, - bucket_name=bucket_name, - object_name=object_name, - logging_payload=logging_payload, - ) - - return vertex_ai_files_transformation.transform_gcs_bucket_response_to_openai_file_object( - create_file_data=create_file_data, - gcs_upload_response=gcs_upload_response, - ) - - def create_file( - self, - _is_async: bool, - create_file_data: CreateFileRequest, - api_base: Optional[str], - vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES], - vertex_project: Optional[str], - vertex_location: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: - """ - Creates a file on VertexAI GCS Bucket - - Only supported for Async litellm.acreate_file - """ - - if _is_async: - return self.async_create_file( - create_file_data=create_file_data, - api_base=api_base, - vertex_credentials=vertex_credentials, - vertex_project=vertex_project, - vertex_location=vertex_location, - timeout=timeout, - max_retries=max_retries, - ) - else: - return asyncio.run( - self.async_create_file( - create_file_data=create_file_data, - api_base=api_base, - vertex_credentials=vertex_credentials, - vertex_project=vertex_project, - vertex_location=vertex_location, - timeout=timeout, - max_retries=max_retries, - ) - ) - def _extract_bucket_and_object_from_file_id( self, file_id: str, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f30518bc7ca..d5164d8c1c2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1,9 +1,21 @@ import base64 +import io +import itertools import json import os import re import time -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import ( + Any, + Callable, + Dict, + Iterable, + Iterator, + List, + Optional, + Tuple, + Union, +) import httpx from httpx import Headers, Response @@ -22,9 +34,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( validate_managed_cloud_file_id, ) from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + extract_file_metadata, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( + BaseFileUploadStream, BaseFilesConfig, LiteLLMLoggingObj, ) @@ -44,8 +60,9 @@ from litellm.types.llms.openai import ( OpenAIFileObject, PathLike, ) +from litellm.types.files import ResumableChunkedUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import ExtractedFileData, LlmProviders, ModelResponse +from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase @@ -137,42 +154,140 @@ def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) -def _openai_batch_jsonl_entries_to_vertex_wrapped_requests( - openai_jsonl_content: List[Dict[str, Any]], +def _openai_batch_jsonl_entry_to_vertex_wrapped_request( + openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> List[Dict[str, Any]]: +) -> Dict[str, Any]: """ - Transforms OpenAI JSONL batch entries to Vertex AI JSONL lines. + Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} - {"request":{"contents": [{"role": "user", "parts": [{"text": "Describe what is happening in this video."}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/another_video.mov", "mimeType": "video/mov"}}]}]}} + """ + openai_request_body = openai_entry.get("body") or {} + vertex_request_body = _transform_request_body( + messages=openai_request_body.get("messages", []), + model=openai_request_body.get("model", ""), + optional_params=map_openai_to_vertex_params(openai_request_body), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + custom_id = openai_entry.get("custom_id") + if custom_id is not None: + if "labels" not in vertex_request_body: + vertex_request_body["labels"] = {} + _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) + + return {"request": vertex_request_body} + + +def _iter_stripped_lines(raw_lines: Iterable[Union[str, bytes]]) -> Iterator[str]: + """Decode (when needed), strip, and drop blank lines from an iterable of lines.""" + for raw in raw_lines: + line = raw.decode("utf-8") if isinstance(raw, (bytes, bytearray)) else raw + line = line.strip() + if line: + yield line + + +def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: + """ + Yield non-empty JSONL lines one at a time without materializing the whole + payload, so peak memory stays bounded regardless of payload size. Mirrors + ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited + JSONL. + """ + content: Any = openai_file_content + if isinstance(content, tuple): + content = content[1] + + if isinstance(content, (bytes, bytearray)): + # Scan for newlines in place so a large in-memory payload is not copied + # into a BytesIO just to iterate it line by line. + newline = ord("\n") + start, length = 0, len(content) + while start < length: + idx = content.find(newline, start) + if idx == -1: + chunk, start = content[start:], length + else: + chunk, start = content[start:idx], idx + 1 + line = chunk.decode("utf-8").strip() + if line: + yield line + return + + if isinstance(content, str): + yield from _iter_stripped_lines(io.StringIO(content)) + return + + if isinstance(content, PathLike): + with open(str(content), "rb") as handle: + yield from _iter_stripped_lines(handle) + return + + if hasattr(content, "read"): + # The handle is read twice per upload (first-row probe for the GCS + # object name, then the body stream), so it must rewind to 0. A + # non-seekable handle would silently resume mid-stream and drop the + # already-consumed first row, so reject it loudly instead. + seek = getattr(content, "seek", None) + if seek is None: + raise ValueError( + "Batch upload file handle must be seekable; got a non-seekable " + "stream. Pass bytes, a path, or a seekable handle." + ) + try: + seek(0) + except (OSError, ValueError) as e: + raise ValueError( + "Batch upload file handle must be seekable so it can be re-read " + "for the GCS object name and the upload body." + ) from e + yield from _iter_stripped_lines(content) + return + + raise ValueError("Unsupported file content type") + + +def _iter_openai_jsonl_entries( + openai_file_content: FileTypes, +) -> Iterator[Dict[str, Any]]: + for line in _iter_openai_jsonl_lines(openai_file_content): + yield json.loads(line) + + +class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): + """Streams an OpenAI batch JSONL upload as Vertex-wrapped JSONL one row at a + time, so the transformed payload is never held in full. + + The transform runs lazily as the HTTP client pulls each chunk, which keeps + peak memory at one row regardless of how large the batch file is. """ - vertex_jsonl_content = [] - for _openai_jsonl_content in openai_jsonl_content: - openai_request_body = _openai_jsonl_content.get("body") or {} - vertex_request_body = _transform_request_body( - messages=openai_request_body.get("messages", []), - model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) + def __init__( + self, + openai_file_content: FileTypes, + map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], + ) -> None: + self._openai_file_content = openai_file_content + self._map_openai_to_vertex_params = map_openai_to_vertex_params - # Add custom_id as a label for correlation in batch outputs - custom_id = _openai_jsonl_content.get("custom_id") - if custom_id is not None: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels( - vertex_request_body["labels"], custom_id + def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: + first = True + for entry in _iter_openai_jsonl_entries(self._openai_file_content): + wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request( + entry, self._map_openai_to_vertex_params ) + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(wrapped).encode("utf-8") - vertex_jsonl_content.append({"request": vertex_request_body}) - return vertex_jsonl_content + def iter_bytes(self) -> Iterator[bytes]: + return self._iter_vertex_jsonl_chunks() class VertexAIFilesConfig(VertexBase, BaseFilesConfig): @@ -181,7 +296,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ def __init__(self): - self.jsonl_transformation = VertexAIJsonlFilesTransformation() super().__init__() @property @@ -208,43 +322,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers["Authorization"] = f"Bearer {api_key}" return headers - def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: - """ - Helper to extract content from various OpenAI file types and return as string. - - Handles: - - Direct content (str, bytes, IO[bytes]) - - Tuple formats: (filename, content, [content_type], [headers]) - - PathLike objects - """ - content: Union[str, bytes] = b"" - # Extract file content from tuple if necessary - if isinstance(openai_file_content, tuple): - # Take the second element which is always the file content - file_content = openai_file_content[1] - else: - file_content = openai_file_content - - # Handle different file content types - if isinstance(file_content, str): - # String content can be used directly - content = file_content - elif isinstance(file_content, bytes): - # Bytes content can be decoded - content = file_content - elif isinstance(file_content, PathLike): # PathLike - with open(str(file_content), "rb") as f: - content = f.read() - elif hasattr(file_content, "read"): # IO[bytes] - # File-like objects need to be read - content = file_content.read() - - # Ensure content is string - if isinstance(content, bytes): - content = content.decode("utf-8") - - return content - def _get_gcs_object_name_from_batch_jsonl( self, openai_jsonl_content: List[Dict[str, Any]], @@ -261,32 +338,21 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): object_name = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name - def get_object_name( - self, extracted_file_data: ExtractedFileData, purpose: str - ) -> str: + def get_object_name(self, file_data: FileTypes, purpose: str) -> str: """ - Get the object name for the request + Get the object name for the request. + + Reads only the first JSONL entry (streamed) for batch files, so a large + upload is never materialized just to derive the GCS object name. """ - extracted_file_data_content = extracted_file_data.get("content") - - if extracted_file_data_content is None: - raise ValueError("file content is required") - if purpose == "batch": - ## 1. If jsonl, check if there's a model name - file_content = self._get_content_from_openai_file( - extracted_file_data_content - ) - - # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - if len(openai_jsonl_content) > 0: - return self._get_gcs_object_name_from_batch_jsonl(openai_jsonl_content) + ## 1. If jsonl, derive the object name from the first entry's model + first_entry = next(_iter_openai_jsonl_entries(file_data), None) + if first_entry is not None: + return self._get_gcs_object_name_from_batch_jsonl([first_entry]) ## 2. If not jsonl, store under a server-generated managed object name - filename = extracted_file_data.get("filename") + filename, _ = extract_file_metadata(file_data) return build_managed_cloud_object_name( prefix=f"{VERTEX_AI_MANAGED_GCS_PREFIX}uploads/", filename=filename, @@ -294,7 +360,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) def _get_configured_bucket_name(self, litellm_params: Dict) -> str: - bucket_name = litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + bucket_name = ( + litellm_params.get("gcs_bucket_name") + or litellm_params.get("bucket_name") + or os.getenv("GCS_BUCKET_NAME") + ) if not bucket_name: raise ValueError("GCS bucket_name is required") return bucket_name @@ -319,12 +389,21 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - extracted_file_data = extract_file_data(file_data) - object_name = self.get_object_name(extracted_file_data, purpose) + _, content_type = extract_file_metadata(file_data) + object_name = self.get_object_name(file_data, purpose) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name = encode_gcs_object_name_for_url(object_name) - endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" + # Batch jsonl is streamed via a resumable session (bounded memory on + # large uploads); everything else is a single simple-media upload. + upload_type = ( + "resumable" + if FilesAPIUtils.is_batch_jsonl_request( + create_file_data=data, content_type=content_type + ) + else "media" + ) + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -366,14 +445,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) return vertex_params - def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - self, openai_jsonl_content: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - return _openai_batch_jsonl_entries_to_vertex_wrapped_requests( - openai_jsonl_content=openai_jsonl_content, - map_openai_to_vertex_params=self._map_openai_to_vertex_params, - ) - def transform_create_file_request( self, model: str, @@ -384,40 +455,34 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ 2 Cases: 1. Handle basic file upload - 2. Handle batch file upload (.jsonl) + 2. Handle batch file upload (.jsonl), streamed to a GCS resumable + session so large uploads stay memory-bounded. """ file_data = create_file_data.get("file") if file_data is None: raise ValueError("file is required") - extracted_file_data = extract_file_data(file_data) - extracted_file_data_content = extracted_file_data.get("content") - if extracted_file_data_content is None: - raise ValueError("file content is required") - - if FilesAPIUtils.is_batch_jsonl_file( + _, content_type = extract_file_metadata(file_data) + if FilesAPIUtils.is_batch_jsonl_request( create_file_data=create_file_data, - extracted_file_data=extracted_file_data, + content_type=content_type, ): - ## 1. If jsonl, check if there's a model name - file_content = self._get_content_from_openai_file( - extracted_file_data_content - ) - - # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - vertex_jsonl_content = ( - self._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content + return { + "resumable_chunked_upload": ResumableChunkedUploadConfig( + body_stream=_OpenAIToVertexBatchUploadStream( + file_data, + self._map_openai_to_vertex_params, + ), + initiate_headers={ + "X-Upload-Content-Type": "application/json", + }, ) - ) - return "\n".join(json.dumps(item) for item in vertex_jsonl_content) - elif isinstance(extracted_file_data_content, bytes): + } + + extracted_file_data_content = extract_file_data(file_data).get("content") + if isinstance(extracted_file_data_content, bytes): return extracted_file_data_content - else: - raise ValueError("Unsupported file content type") + raise ValueError("Unsupported file content type") def transform_create_file_response( self, @@ -642,39 +707,38 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): } """ try: - # Decode content - content_str = content.decode("utf-8") - - # Check if it's JSONL (multiple lines) - lines = content_str.strip().split("\n") - if not lines: + # Read the result file one row at a time. Batch output files can be + # as large as the (multi-GB) input, so splitting into a list of rows + # and building a second list of transformed rows peaks at several full + # copies and OOMs on retrieval. + lines = _iter_openai_jsonl_lines(content) + try: + first_line = next(lines) + except StopIteration: return content - # Try to parse the first line to see if it's Vertex AI batch output - first_line = json.loads(lines[0]) - - # Check if it has Vertex AI batch output structure with discriminating fields - # Must have request, response, and processed_time - # Plus either candidates (success) or status (error) - has_base_structure = ( - "response" in first_line - and "request" in first_line - and "processed_time" in first_line + # Identify a Vertex AI batch output from the first row's + # discriminating fields. Anything else (e.g. a binary file whose + # first line is not valid UTF-8/JSON) raises and falls through to the + # passthrough below, leaving the content untouched. + first_row = json.loads(first_line) + is_vertex_batch_output = ( + "request" in first_row + and "response" in first_row + and "processed_time" in first_row + and ( + "candidates" in first_row.get("response", {}) + or "promptFeedback" in first_row.get("response", {}) + or bool(first_row.get("status")) + ) ) - has_success_or_error = ( - "candidates" in first_line.get("response", {}) - or "promptFeedback" in first_line.get("response", {}) - or bool(first_line.get("status")) - ) - - if not (has_base_structure and has_success_or_error): - # Not a Vertex AI batch output, return as-is + if not is_vertex_batch_output: return content vertex_gemini_config = VertexGeminiConfig() - # Always use a fresh local Logging object for the per-line transformation - # so we never mutate the caller's logging_obj (which already went through - # pre_call and has its own model/start_time/optional_params set). + # Use a fresh Logging object for the per-row transform so we never + # mutate the caller's (which already ran pre_call with its own + # model/start_time/optional_params). batch_transform_logging_obj = Logging( model="", messages=[], @@ -691,29 +755,27 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) - # Transform all lines - transformed_lines = [] - for line in lines: - if not line.strip(): - continue - + # Transform each row straight into the output buffer, so peak memory + # stays at ~one row plus the output. If any row fails, return the + # original content unchanged. + output = bytearray() + for line in itertools.chain([first_line], lines): try: - vertex_output = json.loads(line) openai_output = ( self._transform_single_vertex_batch_output_to_openai( - vertex_output=vertex_output, + vertex_output=json.loads(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, ) ) - transformed_lines.append(json.dumps(openai_output)) except Exception: - # If any line fails, return original content return content + if output: + output += b"\n" + output += json.dumps(openai_output).encode("utf-8") - # Return transformed content - return "\n".join(transformed_lines).encode("utf-8") + return bytes(output) except Exception: # If anything fails, return original content @@ -795,137 +857,3 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "message": f"Failed to transform response: {str(e)}", }, } - - -class VertexAIJsonlFilesTransformation(VertexGeminiConfig): - """ - Transforms OpenAI /v1/files/* requests to VertexAI /v1/files/* requests - """ - - def transform_openai_file_content_to_vertex_ai_file_content( - self, openai_file_content: Optional[FileTypes] = None - ) -> Tuple[str, str]: - """ - Transforms OpenAI FileContentRequest to VertexAI FileContentRequest - """ - - if openai_file_content is None: - raise ValueError("contents of file are None") - # Read the content of the file - file_content = self._get_content_from_openai_file(openai_file_content) - - # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - vertex_jsonl_content = ( - self._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) - vertex_jsonl_string = "\n".join( - json.dumps(item) for item in vertex_jsonl_content - ) - object_name = self._get_gcs_object_name( - openai_jsonl_content=openai_jsonl_content - ) - return vertex_jsonl_string, object_name - - def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - self, openai_jsonl_content: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - return _openai_batch_jsonl_entries_to_vertex_wrapped_requests( - openai_jsonl_content=openai_jsonl_content, - map_openai_to_vertex_params=self._map_openai_to_vertex_params, - ) - - def _get_gcs_object_name( - self, - openai_jsonl_content: List[Dict[str, Any]], - ) -> str: - """ - Gets a unique GCS object name for the VertexAI batch prediction job - - named as: litellm-vertex-{model}-{uuid} - """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path = sanitize_cloud_object_path(_model, fallback="model") - object_name = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" - return object_name - - def _map_openai_to_vertex_params( - self, - openai_request_body: Dict[str, Any], - ) -> Dict[str, Any]: - """ - wrapper to call VertexGeminiConfig.map_openai_params - """ - _model = openai_request_body.get("model", "") - vertex_params = self.map_openai_params( - model=_model, - non_default_params=openai_request_body, - optional_params={}, - drop_params=False, - ) - return vertex_params - - def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: - """ - Helper to extract content from various OpenAI file types and return as string. - - Handles: - - Direct content (str, bytes, IO[bytes]) - - Tuple formats: (filename, content, [content_type], [headers]) - - PathLike objects - """ - content: Union[str, bytes] = b"" - # Extract file content from tuple if necessary - if isinstance(openai_file_content, tuple): - # Take the second element which is always the file content - file_content = openai_file_content[1] - else: - file_content = openai_file_content - - # Handle different file content types - if isinstance(file_content, str): - # String content can be used directly - content = file_content - elif isinstance(file_content, bytes): - # Bytes content can be decoded - content = file_content - elif isinstance(file_content, PathLike): # PathLike - with open(str(file_content), "rb") as f: - content = f.read() - elif hasattr(file_content, "read"): # IO[bytes] - # File-like objects need to be read - content = file_content.read() - - # Ensure content is string - if isinstance(content, bytes): - content = content.decode("utf-8") - - return content - - def transform_gcs_bucket_response_to_openai_file_object( - self, create_file_data: CreateFileRequest, gcs_upload_response: Dict[str, Any] - ) -> OpenAIFileObject: - """ - Transforms GCS Bucket upload file response to OpenAI FileObject - """ - gcs_id = gcs_upload_response.get("id", "") - # Remove the last numeric ID from the path - gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" - - return OpenAIFileObject( - purpose=create_file_data.get("purpose", "batch"), - id=f"gs://{gcs_id}", - filename=gcs_upload_response.get("name", ""), - created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=gcs_upload_response.get("timeCreated", "") - ), - status="uploaded", - bytes=gcs_upload_response.get("size", 0), - object="file", - ) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index f5a2b268263..79f006a83de 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -881,7 +881,9 @@ def _gemini_convert_messages_with_history( ## MERGE CONSECUTIVE ASSISTANT CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "assistant": if isinstance(messages[msg_i], BaseModel): - msg_dict: Union[ChatCompletionAssistantMessage, dict] = messages[msg_i].model_dump() # type: ignore + msg_dict: Union[ChatCompletionAssistantMessage, dict] = messages[ + msg_i + ].model_dump() # type: ignore else: msg_dict = messages[msg_i] # type: ignore assistant_msg = ChatCompletionAssistantMessage(**msg_dict) # type: ignore @@ -945,7 +947,12 @@ def _gemini_convert_messages_with_history( and len(thought_signatures) > 0 ): # Use the first signature for the text part (Gemini expects one signature per part) - assistant_content.append(PartType(text=assistant_text, thoughtSignature=thought_signatures[0])) # type: ignore + assistant_content.append( + PartType( + text=assistant_text, + thoughtSignature=thought_signatures[0], + ) + ) # type: ignore else: assistant_content.append(PartType(text=assistant_text)) # type: ignore @@ -1201,7 +1208,8 @@ def _transform_request_body( ) if supports_response_schema is False: user_response_schema_message = response_schema_prompt( - model=model, response_schema=optional_params.get("response_schema") # type: ignore + model=model, + response_schema=optional_params.get("response_schema"), # type: ignore ) messages.append({"role": "user", "content": user_response_schema_message}) optional_params.pop("response_schema") diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index c171538b9c0..0de34ec257c 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -823,9 +823,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if google_maps_retrieval_config is not None: if "toolConfig" not in optional_params: optional_params["toolConfig"] = {} - optional_params["toolConfig"][ - "retrievalConfig" - ] = google_maps_retrieval_config + optional_params["toolConfig"]["retrievalConfig"] = ( + google_maps_retrieval_config + ) return _tools_list @@ -1271,7 +1271,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): isinstance(value, str) or isinstance(value, dict) ): _tool_choice_value = self.map_tool_choice_values( - model=model, tool_choice=value # type: ignore + model=model, + tool_choice=value, # type: ignore ) if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value @@ -2195,7 +2196,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): from litellm.types.utils import Delta, StreamingChoices annotations = chat_completion_message.get("annotations") # type: ignore - provider_specific_fields = chat_completion_message.get("provider_specific_fields") # type: ignore + provider_specific_fields = chat_completion_message.get( + "provider_specific_fields" + ) # type: ignore # create a streaming choice object choice = StreamingChoices( finish_reason=VertexGeminiConfig._check_finish_reason( @@ -2470,15 +2473,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) if audio_response is not None: - cast(Dict[str, Any], chat_completion_message)[ - "audio" - ] = audio_response + cast(Dict[str, Any], chat_completion_message)["audio"] = ( + audio_response + ) chat_completion_message["content"] = None # OpenAI spec if image_response is not None: # Handle image response - combine with text content into structured format - cast(Dict[str, Any], chat_completion_message)[ - "images" - ] = image_response + cast(Dict[str, Any], chat_completion_message)["images"] = ( + image_response + ) if content is not None: chat_completion_message["content"] = content @@ -2538,13 +2541,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if thought_signatures is not None: if "provider_specific_fields" not in chat_completion_message: chat_completion_message["provider_specific_fields"] = {} - chat_completion_message["provider_specific_fields"]["thought_signatures"] = thought_signatures # type: ignore + chat_completion_message["provider_specific_fields"][ + "thought_signatures" + ] = thought_signatures # type: ignore # Store server-side tool invocations in provider_specific_fields if server_side_tool_invocations is not None: if "provider_specific_fields" not in chat_completion_message: chat_completion_message["provider_specific_fields"] = {} - chat_completion_message["provider_specific_fields"]["server_side_tool_invocations"] = server_side_tool_invocations # type: ignore + chat_completion_message["provider_specific_fields"][ + "server_side_tool_invocations" + ] = server_side_tool_invocations # type: ignore if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( @@ -3313,7 +3320,9 @@ class VertexLLM(VertexBase): client = client try: - response = client.post(url=url, headers=headers, json=data, logging_obj=logging_obj) # type: ignore + response = client.post( + url=url, headers=headers, json=data, logging_obj=logging_obj + ) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index d6441db7856..1fe9f15c9f0 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -32,6 +32,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): self._project = project self._location = location + def _include_function_response_id(self) -> bool: + return False + # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 696341598e5..729cc9c3ead 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -92,13 +92,11 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) - vertex_request: VertexEmbeddingRequest = ( - litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, - optional_params=optional_params, - model=model, - litellm_params=litellm_params, - ) + vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, ) _client_params = {} @@ -192,13 +190,11 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) - vertex_request: VertexEmbeddingRequest = ( - litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, - optional_params=optional_params, - model=model, - litellm_params=litellm_params, - ) + vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, ) _async_client_params = {} diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 990063bb9fb..18a9c98c315 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -179,9 +179,7 @@ class VertexBase: scopes=["https://www.googleapis.com/auth/cloud-platform"], ) if project_id is None: - project_id = ( - creds.quota_project_id - ) # authorized user credentials don't have a project_id, only quota_project_id + project_id = creds.quota_project_id # authorized user credentials don't have a project_id, only quota_project_id else: creds = self._credentials_from_service_account( json_obj, @@ -838,12 +836,13 @@ class VertexBase: self._credentials_project_mapping.pop(credential_cache_key, None) try: - _credentials, credential_project_id = ( - await self._load_and_cache_credentials( - credentials=credentials, - project_id=project_id, - credential_cache_key=credential_cache_key, - ) + ( + _credentials, + credential_project_id, + ) = await self._load_and_cache_credentials( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, ) if project_id is None and isinstance(credential_project_id, str): project_id = credential_project_id @@ -1068,10 +1067,11 @@ class VertexBase: # Load credentials if not cached if _credentials is None: - _credentials, credential_project_id = ( - await self._load_and_cache_credentials( - credentials, project_id, credential_cache_key - ) + ( + _credentials, + credential_project_id, + ) = await self._load_and_cache_credentials( + credentials, project_id, credential_cache_key ) # Resolve project_id from credentials if not provided diff --git a/litellm/llms/volcengine/chat/transformation.py b/litellm/llms/volcengine/chat/transformation.py index 7395f9ce75b..d60f54615fa 100644 --- a/litellm/llms/volcengine/chat/transformation.py +++ b/litellm/llms/volcengine/chat/transformation.py @@ -100,9 +100,9 @@ class VolcEngineChatConfig(OpenAILikeChatConfig): in ["enabled", "disabled", "auto"] # legal values, see docs ): # Add thinking parameter to extra_body for all legal cases - optional_params.setdefault("extra_body", {})[ - "thinking" - ] = thinking_value + optional_params.setdefault("extra_body", {})["thinking"] = ( + thinking_value + ) else: # Skip adding thinking parameter when it's not set or has invalid value pass diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py index 230c9f4cf6e..2aead929c8d 100644 --- a/litellm/llms/watsonx/common_utils.py +++ b/litellm/llms/watsonx/common_utils.py @@ -372,9 +372,7 @@ class IBMWatsonXMixin: def _prepare_payload(self, model: str, api_params: WatsonXAPIParams) -> dict: payload: dict = {} if model.startswith("deployment/"): - return ( - {} - ) # Deployment models do not support 'space_id' or 'project_id' in their payload + return {} # Deployment models do not support 'space_id' or 'project_id' in their payload payload["model_id"] = model if api_params["project_id"] is not None: payload["project_id"] = api_params["project_id"] diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 7180e12162a..8f418567371 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -327,7 +327,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): except ValueError: # datetime.fromisoformat cannot handle 'Z' in Python 3.10 created_datetime = datetime.fromisoformat( - f'{json_resp["created_at"].rstrip("Z")}+00:00' + f"{json_resp['created_at'].rstrip('Z')}+00:00" ) model_response.created = int(created_datetime.timestamp()) else: diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 8019bb67991..87b5757ef35 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -28,7 +28,6 @@ from ...openai.chat.gpt_transformation import ( class XAIChatConfig(OpenAIGPTConfig): - @property def custom_llm_provider(self) -> Optional[str]: return "xai" diff --git a/litellm/llms/you_com/search/transformation.py b/litellm/llms/you_com/search/transformation.py index 3c94b991735..0c7916e4c05 100644 --- a/litellm/llms/you_com/search/transformation.py +++ b/litellm/llms/you_com/search/transformation.py @@ -64,7 +64,13 @@ class YouComSearchConfig(BaseSearchConfig): endpoint with the `X-API-Key` header. Otherwise fall through to the keyless free tier; no auth header is required. """ - api_key = api_key or get_secret_str("YOUCOM_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("YOUCOM_API_KEY",), + base_env_var="YOUCOM_API_BASE", + default_api_base=self.YOU_COM_API_BASE, + ) headers["Content-Type"] = "application/json" # Pin Accept-Encoding to identity: the keyless `api.you.com/v1/agents/search` # endpoint advertises gzip content-encoding but returns body bytes the diff --git a/litellm/main.py b/litellm/main.py index 63c5798e70a..1960c456f80 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -81,11 +81,17 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, +) from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( @@ -118,6 +124,10 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.realtime_api.main import _realtime_health_check from litellm.secret_managers.main import get_secret_bool, get_secret_str +from litellm.types.completion import ( + _CompletionDispatchContext, + _CompletionDispatchResult, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CustomPricingLiteLLMParams, @@ -650,6 +660,39 @@ async def acompletion( response_object=response, model_response_object=litellm.ModelResponse(), ) + # Provider-agnostic dispatch point for the chat-completions agentic loop + # (code-interpreter interception, etc). Chat routing forks per provider + # before this (OpenAI goes through the OpenAI SDK in openai.py, others + # through the shared httpx handler), so a dispatch inside any single + # provider handler would miss the others. Here is where every fork + # reconverges, so the loop runs once for all providers. Responses needs + # no equivalent: every provider already funnels through one shared + # handler where the loop is dispatched. + if isinstance(response, litellm.ModelResponse): + looped = await maybe_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + optional_params={ + k: v + for k, v in completion_kwargs.items() + if v is not None + and k + not in ( + "model", + "messages", + "stream", + "acompletion", + "deployment_id", + ) + }, + kwargs=kwargs, + logging_obj=kwargs.get("litellm_logging_obj"), + custom_llm_provider=custom_llm_provider, + stream=bool(stream), + ) + if looped is not None: + response = looped if isinstance(response, CustomStreamWrapper): response.set_logging_event_loop( loop=loop @@ -1084,6 +1127,3825 @@ def _build_custom_pricing_entry( return entry +def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + _azure_detection_model = ctx._azure_detection_model + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + api_version = ctx.api_version + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + max_retries = ctx.max_retries + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + dynamic_params = False + if client is not None and ( + isinstance(client, openai.AzureOpenAI) + or isinstance(client, openai.AsyncAzureOpenAI) + ): + dynamic_params = _check_dynamic_azure_params( + azure_client_params={"api_version": api_version}, + azure_client=client, + ) + + api_type = get_secret("AZURE_API_TYPE") or "azure" + + api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") + + api_version = ( + api_version + or litellm.api_version + or get_secret_str("AZURE_API_VERSION") + or litellm.AZURE_DEFAULT_API_VERSION + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_OPENAI_API_KEY") + or get_secret_str("AZURE_API_KEY") + ) + + azure_ad_token = optional_params.get("extra_body", {}).pop( + "azure_ad_token", None + ) or get_secret_str("AZURE_AD_TOKEN") + + azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + if max_retries is not None: + optional_params["max_retries"] = max_retries + + if litellm.AzureOpenAIO1Config().is_o_series_model(model=_azure_detection_model): + ## LOAD CONFIG - if set + config = litellm.AzureOpenAIO1Config.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + response = azure_o1_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + api_key=api_key, + api_base=api_base, + api_version=api_version, + dynamic_params=dynamic_params, + azure_ad_token=azure_ad_token, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, # type: ignore + client=client, # pass AsyncAzureOpenAI, AzureOpenAI client + custom_llm_provider=custom_llm_provider, + ) + else: + ## LOAD CONFIG - if set + config = litellm.AzureOpenAIConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ## COMPLETION CALL + response = azure_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + api_key=api_key, + api_base=api_base, + api_version=api_version, + api_type=api_type, + dynamic_params=dynamic_params, + azure_ad_token=azure_ad_token, + azure_ad_token_provider=azure_ad_token_provider, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, # type: ignore + client=client, # pass AsyncAzureOpenAI, AzureOpenAI client + ) + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={ + "headers": headers, + "api_version": api_version, + "api_base": api_base, + }, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + api_version = ctx.api_version + client = ctx.client + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + api_type = get_secret_str("AZURE_API_TYPE") or "azure" + + api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + + if api_base is None: + raise ValueError( + "api_base is required for Azure OpenAI LLM provider. Either set it dynamically or set the AZURE_API_BASE environment variable." + ) + + api_version = ( + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_OPENAI_API_KEY") + or get_secret_str("AZURE_API_KEY") + ) + + azure_ad_token = optional_params.get("extra_body", {}).pop( + "azure_ad_token", None + ) or get_secret_str("AZURE_AD_TOKEN") + + azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + + ## LOAD CONFIG - if set + config = litellm.AzureOpenAIConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ## COMPLETION CALL + response = azure_text_completions.completion( + model=model, + messages=messages, + headers=headers, + api_key=api_key, + api_base=api_base, + api_version=cast(str, api_version), + api_type=api_type, + azure_ad_token=azure_ad_token, + azure_ad_token_provider=azure_ad_token_provider, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + client=client, # pass AsyncAzureOpenAI, AzureOpenAI client + ) + + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={ + "headers": headers, + "api_version": api_version, + "api_base": api_base, + }, + ) + + return response + + +def _complete_deepseek(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_azure_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + azure_ai_route = AzureFoundryModelInfo.get_azure_ai_route(model) + + # Check if this is an agents route - model format: azure_ai/agents/ + if azure_ai_route == "agents": + from litellm.llms.azure_ai.agents import AzureAIAgentsConfig + + api_base = AzureFoundryModelInfo.get_api_base(api_base) + if api_base is None: + raise ValueError( + "Azure AI Agents requests require an api_base. " + "Set `api_base` or the AZURE_AI_API_BASE env var." + ) + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + response = AzureAIAgentsConfig.completion( + model=model, + messages=messages, + api_base=api_base, + api_key=api_key, + model_response=model_response, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, + acompletion=acompletion, + stream=stream, + headers=headers or litellm.headers, + ) + + # Check if this is a Claude model - route to Azure Anthropic handler + elif "claude" in model.lower(): + # Use Azure Anthropic handler for Claude models + api_base = AzureFoundryModelInfo.get_api_base(api_base) + if api_base is None: + raise ValueError( + "Azure Anthropic requests require an api_base. " + "Set `api_base` or the AZURE_AI_API_BASE env var." + ) + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + # Ensure the URL ends with /v1/messages for Anthropic + if api_base: + api_base = api_base.rstrip("/") + if not api_base.endswith("/v1/messages"): + if "/anthropic" in api_base: + parts = api_base.split("/anthropic", 1) + api_base = parts[0] + "/anthropic" + else: + api_base = api_base + "/anthropic" + api_base = api_base + "/v1/messages" + + response = azure_anthropic_chat_completions.completion( + model=model, + messages=messages, + api_base=api_base, + acompletion=acompletion, + custom_prompt_dict=litellm.custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + headers=headers, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + ) + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + ) + response = response + else: + # Non-Claude models use standard Azure AI flow + api_base = AzureFoundryModelInfo.get_api_base(api_base) + # set API KEY + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + + ## FOR COHERE + if "command-r" in model: # make sure tool call in messages are str + messages = stringify_json_tool_call_content(messages=messages) + + ## COMPLETION CALL + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, # pass AsyncOpenAI, OpenAI client + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={"headers": headers}, + ) + + return response + + +def _complete_text_completion_openai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + text_completion = ctx.text_completion + timeout = ctx.timeout + + openai.api_type = "openai" + + api_base = ( + api_base + or litellm.api_base + or get_secret("OPENAI_BASE_URL") + or get_secret("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + + openai.api_version = None + # set API KEY + + api_key = ( + api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") + ) + + headers = headers or litellm.headers + + ## LOAD CONFIG - if set + config = litellm.OpenAITextCompletionConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_text_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + if litellm.organization: + openai.organization = litellm.organization + + ## COMPLETION CALL + _response = openai_text_completions.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + acompletion=acompletion, + client=client, # pass AsyncOpenAI, OpenAI client + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + ) + + if ( + optional_params.get("stream", False) is False + and acompletion is False + and text_completion is False + ): + # convert to chat completion response + _response = ( + litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( + response_object=_response, model_response_object=model_response + ) + ) + + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=_response, + additional_args={"headers": headers}, + ) + return _response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_fireworks_ai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_ragflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_xai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_groq(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or get_secret("GROQ_API_BASE") + or "https://api.groq.com/openai/v1" + ) + + # set API KEY + api_key = ( + api_key + or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there + or litellm.groq_key + or get_secret("GROQ_API_KEY") + ) + + headers = headers or litellm.headers + + ## LOAD CONFIG - if set + config = litellm.GroqChatConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_bedrock_mantle( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = api_base or litellm.api_base or get_secret("BEDROCK_MANTLE_API_BASE") + api_key = api_key or litellm.api_key or get_secret("BEDROCK_MANTLE_API_KEY") + headers = headers or litellm.headers + config = litellm.BedrockMantleChatConfig.get_config() + for k, v in config.items(): + if k not in optional_params: + optional_params[k] = v + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + +def _complete_a2a(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + ( + api_base, + api_key, + headers, + ) = litellm.A2AConfig.resolve_agent_config_from_registry( + model=model, + api_base=api_base, + api_key=api_key, + headers=headers, + optional_params=optional_params, + ) + + # Fall back to environment variables and defaults + api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE") + + if api_base is None: + raise Exception( + "api_base is required for A2A provider. " + "Either provide api_base parameter, set A2A_API_BASE environment variable, " + "or register the agent in the proxy with model='a2a/'." + ) + + headers = headers or litellm.headers + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + provider_config=provider_config, + ) + + +def _complete_gigachat(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.api_key + or litellm.gigachat_key + or get_secret("GIGACHAT_API_KEY") + or get_secret("GIGACHAT_CREDENTIALS") + ) + + headers = headers or litellm.headers or {} + + ## COMPLETION CALL + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_sap(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + headers = headers or litellm.headers + ## LOAD CONFIG - if set + config = litellm.GenAIHubOrchestrationConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + return sap_gen_ai_hub_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + shared_session=shared_session, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + api_key=api_key, + api_base=api_base, + stream=stream, + ) + + +def _complete_aiohttp_openai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or get_secret("OPENAI_BASE_URL") + or get_secret("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + # set API KEY + api_key = ( + api_key + or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or get_secret("OPENAI_API_KEY") + ) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + return base_llm_aiohttp_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + +def _complete_cometapi(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.cometapi_key + or get_secret_str("COMETAPI_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("COMETAPI_API_BASE") + or "https://api.cometapi.com/v1" + ) + + ## COMPLETION CALL + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + + ## LOGGING + logging.post_call(input=messages, api_key=api_key, original_response=response) + + return response + + +def _complete_minimax(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("MINIMAX_API_BASE") + or "https://api.minimax.io/v1" + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + logging.post_call(input=messages, api_key=api_key, original_response=response) + + return response + + +def _complete_hosted_vllm(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE") + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + logging.post_call(input=messages, api_key=api_key, original_response=response) + + return response + + +def _complete_custom_openai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + metadata = ctx.metadata + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + organization = ctx.organization + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or get_secret("OPENAI_BASE_URL") + or get_secret("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + organization + or litellm.organization + or get_secret("OPENAI_ORGANIZATION") + or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 + ) + openai.organization = organization + # set API KEY + api_key = ( + api_key + or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or get_secret("OPENAI_API_KEY") + ) + + headers = headers or litellm.headers + + # Add GitHub Copilot headers (same as /responses endpoint does) + if custom_llm_provider == "github_copilot": + from litellm.llms.github_copilot.authenticator import Authenticator + from litellm.llms.github_copilot.common_utils import ( + get_copilot_default_headers, + ) + + copilot_auth = Authenticator() + copilot_api_key = copilot_auth.get_api_key() + copilot_headers = get_copilot_default_headers(copilot_api_key) + if extra_headers: + copilot_headers.update(extra_headers) + extra_headers = copilot_headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + + if ( + litellm.enable_preview_features and metadata is not None + ): # [PREVIEW] allow metadata to be passed to OPENAI + openai_metadata = get_requester_metadata(metadata) + if openai_metadata is not None: + optional_params["metadata"] = openai_metadata + + ## LOAD CONFIG - if set + config = litellm.OpenAIConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ## COMPLETION CALL + use_base_llm_http_handler = get_secret_bool( + "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER" + ) + + try: + if use_base_llm_http_handler: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + else: + response = openai_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + custom_prompt_dict=custom_prompt_dict, + client=client, # pass AsyncOpenAI, OpenAI client + organization=organization, + custom_llm_provider=custom_llm_provider, + shared_session=shared_session, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={"headers": headers}, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_mistral(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = api_key or litellm.api_key or get_secret("MISTRAL_API_KEY") + api_base = ( + api_base + or litellm.api_base + or get_secret("MISTRAL_API_BASE") + or "https://api.mistral.ai/v1" + ) + + return base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + + +def _complete_replicate(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + replicate_key = ( + api_key + or litellm.replicate_key + or litellm.api_key + or get_secret("REPLICATE_API_KEY") + or get_secret("REPLICATE_API_TOKEN") + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("REPLICATE_API_BASE") + or "https://api.replicate.com/v1" + ) + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + + model_response = replicate_chat_completion( # type: ignore + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), # for calculating input/output tokens + api_key=replicate_key, + logging_obj=logging, + custom_prompt_dict=custom_prompt_dict, + acompletion=acompletion, + headers=headers, + ) + + if optional_params.get("stream", False) is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=replicate_key, + original_response=model_response, + ) + + return model_response + + +def _complete_anthropic_text( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.anthropic_key + or litellm.api_key + or os.environ.get("ANTHROPIC_API_KEY") + ) + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + api_base = cast( + Optional[str], + api_base + or litellm.api_base + or get_secret("ANTHROPIC_API_BASE") + or get_secret("ANTHROPIC_BASE_URL") + or "https://api.anthropic.com/v1/complete", + ) + + # Check if we should disable automatic URL suffix appending + disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") + if ( + api_base is not None + and not disable_url_suffix + and not api_base.endswith("/v1/complete") + ): + api_base += "/v1/complete" + elif disable_url_suffix: + verbose_logger.debug( + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix" + ) + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="anthropic_text", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + +def _complete_anthropic(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.anthropic_key + or litellm.api_key + or os.environ.get("ANTHROPIC_API_KEY") + ) + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + # call /messages + # default route for all anthropic models + api_base = cast( + Optional[str], + api_base + or litellm.api_base + or get_secret("ANTHROPIC_API_BASE") + or get_secret("ANTHROPIC_BASE_URL") + or "https://api.anthropic.com/v1/messages", + ) + + # Check if we should disable automatic URL suffix appending + disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") + if ( + api_base is not None + and not disable_url_suffix + and not api_base.endswith("/v1/messages") + ): + api_base += "/v1/messages" + elif disable_url_suffix: + verbose_logger.debug( + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix" + ) + + response = anthropic_chat_completions.completion( + model=model, + messages=messages, + api_base=api_base, + acompletion=acompletion, + custom_prompt_dict=litellm.custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), # for calculating input/output tokens + api_key=api_key, + logging_obj=logging, + headers=headers, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + ) + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + ) + return response + + +def _complete_nlp_cloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + nlp_cloud_key = ( + api_key + or litellm.nlp_cloud_key + or get_secret("NLP_CLOUD_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("NLP_CLOUD_API_BASE") + or "https://api.nlpcloud.io/v1/gpu/" + ) + + response = nlp_cloud_chat_completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=nlp_cloud_key, + logging_obj=logging, + ) + + if "stream" in optional_params and optional_params["stream"] is True: + # don't try to access stream object, + response = CustomStreamWrapper( + response, + model, + custom_llm_provider="nlp_cloud", + logging_obj=logging, + ) + + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_aleph_alpha(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + api_key = ctx.api_key + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + aleph_alpha_key = ( + api_key + or litellm.aleph_alpha_key + or get_secret("ALEPH_ALPHA_API_KEY") + or get_secret("ALEPHALPHA_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("ALEPH_ALPHA_API_BASE") + or "https://api.aleph-alpha.com/complete" + ) + + model_response = aleph_alpha.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + default_max_tokens_to_sample=litellm.max_tokens, + api_key=aleph_alpha_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + if "stream" in optional_params and optional_params["stream"] is True: + # don't try to access stream object, + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="aleph_alpha", + logging_obj=logging, + ) + return model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_cohere_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + cohere_key = ( + api_key + or litellm.cohere_key + or get_secret_str("COHERE_API_KEY") + or get_secret_str("CO_API_KEY") + or litellm.api_key + ) + + cohere_route = CohereModelInfo.get_cohere_route(model) + verbose_logger.debug(f"Cohere route: {cohere_route}") + # Set API base based on route + if cohere_route == "v2": + api_base = ( + api_base + or litellm.api_base + or get_secret_str("COHERE_API_BASE") + or "https://api.cohere.com/v2/chat" + ) + # Remove v2/ prefix from model name for the actual API call + if "v2/" in model: + model = model.replace("v2/", "") + else: + api_base = ( + api_base + or litellm.api_base + or get_secret_str("COHERE_API_BASE") + or "https://api.cohere.ai/v1/chat" + ) + + headers = headers or litellm.headers or {} + if headers is None: + headers = {} + + if extra_headers is not None: + headers.update(extra_headers) + + verbose_logger.debug(f"Model: {model}, API Base: {api_base}") + verbose_logger.debug(f"Provider Config: {provider_config}") + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="cohere_chat", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=cohere_key, + provider_config=provider_config, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + +def _complete_maritalk(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + maritalk_key = ( + api_key + or litellm.maritalk_key + or get_secret("MARITALK_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("MARITALK_API_BASE") + or "https://chat.maritaca.ai/api" + ) + + return openai_like_chat_completion.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=maritalk_key, + logging_obj=logging, + custom_llm_provider="maritalk", + custom_prompt_dict=custom_prompt_dict, + ) + + +def _complete_amazon_nova(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + api_key = ctx.api_key + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.amazon_nova_api_key + or get_secret_str("AMAZON_NOVA_API_KEY") + or litellm.api_key + ) + api_base = ( + api_base + or litellm.api_base + or get_secret_str("AMAZON_NOVA_API_BASE") + or "https://api.nova.amazon.com/v1" + ) + return openai_like_chat_completion.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + custom_prompt_dict=custom_prompt_dict, + ) + + +def _complete_huggingface(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + huggingface_key = ( + api_key + or litellm.huggingface_key + or os.environ.get("HF_TOKEN") + or os.environ.get("HUGGINGFACE_API_KEY") + or litellm.api_key + ) + hf_headers = headers or litellm.headers + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=hf_headers, + model_response=model_response, + api_key=huggingface_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + +def _complete_oci(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + +def _complete_compactifai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + stream = ctx.stream + timeout = ctx.timeout + + api_key = api_key or get_secret_str("COMPACTIFAI_API_KEY") or litellm.api_key + + api_base = api_base or "https://api.compactif.ai/v1" + + ## COMPLETION CALL + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + + +def _complete_oobabooga(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + model_response = oobabooga.completion( + model=model, + messages=messages, + model_response=model_response, + api_base=api_base, # type: ignore + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=None, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + ) + if "stream" in optional_params and optional_params["stream"] is True: + # don't try to access stream object, + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="oobabooga", + logging_obj=logging, + ) + return model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_databricks(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for databricks we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or os.getenv("DATABRICKS_API_BASE") + ) + + # set API KEY + api_key = ( + api_key + or litellm.api_key # for databricks we check in get_llm_provider and pass in the api key from there + or litellm.databricks_key + or get_secret("DATABRICKS_API_KEY") + ) + + headers = headers or litellm.headers + + ## COMPLETION CALL + try: + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider="databricks", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={"headers": headers}, + ) + + return response + + +def _complete_datarobot(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + stream = ctx.stream + timeout = ctx.timeout + + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + + +def _complete_openrouter(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("OPENROUTER_API_BASE") + or "https://openrouter.ai/api/v1" + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret_str("OPENROUTER_API_KEY") + or get_secret_str("OR_API_KEY") + ) + + openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" + openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" + + openrouter_headers = { + "HTTP-Referer": openrouter_site_url, + "X-Title": openrouter_app_name, + } + + _headers = headers or litellm.headers + if _headers: + openrouter_headers.update(_headers) + + headers = openrouter_headers + + ## Load Config + config = litellm.OpenrouterConfig.get_config() + for k, v in config.items(): + if k == "extra_body": + # we use openai 'extra_body' to pass openrouter specific params - transforms, route, models + if "extra_body" in optional_params: + optional_params[k].update(v) + else: + optional_params[k] = v + elif k not in optional_params: + optional_params[k] = v + + ## COMPLETION CALL + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="openrouter", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + ## LOGGING + logging.post_call( + input=messages, api_key=openai.api_key, original_response=response + ) + + return response + + +def _complete_vercel_ai_gateway( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1" + ) + + api_key = api_key or litellm.api_key or get_secret("VERCEL_AI_GATEWAY_API_KEY") + + vercel_site_url = get_secret("VERCEL_SITE_URL") or "https://litellm.ai" + vercel_app_name = get_secret("VERCEL_APP_NAME") or "liteLLM" + + vercel_headers = { + "http-referer": vercel_site_url, + "x-title": vercel_app_name, + } + + _headers = headers or litellm.headers + if _headers: + vercel_headers.update(_headers) + + headers = vercel_headers + + ## Load Config + config = litellm.VercelAIGatewayConfig.get_config() + for k, v in config.items(): + if k == "extra_body": + # we use openai 'extra_body' to pass vercel specific params - providerOptions + if "extra_body" in optional_params: + optional_params[k].update(v) + else: + optional_params[k] = v + elif k not in optional_params: + optional_params[k] = v + + ## COMPLETION CALL + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="vercel_ai_gateway", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + ## LOGGING + logging.post_call( + input=messages, api_key=openai.api_key, original_response=response + ) + + return response + + +def _complete_vertex_ai_beta( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + vertex_ai_project = ( + optional_params.pop("vertex_project", None) + or optional_params.pop("vertex_ai_project", None) + or litellm.vertex_project + or get_secret("VERTEXAI_PROJECT") + ) + vertex_ai_location = ( + optional_params.pop("vertex_location", None) + or optional_params.pop("vertex_ai_location", None) + or litellm.vertex_location + or get_secret("VERTEXAI_LOCATION") + ) + vertex_credentials = ( + optional_params.pop("vertex_credentials", None) + or optional_params.pop("vertex_ai_credentials", None) + or get_secret("VERTEXAI_CREDENTIALS") + ) + + gemini_api_key = ( + api_key + or get_api_key_from_env() + or get_secret("PALM_API_KEY") # older palm api key should also work + or litellm.api_key + ) + + api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE") + new_params = safe_deep_copy(optional_params or {}) + return vertex_chat_completion.completion( # type: ignore + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + gemini_api_key=gemini_api_key, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + custom_llm_provider=custom_llm_provider, # type: ignore + client=client, + api_base=api_base, + extra_headers=headers, + ) + + +def _complete_vertex_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + vertex_ai_project = ( + optional_params.pop("vertex_project", None) + or optional_params.pop("vertex_ai_project", None) + or litellm.vertex_project + or get_secret("VERTEXAI_PROJECT") + ) + vertex_ai_location = ( + optional_params.pop("vertex_location", None) + or optional_params.pop("vertex_ai_location", None) + or litellm.vertex_location + or get_secret("VERTEXAI_LOCATION") + ) + vertex_credentials = ( + optional_params.pop("vertex_credentials", None) + or optional_params.pop("vertex_ai_credentials", None) + or get_secret("VERTEXAI_CREDENTIALS") + ) + + api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") + + new_params = safe_deep_copy(optional_params or {}) + model_route = get_vertex_ai_model_route(model=model, litellm_params=litellm_params) + + if model_route == VertexAIModelRoute.PARTNER_MODELS: + model_response = vertex_partner_models_chat_completion.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + api_base=api_base, + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + headers=headers, + custom_prompt_dict=custom_prompt_dict, + timeout=timeout, + client=client, + ) + elif model_route == VertexAIModelRoute.GEMINI: + model_response = vertex_chat_completion.completion( # type: ignore + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + gemini_api_key=None, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + custom_llm_provider=custom_llm_provider, # type: ignore + client=client, + api_base=api_base, + extra_headers=headers, + ) + elif model_route == VertexAIModelRoute.GEMMA: + # Vertex Gemma Models with custom prediction endpoint + model_response = vertex_gemma_chat_completion.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + api_base=api_base, + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + headers=headers, + custom_prompt_dict=custom_prompt_dict, + timeout=timeout, + client=client, + ) + elif model_route == VertexAIModelRoute.MODEL_GARDEN: + # Vertex Model Garden - OpenAI compatible models + model_response = vertex_model_garden_chat_completion.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + api_base=api_base, + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + headers=headers, + custom_prompt_dict=custom_prompt_dict, + timeout=timeout, + client=client, + ) + elif model_route == VertexAIModelRoute.AGENT_ENGINE: + # Vertex AI Agent Engine (Reasoning Engines) + from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, + ) + + vertex_agent_engine_config = VertexAgentEngineConfig() + + # Update litellm_params with vertex credentials + litellm_params["vertex_project"] = vertex_ai_project + litellm_params["vertex_location"] = vertex_ai_location + litellm_params["vertex_credentials"] = vertex_credentials + + model_response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + model_response=model_response, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + encoding=_get_encoding(), + api_key=None, + api_base=api_base, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + client=client, + custom_llm_provider="vertex_ai", + provider_config=vertex_agent_engine_config, + headers=headers or {}, + ) + else: # VertexAIModelRoute.NON_GEMINI + model_response = vertex_ai_non_gemini.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + ) + + if ( + "stream" in optional_params + and optional_params["stream"] is True + and acompletion is False + ): + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="vertex_ai", + logging_obj=logging, + ) + return model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_predibase(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + tenant_id = ( + optional_params.pop("tenant_id", None) + or optional_params.pop("predibase_tenant_id", None) + or litellm.predibase_tenant_id + or get_secret("PREDIBASE_TENANT_ID") + ) + + if tenant_id is None: + raise ValueError( + "Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=)`) or in env - `PREDIBASE_TENANT_ID`." + ) + + api_base = ( + api_base + or optional_params.pop("api_base", None) + or optional_params.pop("base_url", None) + or litellm.api_base + or get_secret("PREDIBASE_API_BASE") + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.predibase_key + or get_secret("PREDIBASE_API_KEY") + ) + + _model_response = predibase_chat_completions.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + acompletion=acompletion, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + api_key=api_key, + tenant_id=tenant_id, + timeout=timeout, + ) + + if ( + "stream" in optional_params + and optional_params["stream"] is True + and acompletion is False + ): + return _model_response + return _model_response + + +def _complete_text_completion_codestral( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base + or optional_params.pop("api_base", None) + or optional_params.pop("base_url", None) + or litellm.api_base + or "https://codestral.mistral.ai/v1/fim/completions" + ) + + api_key = api_key or litellm.api_key or get_secret("CODESTRAL_API_KEY") + + text_completion_model_response = litellm.TextCompletionResponse(stream=stream) + + _model_response = codestral_text_completions.completion( # type: ignore + model=model, + messages=messages, + model_response=text_completion_model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + acompletion=acompletion, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + api_key=api_key, + timeout=timeout, + ) + + if ( + "stream" in optional_params + and optional_params["stream"] is True + and acompletion is False + ): + return _model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + return _model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_text_completion_inception( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + text_completion = ctx.text_completion + timeout = ctx.timeout + + passed_api_base = ( + api_base + or optional_params.pop("api_base", None) + or optional_params.pop("base_url", None) + ) + api_base = ( + passed_api_base + or get_secret_str("INCEPTION_API_BASE") + or "https://api.inceptionlabs.ai/v1" + ) + # FIM is served at `/v1/fim/completions`; the OpenAI client appends + # `/completions`, so point it at the `/v1/fim` base. + api_base = api_base.rstrip("/") + if not api_base.endswith("/fim"): + api_base += "/fim" + + # Don't forward the server-managed Inception key to a caller-supplied + # api_base; only resolve it for the default/server base, or when the + # caller passes their own key. + if passed_api_base is None or api_key: + api_key = ( + api_key or litellm.inception_key or get_secret_str("INCEPTION_API_KEY") + ) + + _response = openai_text_completions.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, # type: ignore[arg-type] + custom_llm_provider="text-completion-inception", + api_base=api_base, + acompletion=acompletion, + client=client, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + ) + + if ( + optional_params.get("stream", False) is False + and acompletion is False + and text_completion is False + ): + _response = ( + litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( + response_object=_response, model_response_object=model_response + ) + ) + + if optional_params.get("stream", False) or acompletion is True: + logging.post_call( + input=messages, + api_key=api_key, + original_response=_response, + additional_args={"headers": headers}, + ) + return _response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_sagemaker_chat( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_sagemaker(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + custom_prompt_dict = ctx.custom_prompt_dict + hf_model_name = ctx.hf_model_name + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + return sagemaker_llm.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + custom_prompt_dict=custom_prompt_dict, + hf_model_name=hf_model_name, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + acompletion=acompletion, + ) + + +def _complete_bedrock(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + + if "aws_bedrock_client" in optional_params: + verbose_logger.warning( + "'aws_bedrock_client' is a deprecated param. Please move to another auth method - https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication." + ) + # Extract credentials for legacy boto3 client and pass thru to httpx + aws_bedrock_client = optional_params.pop("aws_bedrock_client") + creds = aws_bedrock_client._get_credentials().get_frozen_credentials() + + if creds.access_key: + optional_params["aws_access_key_id"] = creds.access_key + if creds.secret_key: + optional_params["aws_secret_access_key"] = creds.secret_key + if creds.token: + optional_params["aws_session_token"] = creds.token + if ( + "aws_region_name" not in optional_params + or optional_params["aws_region_name"] is None + ): + optional_params["aws_region_name"] = aws_bedrock_client.meta.region_name + + bedrock_route = BedrockModelInfo.get_bedrock_route(model) + if bedrock_route == "claude_platform": + provider_config = ProviderConfigManager.get_provider_chat_config( + model=model, + provider=LlmProviders.BEDROCK, + ) + model = BedrockModelInfo.get_claude_platform_model(model) + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="bedrock", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + provider_config=provider_config, + ) + elif bedrock_route == "converse": + model = model.replace("converse/", "") + response = bedrock_converse_chat_completion.completion( + model=model, + messages=messages, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + extra_headers=headers, # Use merged headers instead of original extra_headers + timeout=timeout, + acompletion=acompletion, + client=client, + api_base=api_base, + api_key=api_key, + ) + elif bedrock_route == "converse_like": + model = model.replace("converse_like/", "") + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider="bedrock", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + else: + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider="bedrock", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + return response + + +def _complete_watsonx(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + return watsonx_chat_completion.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + custom_prompt_dict=custom_prompt_dict, + client=client, # pass AsyncOpenAI, OpenAI client + encoding=_get_encoding(), + custom_llm_provider="watsonx", + ) + + +def _complete_watsonx_text( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or optional_params.pop("apikey", None) + or get_secret_str("WATSONX_APIKEY") + or get_secret_str("WATSONX_API_KEY") + or get_secret_str("WX_API_KEY") + ) + + api_base = ( + api_base + or optional_params.pop( + "url", + optional_params.pop("api_base", optional_params.pop("base_url", None)), + ) + or get_secret_str("WATSONX_API_BASE") + or get_secret_str("WATSONX_URL") + or get_secret_str("WX_URL") + or get_secret_str("WML_URL") + ) + + wx_credentials = optional_params.pop( + "wx_credentials", + optional_params.pop( + "watsonx_credentials", None + ), # follow {provider}_credentials, same as vertex ai + ) + + token: Optional[str] = None + if wx_credentials is not None: + api_base = wx_credentials.get("url", api_base) + api_key = wx_credentials.get("apikey", wx_credentials.get("api_key", api_key)) + token = wx_credentials.get( + "token", + wx_credentials.get( + "watsonx_token", None + ), # follow format of {provider}_token, same as azure - e.g. 'azure_ad_token=..' + ) + + if token is not None: + optional_params["token"] = token + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="watsonx_text", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_vllm(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + model_response = vllm_handler.completion( + model=model, + messages=messages, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + ) + + if "stream" in optional_params and optional_params["stream"] is True: ## [BETA] + # don't try to access stream object, + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="vllm", + logging_obj=logging, + ) + + ## RESPONSE OBJECT + return model_response + + +def _complete_ollama(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + litellm.api_base + or api_base + or get_secret("OLLAMA_API_BASE") + or "http://localhost:11434" + ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="ollama", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_ollama_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + litellm.api_base + or api_base + or get_secret("OLLAMA_API_BASE") + or "http://localhost:11434" + ) + + api_key = ( + api_key + or litellm.ollama_key + or os.environ.get("OLLAMA_API_KEY") + or litellm.api_key + ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="ollama_chat", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_triton(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = litellm.api_base or api_base + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + ) + + +def _complete_cloudflare(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.cloudflare_api_key + or litellm.api_key + or get_secret("CLOUDFLARE_API_KEY") + ) + api_base = api_base or litellm.api_base or get_secret("CLOUDFLARE_API_BASE") + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="cloudflare", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + +def _complete_petals(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + client = ctx.client + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + + api_base = api_base or litellm.api_base + + stream = optional_params.pop("stream", False) + model_response = petals_handler.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + client=client, + ) + if stream is True: ## [BETA] + # Fake streaming for petals + resp_string = model_response["choices"][0]["message"]["content"] + return CustomStreamWrapper( + resp_string, + model, + custom_llm_provider="petals", + logging_obj=logging, + ) + return model_response + + +def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + client = ( + HTTPHandler(timeout=timeout) if stream is False else None + ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_gradient_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = litellm.api_base or api_base + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="gradient_ai", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + ) + + +def _complete_bytez(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.bytez_key + or get_secret_str("BYTEZ_API_KEY") + or litellm.api_key + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=bytez_transformation, + ) + + pass + + return response + + +def _complete_lemonade(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.lemonade_key + or get_secret_str("LEMONADE_API_KEY") + or litellm.api_key + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=lemonade_transformation, + ) + + pass + + return response + + +def _complete_ovhcloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.ovhcloud_key + or get_secret_str("OVHCLOUD_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("OVHCLOUD_API_BASE") + or "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=ovhcloud_transformation, + ) + + pass + + return response + + +def _complete_custom(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + headers = ctx.headers + kwargs = ctx.kwargs + max_tokens = ctx.max_tokens + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + temperature = ctx.temperature + top_p = ctx.top_p + + url = litellm.api_base or api_base or "" + if url is None or url == "": + raise ValueError( + "api_base not set. Set api_base or litellm.api_base for custom endpoints" + ) + + """ + assume input to custom LLM api bases follow this format: + resp = litellm.module_level_client.post( + api_base, + json={ + 'model': 'meta-llama/Llama-2-13b-hf', # model name + 'params': { + 'prompt': ["The capital of France is P"], + 'max_tokens': 32, + 'temperature': 0.7, + 'top_p': 1.0, + 'top_k': 40, + } + } + ) + + """ + prompt = " ".join([message["content"] for message in messages]) # type: ignore + resp = litellm.module_level_client.post( + url, + headers=headers, + json={ + "model": model, + "params": { + "prompt": [prompt], + "max_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, + "top_k": kwargs.get("top_k"), + }, + **kwargs.get("extra_body", {}), + }, + ) + response_json = resp.json() + """ + assume all responses from custom api_bases of this format: + { + 'data': [ + { + 'prompt': 'The capital of France is P', + 'output': ['The capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France'], + 'params': {'temperature': 0.7, 'top_k': 40, 'top_p': 1}}], + 'message': 'ok' + } + ] + } + """ + string_response = response_json["data"][0]["output"][0] + ## RESPONSE OBJECT + model_response.choices[0].message.content = string_response # type: ignore + model_response.created = int(time.time()) + model_response.model = model + return model_response + + +def _complete_custom_providers( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + custom_handler: Optional[CustomLLM] = None + for item in litellm.custom_provider_map: + if item["provider"] == custom_llm_provider: + custom_handler = item["custom_handler"] + + if custom_handler is None: + raise LiteLLMUnknownProvider( + model=model, custom_llm_provider=custom_llm_provider + ) + + ## ROUTE LLM CALL ## + handler_fn = custom_chat_llm_router( + async_fn=acompletion, stream=stream, custom_llm=custom_handler + ) + + headers = headers or litellm.headers or {} + + ## CALL FUNCTION + response = handler_fn( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + custom_prompt_dict=custom_prompt_dict, + client=client, # pass AsyncOpenAI, OpenAI client + encoding=_get_encoding(), + ) + if stream is True: + return CustomStreamWrapper( + completion_stream=response, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_langgraph(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + from litellm.llms.langgraph.chat.transformation import LangGraphConfig + + ( + api_base, + api_key, + ) = LangGraphConfig()._get_openai_compatible_provider_info( + api_base=api_base or litellm.api_base, + api_key=api_key or litellm.api_key, + ) + + headers = headers or litellm.headers + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + +def _complete_langflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + from litellm.llms.langflow.chat.transformation import LangFlowConfig + + ( + api_base, + api_key, + ) = LangFlowConfig()._get_openai_compatible_provider_info( + api_base=api_base or litellm.api_base, + api_key=api_key or litellm.api_key, + ) + + headers = headers or litellm.headers + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + @tracer.wrap() @client def completion( # type: ignore @@ -1215,9 +5077,7 @@ def completion( # type: ignore if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( tools=tools_for_mcp ): - # Return coroutine - acompletion will await it - # completion() can return a coroutine when MCP tools are present, which acompletion() awaits - return acompletion_with_mcp( # type: ignore[return-value] + return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it model=model, messages=messages, functions=functions, @@ -1389,12 +5249,16 @@ def completion( # type: ignore logging: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, litellm_logging_obj) fallbacks = fallbacks or litellm.model_fallbacks if fallbacks is not None: - return completion_with_fallbacks(**args) + return completion_with_fallbacks( # pyright: ignore[reportReturnType] # fallback runner is untyped; resolves to ModelResponse|CustomStreamWrapper at runtime + **args + ) if model_list is not None: deployments = [ m["litellm_params"] for m in model_list if m["model_name"] == model ] - return litellm.batch_completion_models(deployments=deployments, **args) + return litellm.batch_completion_models( # pyright: ignore[reportReturnType] # batch path returns a list of responses, outside completion()'s single-response return type + deployments=deployments, **args + ) if litellm.model_alias_map and model in litellm.model_alias_map: model = litellm.model_alias_map[ model @@ -1454,7 +5318,7 @@ def completion( # type: ignore timeout, kwargs, custom_llm_provider, - global_timeout=getattr(litellm, "request_timeout", None), + global_timeout=get_configured_request_timeout(), supports_httpx_timeout=supports_httpx_timeout, ) @@ -1584,8 +5448,9 @@ def completion( # type: ignore provider_config=provider_config, ) - if litellm.add_function_to_prompt and optional_params.get( - "functions_unsupported_model", None + if ( + litellm.add_function_to_prompt + and optional_params.get("functions_unsupported_model", None) ): # if user opts to add it to prompt, when API doesn't support function calling functions_unsupported_model = optional_params.pop( "functions_unsupported_model" @@ -1716,7 +5581,7 @@ def completion( # type: ignore else: optional_params["reasoning_effort"] = {"summary": rs_val} - return responses_api_bridge.completion( + return responses_api_bridge.completion( # pyright: ignore[reportReturnType] # bridge returns a coroutine on the acompletion path; awaited by the async caller model=model, messages=messages, headers=headers, @@ -1746,375 +5611,52 @@ def completion( # type: ignore optional_params ) + _dispatch_ctx = _CompletionDispatchContext( + _azure_detection_model=_azure_detection_model, + acompletion=acompletion, + api_base=api_base, + api_key=api_key, + api_version=api_version, + client=client, + custom_llm_provider=custom_llm_provider, + custom_prompt_dict=custom_prompt_dict, + extra_headers=extra_headers, + headers=headers, + hf_model_name=hf_model_name, + kwargs=kwargs, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging=logging, + max_retries=max_retries, + max_tokens=max_tokens, + messages=messages, + metadata=metadata, + model=model, + model_response=model_response, + optional_params=optional_params, + organization=organization, + provider_config=provider_config, + shared_session=shared_session, + stream=stream, + temperature=temperature, + text_completion=text_completion, + timeout=timeout, + top_p=top_p, + ) if custom_llm_provider == "azure": # azure configs ## check dynamic params ## - dynamic_params = False - if client is not None and ( - isinstance(client, openai.AzureOpenAI) - or isinstance(client, openai.AsyncAzureOpenAI) - ): - dynamic_params = _check_dynamic_azure_params( - azure_client_params={"api_version": api_version}, - azure_client=client, - ) - - api_type = get_secret("AZURE_API_TYPE") or "azure" - - api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") - - api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - or litellm.AZURE_DEFAULT_API_VERSION - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) - - azure_ad_token = optional_params.get("extra_body", {}).pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") - - azure_ad_token_provider = litellm_params.get( - "azure_ad_token_provider", None - ) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - if max_retries is not None: - optional_params["max_retries"] = max_retries - - if litellm.AzureOpenAIO1Config().is_o_series_model( - model=_azure_detection_model - ): - ## LOAD CONFIG - if set - config = litellm.AzureOpenAIO1Config.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - response = azure_o1_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - api_key=api_key, - api_base=api_base, - api_version=api_version, - dynamic_params=dynamic_params, - azure_ad_token=azure_ad_token, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, # type: ignore - client=client, # pass AsyncAzureOpenAI, AzureOpenAI client - custom_llm_provider=custom_llm_provider, - ) - else: - ## LOAD CONFIG - if set - config = litellm.AzureOpenAIConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ## COMPLETION CALL - response = azure_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - api_key=api_key, - api_base=api_base, - api_version=api_version, - api_type=api_type, - dynamic_params=dynamic_params, - azure_ad_token=azure_ad_token, - azure_ad_token_provider=azure_ad_token_provider, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, # type: ignore - client=client, # pass AsyncAzureOpenAI, AzureOpenAI client - ) - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={ - "headers": headers, - "api_version": api_version, - "api_base": api_base, - }, - ) + response = _complete_azure(_dispatch_ctx) elif custom_llm_provider == "azure_text": # azure configs - api_type = get_secret_str("AZURE_API_TYPE") or "azure" - - api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - - if api_base is None: - raise ValueError( - "api_base is required for Azure OpenAI LLM provider. Either set it dynamically or set the AZURE_API_BASE environment variable." - ) - - api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) - - azure_ad_token = optional_params.get("extra_body", {}).pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") - - azure_ad_token_provider = litellm_params.get( - "azure_ad_token_provider", None - ) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - - ## LOAD CONFIG - if set - config = litellm.AzureOpenAIConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ## COMPLETION CALL - response = azure_text_completions.completion( - model=model, - messages=messages, - headers=headers, - api_key=api_key, - api_base=api_base, - api_version=cast(str, api_version), - api_type=api_type, - azure_ad_token=azure_ad_token, - azure_ad_token_provider=azure_ad_token_provider, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - client=client, # pass AsyncAzureOpenAI, AzureOpenAI client - ) - - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={ - "headers": headers, - "api_version": api_version, - "api_base": api_base, - }, - ) + response = _complete_azure_text(_dispatch_ctx) elif custom_llm_provider == "deepseek": ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_deepseek(_dispatch_ctx) elif custom_llm_provider == "azure_ai": - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo - - azure_ai_route = AzureFoundryModelInfo.get_azure_ai_route(model) - - # Check if this is an agents route - model format: azure_ai/agents/ - if azure_ai_route == "agents": - from litellm.llms.azure_ai.agents import AzureAIAgentsConfig - - api_base = AzureFoundryModelInfo.get_api_base(api_base) - if api_base is None: - raise ValueError( - "Azure AI Agents requests require an api_base. " - "Set `api_base` or the AZURE_AI_API_BASE env var." - ) - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - response = AzureAIAgentsConfig.completion( - model=model, - messages=messages, - api_base=api_base, - api_key=api_key, - model_response=model_response, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, - acompletion=acompletion, - stream=stream, - headers=headers or litellm.headers, - ) - - # Check if this is a Claude model - route to Azure Anthropic handler - elif "claude" in model.lower(): - # Use Azure Anthropic handler for Claude models - api_base = AzureFoundryModelInfo.get_api_base(api_base) - if api_base is None: - raise ValueError( - "Azure Anthropic requests require an api_base. " - "Set `api_base` or the AZURE_AI_API_BASE env var." - ) - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - # Ensure the URL ends with /v1/messages for Anthropic - if api_base: - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/messages"): - if "/anthropic" in api_base: - parts = api_base.split("/anthropic", 1) - api_base = parts[0] + "/anthropic" - else: - api_base = api_base + "/anthropic" - api_base = api_base + "/v1/messages" - - response = azure_anthropic_chat_completions.completion( - model=model, - messages=messages, - api_base=api_base, - acompletion=acompletion, - custom_prompt_dict=litellm.custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - headers=headers, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - ) - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - ) - response = response - else: - # Non-Claude models use standard Azure AI flow - api_base = AzureFoundryModelInfo.get_api_base(api_base) - # set API KEY - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - - ## FOR COHERE - if "command-r" in model: # make sure tool call in messages are str - messages = stringify_json_tool_call_content(messages=messages) - - ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, # pass AsyncOpenAI, OpenAI client - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={"headers": headers}, - ) + response = _complete_azure_ai(_dispatch_ctx) elif ( custom_llm_provider == "text-completion-openai" or "ft:babbage-002" in model @@ -2123,535 +5665,42 @@ def completion( # type: ignore in litellm.openai_text_completion_compatible_providers and kwargs.get("text_completion") is True ): - openai.api_type = "openai" - - api_base = ( - api_base - or litellm.api_base - or get_secret("OPENAI_BASE_URL") - or get_secret("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - - openai.api_version = None - # set API KEY - - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) - - headers = headers or litellm.headers - - ## LOAD CONFIG - if set - config = litellm.OpenAITextCompletionConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_text_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - if litellm.organization: - openai.organization = litellm.organization - - if ( - len(messages) > 0 - and "content" in messages[0] - and isinstance(messages[0]["content"], list) - ): - # text-davinci-003 can accept a string or array, if it's an array, assume the array is set in messages[0]['content'] - # https://platform.openai.com/docs/api-reference/completions/create - prompt = messages[0]["content"] - else: - prompt = " ".join([message["content"] for message in messages]) # type: ignore - - ## COMPLETION CALL - _response = openai_text_completions.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - acompletion=acompletion, - client=client, # pass AsyncOpenAI, OpenAI client - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - ) - - if ( - optional_params.get("stream", False) is False - and acompletion is False - and text_completion is False - ): - # convert to chat completion response - _response = litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( - response_object=_response, model_response_object=model_response - ) - - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=_response, - additional_args={"headers": headers}, - ) - response = _response + response = _complete_text_completion_openai(_dispatch_ctx) elif custom_llm_provider == "fireworks_ai": ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_fireworks_ai(_dispatch_ctx) elif custom_llm_provider == "heroku": - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_heroku(_dispatch_ctx) elif custom_llm_provider == "ragflow": ## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_ragflow(_dispatch_ctx) elif custom_llm_provider == "xai": ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_xai(_dispatch_ctx) elif custom_llm_provider == "groq": - api_base = ( - api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or get_secret("GROQ_API_BASE") - or "https://api.groq.com/openai/v1" - ) - - # set API KEY - api_key = ( - api_key - or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.groq_key - or get_secret("GROQ_API_KEY") - ) - - headers = headers or litellm.headers - - ## LOAD CONFIG - if set - config = litellm.GroqChatConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_groq(_dispatch_ctx) elif custom_llm_provider == "bedrock_mantle": - api_base = ( - api_base or litellm.api_base or get_secret("BEDROCK_MANTLE_API_BASE") - ) - api_key = api_key or litellm.api_key or get_secret("BEDROCK_MANTLE_API_KEY") - headers = headers or litellm.headers - config = litellm.BedrockMantleChatConfig.get_config() - for k, v in config.items(): - if k not in optional_params: - optional_params[k] = v - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_bedrock_mantle(_dispatch_ctx) elif custom_llm_provider == "a2a": # A2A (Agent-to-Agent) Protocol # Resolve agent configuration from registry if model format is "a2a/" - ( - api_base, - api_key, - headers, - ) = litellm.A2AConfig.resolve_agent_config_from_registry( - model=model, - api_base=api_base, - api_key=api_key, - headers=headers, - optional_params=optional_params, - ) - - # Fall back to environment variables and defaults - api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE") - - if api_base is None: - raise Exception( - "api_base is required for A2A provider. " - "Either provide api_base parameter, set A2A_API_BASE environment variable, " - "or register the agent in the proxy with model='a2a/'." - ) - - headers = headers or litellm.headers - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - provider_config=provider_config, - ) + response = _complete_a2a(_dispatch_ctx) elif custom_llm_provider == "gigachat": # GigaChat - Sber AI's LLM (Russia) - api_key = ( - api_key - or litellm.api_key - or litellm.gigachat_key - or get_secret("GIGACHAT_API_KEY") - or get_secret("GIGACHAT_CREDENTIALS") - ) - - headers = headers or litellm.headers or {} - - ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_gigachat(_dispatch_ctx) elif custom_llm_provider == "sap": - headers = headers or litellm.headers - ## LOAD CONFIG - if set - config = litellm.GenAIHubOrchestrationConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - response = sap_gen_ai_hub_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - shared_session=shared_session, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - api_key=api_key, - api_base=api_base, - stream=stream, - ) + response = _complete_sap(_dispatch_ctx) elif custom_llm_provider == "aiohttp_openai": # NEW aiohttp provider for 10-100x higher RPS - api_base = ( - api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or get_secret("OPENAI_BASE_URL") - or get_secret("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - # set API KEY - api_key = ( - api_key - or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - response = base_llm_aiohttp_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) + response = _complete_aiohttp_openai(_dispatch_ctx) elif custom_llm_provider == "cometapi": - api_key = ( - api_key - or litellm.cometapi_key - or get_secret_str("COMETAPI_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COMETAPI_API_BASE") - or "https://api.cometapi.com/v1" - ) - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - - ## LOGGING - logging.post_call( - input=messages, api_key=api_key, original_response=response - ) + response = _complete_cometapi(_dispatch_ctx) elif custom_llm_provider == "minimax": - api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key - - api_base = ( - api_base - or litellm.api_base - or get_secret_str("MINIMAX_API_BASE") - or "https://api.minimax.io/v1" - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) - logging.post_call( - input=messages, api_key=api_key, original_response=response - ) + response = _complete_minimax(_dispatch_ctx) elif custom_llm_provider == "hosted_vllm": - api_base = ( - api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE") - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) - logging.post_call( - input=messages, api_key=api_key, original_response=response - ) + response = _complete_hosted_vllm(_dispatch_ctx) elif ( model in litellm.open_ai_chat_completion_models or custom_llm_provider == "custom_openai" @@ -2676,205 +5725,17 @@ def completion( # type: ignore ): # allow user to make an openai call with a custom base # note: if a user sets a custom base - we should ensure this works # allow for the setting of dynamic and stateful api-bases - api_base = ( - api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or get_secret("OPENAI_BASE_URL") - or get_secret("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - organization - or litellm.organization - or get_secret("OPENAI_ORGANIZATION") - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - openai.organization = organization - # set API KEY - api_key = ( - api_key - or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) - - headers = headers or litellm.headers - - # Add GitHub Copilot headers (same as /responses endpoint does) - if custom_llm_provider == "github_copilot": - from litellm.llms.github_copilot.authenticator import Authenticator - from litellm.llms.github_copilot.common_utils import ( - get_copilot_default_headers, - ) - - copilot_auth = Authenticator() - copilot_api_key = copilot_auth.get_api_key() - copilot_headers = get_copilot_default_headers(copilot_api_key) - if extra_headers: - copilot_headers.update(extra_headers) - extra_headers = copilot_headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - - if ( - litellm.enable_preview_features and metadata is not None - ): # [PREVIEW] allow metadata to be passed to OPENAI - openai_metadata = get_requester_metadata(metadata) - if openai_metadata is not None: - optional_params["metadata"] = openai_metadata - - ## LOAD CONFIG - if set - config = litellm.OpenAIConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ## COMPLETION CALL - use_base_llm_http_handler = get_secret_bool( - "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER" - ) - - try: - if use_base_llm_http_handler: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) - else: - response = openai_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - custom_prompt_dict=custom_prompt_dict, - client=client, # pass AsyncOpenAI, OpenAI client - organization=organization, - custom_llm_provider=custom_llm_provider, - shared_session=shared_session, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={"headers": headers}, - ) + response = _complete_custom_openai(_dispatch_ctx) elif custom_llm_provider == "mistral": - api_key = api_key or litellm.api_key or get_secret("MISTRAL_API_KEY") - api_base = ( - api_base - or litellm.api_base - or get_secret("MISTRAL_API_BASE") - or "https://api.mistral.ai/v1" - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) + response = _complete_mistral(_dispatch_ctx) elif ( "replicate" in model or custom_llm_provider == "replicate" or model in litellm.replicate_models ): # Setting the relevant API KEY for replicate, replicate defaults to using os.environ.get("REPLICATE_API_TOKEN") - replicate_key = ( - api_key - or litellm.replicate_key - or litellm.api_key - or get_secret("REPLICATE_API_KEY") - or get_secret("REPLICATE_API_TOKEN") - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("REPLICATE_API_BASE") - or "https://api.replicate.com/v1" - ) - - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - - model_response = replicate_chat_completion( # type: ignore - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), # for calculating input/output tokens - api_key=replicate_key, - logging_obj=logging, - custom_prompt_dict=custom_prompt_dict, - acompletion=acompletion, - headers=headers, - ) - - if optional_params.get("stream", False) is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=replicate_key, - original_response=model_response, - ) - - response = model_response + response = _complete_replicate(_dispatch_ctx) elif ( "clarifai" in model or custom_llm_provider == "clarifai" @@ -2882,614 +5743,36 @@ def completion( # type: ignore ): pass # Deprecated - handled in the openai compatible provider section above elif custom_llm_provider == "anthropic_text": - api_key = ( - api_key - or litellm.anthropic_key - or litellm.api_key - or os.environ.get("ANTHROPIC_API_KEY") - ) - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - api_base = ( - api_base - or litellm.api_base - or get_secret("ANTHROPIC_API_BASE") - or get_secret("ANTHROPIC_BASE_URL") - or "https://api.anthropic.com/v1/complete" - ) - - # Check if we should disable automatic URL suffix appending - disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if ( - api_base is not None - and not disable_url_suffix - and not api_base.endswith("/v1/complete") - ): - api_base += "/v1/complete" - elif disable_url_suffix: - verbose_logger.debug( - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix" - ) - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="anthropic_text", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) + response = _complete_anthropic_text(_dispatch_ctx) elif custom_llm_provider == "anthropic": - api_key = ( - api_key - or litellm.anthropic_key - or litellm.api_key - or os.environ.get("ANTHROPIC_API_KEY") - ) - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - # call /messages - # default route for all anthropic models - api_base = ( - api_base - or litellm.api_base - or get_secret("ANTHROPIC_API_BASE") - or get_secret("ANTHROPIC_BASE_URL") - or "https://api.anthropic.com/v1/messages" - ) - - # Check if we should disable automatic URL suffix appending - disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if ( - api_base is not None - and not disable_url_suffix - and not api_base.endswith("/v1/messages") - ): - api_base += "/v1/messages" - elif disable_url_suffix: - verbose_logger.debug( - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix" - ) - - response = anthropic_chat_completions.completion( - model=model, - messages=messages, - api_base=api_base, - acompletion=acompletion, - custom_prompt_dict=litellm.custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), # for calculating input/output tokens - api_key=api_key, - logging_obj=logging, - headers=headers, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - ) - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - ) - response = response + response = _complete_anthropic(_dispatch_ctx) elif custom_llm_provider == "nlp_cloud": - nlp_cloud_key = ( - api_key - or litellm.nlp_cloud_key - or get_secret("NLP_CLOUD_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("NLP_CLOUD_API_BASE") - or "https://api.nlpcloud.io/v1/gpu/" - ) - - response = nlp_cloud_chat_completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=nlp_cloud_key, - logging_obj=logging, - ) - - if "stream" in optional_params and optional_params["stream"] is True: - # don't try to access stream object, - response = CustomStreamWrapper( - response, - model, - custom_llm_provider="nlp_cloud", - logging_obj=logging, - ) - - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - ) - - response = response + response = _complete_nlp_cloud(_dispatch_ctx) elif custom_llm_provider == "aleph_alpha": - aleph_alpha_key = ( - api_key - or litellm.aleph_alpha_key - or get_secret("ALEPH_ALPHA_API_KEY") - or get_secret("ALEPHALPHA_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("ALEPH_ALPHA_API_BASE") - or "https://api.aleph-alpha.com/complete" - ) - - model_response = aleph_alpha.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - default_max_tokens_to_sample=litellm.max_tokens, - api_key=aleph_alpha_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) - - if "stream" in optional_params and optional_params["stream"] is True: - # don't try to access stream object, - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="aleph_alpha", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_aleph_alpha(_dispatch_ctx) elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere": - cohere_key = ( - api_key - or litellm.cohere_key - or get_secret_str("COHERE_API_KEY") - or get_secret_str("CO_API_KEY") - or litellm.api_key - ) - - cohere_route = CohereModelInfo.get_cohere_route(model) - verbose_logger.debug(f"Cohere route: {cohere_route}") - # Set API base based on route - if cohere_route == "v2": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COHERE_API_BASE") - or "https://api.cohere.com/v2/chat" - ) - # Remove v2/ prefix from model name for the actual API call - if "v2/" in model: - model = model.replace("v2/", "") - else: - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COHERE_API_BASE") - or "https://api.cohere.ai/v1/chat" - ) - - headers = headers or litellm.headers or {} - if headers is None: - headers = {} - - if extra_headers is not None: - headers.update(extra_headers) - - verbose_logger.debug(f"Model: {model}, API Base: {api_base}") - verbose_logger.debug(f"Provider Config: {provider_config}") - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="cohere_chat", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=cohere_key, - provider_config=provider_config, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) + response = _complete_cohere_chat(_dispatch_ctx) elif custom_llm_provider == "maritalk": - maritalk_key = ( - api_key - or litellm.maritalk_key - or get_secret("MARITALK_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("MARITALK_API_BASE") - or "https://chat.maritaca.ai/api" - ) - - model_response = openai_like_chat_completion.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=maritalk_key, - logging_obj=logging, - custom_llm_provider="maritalk", - custom_prompt_dict=custom_prompt_dict, - ) - - response = model_response + response = _complete_maritalk(_dispatch_ctx) elif custom_llm_provider == "amazon_nova": - api_key = ( - api_key - or litellm.amazon_nova_api_key - or get_secret_str("AMAZON_NOVA_API_KEY") - or litellm.api_key - ) - api_base = ( - api_base - or litellm.api_base - or get_secret_str("AMAZON_NOVA_API_BASE") - or "https://api.nova.amazon.com/v1" - ) - response = openai_like_chat_completion.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - custom_prompt_dict=custom_prompt_dict, - ) + response = _complete_amazon_nova(_dispatch_ctx) elif custom_llm_provider == "huggingface": - huggingface_key = ( - api_key - or litellm.huggingface_key - or os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_API_KEY") - or litellm.api_key - ) - hf_headers = headers or litellm.headers - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=hf_headers, - model_response=model_response, - api_key=huggingface_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) + response = _complete_huggingface(_dispatch_ctx) elif custom_llm_provider == "oci": - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) + response = _complete_oci(_dispatch_ctx) elif custom_llm_provider == "compactifai": - api_key = ( - api_key or get_secret_str("COMPACTIFAI_API_KEY") or litellm.api_key - ) - - api_base = api_base or "https://api.compactif.ai/v1" - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) + response = _complete_compactifai(_dispatch_ctx) elif custom_llm_provider == "oobabooga": - custom_llm_provider = "oobabooga" - model_response = oobabooga.completion( - model=model, - messages=messages, - model_response=model_response, - api_base=api_base, # type: ignore - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - api_key=None, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - ) - if "stream" in optional_params and optional_params["stream"] is True: - # don't try to access stream object, - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="oobabooga", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_oobabooga(_dispatch_ctx) elif custom_llm_provider == "databricks": - api_base = ( - api_base # for databricks we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or os.getenv("DATABRICKS_API_BASE") - ) - - # set API KEY - api_key = ( - api_key - or litellm.api_key # for databricks we check in get_llm_provider and pass in the api key from there - or litellm.databricks_key - or get_secret("DATABRICKS_API_KEY") - ) - - headers = headers or litellm.headers - - ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="databricks", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={"headers": headers}, - ) + response = _complete_databricks(_dispatch_ctx) elif custom_llm_provider == "datarobot": - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) + response = _complete_datarobot(_dispatch_ctx) elif custom_llm_provider == "openrouter": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("OPENROUTER_API_BASE") - or "https://openrouter.ai/api/v1" - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.openrouter_key - or get_secret_str("OPENROUTER_API_KEY") - or get_secret_str("OR_API_KEY") - ) - - openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" - openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" - - openrouter_headers = { - "HTTP-Referer": openrouter_site_url, - "X-Title": openrouter_app_name, - } - - _headers = headers or litellm.headers - if _headers: - openrouter_headers.update(_headers) - - headers = openrouter_headers - - ## Load Config - config = litellm.OpenrouterConfig.get_config() - for k, v in config.items(): - if k == "extra_body": - # we use openai 'extra_body' to pass openrouter specific params - transforms, route, models - if "extra_body" in optional_params: - optional_params[k].update(v) - else: - optional_params[k] = v - elif k not in optional_params: - optional_params[k] = v - - data = {"model": model, "messages": messages, **optional_params} - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="openrouter", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - ## LOGGING - logging.post_call( - input=messages, api_key=openai.api_key, original_response=response - ) + response = _complete_openrouter(_dispatch_ctx) elif custom_llm_provider == "vercel_ai_gateway": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") - or "https://ai-gateway.vercel.sh/v1" - ) - - api_key = ( - api_key or litellm.api_key or get_secret("VERCEL_AI_GATEWAY_API_KEY") - ) - - vercel_site_url = get_secret("VERCEL_SITE_URL") or "https://litellm.ai" - vercel_app_name = get_secret("VERCEL_APP_NAME") or "liteLLM" - - vercel_headers = { - "http-referer": vercel_site_url, - "x-title": vercel_app_name, - } - - _headers = headers or litellm.headers - if _headers: - vercel_headers.update(_headers) - - headers = vercel_headers - - ## Load Config - config = litellm.VercelAIGatewayConfig.get_config() - for k, v in config.items(): - if k == "extra_body": - # we use openai 'extra_body' to pass vercel specific params - providerOptions - if "extra_body" in optional_params: - optional_params[k].update(v) - else: - optional_params[k] = v - elif k not in optional_params: - optional_params[k] = v - - data = {"model": model, "messages": messages, **optional_params} - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="vercel_ai_gateway", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - ## LOGGING - logging.post_call( - input=messages, api_key=openai.api_key, original_response=response - ) + response = _complete_vercel_ai_gateway(_dispatch_ctx) elif ( custom_llm_provider == "together_ai" or ("togethercomputer" in model) @@ -3504,1114 +5787,75 @@ def completion( # type: ignore "Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en" ) elif custom_llm_provider == "vertex_ai_beta" or custom_llm_provider == "gemini": - vertex_ai_project = ( - optional_params.pop("vertex_project", None) - or optional_params.pop("vertex_ai_project", None) - or litellm.vertex_project - or get_secret("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - optional_params.pop("vertex_location", None) - or optional_params.pop("vertex_ai_location", None) - or litellm.vertex_location - or get_secret("VERTEXAI_LOCATION") - ) - vertex_credentials = ( - optional_params.pop("vertex_credentials", None) - or optional_params.pop("vertex_ai_credentials", None) - or get_secret("VERTEXAI_CREDENTIALS") - ) - - gemini_api_key = ( - api_key - or get_api_key_from_env() - or get_secret("PALM_API_KEY") # older palm api key should also work - or litellm.api_key - ) - - api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE") - new_params = safe_deep_copy(optional_params or {}) - response = vertex_chat_completion.completion( # type: ignore - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - gemini_api_key=gemini_api_key, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - custom_llm_provider=custom_llm_provider, # type: ignore - client=client, - api_base=api_base, - extra_headers=headers, - ) + response = _complete_vertex_ai_beta(_dispatch_ctx) elif custom_llm_provider == "vertex_ai": - vertex_ai_project = ( - optional_params.pop("vertex_project", None) - or optional_params.pop("vertex_ai_project", None) - or litellm.vertex_project - or get_secret("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - optional_params.pop("vertex_location", None) - or optional_params.pop("vertex_ai_location", None) - or litellm.vertex_location - or get_secret("VERTEXAI_LOCATION") - ) - vertex_credentials = ( - optional_params.pop("vertex_credentials", None) - or optional_params.pop("vertex_ai_credentials", None) - or get_secret("VERTEXAI_CREDENTIALS") - ) - - api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") - - new_params = safe_deep_copy(optional_params or {}) - model_route = get_vertex_ai_model_route( - model=model, litellm_params=litellm_params - ) - - if model_route == VertexAIModelRoute.PARTNER_MODELS: - model_response = vertex_partner_models_chat_completion.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - api_base=api_base, - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - headers=headers, - custom_prompt_dict=custom_prompt_dict, - timeout=timeout, - client=client, - ) - elif model_route == VertexAIModelRoute.GEMINI: - model_response = vertex_chat_completion.completion( # type: ignore - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - gemini_api_key=None, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - custom_llm_provider=custom_llm_provider, # type: ignore - client=client, - api_base=api_base, - extra_headers=headers, - ) - elif model_route == VertexAIModelRoute.GEMMA: - # Vertex Gemma Models with custom prediction endpoint - model_response = vertex_gemma_chat_completion.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - api_base=api_base, - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - headers=headers, - custom_prompt_dict=custom_prompt_dict, - timeout=timeout, - client=client, - ) - elif model_route == VertexAIModelRoute.MODEL_GARDEN: - # Vertex Model Garden - OpenAI compatible models - model_response = vertex_model_garden_chat_completion.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - api_base=api_base, - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - headers=headers, - custom_prompt_dict=custom_prompt_dict, - timeout=timeout, - client=client, - ) - elif model_route == VertexAIModelRoute.AGENT_ENGINE: - # Vertex AI Agent Engine (Reasoning Engines) - from litellm.llms.vertex_ai.agent_engine.transformation import ( - VertexAgentEngineConfig, - ) - - vertex_agent_engine_config = VertexAgentEngineConfig() - - # Update litellm_params with vertex credentials - litellm_params["vertex_project"] = vertex_ai_project - litellm_params["vertex_location"] = vertex_ai_location - litellm_params["vertex_credentials"] = vertex_credentials - - model_response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - model_response=model_response, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - encoding=_get_encoding(), - api_key=None, - api_base=api_base, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - client=client, - custom_llm_provider="vertex_ai", - provider_config=vertex_agent_engine_config, - headers=headers or {}, - ) - else: # VertexAIModelRoute.NON_GEMINI - model_response = vertex_ai_non_gemini.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - ) - - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="vertex_ai", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_vertex_ai(_dispatch_ctx) elif custom_llm_provider == "predibase": - tenant_id = ( - optional_params.pop("tenant_id", None) - or optional_params.pop("predibase_tenant_id", None) - or litellm.predibase_tenant_id - or get_secret("PREDIBASE_TENANT_ID") - ) - - if tenant_id is None: - raise ValueError( - "Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=)`) or in env - `PREDIBASE_TENANT_ID`." - ) - - api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - or litellm.api_base - or get_secret("PREDIBASE_API_BASE") - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.predibase_key - or get_secret("PREDIBASE_API_KEY") - ) - - _model_response = predibase_chat_completions.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - acompletion=acompletion, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - api_key=api_key, - tenant_id=tenant_id, - timeout=timeout, - ) - - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): - return _model_response - response = _model_response + response = _complete_predibase(_dispatch_ctx) elif custom_llm_provider == "text-completion-codestral": - api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - or litellm.api_base - or "https://codestral.mistral.ai/v1/fim/completions" - ) - - api_key = api_key or litellm.api_key or get_secret("CODESTRAL_API_KEY") - - text_completion_model_response = litellm.TextCompletionResponse( - stream=stream - ) - - _model_response = codestral_text_completions.completion( # type: ignore - model=model, - messages=messages, - model_response=text_completion_model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - acompletion=acompletion, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - api_key=api_key, - timeout=timeout, - ) - - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): - return _model_response - response = _model_response + response = _complete_text_completion_codestral(_dispatch_ctx) elif custom_llm_provider == "text-completion-inception": - passed_api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - ) - api_base = ( - passed_api_base - or get_secret_str("INCEPTION_API_BASE") - or "https://api.inceptionlabs.ai/v1" - ) - # FIM is served at `/v1/fim/completions`; the OpenAI client appends - # `/completions`, so point it at the `/v1/fim` base. - api_base = api_base.rstrip("/") - if not api_base.endswith("/fim"): - api_base += "/fim" - - # Don't forward the server-managed Inception key to a caller-supplied - # api_base; only resolve it for the default/server base, or when the - # caller passes their own key. - if passed_api_base is None or api_key: - api_key = ( - api_key - or litellm.inception_key - or get_secret_str("INCEPTION_API_KEY") - ) - - _response = openai_text_completions.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, # type: ignore[arg-type] - custom_llm_provider="text-completion-inception", - api_base=api_base, - acompletion=acompletion, - client=client, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - ) - - if ( - optional_params.get("stream", False) is False - and acompletion is False - and text_completion is False - ): - _response = litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( - response_object=_response, model_response_object=model_response - ) - - if optional_params.get("stream", False) or acompletion is True: - logging.post_call( - input=messages, - api_key=api_key, - original_response=_response, - additional_args={"headers": headers}, - ) - response = _response + response = _complete_text_completion_inception(_dispatch_ctx) elif custom_llm_provider in ("sagemaker_chat", "sagemaker_nova"): # boto3 reads keys from .env # sagemaker_chat: HF Messages API endpoints # sagemaker_nova: Nova models on SageMaker (OpenAI-compatible) - model_response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - - ## RESPONSE OBJECT - response = model_response + response = _complete_sagemaker_chat(_dispatch_ctx) elif custom_llm_provider == "sagemaker": # boto3 reads keys from .env - model_response = sagemaker_llm.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - custom_prompt_dict=custom_prompt_dict, - hf_model_name=hf_model_name, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - acompletion=acompletion, - ) - - ## RESPONSE OBJECT - response = model_response + response = _complete_sagemaker(_dispatch_ctx) elif custom_llm_provider == "bedrock": # boto3 reads keys from .env - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - - if "aws_bedrock_client" in optional_params: - verbose_logger.warning( - "'aws_bedrock_client' is a deprecated param. Please move to another auth method - https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication." - ) - # Extract credentials for legacy boto3 client and pass thru to httpx - aws_bedrock_client = optional_params.pop("aws_bedrock_client") - creds = aws_bedrock_client._get_credentials().get_frozen_credentials() - - if creds.access_key: - optional_params["aws_access_key_id"] = creds.access_key - if creds.secret_key: - optional_params["aws_secret_access_key"] = creds.secret_key - if creds.token: - optional_params["aws_session_token"] = creds.token - if ( - "aws_region_name" not in optional_params - or optional_params["aws_region_name"] is None - ): - optional_params["aws_region_name"] = ( - aws_bedrock_client.meta.region_name - ) - - bedrock_route = BedrockModelInfo.get_bedrock_route(model) - if bedrock_route == "claude_platform": - provider_config = ProviderConfigManager.get_provider_chat_config( - model=model, - provider=LlmProviders.BEDROCK, - ) - model = BedrockModelInfo.get_claude_platform_model(model) - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - provider_config=provider_config, - ) - return response - elif bedrock_route == "converse": - model = model.replace("converse/", "") - response = bedrock_converse_chat_completion.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - extra_headers=headers, # Use merged headers instead of original extra_headers - timeout=timeout, - acompletion=acompletion, - client=client, - api_base=api_base, - api_key=api_key, - ) - elif bedrock_route == "converse_like": - model = model.replace("converse_like/", "") - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - else: - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_bedrock(_dispatch_ctx) elif custom_llm_provider == "watsonx": - response = watsonx_chat_completion.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - custom_prompt_dict=custom_prompt_dict, - client=client, # pass AsyncOpenAI, OpenAI client - encoding=_get_encoding(), - custom_llm_provider="watsonx", - ) + response = _complete_watsonx(_dispatch_ctx) elif custom_llm_provider == "watsonx_text": - api_key = ( - api_key - or optional_params.pop("apikey", None) - or get_secret_str("WATSONX_APIKEY") - or get_secret_str("WATSONX_API_KEY") - or get_secret_str("WX_API_KEY") - ) - - api_base = ( - api_base - or optional_params.pop( - "url", - optional_params.pop( - "api_base", optional_params.pop("base_url", None) - ), - ) - or get_secret_str("WATSONX_API_BASE") - or get_secret_str("WATSONX_URL") - or get_secret_str("WX_URL") - or get_secret_str("WML_URL") - ) - - wx_credentials = optional_params.pop( - "wx_credentials", - optional_params.pop( - "watsonx_credentials", None - ), # follow {provider}_credentials, same as vertex ai - ) - - token: Optional[str] = None - if wx_credentials is not None: - api_base = wx_credentials.get("url", api_base) - api_key = wx_credentials.get( - "apikey", wx_credentials.get("api_key", api_key) - ) - token = wx_credentials.get( - "token", - wx_credentials.get( - "watsonx_token", None - ), # follow format of {provider}_token, same as azure - e.g. 'azure_ad_token=..' - ) - - if token is not None: - optional_params["token"] = token - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="watsonx_text", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_watsonx_text(_dispatch_ctx) elif custom_llm_provider == "vllm": - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - model_response = vllm_handler.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - ) - - if ( - "stream" in optional_params and optional_params["stream"] is True - ): ## [BETA] - # don't try to access stream object, - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="vllm", - logging_obj=logging, - ) - return response - - ## RESPONSE OBJECT - response = model_response + response = _complete_vllm(_dispatch_ctx) elif custom_llm_provider == "ollama": - api_base = ( - litellm.api_base - or api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" - ) - if api_key is not None and "Authorization" not in headers: - headers["Authorization"] = f"Bearer {api_key}" - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="ollama", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_ollama(_dispatch_ctx) elif custom_llm_provider == "ollama_chat": - api_base = ( - litellm.api_base - or api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" - ) - - api_key = ( - api_key - or litellm.ollama_key - or os.environ.get("OLLAMA_API_KEY") - or litellm.api_key - ) - if api_key is not None and "Authorization" not in headers: - headers["Authorization"] = f"Bearer {api_key}" - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="ollama_chat", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_ollama_chat(_dispatch_ctx) elif custom_llm_provider == "triton": - api_base = litellm.api_base or api_base - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - ) + response = _complete_triton(_dispatch_ctx) elif custom_llm_provider == "cloudflare": - api_key = ( - api_key - or litellm.cloudflare_api_key - or litellm.api_key - or get_secret("CLOUDFLARE_API_KEY") - ) - account_id = get_secret("CLOUDFLARE_ACCOUNT_ID") - api_base = ( - api_base - or litellm.api_base - or get_secret("CLOUDFLARE_API_BASE") - or f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/" - ) - - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="cloudflare", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) + response = _complete_cloudflare(_dispatch_ctx) elif custom_llm_provider == "petals" or model in litellm.petals_models: - api_base = api_base or litellm.api_base - - custom_llm_provider = "petals" - stream = optional_params.pop("stream", False) - model_response = petals_handler.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - client=client, - ) - if stream is True: ## [BETA] - # Fake streaming for petals - resp_string = model_response["choices"][0]["message"]["content"] - response = CustomStreamWrapper( - resp_string, - model, - custom_llm_provider="petals", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_petals(_dispatch_ctx) elif custom_llm_provider == "snowflake" or model in litellm.snowflake_models: - try: - client = ( - HTTPHandler(timeout=timeout) if stream is False else None - ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) - - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_snowflake(_dispatch_ctx) elif custom_llm_provider == "gradient_ai": - api_base = litellm.api_base or api_base - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="gradient_ai", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - ) + response = _complete_gradient_ai(_dispatch_ctx) elif custom_llm_provider == "bytez": - api_key = ( - api_key - or litellm.bytez_key - or get_secret_str("BYTEZ_API_KEY") - or litellm.api_key - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=bytez_transformation, - ) - - pass + response = _complete_bytez(_dispatch_ctx) elif custom_llm_provider == "lemonade": - api_key = ( - api_key - or litellm.lemonade_key - or get_secret_str("LEMONADE_API_KEY") - or litellm.api_key - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=lemonade_transformation, - ) - - pass + response = _complete_lemonade(_dispatch_ctx) elif custom_llm_provider == "ovhcloud" or model in litellm.ovhcloud_models: - api_key = ( - api_key - or litellm.ovhcloud_key - or get_secret_str("OVHCLOUD_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret_str("OVHCLOUD_API_BASE") - or "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=ovhcloud_transformation, - ) - - pass + response = _complete_ovhcloud(_dispatch_ctx) elif custom_llm_provider == "custom": - url = litellm.api_base or api_base or "" - if url is None or url == "": - raise ValueError( - "api_base not set. Set api_base or litellm.api_base for custom endpoints" - ) - - """ - assume input to custom LLM api bases follow this format: - resp = litellm.module_level_client.post( - api_base, - json={ - 'model': 'meta-llama/Llama-2-13b-hf', # model name - 'params': { - 'prompt': ["The capital of France is P"], - 'max_tokens': 32, - 'temperature': 0.7, - 'top_p': 1.0, - 'top_k': 40, - } - } - ) - - """ - prompt = " ".join([message["content"] for message in messages]) # type: ignore - resp = litellm.module_level_client.post( - url, - headers=headers, - json={ - "model": model, - "params": { - "prompt": [prompt], - "max_tokens": max_tokens, - "temperature": temperature, - "top_p": top_p, - "top_k": kwargs.get("top_k"), - }, - **kwargs.get("extra_body", {}), - }, - ) - response_json = resp.json() - """ - assume all responses from custom api_bases of this format: - { - 'data': [ - { - 'prompt': 'The capital of France is P', - 'output': ['The capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France'], - 'params': {'temperature': 0.7, 'top_k': 40, 'top_p': 1}}], - 'message': 'ok' - } - ] - } - """ - string_response = response_json["data"][0]["output"][0] - ## RESPONSE OBJECT - model_response.choices[0].message.content = string_response # type: ignore - model_response.created = int(time.time()) - model_response.model = model - response = model_response + response = _complete_custom(_dispatch_ctx) elif ( custom_llm_provider in litellm._custom_providers ): # Assume custom LLM provider # Get the Custom Handler - custom_handler: Optional[CustomLLM] = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] - - if custom_handler is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) - - ## ROUTE LLM CALL ## - handler_fn = custom_chat_llm_router( - async_fn=acompletion, stream=stream, custom_llm=custom_handler - ) - - headers = headers or litellm.headers or {} - - ## CALL FUNCTION - response = handler_fn( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - custom_prompt_dict=custom_prompt_dict, - client=client, # pass AsyncOpenAI, OpenAI client - encoding=_get_encoding(), - ) - if stream is True: - return CustomStreamWrapper( - completion_stream=response, - model=model, - custom_llm_provider=custom_llm_provider, - logging_obj=logging, - ) + response = _complete_custom_providers(_dispatch_ctx) elif custom_llm_provider == "langgraph": # LangGraph - Agent Runtime Provider - from litellm.llms.langgraph.chat.transformation import LangGraphConfig - - ( - api_base, - api_key, - ) = LangGraphConfig()._get_openai_compatible_provider_info( - api_base=api_base or litellm.api_base, - api_key=api_key or litellm.api_key, - ) - - headers = headers or litellm.headers - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_langgraph(_dispatch_ctx) elif custom_llm_provider == "langflow": # LangFlow - Visual AI Agent Platform - from litellm.llms.langflow.chat.transformation import LangFlowConfig - - ( - api_base, - api_key, - ) = LangFlowConfig()._get_openai_compatible_provider_info( - api_base=api_base or litellm.api_base, - api_key=api_key or litellm.api_key, - ) - - headers = headers or litellm.headers - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_langflow(_dispatch_ctx) else: raise LiteLLMUnknownProvider( @@ -6341,8 +7585,8 @@ def text_completion( kwargs.pop("prompt", None) - if _model is not None and ( - custom_llm_provider == "openai" + if ( + _model is not None and (custom_llm_provider == "openai") ): # for openai compatible endpoints - e.g. vllm, call the native /v1/completions endpoint for text completion calls if _model not in litellm.open_ai_chat_completion_models: model = "text-completion-openai/" + _model @@ -6413,7 +7657,9 @@ async def aadapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) - response: Union[ModelResponse, CustomStreamWrapper] = await acompletion(**new_kwargs) # type: ignore + response: Union[ModelResponse, CustomStreamWrapper] = await acompletion( + **new_kwargs + ) # type: ignore translated_response: Optional[ Union[BaseModel, AdapterCompletionStreamWrapper] ] = None @@ -6819,7 +8065,12 @@ def transcription( ) # set API KEY - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") # type: ignore + api_key = ( + api_key + or litellm.api_key + or litellm.openai_key + or get_secret("OPENAI_API_KEY") + ) # type: ignore response = openai_audio_transcriptions.audio_transcriptions( model=model, audio_file=file, @@ -7160,7 +8411,9 @@ def speech( ) api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore + api_version = ( + api_version or litellm.api_version or get_secret("AZURE_API_VERSION") + ) # type: ignore api_key = ( api_key @@ -7172,9 +8425,7 @@ def speech( azure_ad_token: Optional[str] = optional_params.get("extra_body", {}).pop( # type: ignore "azure_ad_token", None - ) or get_secret( - "AZURE_AD_TOKEN" - ) + ) or get_secret("AZURE_AD_TOKEN") azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None) if extra_headers: @@ -7631,9 +8882,7 @@ def stream_chunk_builder_text_completion( response["usage"]["prompt_tokens"] = token_counter( model=model, messages=messages ) - except ( - Exception - ): # don't allow this failing to block a complete streaming response from being returned + except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") response["usage"]["prompt_tokens"] = 0 response["usage"]["completion_tokens"] = token_counter( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5c962cf8440..6ebac7efc8d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -571,7 +571,7 @@ "output_vector_size": 1536 }, "amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 2e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "max_tokens": 8192, @@ -10443,7 +10443,8 @@ "fast": 6.0 }, "supports_output_config": true, - "supports_max_reasoning_effort": true + "supports_max_reasoning_effort": true, + "supports_speed": true }, "claude-opus-4-6-20260205": { "cache_creation_input_token_cost": 6.25e-06, @@ -10476,7 +10477,8 @@ "fast": 6.0 }, "supports_max_reasoning_effort": true, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -10511,7 +10513,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-opus-4-7-20260416": { "cache_creation_input_token_cost": 6.25e-06, @@ -10546,7 +10549,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, @@ -10615,7 +10619,8 @@ "us": 1.1, "fast": 2.0 }, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-05-14", @@ -10684,6 +10689,268 @@ "mode": "chat", "output_cost_per_token": 1.923e-06 }, + "cloudflare/@cf/openai/gpt-oss-120b": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-2b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.2-3b-instruct": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 3.35e-07 + }, + "cloudflare/@cf/meta/llama-guard-3-8b": { + "input_cost_per_token": 4.84e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3e-08 + }, + "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 15000, + "max_output_tokens": 15000, + "max_tokens": 15000, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/deepseek-ai/deepseek-r1-distill-qwen-32b": { + "input_cost_per_token": 4.97e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 4.881e-06, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { + "input_cost_per_token": 1.52e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "output_cost_per_token": 2.87e-07 + }, + "cloudflare/@cf/meta/llama-3.2-1b-instruct": { + "input_cost_per_token": 2.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 60000, + "max_output_tokens": 60000, + "max_tokens": 60000, + "mode": "chat", + "output_cost_per_token": 2.01e-07 + }, + "cloudflare/@cf/moonshotai/kimi-k2.6": { + "cache_read_input_token_cost": 1.6e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/zai-org/glm-4.7-flash": { + "input_cost_per_token": 6.05e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta-llama/llama-2-7b-chat-hf-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { + "input_cost_per_token": 2.93e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 2.253e-06, + "supports_function_calling": true + }, + "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { + "input_cost_per_token": 1.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.12e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-06 + }, + "cloudflare/@cf/zai-org/glm-5.2": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/nvidia/nemotron-3-120b-a12b": { + "input_cost_per_token": 5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/aisingapore/gemma-sea-lion-v4-27b-it": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07 + }, + "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.35e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-7b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 3500, + "max_output_tokens": 3500, + "max_tokens": 3500, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/mistralai/mistral-small-3.1-24b-instruct": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { + "input_cost_per_token": 4.85e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.76e-07, + "supports_vision": true + }, + "cloudflare/@cf/openai/gpt-oss-20b": { + "input_cost_per_token": 2e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-4-scout-17b-16e-instruct": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 8.5e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwq-32b": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 1e-06, + "supports_reasoning": true + }, "codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "codestral", @@ -20088,8 +20355,6 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20163,8 +20428,6 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20238,8 +20501,6 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20311,8 +20572,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20354,8 +20613,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20377,8 +20634,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20667,8 +20922,6 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_priority": 1e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -21372,8 +21625,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21767,6 +22018,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21815,6 +22068,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21859,6 +22114,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21903,6 +22160,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21951,6 +22210,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21998,6 +22259,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22038,6 +22301,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22081,6 +22346,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22126,6 +22393,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22172,6 +22441,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22215,6 +22486,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22258,6 +22531,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22296,8 +22571,6 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -22704,8 +22977,6 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22787,8 +23058,6 @@ "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, @@ -39908,24 +40177,6 @@ "litellm_provider": "fireworks_ai", "mode": "chat" }, - "fireworks_ai/accounts/fireworks/models/whisper-v3": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, - "fireworks_ai/accounts/fireworks/models/whisper-v3-turbo": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, "fireworks_ai/accounts/fireworks/models/yi-34b": { "max_tokens": 4096, "max_input_tokens": 4096, @@ -43061,6 +43312,40 @@ "supports_tool_choice": true, "supports_vision": false }, + "darkbloom/gemma-4-26b": { + "input_cost_per_token": 3e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.65e-07, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "darkbloom/gpt-oss-20b": { + "input_cost_per_token": 1.45e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7e-08, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "deepseek/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 3.625e-09, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b27082c361a..3a9ef8db804 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -10,7 +10,7 @@ import os import re from functools import partial from io import IOBase -from typing import Any, Coroutine, Dict, Optional, Union +from typing import Any, Callable, Coroutine, Dict, Optional, Union, cast import httpx @@ -20,6 +20,7 @@ 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 BaseOCRConfig, OCRResponse from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.ocr.rust_bridge import RustOcr, load_rust_ocr, rust_ocr_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -28,6 +29,82 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +def _timeout_to_seconds( + timeout: Optional[Union[float, httpx.Timeout]], +) -> Optional[float]: + """Convert the Python OCR timeout to a single seconds value for the Rust bridge. + + The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate + connect/read/write/pool values, so pick the read deadline as the closest + analog to a total-request timeout. + """ + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) + + +def _run_rust_ocr( + rust_ocr: RustOcr, + logging_obj: LiteLLMLoggingObj, + provider_config: BaseOCRConfig, + resolve_api_key: Callable[[str], Optional[str]], + model: str, + document: dict[str, object], + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict[str, object], + litellm_params: dict[str, object], + timeout_seconds: Optional[float], +) -> OCRResponse: + """Run the Mistral OCR call through the Rust bridge and wrap the result. + + Resolves the key the same way the Python path does so secret-manager backends + (AWS/Azure/GCP/Vault) work; the Rust bridge's own fallback only reads the + process environment. The request that Rust actually sends (resolved URL and + headers) is mirrored into pre_call so logs match the wire. Dependencies are + injected so this stays unit-testable without patching module globals. + """ + resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY") + resolved_headers = provider_config.validate_environment( + headers={}, + model=model, + api_key=resolved_api_key, + api_base=api_base, + litellm_params=litellm_params, + ) + resolved_complete_url = provider_config.get_complete_url( + api_base=api_base, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + ) + logging_obj.pre_call( + input="OCR document processing", + api_key=resolved_api_key, + additional_args={ + "complete_input_dict": { + "model": model, + "document": document, + **optional_params, + }, + "api_base": resolved_complete_url, + "headers": resolved_headers, + }, + ) + return OCRResponse.model_validate( + rust_ocr( + model=model, + document=document, + api_key=resolved_api_key, + api_base=api_base, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ) + ) + + @client async def aocr( model: str, @@ -220,7 +297,7 @@ def ocr( """ local_vars = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("aocr", False) is True @@ -261,7 +338,6 @@ def ocr( if dynamic_api_base: api_base = dynamic_api_base - # Get provider config ocr_provider_config: Optional[BaseOCRConfig] = ( ProviderConfigManager.get_provider_ocr_config( model=model, @@ -278,17 +354,14 @@ def ocr( f"OCR call - model: {model}, provider: {custom_llm_provider}" ) - # Get litellm params using GenericLiteLLMParams (same as responses API) litellm_params = GenericLiteLLMParams(**kwargs) - # Extract OCR-specific parameters from kwargs supported_params = ocr_provider_config.get_supported_ocr_params(model=model) non_default_params = {} for param in supported_params: if param in kwargs: non_default_params[param] = kwargs.pop(param) - # Map parameters to provider-specific format optional_params = ocr_provider_config.map_ocr_params( non_default_params=non_default_params, optional_params={}, @@ -297,7 +370,8 @@ def ocr( verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") - # Pre Call logging + effective_timeout = timeout or request_timeout + litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, @@ -309,12 +383,35 @@ def ocr( custom_llm_provider=custom_llm_provider, ) - # Call the handler - pass document dict directly + # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. + if custom_llm_provider == "mistral" and rust_ocr_enabled(): + rust_ocr = load_rust_ocr() + if rust_ocr is None: + verbose_logger.debug( + "Rust OCR bridge unavailable; falling back to Python path" + ) + else: + from litellm.secret_managers.main import get_secret_str + + return _run_rust_ocr( + rust_ocr=rust_ocr, + logging_obj=litellm_logging_obj, + provider_config=ocr_provider_config, + resolve_api_key=get_secret_str, + model=model, + document=document, + api_key=api_key, + api_base=api_base, + optional_params=optional_params, + litellm_params=dict(litellm_params), + timeout_seconds=_timeout_to_seconds(effective_timeout), + ) + response = base_llm_http_handler.ocr( model=model, - document=document, # Pass the entire document dict + document=document, optional_params=optional_params, - timeout=timeout or request_timeout, + timeout=effective_timeout, logging_obj=litellm_logging_obj, api_key=api_key, api_base=api_base, diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py new file mode 100644 index 00000000000..61f9e8ca69a --- /dev/null +++ b/litellm/ocr/rust_bridge.py @@ -0,0 +1,74 @@ +""" +Optional Rust-backed OCR path. + +Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint +then routes supported Mistral calls through the compiled ``litellm_python_bridge`` +extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. + +No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` +can import it statically without forming an import cycle. +""" + +from __future__ import annotations + +from typing import Final, Protocol, cast + + +class RustOcr(Protocol): + """Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint.""" + + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: ... + + +class _Unset: + """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it.""" + + +_UNSET: Final[_Unset] = _Unset() + +_rust_ocr_enabled = False +_rust_ocr_impl: RustOcr | None = None + + +def use_litellm_rust( + enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET +) -> None: + """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension. + + ``ocr`` injects the bridge callable; when omitted the compiled extension is + loaded on demand and any previously injected bridge is preserved. Pass + ``ocr=None`` explicitly to clear a prior injection. + """ + global _rust_ocr_enabled, _rust_ocr_impl + _rust_ocr_enabled = enabled + if not isinstance(ocr, _Unset): + _rust_ocr_impl = ocr + + +def rust_ocr_enabled() -> bool: + """Whether the Rust OCR path has been turned on via ``use_litellm_rust()``.""" + return _rust_ocr_enabled + + +def load_rust_ocr() -> RustOcr | None: + """Return the Rust OCR callable, or ``None`` when no bridge is available. + + Prefers an injected implementation, otherwise loads the compiled + ``litellm_python_bridge`` extension; a missing extension yields ``None`` so + the caller can fall back to the Python path instead of hard-failing. + """ + if _rust_ocr_impl is not None: + return _rust_ocr_impl + try: + import litellm_python_bridge + except ImportError: + return None + return cast(RustOcr, litellm_python_bridge.ocr) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 3e60988b9e7..b1d4864cea7 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -345,7 +345,9 @@ def llm_passthrough_route( ) else: # Sync path - client.client.send returns Response directly - response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore + response: httpx.Response = client.client.send( + request=request, stream=is_streaming_request + ) # type: ignore response.raise_for_status() if ( diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index db6183edaa0..dd7712aabca 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1835,6 +1835,23 @@ "interactions": true } }, + "darkbloom": { + "display_name": "Darkbloom (`darkbloom`)", + "url": "https://docs.litellm.ai/docs/providers/darkbloom", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "predibase": { "display_name": "Predibase (`predibase`)", "url": "https://docs.litellm.ai/docs/providers/predibase", 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 90108de25c3..19f508e6af9 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 @@ -521,9 +521,9 @@ class MCPRequestHandler: if server_alias not in server_auth_headers: server_auth_headers[server_alias] = {} - server_auth_headers[server_alias][ - auth_header_name - ] = header_value + server_auth_headers[server_alias][auth_header_name] = ( + header_value + ) verbose_logger.debug( f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}..." ) @@ -1383,9 +1383,9 @@ class MCPRequestHandler: cache_key = f"agent_object_permission_id:{agent_id}" try: - object_permission_id: Optional[str] = ( - await user_api_key_cache.async_get_cache(key=cache_key) - ) + object_permission_id: Optional[ + str + ] = await user_api_key_cache.async_get_cache(key=cache_key) if object_permission_id == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 8edb831a9df..bc1068d77d0 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -695,7 +695,8 @@ async def update_mcp_server( data_dict["updated_by"] = touched_by updated_mcp_server = await MCPServerRepository(prisma_client).table.update( - where={"server_id": data.server_id}, data=data_dict # type: ignore + where={"server_id": data.server_id}, + data=data_dict, # type: ignore ) _decrypt_env_vars_on_returned_row(updated_mcp_server) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 5e704b889ae..b9e0379e445 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -56,6 +56,16 @@ from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + Error, + Ok, + UpstreamCredentialProvider, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_public, + to_server_spec, + to_subject, +) from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -72,6 +82,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( normalize_server_name, parse_admin_env_vars, split_server_prefix_from_name, + strip_known_server_prefix, validate_mcp_server_name, ) from litellm.proxy._types import ( @@ -511,7 +522,8 @@ class MCPServerManager: return "client_credentials" return None - def __init__(self): + def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None): + self._cred_provider = cred_provider or UpstreamCredentialProvider() self.registry: Dict[str, MCPServer] = {} self.config_mcp_servers: Dict[str, MCPServer] = {} """ @@ -1470,7 +1482,8 @@ class MCPServerManager: for toolset in toolsets: for tool in toolset.tools: raw_name = tool["tool_name"] - unprefixed, _ = split_server_prefix_from_name(raw_name) + server = self.get_mcp_server_by_id(tool["server_id"]) + unprefixed = strip_known_server_prefix(raw_name, server) tool_permissions.setdefault(tool["server_id"], []) if unprefixed not in tool_permissions[tool["server_id"]]: tool_permissions[tool["server_id"]].append(unprefixed) @@ -1942,11 +1955,19 @@ class MCPServerManager: Returns: Configured MCP client instance. """ - auth_value = await resolve_mcp_auth( - server, mcp_auth_header, subject_token=subject_token - ) - transport = server.transport or MCPTransport.sse + spec = None if transport == MCPTransport.stdio else to_server_spec(server) + # A per-request override is the caller-supplied credential v1 turns into the upstream + # auth, so it must win; defer those to v1 (this defer falls away once the per-user modes + # stop writing mcp_auth_header). An inbound header already in extra_headers is handled on + # the v2 path below, not here. + if spec is not None and mcp_auth_header: + spec = None + auth_value = ( + await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) + if spec is None + else None + ) # Create sampling and elicitation callbacks for this client sampling_cb = ( @@ -2017,6 +2038,43 @@ class MCPServerManager: # For HTTP/SSE transports server_url = server.url or "" + if spec is not None: + match await self._cred_provider.resolve_credentials( + to_subject(user_api_key_auth, subject_token), spec + ): + case Ok(auth): + resolved_auth = auth + # Do not override an Authorization already supplied via extra_headers + # (a guardrail hook such as the JWT signer, static_headers, or a + # forwarded caller header): v1 applies those last, so they win. NoOpAuth + # has no header_name and so never skips. + header_name = getattr(resolved_auth, "header_name", None) + if ( + header_name + and extra_headers + and any( + key.lower() == header_name.lower() + for key in extra_headers + ) + ): + resolved_auth = None + case Error(err): + raise_public(err) + return MCPClient( + server_url=server_url, + transport_type=transport, + auth_type=server.auth_type, + timeout=( + server.timeout + if server.timeout is not None + else MCP_CLIENT_TIMEOUT + ), + extra_headers=extra_headers, + resolved_auth=resolved_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ) + # Create SigV4 auth if configured aws_auth = None if server.auth_type == MCPAuth.aws_sigv4: @@ -3643,7 +3701,7 @@ class MCPServerManager: return stored_headers except Exception as _lookup_exc: verbose_logger.debug( - "call_tool: per-user token lookup failed for " "user=%s server=%s: %s", + "call_tool: per-user token lookup failed for user=%s server=%s: %s", user_id, mcp_server.server_id, _lookup_exc, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py index b357504979d..73166a45d6e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py @@ -1,16 +1,20 @@ """Typed upstream-credential resolution for MCP servers. -This subpackage houses the typed credential vocabulary and (in a later PR) the -``resolve_credentials`` dispatch. A server declares one per-mode config from the -``AuthConfig`` discriminated union; failures are modeled as values via :mod:`.result` -(``Result[T, CredError]``) rather than raised, so every seam is total. Nothing here is -wired onto a live request path yet. +This subpackage houses the typed credential vocabulary and the ``resolve_credentials`` +dispatch. A server declares one per-mode config from the ``AuthConfig`` discriminated union; +``UpstreamCredentialProvider.resolve_credentials`` selects one arm and returns an ``httpx.Auth`` +or a typed ``CredError``. Failures are modeled as values via :mod:`.result` (``Result[T, +CredError]``) rather than raised, so every seam is total. Nothing here is wired onto a live +request path yet. """ from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( NoOpAuth, StaticHeaderAuth, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import ( + UpstreamCredentialProvider, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Error, Ok, @@ -45,6 +49,7 @@ __all__ = [ "Result", "NoOpAuth", "StaticHeaderAuth", + "UpstreamCredentialProvider", "AuthSpecKind", "CredError", "Subject", diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py new file mode 100644 index 00000000000..39db2314aee --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -0,0 +1,140 @@ +"""The v1 <-> v2 bridge for the credential resolver. + +These edge functions translate v1's request objects into the resolver's typed inputs and map +its typed errors onto the proxy's public exception contract. They import v1 and live outside the +package's public surface so the resolver core (``resolver.py`` / ``types.py``) stays v1-free. +Nothing wires them into ``_create_mcp_client`` yet. + +``to_server_spec`` maps only the modes the resolver has gone live for, returning ``None`` for +every other mode so the caller defers to v1 (parity-safe); it grows one branch per migrated mode. +""" + +from __future__ import annotations + +import base64 +from typing import TYPE_CHECKING, NoReturn, Optional + +from fastapi import HTTPException +from pydantic import SecretStr +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ApiKeyConfig, + CredError, + NoneConfig, + ServerSpec, + SharedKey, + Subject, +) +from litellm.types.mcp import MCPAuth + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def to_subject( + user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str] +) -> Subject: + """Map v1's authenticated principal onto the resolver's Subject. + + tenant_id / subject_id are empty for an unauthenticated caller; the per-user arms must reject + an empty subject rather than share one credential slot across callers. + """ + inbound = SecretStr(subject_token) if subject_token else None + if user_api_key_auth is None: + return Subject(tenant_id="", subject_id="", inbound_token=inbound) + return Subject( + tenant_id=user_api_key_auth.org_id or user_api_key_auth.team_id or "", + subject_id=user_api_key_auth.user_id or "", + inbound_token=inbound, + ) + + +def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: + """Map a v1 server onto a ServerSpec for a migrated mode, or None to defer to v1. + + BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just + like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers + to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later). + + Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with + an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is + explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live + modes: ``none`` and the static-header family (``api_key`` plus the Authorization schemes), + all shared-key; every other mode returns None and stays on v1. + """ + if server.is_byok: + return ( + None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) + ) + resource = server.url or server.server_id + auth_type = server.auth_type + match auth_type: + case None | MCPAuth.none: + if server.is_oauth_passthrough: + return None # passthrough is not migrated yet -> defer to v1 + return ServerSpec( + server_id=server.server_id, resource=resource, config=NoneConfig() + ) + case MCPAuth.api_key: + return _shared_key_spec(server, resource, "X-API-Key", "") + case MCPAuth.bearer_token: + return _shared_key_spec(server, resource, "Authorization", "Bearer") + case MCPAuth.token: + return _shared_key_spec(server, resource, "Authorization", "token") + case MCPAuth.authorization: + return _shared_key_spec(server, resource, "Authorization", "") + case MCPAuth.basic: + return _shared_key_spec( + server, resource, "Authorization", "Basic", encode=True + ) + case MCPAuth.oauth2 | MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4: + return None # OAuth grants and SigV4 are not migrated yet -> defer to v1 + assert_never(auth_type) + + +def _shared_key_spec( + server: MCPServer, + resource: str, + header_name: str, + value_prefix: str, + *, + encode: bool = False, +) -> Optional[ServerSpec]: + """Build an api_key spec from the server's static token, or defer (None) if it is absent. + + Covers the whole shared-key static-header family: ``api_key`` on ``X-API-Key`` and the + Authorization schemes (bearer / token / authorization sent verbatim, basic base64-encoded). + """ + token = server.authentication_token + if not token: + return None # no key configured -> defer to v1 (parity-safe) + value = base64.b64encode(token.encode("utf-8")).decode() if encode else token + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=ApiKeyConfig( + header_name=header_name, + value_prefix=value_prefix, + key_source=SharedKey(value=SecretStr(value)), + ), + ) + + +def raise_public(error: CredError) -> NoReturn: + """Map a resolver CredError onto the proxy's public HTTP contract. The one edge that raises.""" + match error.tag: + case "unauthorized": + raise HTTPException(status_code=401, detail=error.summary) + case "misconfigured": + raise HTTPException(status_code=500, detail=error.summary) + case "upstream_unavailable": + raise HTTPException(status_code=503, detail=error.summary) + case "unsupported_mode": + raise HTTPException(status_code=500, detail=error.summary) + case "precondition_required": + raise HTTPException(status_code=412, detail=error.summary) + case "not_implemented": + raise HTTPException(status_code=501, detail=error.summary) + assert_never(error.tag) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py new file mode 100644 index 00000000000..969bbf01ec8 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -0,0 +1,93 @@ +"""The one credential resolver: dispatch on the declared mode, fail closed. + +`resolve_credentials` selects exactly one arm off the server's typed `config` and either +produces an `httpx.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig` +variant, so each arm receives its own fully-typed config with no field-presence inference and +no precedence cascade. It is wildcard-free with an `assert_never` tail, so adding a mode without +an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly +at runtime instead of returning `None`. + +`none` and `api_key` (shared-key source) are live; the remaining arms are `not_implemented` +stubs that each land in a follow-up PR with their injected seam. The self-contained arms read +straight from the config and need no collaborator. Pure v2: no imports from v1. +""" + +from __future__ import annotations + +import httpx +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + NoOpAuth, + StaticHeaderAuth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ApiKeyConfig, + AuthorizationCodeConfig, + AuthSpecKind, + AwsSigV4Config, + Byok, + ClientCredentialsConfig, + CredError, + NoneConfig, + PassthroughConfig, + ServerSpec, + SharedKey, + Subject, + TokenExchangeConfig, +) + + +class UpstreamCredentialProvider: + """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode. + + Collaborators (the per-mode credential stores and token fetchers) are injected as each arm + is built; the live `none` and `api_key`-shared arms read from the config and need none. + """ + + async def resolve_credentials( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + match server.config: + case NoneConfig(): + return Ok(NoOpAuth()) + case ApiKeyConfig() as config: + return self._api_key(config) + case PassthroughConfig(): + return _not_implemented(AuthSpecKind.passthrough) + case ClientCredentialsConfig(): + return _not_implemented(AuthSpecKind.client_credentials) + case TokenExchangeConfig(): + return _not_implemented(AuthSpecKind.token_exchange) + case AuthorizationCodeConfig(): + return _not_implemented(AuthSpecKind.authorization_code) + case AwsSigV4Config(): + return _not_implemented(AuthSpecKind.aws_sigv4) + assert_never(server.config) + + def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]: + match config.key_source: + case SharedKey() as source: + header_name, header_value = config.header( + source.value.get_secret_value() + ) + return Ok(StaticHeaderAuth(header_value, header_name=header_name)) + case Byok(): + # Per-user key pulled from the credential store; lands with that seam. + return Error( + CredError.of_not_implemented( + "api_key BYOK source not implemented yet" + ) + ) + assert_never(config.key_source) + + +def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: + return Error( + CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet") + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e891425274f..d7975303802 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -138,9 +138,7 @@ try: import weakref # Robust auth lookup keyed by session_object. - _session_obj_auth_storage: ( - "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" - ) = weakref.WeakKeyDictionary() + _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = ( contextvars.ContextVar("active_mcp_session", default=None) @@ -300,6 +298,7 @@ if MCP_AVAILABLE: is_tool_name_prefixed, normalize_server_name, split_server_prefix_from_name, + strip_known_server_prefix, ) ###################################################### @@ -517,7 +516,12 @@ if MCP_AVAILABLE: async def initialize_session_managers(): """Initialize the session managers. Can be called from main app lifespan.""" - global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _session_manager_stateful_cm, _sse_session_manager_cm, _stateful_auth_context_cleanup_task + global \ + _SESSION_MANAGERS_INITIALIZED, \ + _session_manager_cm, \ + _session_manager_stateful_cm, \ + _sse_session_manager_cm, \ + _stateful_auth_context_cleanup_task # Use async lock to prevent concurrent initialization async with _INITIALIZATION_LOCK: @@ -546,7 +550,12 @@ if MCP_AVAILABLE: async def shutdown_session_managers(): """Shutdown the session managers.""" - global _SESSION_MANAGERS_INITIALIZED, _session_manager_cm, _session_manager_stateful_cm, _sse_session_manager_cm, _stateful_auth_context_cleanup_task + global \ + _SESSION_MANAGERS_INITIALIZED, \ + _session_manager_cm, \ + _session_manager_stateful_cm, \ + _sse_session_manager_cm, \ + _stateful_auth_context_cleanup_task if _SESSION_MANAGERS_INITIALIZED: verbose_logger.info("Shutting down MCP session managers...") @@ -2109,20 +2118,18 @@ if MCP_AVAILABLE: server_id=server_id, user_api_key_auth=user_api_key_auth, ) - if allowed_tool_names is not None: - # Strip prefix from tool names before comparing - # Tools are stored in DB without prefix, but come from MCP server with prefix - filtered_tools = [] - for t in tools: - # Get tool name without server prefix - unprefixed_tool_name, _ = split_server_prefix_from_name(t.name) - if unprefixed_tool_name in allowed_tool_names: - filtered_tools.append(t) - else: - # No restrictions, return all tools - filtered_tools = tools + if allowed_tool_names is None: + return tools - return filtered_tools + # Tools arrive prefixed with the server's own prefix; strip exactly that + # prefix (resolved from the server) rather than the first separator, so a + # prefix containing the separator still reduces to the stored bare name. + server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + return [ + t + for t in tools + if strip_known_server_prefix(t.name, server) in allowed_tool_names + ] async def _merge_toolset_permissions( user_api_key_auth: Optional[UserAPIKeyAuth], @@ -3501,7 +3508,19 @@ if MCP_AVAILABLE: if stored_oauth_headers: continue if getattr(server, "delegate_auth_to_upstream", False) is True: - continue + # Delegate-auth servers run upstream PKCE: challenge with + # the proxied resource_metadata (RFC 9728), not the + # gateway authorization_uri below which would authorize + # against the gateway instead of the upstream IdP. + www_authenticate = _get_passthrough_www_authenticate( + scope=scope, + server_name=server_name, + ) + raise HTTPException( + status_code=401, + detail="Unauthorized", + headers={"www-authenticate": www_authenticate}, + ) request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index b0141d3207c..3418417a8f8 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -332,6 +332,30 @@ def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: return prefixed_name, "" +def strip_known_server_prefix(name: str, server: Optional[Any]) -> str: + """Strip ``server``'s registered prefix from a prefixed tool/resource name. + + Unlike :func:`split_server_prefix_from_name`, which guesses the boundary at + the first separator, this removes exactly ``{known_prefix}{separator}`` for + one of the server's actual registered prefixes. It therefore stays correct + when a prefix itself contains the separator (e.g. the UUID ``server_id`` + used as the fallback prefix when a server has no alias, or a legacy + hyphenated alias), where the first-separator split would cut inside the + prefix and never match the stored bare tool name. + + Returns ``name`` unchanged when ``server`` is known but none of its prefixes + match (the name is already unprefixed). Falls back to the legacy split only + when ``server`` is ``None``. + """ + if server is None: + return split_server_prefix_from_name(name)[0] + for prefix in iter_known_server_prefixes(server): + candidate = normalize_server_name(prefix) + MCP_TOOL_PREFIX_SEPARATOR + if name.startswith(candidate): + return name[len(candidate) :] + return name + + def is_tool_name_prefixed( tool_name: str, known_server_prefixes: Optional[set] = None, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ac90302fdaa..7db0156bbe0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1052,9 +1052,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -2361,6 +2361,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs.", ) + mcp_xff_num_trusted_hops: Optional[int] = Field( + None, + ge=1, + description="Number of trusted reverse proxies/load balancers in front of the gateway that append to X-Forwarded-For. When set (and mcp_trusted_proxy_ranges validates the direct peer), the client IP for MCP access control is read this many entries from the right of the chain instead of the spoofable leftmost value, defeating append-style X-Forwarded-For forgery.", + ) trusted_proxy_ranges: Optional[List[str]] = Field( None, description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler.", @@ -3358,7 +3363,9 @@ class ProxyException(Exception): class CommonProxyErrors(str, enum.Enum): db_not_connected_error = ( - "DB not connected. See https://docs.litellm.ai/docs/proxy/virtual_keys" + "DB not connected. This endpoint needs a database; set DATABASE_URL to a " + "PostgreSQL connection string (postgresql://...) to enable it. " + "See https://docs.litellm.ai/docs/proxy/virtual_keys" ) no_llm_router = "No models configured on proxy" not_allowed_access = "Admin-only endpoint. Not allowed to access this." @@ -3992,9 +3999,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 7446f61ad1c..b1fb72619b1 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -305,16 +305,19 @@ async def _handle_stream_message( if not A2A_SDK_AVAILABLE: async def _error_stream(): - yield json.dumps( - { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32603, - "message": "Server error: 'a2a' package not installed", - }, - } - ) + "\n" + yield ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": "Server error: 'a2a' package not installed", + }, + } + ) + + "\n" + ) return StreamingResponse(_error_stream(), media_type="application/x-ndjson") @@ -392,9 +395,10 @@ async def _handle_stream_message( else: async for chunk in a2a_stream: if hasattr(chunk, "model_dump"): - yield json.dumps( - chunk.model_dump(mode="json", exclude_none=True) - ) + "\n" + yield ( + json.dumps(chunk.model_dump(mode="json", exclude_none=True)) + + "\n" + ) else: yield json.dumps(chunk) + "\n" except Exception as e: @@ -414,13 +418,19 @@ async def _handle_stream_message( e = transformed_exception if isinstance(e, HTTPException): raise - yield json.dumps( - { - "jsonrpc": "2.0", - "id": request_id, - "error": {"code": -32603, "message": f"Streaming error: {str(e)}"}, - } - ) + "\n" + yield ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": f"Streaming error: {str(e)}", + }, + } + ) + + "\n" + ) return StreamingResponse(stream_response(), media_type="application/x-ndjson") @@ -826,9 +836,9 @@ async def invoke_agent_a2a( ) if method == "agent/getAuthenticatedExtendedCard": if isinstance(result.get("result"), dict) and "url" in result["result"]: - result["result"][ - "url" - ] = f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}" + result["result"]["url"] = ( + f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}" + ) from litellm.types.agents import LiteLLMSendMessageResponse response = LiteLLMSendMessageResponse.from_dict( diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 11fd01e2369..c1f9c89529b 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -65,7 +65,9 @@ class AgentRegistry: # create a stable hash id for config item config_hash = self._create_agent_id(agent_config_item) - self.register_agent(agent_config=AgentResponse(agent_id=config_hash, **agent_config_item)) # type: ignore + self.register_agent( + agent_config=AgentResponse(agent_id=config_hash, **agent_config_item) + ) # type: ignore def load_agents_from_db_and_config( self, @@ -79,7 +81,12 @@ class AgentRegistry: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") - self.register_agent(agent_config=AgentResponse(agent_id=self._create_agent_id(agent_config_item), **agent_config_item)) # type: ignore + self.register_agent( + agent_config=AgentResponse( + agent_id=self._create_agent_id(agent_config_item), + **agent_config_item, + ) + ) # type: ignore if db_agents: for db_agent in db_agents: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 52788ed9238..88db2a2b7ea 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2954,6 +2954,26 @@ async def _get_agent_ids_from_access_groups( ) +def _resolve_all_team_model_sentinel_for_auth_check( + models: List[str], + llm_router: Optional[Router], + team_id: Optional[str], +) -> List[str]: + if ( + SpecialModelNames.all_team_models.value not in models + or team_id is None + or llm_router is None + ): + return models + proxy_models = llm_router.get_model_names() + non_sentinel_models = [ + model for model in models if model != SpecialModelNames.all_team_models.value + ] + if not proxy_models: + return non_sentinel_models or models + return list(dict.fromkeys(non_sentinel_models + proxy_models)) + + def _check_model_access_helper( model: str, llm_router: Optional[Router], @@ -2971,6 +2991,12 @@ def _check_model_access_helper( model_name=model, team_id=team_id ) + models = _resolve_all_team_model_sentinel_for_auth_check( + models=models, + llm_router=llm_router, + team_id=team_id, + ) + if ( len(access_groups) > 0 and llm_router is not None ): # check if token contains any model access groups diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index d89afcffa9a..00aac0d48f9 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/litellm/proxy/auth/auth_checks_organization.py @@ -134,7 +134,9 @@ def get_user_organization_info( for _membership in user_object.organization_memberships: if _membership.organization_id is not None: _user_organizations.append(_membership.organization_id) - _user_organization_role_mapping[_membership.organization_id] = _membership.user_role # type: ignore + _user_organization_role_mapping[_membership.organization_id] = ( + _membership.user_role + ) # type: ignore return _user_organizations, _user_organization_role_mapping diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 90845dfd824..e73c09719f5 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1970,9 +1970,8 @@ class JWTAuthManager: """Main authentication and authorization builder""" # Check if OIDC UserInfo endpoint is enabled, but fall back to standard # JWT auth if the token itself is a well-formed JWT (3-part structure). - if ( - jwt_handler.litellm_jwtauth.oidc_userinfo_enabled - and not jwt_handler.is_jwt(token=api_key) + if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt( + token=api_key ): verbose_proxy_logger.debug( "OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint." @@ -2173,16 +2172,18 @@ class JWTAuthManager: # If JWT did not resolve team_id, attempt single-team DB fallback. if team_id is None: - team_id, team_object, team_membership_object = ( - await JWTAuthManager._resolve_single_team_fallback( - user_object=user_object, - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, - ) + ( + team_id, + team_object, + team_membership_object, + ) = await JWTAuthManager._resolve_single_team_fallback( + user_object=user_object, + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, ) ## MAP USER TO TEAMS diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index be0d83dfcdc..f1ef05c7ad4 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -6,9 +6,11 @@ External callers (public IPs) only see servers with available_on_public_internet """ import ipaddress +from dataclasses import dataclass from typing import Any, Dict, List, Optional, Union from fastapi import Request +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.auth_utils import _get_request_ip_address @@ -17,6 +19,34 @@ from litellm.proxy.auth.auth_utils import _get_request_ip_address # behaviour see an actionable message in their logs the first time it triggers. _warned_xff_without_trusted_ranges = False +# Error for the inverse footgun: requests arrive with an X-Forwarded-For header +# but use_x_forwarded_for is off, so the real client IP is silently dropped and +# "internal network only" access control trusts the load balancer's IP instead. +# Logged once per misconfiguration window (not per-request) so a flood of crafted +# XFF headers can't spam the logs; re-arms whenever use_x_forwarded_for is observed +# enabled, so a later rollback to disabled warns again. +_warned_xff_present_but_disabled = False + +_NUM_TRUSTED_HOPS_ADAPTER = TypeAdapter(int) + + +@dataclass(frozen=True, slots=True) +class _HopCountUnset: + """mcp_xff_num_trusted_hops is absent: keep the legacy leftmost-XFF path.""" + + +@dataclass(frozen=True, slots=True) +class _HopCountInvalid: + """mcp_xff_num_trusted_hops is present but unusable: fail closed, never legacy.""" + + +@dataclass(frozen=True, slots=True) +class _HopCount: + value: int + + +_HopCountSetting = Union[_HopCountUnset, _HopCountInvalid, _HopCount] + class IPAddressUtils: """Static utilities for IP-based MCP access control.""" @@ -152,12 +182,12 @@ class IPAddressUtils: if not _warned_xff_without_trusted_ranges: verbose_proxy_logger.warning( "use_x_forwarded_for is enabled but mcp_trusted_proxy_ranges " - "is not configured. X-Forwarded-* headers will NOT be " - "trusted, so MCP OAuth discovery URLs and access-control " - "client IPs will use the proxy's literal request values. " - "Set mcp_trusted_proxy_ranges in " - "general_settings to your reverse-proxy CIDR(s) to allow " - "X-Forwarded-* through." + "is not configured, so X-Forwarded-* headers will NOT be trusted. " + "MCP OAuth discovery URLs fall back to the proxy's literal request " + "URL, and MCP access-control client-IP resolution fails closed " + "(callers are treated as external). Set mcp_trusted_proxy_ranges in " + "general_settings to your reverse-proxy CIDR(s) to trust " + "X-Forwarded-*." ) _warned_xff_without_trusted_ranges = True return False @@ -166,6 +196,60 @@ class IPAddressUtils: trusted_networks = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges) return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks) + @staticmethod + def extract_client_ip_from_xff_hops( + xff_header: str, + num_trusted_hops: int, + ) -> Optional[str]: + """ + Resolve the originating client IP from an X-Forwarded-For chain by + counting ``num_trusted_hops`` entries from the right. + + Each trusted proxy appends the address it received the connection from, + so the right end of the chain is written by infrastructure while the + left end is attacker-controllable. Selecting the Nth entry from the + right, where N is the number of trusted appending proxies in front of + the gateway, yields the real client IP and discards any values a client + prepended to spoof an allowed address. + + Returns None when the chain has fewer than ``num_trusted_hops`` entries + or the selected entry is not a valid IP, so callers can fail closed. + """ + entries = tuple(part.strip() for part in xff_header.split(",") if part.strip()) + if num_trusted_hops < 1 or len(entries) < num_trusted_hops: + return None + candidate = entries[-num_trusted_hops] + try: + ipaddress.ip_address(candidate) + except ValueError: + return None + return candidate + + @staticmethod + def _resolve_num_trusted_hops(raw_num_trusted_hops: object) -> _HopCountSetting: + if raw_num_trusted_hops is None: + return _HopCountUnset() + try: + num_hops = _NUM_TRUSTED_HOPS_ADAPTER.validate_python(raw_num_trusted_hops) + except ValidationError: + verbose_proxy_logger.warning( + "Invalid mcp_xff_num_trusted_hops value %r; failing closed for " + "MCP client IP resolution. Set it to a positive integer, or " + "remove the setting to restore the legacy X-Forwarded-For path", + raw_num_trusted_hops, + ) + return _HopCountInvalid() + if num_hops < 1: + verbose_proxy_logger.warning( + "mcp_xff_num_trusted_hops=%s is below the minimum of 1; failing " + "closed for MCP client IP resolution. Set it to a positive " + "integer, or remove the setting to restore the legacy " + "X-Forwarded-For path", + num_hops, + ) + return _HopCountInvalid() + return _HopCount(num_hops) + @staticmethod def get_mcp_client_ip( request: Request, @@ -178,6 +262,12 @@ class IPAddressUtils: 1. use_x_forwarded_for is enabled in settings 2. The direct connection is from a trusted proxy (if mcp_trusted_proxy_ranges configured) + When ``mcp_xff_num_trusted_hops`` is set, the client IP is read that many + entries from the right of the chain instead of the spoofable leftmost + value, defeating append-style X-Forwarded-For forgery. A present-but-invalid + value (non-integer or below 1) fails closed rather than silently reverting + to the legacy path, so a config typo cannot quietly weaken access control. + Args: request: FastAPI request object general_settings: Optional settings dict. If not provided, imports from proxy_server. @@ -198,6 +288,28 @@ class IPAddressUtils: use_xff = general_settings.get("use_x_forwarded_for", False) + global _warned_xff_present_but_disabled + if use_xff: + _warned_xff_present_but_disabled = False + elif "x-forwarded-for" in request.headers: + if not _warned_xff_present_but_disabled: + verbose_proxy_logger.error( + "Received a request with an X-Forwarded-For header but " + "use_x_forwarded_for is not enabled. The real client IP is " + "being ignored and the direct peer's IP (typically your load " + "balancer / reverse proxy) is used for MCP access control. " + "Because that peer almost always falls within " + "general_settings.mcp_internal_ip_ranges, every external caller " + "is treated as internal and 'available_on_public_internet: " + "false' MCP servers are effectively exposed. Set " + "use_x_forwarded_for: true (and mcp_trusted_proxy_ranges to " + "your proxy CIDRs) in general_settings to honor the real " + "client IP. Not failing the request: if there is no load " + "balancer, a crafted X-Forwarded-For header must not be able " + "to take down the service." + ) + _warned_xff_present_but_disabled = True + # If XFF is enabled, validate the request comes from a trusted proxy if use_xff and "x-forwarded-for" in request.headers: if not IPAddressUtils.is_request_from_trusted_proxy( @@ -215,4 +327,24 @@ class IPAddressUtils: # returning it would mis-classify external callers as internal. # Fail closed for access control. return "" + match IPAddressUtils._resolve_num_trusted_hops( + general_settings.get("mcp_xff_num_trusted_hops") + ): + case _HopCountInvalid(): + return "" + case _HopCount(value=num_trusted_hops): + client_ip = IPAddressUtils.extract_client_ip_from_xff_hops( + request.headers["x-forwarded-for"], num_trusted_hops + ) + if client_ip is None: + verbose_proxy_logger.warning( + "X-Forwarded-For chain has fewer than " + "mcp_xff_num_trusted_hops=%s entries or an invalid " + "address; failing closed", + num_trusted_hops, + ) + return "" + return client_ip + case _HopCountUnset(): + pass return _get_request_ip_address(request, use_x_forwarded_for=use_xff) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index b89db51c6f1..aa53954da8f 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -122,9 +122,16 @@ def get_key_models( SpecialModelNames.all_team_models.value in all_models and user_api_key_dict.team_id is not None ): - all_models = list( - user_api_key_dict.team_models - ) # copy to avoid mutating cached objects + all_models = list(user_api_key_dict.team_models) + if SpecialModelNames.all_team_models.value in all_models: + all_models = [ + model + for model in all_models + if model != SpecialModelNames.all_team_models.value + ] + all_models.extend(proxy_model_list) + if include_model_access_groups: + all_models.extend(model_access_groups.keys()) if SpecialModelNames.all_proxy_models.value in all_models: all_models = list(proxy_model_list) # copy to avoid mutating caller's list if include_model_access_groups: @@ -160,6 +167,12 @@ def get_team_models( all_models_set.update(team_models) if SpecialModelNames.all_team_models.value in all_models_set: all_models_set.update(team_models) + # GH#30619: expand all-team-models sentinel + # to the actual proxy model list + all_models_set.discard(SpecialModelNames.all_team_models.value) + all_models_set.update(proxy_model_list) + if include_model_access_groups: + all_models_set.update(model_access_groups.keys()) if SpecialModelNames.all_proxy_models.value in all_models_set: all_models_set.update(proxy_model_list) if include_model_access_groups: diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 344f90aa144..6db75eeb3d9 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -282,7 +282,8 @@ async def create_batch( else: # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) response = await litellm.acreate_batch( - custom_llm_provider=custom_llm_provider, **_create_batch_data # type: ignore + custom_llm_provider=custom_llm_provider, + **_create_batch_data, # type: ignore ) ### CALL HOOKS ### - modify outgoing data @@ -523,7 +524,8 @@ async def retrieve_batch( or "openai" ) response = await litellm.aretrieve_batch( - custom_llm_provider=custom_llm_provider, **data # type: ignore + custom_llm_provider=custom_llm_provider, + **data, # type: ignore ) # FIX: Update the database with the latest state from provider @@ -735,7 +737,9 @@ async def list_batches( ## POST CALL HOOKS ### _response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response # type: ignore + data=data, + user_api_key_dict=user_api_key_dict, + response=response, # type: ignore ) if _response is not None and type(response) is type(_response): response = _response diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8ef931e8d25..92f8dd8df98 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1037,6 +1037,8 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + if not general_settings.get("expose_fallback_errors_to_caller"): + self.data.pop("include_fallback_errors", None) if route_type in {"aresponses", "_aresponses_websocket"}: await _authorize_response_file_search_vector_stores( data=self.data, @@ -1061,9 +1063,9 @@ class ProxyBaseLLMRequestProcessing: self.data[_metadata_variable_name] = {} if not isinstance(self.data[_metadata_variable_name], dict): self.data[_metadata_variable_name] = {} - self.data[_metadata_variable_name][ - "queue_time_seconds" - ] = queue_time_seconds + self.data[_metadata_variable_name]["queue_time_seconds"] = ( + queue_time_seconds + ) self.data["model"] = ( general_settings.get("completion_model", None) # server default @@ -1506,7 +1508,9 @@ class ProxyBaseLLMRequestProcessing: cache_hit=cache_hit, ) - logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[union-attr] + logging_obj._on_deferred_stream_complete = ( + _on_deferred_stream_complete # type: ignore[union-attr] + ) if route_type == "allm_passthrough_route": # Check if response is an async generator diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d48499af6f0..cd41bce97d6 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -186,11 +186,8 @@ def initialize_callbacks_on_proxy( ) init_params = {} - if ( - "lakera_prompt_injection" in callback_specific_params - and isinstance( - callback_specific_params["lakera_prompt_injection"], dict - ) + if "lakera_prompt_injection" in callback_specific_params and isinstance( + callback_specific_params["lakera_prompt_injection"], dict ): init_params = callback_specific_params["lakera_prompt_injection"] lakera_moderations_object = lakeraAI_Moderation(**init_params) @@ -343,11 +340,8 @@ def initialize_callbacks_on_proxy( ) init_params = {} - if ( - "datadog_cost_management" in callback_specific_params - and isinstance( - callback_specific_params["datadog_cost_management"], dict - ) + if "datadog_cost_management" in callback_specific_params and isinstance( + callback_specific_params["datadog_cost_management"], dict ): init_params = callback_specific_params["datadog_cost_management"] datadog_cost_management_obj = DatadogCostManagementLogger(**init_params) diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index fa3cb02195b..ce5d2539fb3 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -156,9 +156,9 @@ class CustomOpenAPISpec: filtered_params = [ param for param in existing_params if param.get("in") == "path" ] - openapi_schema["paths"][path]["post"][ - "parameters" - ] = filtered_params + openapi_schema["paths"][path]["post"]["parameters"] = ( + filtered_params + ) @staticmethod def _move_defs_to_components( diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 7eeb11fc372..52bd9d66b7d 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -390,12 +390,13 @@ async def _process_request( # Validate container_id ownership if present in path_params. if "container_id" in path_params: - original_container_id, resolved_provider = ( - await assert_user_can_access_container( - container_id=path_params["container_id"], - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, - ) + ( + original_container_id, + resolved_provider, + ) = await assert_user_can_access_container( + container_id=path_params["container_id"], + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, ) data.update( await get_container_forwarding_params( diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index aab92a54577..57a062b509a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -843,9 +843,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_transactions, daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions, - ) = ( - await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() - ) + ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() if db_spend_update_transactions is not None: verbose_proxy_logger.info( @@ -960,9 +958,7 @@ class DBSpendUpdateWriter: # Aggregate all in memory spend updates (key, user, end_user, team, team_member, org) and commit to db ################## Spend Update Transactions ################## - db_spend_update_transactions = ( - await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() - ) + db_spend_update_transactions = await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, @@ -1089,9 +1085,7 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for daily tag spend updates") try: - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) + daily_tag_spend_update_transactions = await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() if daily_tag_spend_update_transactions: await DBSpendUpdateWriter.update_daily_tag_spend( diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index f47b694d44e..a72ffad7e9b 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -73,9 +73,9 @@ class DailySpendUpdateQueue(BaseUpdateQueue): Combine all updates in the queue into a single update. This is used to reduce the size of the in-memory queue. """ - updates: List[Dict[str, BaseDailySpendTransaction]] = ( - await self.flush_all_updates_from_in_memory_queue() - ) + updates: List[ + Dict[str, BaseDailySpendTransaction] + ] = await self.flush_all_updates_from_in_memory_queue() aggregated_updates = self.get_aggregated_daily_spend_update_transactions( updates ) diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 1e3014dbf3c..6cbfb37396c 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -186,24 +186,12 @@ class RedisUpdateBuffer: return # Get all transactions - db_spend_update_transactions = ( - await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() - ) - daily_spend_update_transactions = ( - await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) - daily_team_spend_update_transactions = ( - await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) - daily_org_spend_update_transactions = ( - await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) - daily_end_user_spend_update_transactions = ( - await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) - daily_agent_spend_update_transactions = ( - await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) + db_spend_update_transactions = await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + daily_spend_update_transactions = await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + daily_team_spend_update_transactions = await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + daily_org_spend_update_transactions = await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + daily_end_user_spend_update_transactions = await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + daily_agent_spend_update_transactions = await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() verbose_proxy_logger.debug( "ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions @@ -576,9 +564,7 @@ class RedisUpdateBuffer: """ Flush in-memory daily tag spend updates and append them to Redis. """ - daily_tag_spend_update_transactions = ( - await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) + daily_tag_spend_update_transactions = await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() await self._store_transactions_in_redis( transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 8100a1e8a12..727e8dc1d5a 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -53,9 +53,9 @@ class SpendUpdateQueue(BaseUpdateQueue): async def aggregate_queue_updates(self): """Concatenate all updates in the queue to reduce the size of in-memory queue""" - updates: List[SpendUpdateQueueItem] = ( - await self.flush_all_updates_from_in_memory_queue() - ) + updates: List[ + SpendUpdateQueueItem + ] = await self.flush_all_updates_from_in_memory_queue() aggregated_updates = self._get_aggregated_spend_update_queue_item(updates) for update in aggregated_updates: await self.update_queue.put(update) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 58478db5e2e..ae2307658dd 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -32,7 +32,7 @@ password when their ``*_READ_REPLICA`` counterpart is unset. import os import urllib.parse -from typing import Optional, cast +from typing import Final, cast from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -44,6 +44,41 @@ from litellm.proxy.auth import rds_iam_token _IAM_ENV_KEY = "IAM_TOKEN_DB_AUTH" _DEFAULT_PG_PORT = "5432" +# schema.prisma pins `provider = "postgresql"`, so these are the only schemes +# Prisma can actually connect with. +SUPPORTED_DB_SCHEMES: Final[frozenset[str]] = frozenset({"postgresql", "postgres"}) +_MISSING_SCHEME = "" + + +def unsupported_db_scheme(database_url: str) -> str | None: + """Return the connection URL scheme when it is not PostgreSQL, else None. + + A `sqlite://` / `mysql://` URL can never connect against the + postgresql-only datasource, but the resulting Prisma failure is opaque and + version-dependent (a confusing migration error, or a startup that never + binds). Callers use this to reject the URL up front with an actionable + error instead. + + A schemeless value (e.g. a malformed DSN like ``user:pass@host/db``) yields + the ``_MISSING_SCHEME`` placeholder rather than the raw URL, so callers that + log the return value never echo embedded credentials. + """ + scheme = urllib.parse.urlsplit(database_url).scheme.lower() + if scheme in SUPPORTED_DB_SCHEMES: + return None + return scheme or _MISSING_SCHEME + + +def unsupported_db_scheme_message(env_var: str, scheme: str) -> str: + """Operator-facing message naming the offending env var and scheme.""" + return ( + f"{env_var} uses unsupported scheme '{scheme}'. LiteLLM's database " + "features (virtual keys, store_model_in_db, spend tracking) require " + "PostgreSQL; use a 'postgresql://' connection string. SQLite and other " + "engines are not supported. " + "See https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + class DatabaseURLSettings(BaseSettings): """Discrete ``DATABASE_*`` env vars, loaded once at process start. @@ -58,46 +93,47 @@ class DatabaseURLSettings(BaseSettings): iam_token_db_auth: bool = Field(default=False, validation_alias=_IAM_ENV_KEY) # Writer - database_url: Optional[str] = Field(default=None, validation_alias="DATABASE_URL") - database_host: Optional[str] = Field(default=None, validation_alias="DATABASE_HOST") + database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") + direct_url: str | None = Field(default=None, validation_alias="DIRECT_URL") + database_host: str | None = Field(default=None, validation_alias="DATABASE_HOST") database_port: str = Field( default=_DEFAULT_PG_PORT, validation_alias="DATABASE_PORT" ) - database_user: Optional[str] = Field( + database_user: str | None = Field( default=None, validation_alias=AliasChoices("DATABASE_USER", "DATABASE_USERNAME"), ) - database_name: Optional[str] = Field(default=None, validation_alias="DATABASE_NAME") - database_schema: Optional[str] = Field( + database_name: str | None = Field(default=None, validation_alias="DATABASE_NAME") + database_schema: str | None = Field( default=None, validation_alias="DATABASE_SCHEMA" ) - database_password: Optional[str] = Field( + database_password: str | None = Field( default=None, validation_alias="DATABASE_PASSWORD" ) # Read replica - database_url_read_replica: Optional[str] = Field( + database_url_read_replica: str | None = Field( default=None, validation_alias="DATABASE_URL_READ_REPLICA" ) - database_host_read_replica: Optional[str] = Field( + database_host_read_replica: str | None = Field( default=None, validation_alias="DATABASE_HOST_READ_REPLICA" ) - database_port_read_replica: Optional[str] = Field( + database_port_read_replica: str | None = Field( default=None, validation_alias="DATABASE_PORT_READ_REPLICA" ) - database_user_read_replica: Optional[str] = Field( + database_user_read_replica: str | None = Field( default=None, validation_alias=AliasChoices( "DATABASE_USER_READ_REPLICA", "DATABASE_USERNAME_READ_REPLICA" ), ) - database_name_read_replica: Optional[str] = Field( + database_name_read_replica: str | None = Field( default=None, validation_alias="DATABASE_NAME_READ_REPLICA" ) - database_schema_read_replica: Optional[str] = Field( + database_schema_read_replica: str | None = Field( default=None, validation_alias="DATABASE_SCHEMA_READ_REPLICA" ) - database_password_read_replica: Optional[str] = Field( + database_password_read_replica: str | None = Field( default=None, validation_alias="DATABASE_PASSWORD_READ_REPLICA" ) @@ -106,7 +142,7 @@ class DatabaseURLSettings(BaseSettings): """Load the settings from ``os.environ`` (read at call time).""" return cls() - def build_writer_url(self) -> Optional[str]: + def build_writer_url(self) -> str | None: """Return the writer URL to set, or ``None`` to leave it as-is. Raises ``RuntimeError`` (naming the offending vars) when IAM auth is @@ -156,7 +192,7 @@ class DatabaseURLSettings(BaseSettings): ) return None - def build_reader_url(self) -> Optional[str]: + def build_reader_url(self) -> str | None: """Return the read-replica URL to set, or ``None`` to leave it as-is. Opt-in via ``DATABASE_HOST_READ_REPLICA``; never clobbers a @@ -217,11 +253,11 @@ class DatabaseURLSettings(BaseSettings): def _password_url( *, user: str, - password: Optional[str], + password: str | None, host: str, port: str, name: str, - schema: Optional[str], + schema: str | None, ) -> str: """Percent-encode credentials into a ``postgresql://`` URL. @@ -239,6 +275,26 @@ class DatabaseURLSettings(BaseSettings): url += f"?schema={schema}" return url + def _raise_for_unsupported_scheme(self) -> None: + """Reject an operator-pinned non-PostgreSQL writer / direct / reader URL. + + The componentized entrypoints (gateway / backend / migrations) call + ``apply_to_env`` and then hand the URL straight to Prisma, bypassing + the CLI's own guard. A pinned URL flows through untouched, so validate + the same three vars the CLI guard checks (DATABASE_URL, DIRECT_URL, and + the read replica) rather than letting Prisma stall on an unusable scheme. + """ + for env_var, url in ( + ("DATABASE_URL", self.database_url), + ("DIRECT_URL", self.direct_url), + ("DATABASE_URL_READ_REPLICA", self.database_url_read_replica), + ): + if not url: + continue + bad_scheme = unsupported_db_scheme(url) + if bad_scheme is not None: + raise RuntimeError(unsupported_db_scheme_message(env_var, bad_scheme)) + def apply_to_env(self) -> bool: """Write the assembled URL(s) into ``os.environ``. @@ -246,6 +302,7 @@ class DatabaseURLSettings(BaseSettings): password auth that assembled a fresh URL). False means there was nothing to do — an operator-pinned URL, or no discrete fields. """ + self._raise_for_unsupported_scheme() wrote_writer = False writer_url = self.build_writer_url() if writer_url is not None: diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 628509d9c36..57ebb9678cb 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -25,7 +25,10 @@ class DynamoDBWrapper(CustomDB): and database_arguments.write_capacity_units is not None and isinstance(database_arguments.write_capacity_units, int) ): - self.throughput_type = Throughput(read=database_arguments.read_capacity_units, write=database_arguments.write_capacity_units) # type: ignore + self.throughput_type = Throughput( + read=database_arguments.read_capacity_units, + write=database_arguments.write_capacity_units, + ) # type: ignore else: raise Exception( f"Invalid args passed in. Need to set both read_capacity_units and write_capacity_units. Args passed in - {database_arguments}" diff --git a/litellm/proxy/db/log_db_metrics.py b/litellm/proxy/db/log_db_metrics.py index eb4961062df..837e94f1a85 100644 --- a/litellm/proxy/db/log_db_metrics.py +++ b/litellm/proxy/db/log_db_metrics.py @@ -66,9 +66,7 @@ def log_db_metrics(func): elif ( # in litellm custom callbacks kwargs is passed as arg[0] # https://docs.litellm.ai/docs/observability/custom_callback#callback-functions - args is not None - and len(args) > 1 - and isinstance(args[1], dict) + args is not None and len(args) > 1 and isinstance(args[1], dict) ): passed_kwargs = args[1] parent_otel_span = _get_parent_otel_span_from_kwargs( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 9f8ea584103..c2175ce95e0 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -617,9 +617,7 @@ class GuardrailSubmissionItem(BaseModel): guardrail_name: str status: str # pending_review | active | rejected team_id: Optional[str] = None - team_guardrail: bool = ( - False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails - ) + team_guardrail: bool = False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails litellm_params: Optional[Dict[str, Any]] = None guardrail_info: Optional[Dict[str, Any]] = None submitted_by_user_id: Optional[str] = None @@ -2322,16 +2320,17 @@ async def apply_guardrail( ) request_processor = ProxyBaseLLMRequestProcessing(data=data) - data, litellm_logging_obj = ( - await request_processor.common_processing_pre_call_logic( - request=fastapi_request, - general_settings=general_settings, - user_api_key_dict=user_api_key_dict, - version=version, - proxy_logging_obj=proxy_logging_obj, - proxy_config=proxy_config, - route_type="apply_guardrail", - ) + ( + data, + litellm_logging_obj, + ) = await request_processor.common_processing_pre_call_logic( + request=fastapi_request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type="apply_guardrail", ) if litellm_logging_obj is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 55b446eb529..e21be6ffdbe 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -65,9 +65,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr **kwargs, ) - self.optional_params_request_body: ( - AzureTextModerationRequestBodyOptionalParams - ) = { + self.optional_params_request_body: AzureTextModerationRequestBodyOptionalParams = { "categories": kwargs.get("categories") or [ "Hate", diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index ba2f531f26e..79e81b6c4c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self.guardrail_name, type(all_chunks[0]).__name__, ) - yield f'data: {json.dumps({"error": {"message": "Cisco AI Defense: unsupported streaming format — response withheld for safety", "type": "guardrail_unsupported_stream", "code": 400, "guardrail": self.guardrail_name}})}\n\n' + yield f"data: {json.dumps({'error': {'message': 'Cisco AI Defense: unsupported streaming format — response withheld for safety', 'type': 'guardrail_unsupported_stream', 'code': 400, 'guardrail': self.guardrail_name}})}\n\n" return assembled = stream_chunk_builder(chunks=all_chunks) @@ -546,7 +546,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self.guardrail_name, type(assembled).__name__, ) - yield f'data: {json.dumps({"error": {"message": "Cisco AI Defense: unsupported streaming format — response withheld for safety", "type": "guardrail_unsupported_stream", "code": 400, "guardrail": self.guardrail_name}})}\n\n' + yield f"data: {json.dumps({'error': {'message': 'Cisco AI Defense: unsupported streaming format — response withheld for safety', 'type': 'guardrail_unsupported_stream', 'code': 400, 'guardrail': self.guardrail_name}})}\n\n" return response_messages = self._extract_response_messages(assembled) @@ -586,14 +586,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return except Exception as exc: verbose_proxy_logger.error( - "Cisco AI Defense guardrail (%s): streaming response " - "scan failed: %s", + "Cisco AI Defense guardrail (%s): streaming response scan failed: %s", self.guardrail_name, exc, ) error_obj = { "message": ( - "Cisco AI Defense streaming scan failed — response " "withheld." + "Cisco AI Defense streaming scan failed — response withheld." ), "type": "guardrail_scan_error", "code": 500, @@ -932,8 +931,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ) from exc except httpx.TimeoutException as exc: raise CiscoAIDefenseGuardrailAPIError( - f"Cisco AI Defense {surface} API call timed out after " - f"{self.timeout}s" + f"Cisco AI Defense {surface} API call timed out after {self.timeout}s" ) from exc except httpx.RequestError as exc: raise CiscoAIDefenseGuardrailAPIError( @@ -1178,8 +1176,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ) if redacted: verbose_proxy_logger.info( - "Cisco AI Defense guardrail (%s): redaction applied " - "(event_id=%s)", + "Cisco AI Defense guardrail (%s): redaction applied (event_id=%s)", context.surface, verdict.event_id, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index c6dfe141ab5..cc944fb46dc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -212,17 +212,17 @@ class ContentFilterGuardrail(CustomGuardrail): self.image_model = image_model # Store loaded categories self.loaded_categories: Dict[str, CategoryConfig] = {} - self.category_keywords: Dict[str, Tuple[str, str, ContentFilterAction]] = ( - {} - ) # keyword -> (category, severity, action) + self.category_keywords: Dict[ + str, Tuple[str, str, ContentFilterAction] + ] = {} # keyword -> (category, severity, action) # Always-block keywords are checked after exceptions (exceptions take precedence) self.always_block_category_keywords: Dict[ str, Tuple[str, str, ContentFilterAction] ] = {} # Store conditional categories (identifier_words + block_words) - self.conditional_categories: Dict[str, Dict[str, Any]] = ( - {} - ) # category_name -> {identifier_words, block_words, action, severity} + self.conditional_categories: Dict[ + str, Dict[str, Any] + ] = {} # category_name -> {identifier_words, block_words, action, severity} # Competitor intent checker (optional; airline uses major_airlines.json, generic requires competitors) self._competitor_intent_checker: Optional[BaseCompetitorIntentChecker] = None diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index 45bdd6dd09f..bc234c0a6f8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -68,11 +68,11 @@ def _build_judge_prompt( response_text: str, ) -> str: criteria_block = "\n".join( - f'- {c.get("name", "")} (weight {c.get("weight", 0)}%): {c.get("description", "")}' + f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" for c in criteria ) conversation = "\n".join( - f'{m.get("role", "user").upper()}: {_extract_text_from_content(m.get("content", ""))}' + f"{m.get('role', 'user').upper()}: {_extract_text_from_content(m.get('content', ''))}" for m in messages if m.get("content") is not None ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index d00e77daa07..7334e82549b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -297,7 +297,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): filters = ( list(filter_results.values()) if isinstance(filter_results, dict) - else filter_results if isinstance(filter_results, list) else [] + else filter_results + if isinstance(filter_results, list) + else [] ) # Prefer sanitized text from deidentifyResult if present diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index d80a1a299d5..aadfdf66531 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -185,8 +185,10 @@ class NomaGuardrail(CustomGuardrail): if not messages: return None - input_items, instructions = self._responses_transform_handler.convert_chat_completion_messages_to_responses_api( # type: ignore[arg-type] - messages + input_items, instructions = ( + self._responses_transform_handler.convert_chat_completion_messages_to_responses_api( # type: ignore[arg-type] + messages + ) ) if instructions: diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index e8887fa712a..04d68edeb80 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -1502,7 +1502,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): yield f"data: {json.dumps({'error': error_obj})}\n\n" except Exception as e: verbose_proxy_logger.error(f"PANW Prisma AIRS streaming error: {str(e)}") - yield f'data: {json.dumps({"error": {"message": "Security scan failed - streaming response blocked for safety", "type": "guardrail_scan_error", "code": 500, "guardrail": self.guardrail_name}})}\n\n' + yield f"data: {json.dumps({'error': {'message': 'Security scan failed - streaming response blocked for safety', 'type': 'guardrail_scan_error', 'code': 500, 'guardrail': self.guardrail_name}})}\n\n" async def _scan_tool_calls_for_guardrail( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index a8afe4efe2a..1367f95c5c8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -93,9 +93,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): kwargs["event_hook"] = GuardrailEventHooks.logging_only super().__init__(**kwargs) self.guardrail_provider = "presidio" - self.pii_tokens: dict = ( - {} - ) # mapping of PII token to original text - only used with Presidio `replace` operation + self.pii_tokens: dict = {} # mapping of PII token to original text - only used with Presidio `replace` operation self.mock_redacted_text = mock_redacted_text self.output_parse_pii = output_parse_pii or False self.apply_to_output = apply_to_output @@ -164,15 +162,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): presidio_analyzer_api_base: Optional[str] = None, presidio_anonymizer_api_base: Optional[str] = None, ): - self.presidio_analyzer_api_base: Optional[ - str - ] = presidio_analyzer_api_base or get_secret( - "PRESIDIO_ANALYZER_API_BASE", None + self.presidio_analyzer_api_base: Optional[str] = ( + presidio_analyzer_api_base or get_secret("PRESIDIO_ANALYZER_API_BASE", None) ) # type: ignore - self.presidio_anonymizer_api_base: Optional[ - str - ] = presidio_anonymizer_api_base or litellm.get_secret( - "PRESIDIO_ANONYMIZER_API_BASE", None + self.presidio_anonymizer_api_base: Optional[str] = ( + presidio_anonymizer_api_base + or litellm.get_secret("PRESIDIO_ANONYMIZER_API_BASE", None) ) # type: ignore if self.presidio_analyzer_api_base is None: @@ -762,9 +757,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if messages is None: return data tasks = [] - task_mappings: List[Tuple[int, Optional[int]]] = ( - [] - ) # Track (message_index, content_index) for each task + task_mappings: List[ + Tuple[int, Optional[int]] + ] = [] # Track (message_index, content_index) for each task for msg_idx, m in enumerate(messages): content = m.get("content", None) @@ -808,9 +803,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if content is None: continue if isinstance(content, str) and content_idx_optional is None: - messages[msg_idx][ - "content" - ] = r # replace content with redacted string + messages[msg_idx]["content"] = ( + r # replace content with redacted string + ) elif isinstance(content, list) and content_idx_optional is not None: messages[msg_idx]["content"][content_idx_optional]["text"] = r @@ -865,9 +860,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ): # /chat/completions requests messages: Optional[List] = kwargs.get("messages", None) tasks = [] - task_mappings: List[Tuple[int, Optional[int]]] = ( - [] - ) # Track (message_index, content_index) for each task + task_mappings: List[ + Tuple[int, Optional[int]] + ] = [] # Track (message_index, content_index) for each task if messages is None: return kwargs, result @@ -916,9 +911,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if content is None: continue if isinstance(content, str) and content_idx_optional is None: - messages[msg_idx][ - "content" - ] = r # replace content with redacted string + messages[msg_idx]["content"] = ( + r # replace content with redacted string + ) elif isinstance(content, list) and content_idx_optional is not None: messages[msg_idx]["content"][content_idx_optional]["text"] = r diff --git a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py index 34f38036265..a796fedbe27 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py @@ -203,12 +203,10 @@ class RepelloAIGuardrail(CustomGuardrail): repelloai_response: RepelloAIAnalyzeResponse | None = None try: verbose_proxy_logger.debug("RepelloAI Argus request: %s", request) - raw_response: HttpxResponse | None = ( - await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] - url=endpoint, - headers={"X-API-Key": self.repelloai_api_key}, - json=request, - ) + raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] + url=endpoint, + headers={"X-API-Key": self.repelloai_api_key}, + json=request, ) if raw_response is None: raise ValueError("RepelloAI Argus returned no response") @@ -235,7 +233,9 @@ class RepelloAIGuardrail(CustomGuardrail): return repelloai_response except HTTPException as e: status = "guardrail_failed_to_respond" - guardrail_json_response = str(e.detail) if not isinstance(e.detail, (dict, list)) else e.detail # type: ignore[assignment] + guardrail_json_response = ( + str(e.detail) if not isinstance(e.detail, (dict, list)) else e.detail + ) # type: ignore[assignment] raise except HTTPError as e: status = "guardrail_failed_to_respond" diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index b99ea8f14a0..734330a2167 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -479,7 +479,9 @@ class InMemoryGuardrailHandler: sig = inspect.signature(initializer) if "llm_router" in sig.parameters: custom_guardrail_callback = initializer( - litellm_params, guardrail, llm_router # type: ignore + litellm_params, + guardrail, + llm_router, # type: ignore ) else: custom_guardrail_callback = initializer(litellm_params, guardrail) diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 83f1281dc02..9e45f18232c 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -137,7 +137,9 @@ def initialize_guardrails( if guardrail.logging_only is True: if callback == "presidio": - callback_specific_params["presidio"] = {"logging_only": True} # type: ignore + callback_specific_params["presidio"] = { + "logging_only": True + } # type: ignore default_on_callbacks_list = list(default_on_callbacks) if len(default_on_callbacks_list) > 0: diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 488467e1b99..ed91645eb43 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -194,7 +194,8 @@ async def _run_model_health_check(model: dict): litellm_params = model["litellm_params"] model_info = model.get("model_info", {}) mode = _resolve_health_check_mode( - model_info, litellm_params # any-ok: untyped router config dict + model_info, + litellm_params, # any-ok: untyped router config dict ) litellm_params = _update_litellm_params_for_health_check(model_info, litellm_params) timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS @@ -454,11 +455,13 @@ def _update_litellm_params_for_health_check( - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID, and pins `custom_llm_provider` to `bedrock` (only when the deployment hasn't already set one, so an explicit `bedrock_converse` survives) so the bare model id still resolves to the provider (e.g. cross-region ids like `us.cohere.embed-v4:0`) """ mode = _resolve_health_check_mode( - model_info, litellm_params # any-ok: untyped router config dict + model_info, + litellm_params, # any-ok: untyped router config dict ) litellm_params["messages"] = _get_random_llm_message() if _should_inject_health_check_max_tokens( - model_info, mode # any-ok: untyped router config dict + model_info, + mode, # any-ok: untyped router config dict ): _resolved_max_tokens = _resolve_health_check_max_tokens( model_info, litellm_params diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 5b691beccbf..91c604e6204 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -21,6 +21,7 @@ from typing import ( TYPE_CHECKING, Any, Dict, + Iterable, List, Literal, NoReturn, @@ -32,13 +33,15 @@ from typing import ( from fastapi import HTTPException from pydantic import BaseModel +import json + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import ( + _count_entry_tokens, + _estimate_batch_entry_tokens, _extract_file_access_credentials, - _get_batch_job_input_file_usage, - _get_file_content_as_dictionary, - _get_models_from_batch_input_file_content, + _iter_batch_input_lines, ) from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger @@ -537,6 +540,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) + # For managed files the unified file id encodes the proxy model + # alias(es) the file was uploaded for; auth validates against those. target_model_names = ( get_models_from_unified_file_id(is_managed_file) if is_managed_file @@ -568,7 +573,38 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Expected bytes content from file retrieval for {file_id}, " f"got {type(file_content_bytes)}" ) - file_content_as_dict = _get_file_content_as_dictionary(file_content_bytes) + + # Single streaming pass over the JSONL lines, accounting each row + # independently. One bad row can never abort the pass: a malformed + # line is skipped (its request can't run upstream anyway) and a row + # the token counter can't measure falls back to a conservative + # size-based estimate. This guarantees two things a restricted caller + # must not be able to break by crafting a row that raises: + # 1. The allowlist check below always sees every parseable + # ``body.model`` (the loop never stops early), so models can't be + # smuggled in after a bad row. + # 2. The token total is never silently zeroed, so the TPM limit + # can't be evaded by sending uncountable rows. + # Counting stays best-effort, so a legitimate (e.g. multimodal) row + # the counter can't measure is estimated, not hard-rejected. + models: set = set() + total_tokens = 0 + request_count = 0 + for raw_line in _iter_batch_input_lines(file_content_bytes): + request_count += 1 + try: + entry = json.loads(raw_line) + except Exception: + total_tokens += _estimate_batch_entry_tokens(raw_line) + continue + if isinstance(entry, dict): + model = (entry.get("body") or {}).get("model") + if model: + models.add(model) + try: + total_tokens += _count_entry_tokens(entry) + except Exception: + total_tokens += _estimate_batch_entry_tokens(raw_line) # Validate every model named in the batch JSONL against the # caller's per-key model allowlist. Without this, a caller @@ -578,17 +614,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): if user_api_key_dict is not None: await self._enforce_batch_file_model_access( user_api_key_dict=user_api_key_dict, - file_content_as_dict=file_content_as_dict, + models=models, target_model_names=target_model_names or None, ) - input_file_usage = _get_batch_job_input_file_usage( - file_content_dictionary=file_content_as_dict, - custom_llm_provider=custom_llm_provider, - ) - request_count = len(file_content_as_dict) return BatchFileUsage( - total_tokens=input_file_usage.total_tokens, + total_tokens=total_tokens, request_count=request_count, ) @@ -614,14 +645,15 @@ class _PROXY_BatchRateLimiter(CustomLogger): async def _enforce_batch_file_model_access( self, user_api_key_dict: UserAPIKeyAuth, - file_content_as_dict: List[dict], + models: Optional[Iterable[str]] = None, target_model_names: Optional[List[str]] = None, ) -> None: """Reject the batch if the caller is not authorized for the upload target. For managed files, ``target_model_names`` (from the unified file id) is - the proxy alias the file was uploaded for and is used directly for auth. - For legacy/non-managed files, falls back to ``body.model`` values in the JSONL. + the proxy alias the file was uploaded for and is checked directly. + Otherwise the ``body.model`` values collected from the JSONL (``models``) + are checked. Reuses standard auth helpers so the same model access rules the proxy enforces on `/chat/completions` apply here. @@ -640,10 +672,9 @@ class _PROXY_BatchRateLimiter(CustomLogger): if target_model_names: models = target_model_names - else: - models = _get_models_from_batch_input_file_content(file_content_as_dict) - if not models: - return + + if not models: + return team_object = None if ( diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index b9e2bd12ecf..7edc5f4698c 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -272,10 +272,10 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): model_info = self.llm_router.get_model_info( id=response._hidden_params["model_id"] ) - assert ( - model_info is not None - ), "Model info for model with id={} is None".format( - response._hidden_params["model_id"] + assert model_info is not None, ( + "Model info for model with id={} is None".format( + response._hidden_params["model_id"] + ) ) key_priority: Optional[str] = user_api_key_dict.metadata.get( "priority", None @@ -289,16 +289,16 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) = await self.check_available_usage( model=model_info["model_name"], priority=key_priority ) - response._hidden_params["additional_headers"] = ( - { # Add additional response headers - easier debugging - "x-litellm-model_group": model_info["model_name"], - "x-ratelimit-remaining-litellm-project-tokens": available_tpm, - "x-ratelimit-remaining-litellm-project-requests": available_rpm, - "x-ratelimit-remaining-model-tokens": model_tpm, - "x-ratelimit-remaining-model-requests": model_rpm, - "x-ratelimit-current-active-projects": active_projects, - } - ) + response._hidden_params[ + "additional_headers" + ] = { # Add additional response headers - easier debugging + "x-litellm-model_group": model_info["model_name"], + "x-ratelimit-remaining-litellm-project-tokens": available_tpm, + "x-ratelimit-remaining-litellm-project-requests": available_rpm, + "x-ratelimit-remaining-model-tokens": model_tpm, + "x-ratelimit-remaining-model-requests": model_rpm, + "x-ratelimit-current-active-projects": active_projects, + } return response return await super().async_post_call_success_hook( diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 6bb9c1f507c..ee6c063e8e7 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -376,10 +376,10 @@ class KeyManagementEventHooks: if key.key_alias is not None: team_id = getattr(key, "team_id", None) if team_id not in team_settings_cache: - team_settings_cache[team_id] = ( - await KeyManagementEventHooks._get_secret_manager_optional_params( - team_id - ) + team_settings_cache[ + team_id + ] = await KeyManagementEventHooks._get_secret_manager_optional_params( + team_id ) optional_params = team_settings_cache[team_id] await litellm.secret_manager_client.async_delete_secret( diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 6343faaa965..9888baf897e 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -123,11 +123,70 @@ class SemanticToolFilterHook(CustomLogger): return openai_tools_as_dicts + def _is_mcp_tool(self, tool: object) -> bool: + """ + Check whether *tool* is registered in the MCP semantic router. + + Classification strategy (shape-first, lookup-second): + 1. Chat Completions format dicts are always native. + 2. Responses API function tools are always native. + 3. Everything else is looked up by name in the MCP registry. + """ + if ( + isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("function"), dict) + ): + return False + if ( + isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("name"), str) + ): + return False + name, _ = self.filter._extract_tool_info(tool) + return bool(name) and name in self.filter._tool_map + def _get_metadata_variable_name(self, data: dict) -> str: if "litellm_metadata" in data: return "litellm_metadata" return "metadata" + def _emit_filter_metadata( + self, + data: dict, + mcp_tools: list[object], + filtered_mcp_tools: list[object], + native_tools: list[object], + filtered_tools: list[object], + ) -> None: + """ + Emit response-header metadata when MCP tools were filtered. + + Stats report MCP-only counts so downstream consumers see accurate + semantic filter metrics. Skips metadata entirely for purely-native + requests to avoid spurious headers. + """ + if mcp_tools: + filter_stats = f"{len(mcp_tools)}->{len(filtered_mcp_tools)}" + tool_names_csv = self._get_tool_names_csv(filtered_mcp_tools) + + _metadata_variable_name = self._get_metadata_variable_name(data) + metadata = data.setdefault(_metadata_variable_name, {}) + metadata["litellm_semantic_filter_stats"] = filter_stats + metadata["litellm_semantic_filter_tools"] = tool_names_csv + + verbose_proxy_logger.info( + f"Semantic tool filter: {filter_stats} MCP tools " + f"({len(native_tools)} native preserved, " + f"{len(filtered_tools)} total)" + ) + else: + verbose_proxy_logger.info( + f"Semantic tool filter: all {len(native_tools)} tools " + f"are native, no MCP filtering applied" + ) + async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", @@ -140,53 +199,55 @@ class SemanticToolFilterHook(CustomLogger): This hook is called before the LLM request is made. It filters the tools list to only include semantically relevant tools. - - Args: - user_api_key_dict: User authentication - cache: Cache instance - data: Request data containing messages and tools - call_type: Type of call (completion, acompletion, etc.) - - Returns: - Modified data dict with filtered tools, or None if no changes """ - # Only filter endpoints that support tools if call_type not in ("completion", "acompletion", "aresponses"): verbose_proxy_logger.debug( f"Skipping semantic filter for call_type={call_type}" ) return None - # Check if tools are present tools = data.get("tools") if not tools: verbose_proxy_logger.debug("No tools in request, skipping semantic filter") return None - original_tool_count = len(tools) - - # Check for MCP references (server_url="litellm_proxy") and expand them + # Expanded MCP tools are in OpenAI nested format which + # filter_tools/_extract_tool_info cannot name-match, so we skip + # semantic filtering and return early. if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug( "Detected litellm_proxy MCP references, expanding before semantic filtering" ) try: + native_tools_before_expand = [ + t + for t in tools + if not (isinstance(t, dict) and t.get("type") == "mcp") + ] + expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict) if not expanded_tools: + if native_tools_before_expand: + data["tools"] = native_tools_before_expand + verbose_proxy_logger.warning( + "No MCP tools expanded, preserving " + f"{len(native_tools_before_expand)} native tools" + ) + return data verbose_proxy_logger.warning( "No tools expanded from MCP references" ) return None + data["tools"] = native_tools_before_expand + expanded_tools verbose_proxy_logger.info( - f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools" + f"Expanded MCP references to {len(expanded_tools)} tools " + f"({len(native_tools_before_expand)} native preserved), " + f"skipping semantic filter (OpenAI nested format)" ) - - # Update tools for filtering - tools = expanded_tools - original_tool_count = len(tools) + return data except Exception as e: verbose_proxy_logger.error( @@ -194,7 +255,6 @@ class SemanticToolFilterHook(CustomLogger): ) return None - # Check if messages are present (try both "messages" and "input" for responses API) messages = data.get("messages", []) if not messages: messages = data.get("input", []) @@ -204,13 +264,11 @@ class SemanticToolFilterHook(CustomLogger): ) return None - # Check if filter is enabled if not self.filter.enabled: verbose_proxy_logger.debug("Semantic filter disabled, skipping") return None try: - # Extract user query from messages user_query = self.filter.extract_user_query(messages) if not user_query: verbose_proxy_logger.debug( @@ -218,33 +276,60 @@ class SemanticToolFilterHook(CustomLogger): ) return None + native_tools: list[object] = [] + mcp_tools: list[object] = [] + mcp_indices: set[int] = set() + for i, t in enumerate(tools): + if self._is_mcp_tool(t): + mcp_tools.append(t) + mcp_indices.add(i) + else: + native_tools.append(t) + verbose_proxy_logger.debug( - f"Applying semantic filter to {len(tools)} tools " - f"with query: '{user_query[:50]}...'" + f"Applying semantic filter: {len(mcp_tools)} MCP tools, " + f"{len(native_tools)} native tools, " + f"query: '{user_query[:50]}...'" ) - # Filter tools semantically - filtered_tools = await self.filter.filter_tools( - query=user_query, - available_tools=tools, # type: ignore - ) + if mcp_tools: + filtered_mcp_tools = await self.filter.filter_tools( + query=user_query, + available_tools=mcp_tools, # type: ignore + ) + else: + filtered_mcp_tools = [] + + filtered_mcp_names: set[str] = set() + for t in filtered_mcp_tools: + name, _ = self.filter._extract_tool_info(t) + if name: + filtered_mcp_names.add(name) + + filtered_tools: list[object] = [] + for i, t in enumerate(tools): + if i in mcp_indices: + name, _ = self.filter._extract_tool_info(t) + if name in filtered_mcp_names: + filtered_tools.append(t) + else: + filtered_tools.append(t) - # Always update tools and emit header (even if count unchanged) data["tools"] = filtered_tools - # Store filter stats and tool names for response header - filter_stats = f"{original_tool_count}->{len(filtered_tools)}" - tool_names_csv = self._get_tool_names_csv(filtered_tools) - - _metadata_variable_name = self._get_metadata_variable_name(data) - data[_metadata_variable_name][ - "litellm_semantic_filter_stats" - ] = filter_stats - data[_metadata_variable_name][ - "litellm_semantic_filter_tools" - ] = tool_names_csv - - verbose_proxy_logger.info(f"Semantic tool filter: {filter_stats} tools") + try: + self._emit_filter_metadata( + data=data, + mcp_tools=mcp_tools, + filtered_mcp_tools=filtered_mcp_tools, + native_tools=native_tools, + filtered_tools=filtered_tools, + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to emit semantic filter metadata: {e}", + exc_info=True, + ) return data @@ -266,7 +351,7 @@ class SemanticToolFilterHook(CustomLogger): from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH _metadata_variable_name = self._get_metadata_variable_name(data) - metadata = data[_metadata_variable_name] + metadata = data.get(_metadata_variable_name, {}) filter_stats = metadata.get("litellm_semantic_filter_stats") if not filter_stats: diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index d36e9858b5a..37da5671b64 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -263,9 +263,9 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if rpm_limit is None: rpm_limit = sys.maxsize - values_to_update_in_cache: List[Tuple[Any, Any]] = ( - [] - ) # values that need to get updated in cache, will run a batch_set_cache after this function + values_to_update_in_cache: List[ + Tuple[Any, Any] + ] = [] # values that need to get updated in cache, will run a batch_set_cache after this function # ------------ # Setup values @@ -901,11 +901,11 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): current_minute = datetime.now().strftime("%M") precise_minute = f"{current_date}-{current_hour}-{current_minute}" request_count_api_key = f"{api_key}::{precise_minute}::request_count" - current: Optional[CurrentItemRateLimit] = ( - await self.internal_usage_cache.async_get_cache( - key=request_count_api_key, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - ) + current: Optional[ + CurrentItemRateLimit + ] = await self.internal_usage_cache.async_get_cache( + key=request_count_api_key, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) key_remaining_rpm_limit: Optional[int] = None diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 85d034b7a41..5ca1b2caccf 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -2943,9 +2943,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if pipeline_operations: - await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - litellm_parent_otel_span=litellm_parent_otel_span, + await ( + self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + litellm_parent_otel_span=litellm_parent_otel_span, + ) ) if reserved_tokens > 0: self._mark_reservation_released(kwargs) @@ -3104,9 +3106,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"Releasing reserved TPM tokens on proxy-level " f"rejection: {reserved_tokens}" ) - await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=ops, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + await ( + self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=ops, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) ) self._mark_reservation_released(request_data) except Exception as e: diff --git a/litellm/proxy/hooks/prompt_injection_detection.py b/litellm/proxy/hooks/prompt_injection_detection.py index 6678ccd7e0b..bfe85edf4cb 100644 --- a/litellm/proxy/hooks/prompt_injection_detection.py +++ b/litellm/proxy/hooks/prompt_injection_detection.py @@ -270,7 +270,10 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): if isinstance(response, litellm.ModelResponse) and isinstance( response.choices[0], litellm.Choices ): - if self.prompt_injection_params.llm_api_fail_call_string in response.choices[0].message.content: # type: ignore + if ( + self.prompt_injection_params.llm_api_fail_call_string + in response.choices[0].message.content + ): # type: ignore is_prompt_attack = True if is_prompt_attack is True: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 8fc9d009e67..e3506c78096 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -281,12 +281,14 @@ class _ProxyDBLogger(CustomLogger): ) ) - await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( - token=user_api_key, - key_alias=key_alias, - end_user_id=end_user_id, - response_cost=response_cost, - max_budget=end_user_max_budget, + await ( + proxy_logging_obj.slack_alerting_instance.customer_spend_alert( + token=user_api_key, + key_alias=key_alias, + end_user_id=end_user_id, + response_cost=response_cost, + max_budget=end_user_max_budget, + ) ) elif budget_reservation is not None: await _release_budget_reservation( diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index c22fd1d6579..8122c5e68c6 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -123,8 +123,10 @@ class UserManagementEventHooks: use_enterprise_email_hooks = False if use_enterprise_email_hooks and (data.send_invite_email is True): - initialized_email_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=BaseEmailLogger # type: ignore + initialized_email_loggers = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger # type: ignore + ) ) if len(initialized_email_loggers) > 0: for email_logger in initialized_email_loggers: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c0cdf84dfb6..36881765596 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1049,9 +1049,9 @@ class LiteLLMProxyRequestSetup: ) ) data[_metadata_variable_name].update(user_api_key_logged_metadata) - data[_metadata_variable_name][ - "user_api_key" - ] = user_api_key_dict.api_key # this is just the hashed token + data[_metadata_variable_name]["user_api_key"] = ( + user_api_key_dict.api_key + ) # this is just the hashed token # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none _key_agent_id = getattr(user_api_key_dict, "agent_id", None) @@ -1063,9 +1063,9 @@ class LiteLLMProxyRequestSetup: user_api_key_dict, "end_user_max_budget", None ) if user_api_key_dict.budget_reservation is not None: - data[_metadata_variable_name][ - "user_api_key_budget_reservation" - ] = user_api_key_dict.budget_reservation + data[_metadata_variable_name]["user_api_key_budget_reservation"] = ( + user_api_key_dict.budget_reservation + ) # Add the full UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict return data @@ -1140,9 +1140,9 @@ class LiteLLMProxyRequestSetup: if ( key not in data[_metadata_variable_name]["spend_logs_metadata"] ): # don't override k-v pair sent by request (user request) - data[_metadata_variable_name]["spend_logs_metadata"][ - key - ] = value + data[_metadata_variable_name]["spend_logs_metadata"][key] = ( + value + ) else: data[_metadata_variable_name]["spend_logs_metadata"] = key_metadata[ "spend_logs_metadata" @@ -1712,41 +1712,41 @@ async def add_litellm_data_to_request( ) # Team spend, budget - used by prometheus.py - data[_metadata_variable_name][ - "user_api_key_team_max_budget" - ] = user_api_key_dict.team_max_budget - data[_metadata_variable_name][ - "user_api_key_team_spend" - ] = user_api_key_dict.team_spend - data[_metadata_variable_name][ - "user_api_key_request_route" - ] = user_api_key_dict.request_route + data[_metadata_variable_name]["user_api_key_team_max_budget"] = ( + user_api_key_dict.team_max_budget + ) + data[_metadata_variable_name]["user_api_key_team_spend"] = ( + user_api_key_dict.team_spend + ) + data[_metadata_variable_name]["user_api_key_request_route"] = ( + user_api_key_dict.request_route + ) # API Key spend, budget - used by prometheus.py data[_metadata_variable_name]["user_api_key_spend"] = user_api_key_dict.spend - data[_metadata_variable_name][ - "user_api_key_max_budget" - ] = user_api_key_dict.max_budget - data[_metadata_variable_name][ - "user_api_key_model_max_budget" - ] = user_api_key_dict.model_max_budget - data[_metadata_variable_name][ - "user_api_key_end_user_model_max_budget" - ] = user_api_key_dict.end_user_model_max_budget + data[_metadata_variable_name]["user_api_key_max_budget"] = ( + user_api_key_dict.max_budget + ) + data[_metadata_variable_name]["user_api_key_model_max_budget"] = ( + user_api_key_dict.model_max_budget + ) + data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = ( + user_api_key_dict.end_user_model_max_budget + ) # User spend, budget - used by prometheus.py # Follow same pattern as team and API key budgets - data[_metadata_variable_name][ - "user_api_key_user_spend" - ] = user_api_key_dict.user_spend - data[_metadata_variable_name][ - "user_api_key_user_max_budget" - ] = user_api_key_dict.user_max_budget + data[_metadata_variable_name]["user_api_key_user_spend"] = ( + user_api_key_dict.user_spend + ) + data[_metadata_variable_name]["user_api_key_user_max_budget"] = ( + user_api_key_dict.user_max_budget + ) data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata - data[_metadata_variable_name][ - "user_api_key_team_metadata" - ] = user_api_key_dict.team_metadata + data[_metadata_variable_name]["user_api_key_team_metadata"] = ( + user_api_key_dict.team_metadata + ) data[_metadata_variable_name]["user_api_key_object_permission_id"] = getattr( user_api_key_dict, "object_permission_id", None ) @@ -1764,9 +1764,9 @@ async def add_litellm_data_to_request( # OTEL Controls / Tracing # Add the OTEL Parent Trace before sending it LiteLLM - data[_metadata_variable_name][ - "litellm_parent_otel_span" - ] = user_api_key_dict.parent_otel_span + data[_metadata_variable_name]["litellm_parent_otel_span"] = ( + user_api_key_dict.parent_otel_span + ) _add_otel_traceparent_to_data(data, request=request) ### END-USER SPECIFIC PARAMS ### @@ -2578,9 +2578,9 @@ async def move_guardrails_to_metadata( request_body_guardrail_config ) else: - data[_metadata_variable_name][ - "guardrail_config" - ] = request_body_guardrail_config + data[_metadata_variable_name]["guardrail_config"] = ( + request_body_guardrail_config + ) def _is_policy_version_id(s: str) -> bool: @@ -2710,9 +2710,9 @@ def _apply_resolved_guardrails_to_metadata( pipelines ) data[metadata_variable_name]["_guardrail_pipelines"] = pipelines - data[metadata_variable_name][ - "_pipeline_managed_guardrails" - ] = pipeline_managed_guardrails + data[metadata_variable_name]["_pipeline_managed_guardrails"] = ( + pipeline_managed_guardrails + ) verbose_proxy_logger.debug( f"Policy engine: resolved {len(pipelines)} pipeline(s), " f"managed guardrails: {pipeline_managed_guardrails}" diff --git a/litellm/proxy/logging_endpoints/__init__.py b/litellm/proxy/logging_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py new file mode 100644 index 00000000000..a96a5431294 --- /dev/null +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -0,0 +1,208 @@ +""" +Ingest pre-built logging payloads from external producers and replay them +through LiteLLM's standard success/failure callback fan-out. + +This exists for hosts that own a request outside the Python process — e.g. the +`litellm-rust` gateway proxying realtime websockets. Those hosts can't use the +in-process logging object, so they POST a finished `StandardLoggingPayload` here +and Python replays it through the exact same path a normal completion uses: +`Logging.async_success_handler` / `async_failure_handler`. Every registered +callback (spend logs, Langfuse, Datadog, ...) fires unchanged — there is no +spend-logs-specific or callback-specific code here, only the replay. + +The endpoint is generic: realtime is the first producer, but the contract is the +self-describing `StandardLoggingPayload`, so completions/responses can use it too. +""" + +import uuid +from datetime import datetime, timezone +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.proxy.callback_logs_endpoints import ( + CallbackLogFailure, + CallbackLogRecord, + CallbackLogsRequest, + CallbackLogsResponse, +) + +# Routes the Python proxy exposes for the Rust data-plane gateway to call into +# (logging today; auth/budgets later). Namespaced under /v1/rust_control_plane so +# they're clearly distinct from the proxy's own control-plane/management routes. +rust_control_plane_router = APIRouter( + prefix="/v1/rust_control_plane", tags=["rust control plane"] +) + + +class CallbackLogsReplayer: + """ + Replays finished logging payloads through LiteLLM's callback fan-out. + + Each helper is small and pure so the replay path is easy to read and test: + rebuild a `Logging` object from the payload, seed `model_call_details` with + exactly what the callbacks read, then dispatch to the success/failure + handler. No spend/callback logic lives here — only the replay. + """ + + @staticmethod + def _epoch_to_datetime(value: Any) -> datetime: + """`StandardLoggingPayload` stores startTime/endTime as float epoch seconds.""" + if isinstance(value, (int, float)): + return datetime.fromtimestamp(float(value), tz=timezone.utc) + if isinstance(value, datetime): + return value + return datetime.now(tz=timezone.utc) + + @staticmethod + def _build_logging_obj(payload: dict[str, Any]) -> LiteLLMLogging: + """ + Reconstruct a `Logging` object from a finished payload and seed + `model_call_details` with exactly what the success/failure callbacks + read: the prebuilt `standard_logging_object`, the resolved + `response_cost`, and the `litellm_params.metadata` keys used for cost + attribution. Setting `standard_logging_object` up front makes the handler + skip rebuilding it. + """ + model = payload.get("model") or "" + call_type = payload.get("call_type") or "acompletion" + start_time = CallbackLogsReplayer._epoch_to_datetime(payload.get("startTime")) + call_id = ( + payload.get("litellm_call_id") or payload.get("id") or str(uuid.uuid4()) + ) + + logging_obj = LiteLLMLogging( + model=model, + messages=payload.get("messages") or [], + # A replayed payload is always a *terminal*, fully-aggregated event — + # the producer (e.g. the rust gateway) already collected the whole + # session before POSTing. Never mark it streaming: a streaming + # Logging object makes async_success_handler wait for a + # complete_streaming_response that will never arrive, so the spend + # log is never written. + stream=False, + call_type=call_type, + start_time=start_time, + litellm_call_id=call_id, + function_id="", + ) + + metadata: dict[str, Any] = payload.get("metadata") or {} + litellm_metadata: dict[str, Any] = { + "user_api_key": metadata.get("user_api_key_hash"), + "user_api_key_alias": metadata.get("user_api_key_alias"), + "user_api_key_user_id": metadata.get("user_api_key_user_id"), + "user_api_key_team_id": metadata.get("user_api_key_team_id"), + "user_api_key_org_id": metadata.get("user_api_key_org_id"), + "user_api_key_end_user_id": metadata.get("user_api_key_end_user_id"), + "spend_logs_metadata": metadata.get("spend_logs_metadata"), + } + + logging_obj.model_call_details.update( + { + "model": model, + "call_type": call_type, + "custom_llm_provider": payload.get("custom_llm_provider"), + "response_cost": payload.get("response_cost") or 0.0, + "standard_logging_object": payload, + "litellm_params": {"metadata": litellm_metadata}, + "cache_hit": payload.get("cache_hit") or False, + } + ) + return logging_obj + + @staticmethod + def _response_obj_from_payload(payload: dict[str, Any]) -> dict[str, Any]: + """Minimal response object so usage-derived spend-log fields resolve.""" + return { + "id": payload.get("id"), + "usage": { + "prompt_tokens": payload.get("prompt_tokens", 0), + "completion_tokens": payload.get("completion_tokens", 0), + "total_tokens": payload.get("total_tokens", 0), + }, + } + + async def replay(self, record: CallbackLogRecord) -> None: + """Replay one record through the matching success/failure handler.""" + payload = record.standard_logging_payload + verbose_proxy_logger.debug( + "CallbackLogsReplayer: replaying %s record id=%s model=%s call_type=%s", + record.status, + payload.get("id"), + payload.get("model"), + payload.get("call_type"), + ) + + logging_obj = self._build_logging_obj(payload) + start_time = self._epoch_to_datetime(payload.get("startTime")) + end_time = self._epoch_to_datetime(payload.get("endTime")) + + if record.status == "success": + await logging_obj.async_success_handler( + result=self._response_obj_from_payload(payload), + start_time=start_time, + end_time=end_time, + ) + else: + error_str = record.error or payload.get("error_str") or "replayed failure" + await logging_obj.async_failure_handler( + Exception(error_str), + traceback_exception="", + start_time=start_time, + end_time=end_time, + ) + + async def replay_batch( + self, records: list[CallbackLogRecord] + ) -> CallbackLogsResponse: + """Replay a batch; a single bad record never sinks the rest. Each failure + is reported back with its batch index so the caller can retry/triage it.""" + processed = 0 + failures: list[CallbackLogFailure] = [] + for index, record in enumerate(records): + try: + await self.replay(record) + processed += 1 + except Exception as e: + failures.append(CallbackLogFailure(index=index, error=str(e))) + verbose_proxy_logger.exception( + "CallbackLogsReplayer: failed to replay record %s: %s", + index, + str(e), + ) + verbose_proxy_logger.debug( + "CallbackLogsReplayer: batch done processed=%s failed=%s", + processed, + len(failures), + ) + return CallbackLogsResponse( + processed=processed, failed=len(failures), failures=failures + ) + + +@rust_control_plane_router.post( + "/logs", + dependencies=[Depends(user_api_key_auth)], +) +async def ingest_callback_logs( + body: CallbackLogsRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> CallbackLogsResponse: + """ + Replay a batch of finished logging payloads through the callback fan-out. + + Admin-only: the payloads write spend logs and trigger every callback, so this + is a trusted internal route, not a public surface. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="/v1/rust_control_plane/logs is admin-only (proxy admin key required).", + ) + + return await CallbackLogsReplayer().replay_batch(body.records) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 79882909c23..ba41852f6d8 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -132,13 +132,11 @@ def update_breakdown_metrics( ), ) ) - breakdown.models[record.model].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.models[record.model] - .api_key_breakdown[record.api_key] - .metrics, - record, - ) + breakdown.models[record.model].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.models[record.model].api_key_breakdown[record.api_key].metrics, + record, ) # Update model group breakdown @@ -247,11 +245,11 @@ def update_breakdown_metrics( ), ) ) - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, - record, - ) + breakdown.providers[provider].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, + record, ) # Update endpoint breakdown @@ -338,13 +336,11 @@ def update_breakdown_metrics( ), ) ) - breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.entities[entity_value] - .api_key_breakdown[record.api_key] - .metrics, - record, - ) + breakdown.entities[entity_value].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, + record, ) return breakdown diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 50a1bc23a6d..ecc338c2052 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -342,7 +342,8 @@ async def new_end_user( budget_record = await BudgetRepository(prisma_client).table.create( data={ **_new_budget.model_dump(exclude_unset=True), - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore + "created_by": user_api_key_dict.user_id + or litellm_proxy_admin_name, # type: ignore "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, } @@ -563,10 +564,14 @@ async def update_end_user( # get non default values for key non_default_values = {} for k, v in data_json.items(): - if v is not None and v not in ( - [], - {}, - 0, + if ( + v is not None + and v + not in ( + [], + {}, + 0, + ) ): # models default to [], spend defaults to 0, we should not reset these values non_default_values[k] = v @@ -655,7 +660,9 @@ async def update_end_user( update_end_user_table_data["user_id"] = data.user_id # type: ignore verbose_proxy_logger.debug("In update customer, user_id condition block.") response = await EndUserRepository(prisma_client).table.update( - where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True, "object_permission": True} # type: ignore + where={"user_id": data.user_id}, + data=update_end_user_table_data, + include={"litellm_budget_table": True, "object_permission": True}, # type: ignore ) if response is None: raise ValueError( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6d7f565fb85..8548d1c672c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -746,7 +746,7 @@ def _build_user_info_response( user_info = {"spend": spend} returned_keys = _process_keys_for_user_info(keys=keys, all_teams=teams_1) - team_list.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) + team_list.sort(key=lambda x: getattr(x, "team_alias", "") or "") _user_info = ( user_info.model_dump() if isinstance(user_info, BaseModel) else user_info @@ -1053,7 +1053,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): # cast all teams to LiteLLM_TeamTable _teams_in_db: List = results[0]["teams"] or [] _teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db] - _teams_in_db.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) + _teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "") returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) # Get admin's own user_id and user_info @@ -1705,7 +1705,8 @@ async def bulk_user_update( try: # Perform bulk database update await UserRepository(prisma_client).table.update_many( - where={}, data=non_default_values # Update all users + where={}, + data=non_default_values, # Update all users ) # Create individual success results diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2d49297c8e9..db205803a47 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -229,10 +229,10 @@ def _is_allowed_to_make_key_request( return True if user_id is not None: - assert ( - user_id == user_api_key_dict.user_id - ), "User can only create keys for themselves. Got user_id={}, Your ID={}".format( - user_id, user_api_key_dict.user_id + assert user_id == user_api_key_dict.user_id, ( + "User can only create keys for themselves. Got user_id={}, Your ID={}".format( + user_id, user_api_key_dict.user_id + ) ) if team_id is not None: @@ -382,7 +382,9 @@ def _personal_key_generation_check( ): return True - _personal_key_generation = litellm.key_generation_settings["personal_key_generation"] # type: ignore + _personal_key_generation = litellm.key_generation_settings[ + "personal_key_generation" + ] # type: ignore _personal_key_membership_check( user_api_key_dict, @@ -874,6 +876,8 @@ async def _common_key_generation_helper( object_permission=data_json.get("object_permission"), team_obj=team_table, prisma_client=prisma_client, + is_proxy_admin=user_api_key_dict.user_role + == LitellmUserRoles.PROXY_ADMIN.value, ) if normalized_object_permission is not None: data_json["object_permission"] = normalized_object_permission @@ -961,9 +965,7 @@ async def _common_key_generation_helper( response = GenerateKeyResponse(**response) - response.token = ( - response.token_id - ) # remap token to use the hash, and leave the key in the `key` field [TODO]: clean up generate_key_helper_fn to do this + response.token = response.token_id # remap token to use the hash, and leave the key in the `key` field [TODO]: clean up generate_key_helper_fn to do this asyncio.create_task( KeyManagementEventHooks.async_key_generated_hook( @@ -2050,12 +2052,14 @@ async def _process_single_key_update( # Check team member permissions if prisma_client is not None: - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=existing_key_row, - user_api_key_cache=user_api_key_cache, + await ( + TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=existing_key_row, + user_api_key_cache=user_api_key_cache, + ) ) # Custom key update hook @@ -2164,6 +2168,7 @@ async def _validate_mcp_servers_for_key_update( existing_key_row: Any, prisma_client: Any, user_api_key_cache: Any, + is_proxy_admin: bool, ) -> Optional[dict]: """Validate MCP servers in object_permission against the effective team.""" effective_team_obj = team_obj @@ -2186,6 +2191,7 @@ async def _validate_mcp_servers_for_key_update( object_permission=object_permission_dict, team_obj=effective_team_obj, prisma_client=prisma_client, + is_proxy_admin=is_proxy_admin, ) await validate_key_search_tools_against_team( object_permission=object_permission_dict, @@ -2422,6 +2428,7 @@ async def _validate_update_key_data( existing_key_row=existing_key_row, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + is_proxy_admin=_is_proxy_admin, ) if normalized_object_permission is not None: data.object_permission = LiteLLM_ObjectPermissionBase( @@ -2963,12 +2970,14 @@ async def bulk_update_team_keys( models=[], ) ) - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=auth_anchor, - user_api_key_cache=user_api_key_cache, + await ( + TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=auth_anchor, + user_api_key_cache=user_api_key_cache, + ) ) # Block metadata.allowed_passthrough_routes for non-admins — the runtime @@ -3975,10 +3984,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[LiteLLM_VerificationToken] = ( - await VerificationTokenRepository(prisma_client).table.find_many( - where={"token": {"in": tokens}} - ) + _keys_being_deleted: List[ + LiteLLM_VerificationToken + ] = await VerificationTokenRepository(prisma_client).table.find_many( + where={"token": {"in": tokens}} ) if len(_keys_being_deleted) == 0: @@ -4665,12 +4674,14 @@ async def regenerate_key_fn( ) # check if user has permission to regenerate key - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_REGENERATE, - prisma_client=prisma_client, - existing_key_row=_key_in_db, - user_api_key_cache=user_api_key_cache, + await ( + TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_REGENERATE, + prisma_client=prisma_client, + existing_key_row=_key_in_db, + user_api_key_cache=user_api_key_cache, + ) ) # check if user has ownership permission to regenerate key @@ -5987,7 +5998,8 @@ async def block_key( ) record = await VerificationTokenRepository(prisma_client).table.update( - where={"token": hashed_token}, data={"blocked": True} # type: ignore + where={"token": hashed_token}, + data={"blocked": True}, # type: ignore ) ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB @@ -6101,7 +6113,8 @@ async def unblock_key( ) record = await VerificationTokenRepository(prisma_client).table.update( - where={"token": hashed_token}, data={"blocked": False} # type: ignore + where={"token": hashed_token}, + data={"blocked": False}, # type: ignore ) ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e86982307e7..4ab308990b5 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1016,15 +1016,10 @@ if MCP_AVAILABLE: if is_restricted_virtual_key: return _sanitize_mcp_server_list_for_virtual_key(redacted_mcp_servers) - # Non-admin authenticated users may see the server inventory but - # not credential-bearing fields like `url` (often contains bearer - # tokens) or headers/env (often contain Authorization). - if not _user_has_admin_view(user_api_key_dict): - return _sanitize_mcp_server_list_for_non_admin(redacted_mcp_servers) - + # only a full PROXY_ADMIN sees credential-bearing fields; everyone else + # goes through the non-admin sanitizer if not _user_is_full_admin(user_api_key_dict): - for server in redacted_mcp_servers: - _redact_global_env_var_values(server) + return _sanitize_mcp_server_list_for_non_admin(redacted_mcp_servers) return redacted_mcp_servers @@ -1415,10 +1410,10 @@ if MCP_AVAILABLE: redacted = _redact_mcp_credentials(mcp_server) if is_restricted_virtual_key: return _sanitize_mcp_server_for_virtual_key(redacted) - if not _user_has_admin_view(user_api_key_dict): - return _sanitize_mcp_server_for_non_admin(redacted) + # only a full PROXY_ADMIN sees credential-bearing fields; everyone else + # goes through the non-admin sanitizer if not _user_is_full_admin(user_api_key_dict): - _redact_global_env_var_values(redacted) + return _sanitize_mcp_server_for_non_admin(redacted) return redacted @router.post( @@ -1935,12 +1930,9 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) - mcp_server = await get_mcp_server(prisma_client, server_id) - if mcp_server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server {server_id} not found"}, - ) + mcp_server = await _authorize_and_fetch_mcp_server( + prisma_client, user_api_key_dict, server_id + ) if not getattr(mcp_server, "is_byok", False): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -2015,12 +2007,9 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) - mcp_server = await get_mcp_server(prisma_client, server_id) - if mcp_server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server {server_id} not found"}, - ) + await _authorize_and_fetch_mcp_server( + prisma_client, user_api_key_dict, server_id + ) user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2182,37 +2171,47 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, server_id: str, ) -> LiteLLM_MCPServerTable: - """Return the MCP server the caller may manage env vars for. + """Resolve the MCP server a caller may manage their own per-user state for. - Admins look the server up directly. Non-admins reuse the access-scoped - listing that already loads every server they can see, so we don't issue - a second per-server query just to re-fetch a record the authorization - check produced. A non-admin who can't see the server gets 403 (never - 404) so server ids can't be enumerated. + Looks the server up in the DB, then the in-memory registry, so a + config-defined server (which never gets a DB row) resolves too. Admins + may reach any server and get a 404 for an unknown id. A non-admin may + only reach a server in their allowed set and otherwise gets 403 (never + 404, so server ids can't be enumerated), using the same allowed-server + resolution the MCP gateway enforces on tool calls. """ + server = await get_mcp_server(prisma_client, server_id) + if server is None: + registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + if registry_server is not None: + server = global_mcp_server_manager._build_mcp_server_table( + registry_server + ) + if _user_has_admin_view(user_api_key_dict): - server = await get_mcp_server(prisma_client, server_id) if server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={"error": f"MCP Server {server_id} not found"}, ) return server - accessible = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) - for server in accessible: - if server.server_id == server_id: - return server - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": ( - f"User does not have permission to access mcp server with id {server_id}. " - "You can only manage env vars for mcp servers that you have access to." - ) - }, - ) + + allowed_server_ids: set[str] = set() + for auth_context in await build_effective_auth_contexts(user_api_key_dict): + allowed_server_ids.update( + await global_mcp_server_manager.get_allowed_mcp_servers(auth_context) + ) + if server is None or server.server_id not in allowed_server_ids: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"User does not have permission to access mcp server with id {server_id}. " + "You can only manage mcp servers that you have access to." + ) + }, + ) + return server def _compute_user_env_var_status( *, @@ -2541,9 +2540,9 @@ if MCP_AVAILABLE: if "litellm_settings" not in config or config["litellm_settings"] is None: config["litellm_settings"] = {} - config["litellm_settings"][ - "public_mcp_servers" - ] = litellm.public_mcp_servers + config["litellm_settings"]["public_mcp_servers"] = ( + litellm.public_mcp_servers + ) # Save the updated config await proxy_config.save_config(new_config=config) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index def6e271635..df4d46b9098 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -518,7 +518,9 @@ async def _add_model_to_db( _data: dict = { "model_id": model_params.model_info.id, "model_name": model_params.model_name, - "litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), # type: ignore + "litellm_params": model_params.litellm_params.model_dump_json( + exclude_none=True + ), # type: ignore "model_info": model_params.model_info.model_dump_json( # type: ignore exclude_none=True ), diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 99659121b27..a45382b54d0 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1299,7 +1299,9 @@ async def add_member_to_organization( user_email=member.user_email, ) - _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore + _returned_user = await prisma_client.insert_data( + data=new_user_defaults, table_name="user" + ) # type: ignore if _returned_user is not None: user_object = LiteLLM_UserTable(**_returned_user.model_dump()) elif existing_user_email_row is not None and len(existing_user_email_row) > 1: diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 0c11507697d..94c84387be3 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -251,7 +251,8 @@ async def add_team_callbacks( team_metadata_json = json.dumps(team_metadata) # update team_metadata new_team_row = await TeamRepository(prisma_client).table.update( - where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore + where={"team_id": team_id}, + data={"metadata": team_metadata_json}, # type: ignore ) await _emit_team_callback_audit_log( @@ -355,7 +356,8 @@ async def disable_team_logging( # Update team in database updated_team = await TeamRepository(prisma_client).table.update( - where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore + where={"team_id": team_id}, + data={"metadata": team_metadata_json}, # type: ignore ) if updated_team is None: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3d90e7b5ab9..249cfdae4eb 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -281,9 +281,9 @@ class TeamMemberBudgetHandler: # Add team_member_budget_id as metadata field to team table if new_team_data_json.get("metadata") is None: new_team_data_json["metadata"] = {} - new_team_data_json["metadata"][ - "team_member_budget_id" - ] = team_member_budget_table.budget_id + new_team_data_json["metadata"]["team_member_budget_id"] = ( + team_member_budget_table.budget_id + ) # Remove team member fields from new_team_data_json TeamMemberBudgetHandler._clean_team_member_fields(new_team_data_json) @@ -3889,7 +3889,8 @@ async def block_team( ) record = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, data={"blocked": True} # type: ignore + where={"team_id": data.team_id}, + data={"blocked": True}, # type: ignore ) return record @@ -3941,7 +3942,8 @@ async def unblock_team( ) record = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, data={"blocked": False} # type: ignore + where={"team_id": data.team_id}, + data={"blocked": False}, # type: ignore ) return record @@ -4651,7 +4653,7 @@ async def list_team( verbose_proxy_logger.exception(team_exception) continue # Sort the responses by team_alias - returned_responses.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) + returned_responses.sort(key=lambda x: getattr(x, "team_alias", "") or "") if organization_id is not None: if organization_id == SpecialManagementEndpointEnums.DEFAULT_ORGANIZATION.value: @@ -4692,7 +4694,9 @@ async def get_paginated_teams( # Get paginated teams teams = await TeamRepository(prisma_client).table.find_many( - skip=skip, take=page_size, order={"team_alias": "asc"} # Sort by team_alias + skip=skip, + take=page_size, + order={"team_alias": "asc"}, # Sort by team_alias ) return teams, total_count except Exception as e: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 199de54ff09..480861fb517 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2777,7 +2777,9 @@ class SSOAuthenticationHandler: state_only_params[key] = value # Get the redirect response from fastapi-sso with only state param - redirect_response = await generic_sso.get_login_redirect(**state_only_params) # type: ignore + redirect_response = await generic_sso.get_login_redirect( + **state_only_params + ) # type: ignore # If PKCE is enabled, add PKCE parameters to the redirect URL if code_verifier and "state" in redirect_params: @@ -3204,7 +3206,9 @@ class SSOAuthenticationHandler: user_id = getattr(result, "id", None) user_email = normalize_email(getattr(result, "email", None)) if user_role is None: - _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore + _role_from_attr = getattr( + result, generic_user_role_attribute_name, None + ) # type: ignore if _role_from_attr is not None: # Convert enum to string if needed user_role = ( @@ -4539,14 +4543,16 @@ async def debug_sso_callback(request: Request): ) elif generic_client_id is not None: - result, received_response, access_token_payload = ( - await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) + ( + result, + received_response, + access_token_payload, + ) = await get_generic_sso_response( + request=request, + jwt_handler=jwt_handler, + generic_client_id=generic_client_id, + redirect_url=redirect_url, + sso_jwt_handler=sso_jwt_handler, ) # If result is None, return a basic error message diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index a50ce1d3c48..c9884c02e14 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -197,8 +197,7 @@ def _build_system_prompt(is_admin: bool) -> str: """Build role-appropriate system prompt with today's date.""" tool_desc = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE return ( - f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}" - f"Today's date: {date.today().isoformat()}" + f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}Today's date: {date.today().isoformat()}" ) @@ -352,7 +351,9 @@ def _summarise_usage_data(data: Dict[str, Any]) -> str: model_lines = _ranked_lines( models, - lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", + lambda n, d: ( + f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)" + ), TOP_N_MODELS, ) provider_lines = _ranked_lines( diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 07c355f2cd9..0a61fc5d3dd 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -361,10 +361,10 @@ async def _resolve_team_allowed_mcp_servers( ) direct_servers: List[str] = team_object_permission.mcp_servers or [] - access_group_servers: List[str] = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - team_object_permission.mcp_access_groups or [] - ) + access_group_servers: List[ + str + ] = await MCPRequestHandler._get_mcp_servers_from_access_groups( + team_object_permission.mcp_access_groups or [] ) raw_tool_perms = team_object_permission.mcp_tool_permissions or {} if isinstance(raw_tool_perms, str): @@ -469,6 +469,7 @@ async def validate_key_mcp_servers_against_team( object_permission: Optional[dict], team_obj: Optional["LiteLLM_TeamTableCachedObj"], prisma_client: Optional[PrismaClient] = None, + is_proxy_admin: bool = False, ) -> Optional[dict]: """ Validate that MCP servers requested on a key are within the allowed scope. @@ -476,12 +477,17 @@ async def validate_key_mcp_servers_against_team( Rules: - If key is in a team: key's mcp_servers must be a subset of (team's allowed servers + allow_all_keys servers) - - If key is NOT in a team: key's mcp_servers must only contain - allow_all_keys servers + - If key is NOT in a team and the caller is a proxy admin: any server or + access group may be assigned. A proxy admin can already reach every MCP + server, and runtime access is granted directly from the key's own + object_permission, so the key is scoped to exactly what the admin selected + - If key is NOT in a team and the caller is not a proxy admin: key's + mcp_servers must only contain allow_all_keys servers - If team has no MCP config: key can only use allow_all_keys servers Raises HTTPException(403) if validation fails. """ + teamless_admin_assignment = team_obj is None and is_proxy_admin requested_servers = _extract_requested_mcp_server_ids(object_permission) requested_access_groups = _extract_requested_mcp_access_groups(object_permission) @@ -526,7 +532,11 @@ async def validate_key_mcp_servers_against_team( identifier_to_server_ids ) - disallowed_servers = active_requested_servers - all_allowed_servers + allowed_servers = all_allowed_servers + if teamless_admin_assignment: + allowed_servers = all_allowed_servers | active_requested_servers + + disallowed_servers = active_requested_servers - allowed_servers if disallowed_servers: if team_obj is not None: team_id = team_obj.team_id @@ -557,7 +567,11 @@ async def validate_key_mcp_servers_against_team( ): team_access_groups = set(team_obj.object_permission.mcp_access_groups) - disallowed_groups = requested_access_groups - team_access_groups + allowed_access_groups = team_access_groups + if teamless_admin_assignment: + allowed_access_groups = team_access_groups | requested_access_groups + + disallowed_groups = requested_access_groups - allowed_access_groups if disallowed_groups: if team_obj is not None: team_id = team_obj.team_id diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 830d6f84b85..4ce8633ed57 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -256,7 +256,9 @@ async def add_new_member( isinstance(existing_user_row, list) and len(existing_user_row) == 0 ): new_user_defaults["teams"] = [team_id] - _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore + _returned_user = await prisma_client.insert_data( + data=new_user_defaults, table_name="user" + ) # type: ignore if _returned_user is not None: returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) @@ -433,10 +435,12 @@ async def send_management_endpoint_alert( # replace all "_" with " " and capitalize event_name = _event_name.replace("_", " ").title() - await proxy_logging_obj.slack_alerting_instance.send_virtual_key_event_slack( - key_event=key_event, - event_name=event_name, - alert_type=_event_name, + await ( + proxy_logging_obj.slack_alerting_instance.send_virtual_key_event_slack( + key_event=key_event, + event_name=event_name, + alert_type=_event_name, + ) ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 944423632ef..d8f8726b6b6 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,7 +7,7 @@ import asyncio import traceback -from typing import Any, Optional, cast, get_args +from typing import Any, BinaryIO, Optional, Union, cast, get_args import httpx from fastapi import ( @@ -97,16 +97,18 @@ def get_files_provider_config( return None -def get_first_json_object(file_content_bytes: bytes) -> Optional[dict]: +def get_first_json_object(file_source: Union[bytes, BinaryIO]) -> Optional[dict]: try: - # Decode the bytes to a string and split into lines - file_content = file_content_bytes.decode("utf-8") - first_line = file_content.splitlines()[0].strip() - - # Parse the JSON object from the first line - json_object = json.loads(first_line) - return json_object - except (json.JSONDecodeError, UnicodeDecodeError): + if isinstance(file_source, (bytes, bytearray)): + newline = file_source.find(b"\n") + raw = file_source if newline == -1 else file_source[:newline] + first_line = raw.decode("utf-8") + else: + file_source.seek(0) + first_line = file_source.readline().decode("utf-8") + file_source.seek(0) + return json.loads(first_line.strip()) + except (json.JSONDecodeError, UnicodeDecodeError, OSError, ValueError): return None @@ -268,7 +270,9 @@ async def route_create_file( _create_file_request.update(llm_provider_config) _create_file_request.pop("custom_llm_provider", None) # type: ignore # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch - response = await litellm.acreate_file(**_create_file_request, custom_llm_provider=custom_llm_provider) # type: ignore + response = await litellm.acreate_file( + **_create_file_request, custom_llm_provider=custom_llm_provider + ) # type: ignore return response @@ -327,9 +331,15 @@ async def create_file( data: Dict = {} try: - # Use orjson to parse JSON data, orjson speeds up requests significantly - # Read the file content - file_content = await file.read() + # Batch uploads can be gigabytes. Starlette has already spooled the upload + # to disk, so stream from that handle instead of reading it into memory. + # Other uploads are small and stay in-memory bytes. + file_source: Union[bytes, BinaryIO] + if purpose == "batch": + await file.seek(0) + file_source = file.file + else: + file_source = await file.read() custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -454,13 +464,13 @@ async def create_file( ) # Prepare the file data according to FileTypes - file_data = (file.filename, file_content, file.content_type) + file_data = (file.filename, file_source, file.content_type) ## check if model is a loadbalanced model router_model: Optional[str] = None is_router_model = False if litellm.enable_loadbalancing_on_batch_endpoints is True: - json_obj = get_first_json_object(file_content_bytes=file_content) + json_obj = get_first_json_object(file_source) if json_obj: router_model = get_model_from_json_obj(json_object=json_obj) is_router_model = is_known_model( @@ -1016,7 +1026,9 @@ async def get_file( # data was initialized with {"file_id": file_id} data.pop("file_id", None) response = await litellm.afile_retrieve( - custom_llm_provider=custom_llm_provider, file_id=file_id, **data # type: ignore + custom_llm_provider=custom_llm_provider, + file_id=file_id, + **data, # type: ignore ) ### ALERTING ### @@ -1220,7 +1232,9 @@ async def delete_file( else: data.pop("file_id", None) response = await litellm.afile_delete( - custom_llm_provider=custom_llm_provider, file_id=file_id, **data # type: ignore + custom_llm_provider=custom_llm_provider, + file_id=file_id, + **data, # type: ignore ) ### ALERTING ### @@ -1402,7 +1416,9 @@ async def list_files( prepare_data_with_credentials(data=data, credentials=team_credentials) response = await litellm.afile_list( - custom_llm_provider=custom_llm_provider, purpose=purpose, **data # type: ignore + custom_llm_provider=custom_llm_provider, + purpose=purpose, + **data, # type: ignore ) if response is None: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c7db818a07e..5e688702f13 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2089,7 +2089,9 @@ class BaseOpenAIPassThroughHandler: custom_llm_provider=( custom_llm_provider.value if hasattr(custom_llm_provider, "value") - else str(custom_llm_provider) if custom_llm_provider else None + else str(custom_llm_provider) + if custom_llm_provider + else None ), ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index b77c6e2f655..1d76165c30c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -478,7 +478,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): endpoint_type = ( "chat_completions" if is_chat_completions - else "image_generation" if is_image_generation else "image_editing" + else "image_generation" + if is_image_generation + else "image_editing" ) verbose_proxy_logger.debug( f"OpenAI passthrough cost tracking - Endpoint: {endpoint_type}, Model: {model}, Cost: ${response_cost:.6f}" diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 73d4245670a..538267e9c84 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -219,9 +219,9 @@ class VertexPassthroughLoggingHandler: kwargs["response_cost"] = response_cost kwargs["model"] = "vertex_ai/search_api" logging_obj.model_call_details.setdefault("litellm_params", {}) - logging_obj.model_call_details["litellm_params"][ - "base_model" - ] = "vertex_ai/search_api" + logging_obj.model_call_details["litellm_params"]["base_model"] = ( + "vertex_ai/search_api" + ) logging_obj.model_call_details["response_cost"] = response_cost return { diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index b84746758fb..a7c7786aa91 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -484,10 +484,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): for field_name, field_value in form_data.items(): if isinstance(field_value, (StarletteUploadFile, UploadFile)): - files[field_name] = ( - await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file=field_value - ) + files[ + field_name + ] = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( + upload_file=field_value ) else: form_data_dict[field_name] = field_value diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index 25368c2a834..ae5ce177853 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -97,9 +97,9 @@ class InMemoryPromptRegistry: Prompt id to Prompt object mapping """ - self.prompt_id_to_custom_prompt: Dict[str, Optional[CustomPromptManagement]] = ( - {} - ) + self.prompt_id_to_custom_prompt: Dict[ + str, Optional[CustomPromptManagement] + ] = {} """ Guardrail id to CustomGuardrail object mapping """ @@ -142,7 +142,9 @@ class InMemoryPromptRegistry: raise ValueError( f"CustomPromptManagement is required, got {type(custom_prompt_callback)}" ) - litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback) # type: ignore + litellm.logging_callback_manager.add_litellm_callback( + custom_prompt_callback + ) # type: ignore else: raise ValueError(f"Unsupported prompt: {prompt_integration}") diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 9c4d7b1bb5d..a8efbadb253 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -438,7 +438,9 @@ class ProxyInitializationHelpers: _endpoint_str = ( f"curl --location 'http://0.0.0.0:{port}/chat/completions' \\" ) - curl_command = _endpoint_str + """ + curl_command = ( + _endpoint_str + + """ --header 'Content-Type: application/json' \\ --data ' { "model": "gpt-3.5-turbo", @@ -451,6 +453,7 @@ class ProxyInitializationHelpers: }' \n """ + ) print() print( '\033[1;34mLiteLLM: Test your local proxy with: "litellm --test" This runs an openai.ChatCompletion request to your proxy [In a new terminal tab]\033[0m\n' @@ -1195,6 +1198,25 @@ def run_server( os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None ): + from litellm.proxy.db.db_url_settings import ( + unsupported_db_scheme, + unsupported_db_scheme_message, + ) + + for _db_env in ("DATABASE_URL", "DIRECT_URL"): + _candidate_url = os.getenv(_db_env) + if _candidate_url is None: + continue + _bad_scheme = unsupported_db_scheme(_candidate_url) + if _bad_scheme is not None: + print( + f"\033[1;31mLiteLLM Proxy: " + f"{unsupported_db_scheme_message(_db_env, _bad_scheme)}" + "\033[0m", + file=sys.stderr, + flush=True, + ) + sys.exit(1) try: from litellm.secret_managers.main import get_secret diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5676fbb32f..b0beaae02cf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -38,7 +38,7 @@ from typing import ( import anyio import websockets import websockets.exceptions -from pydantic import BaseModel, Json +from pydantic import BaseModel, Json, JsonValue from litellm._uuid import uuid from litellm.constants import ( @@ -106,6 +106,10 @@ from litellm.proxy.common_utils.callback_utils import ( process_callback, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.router_utils.add_retry_fallback_headers import ( + get_fallback_errors_from_headers, + get_hidden_params_dict, +) from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -342,6 +346,9 @@ from litellm.proxy.hooks.prompt_injection_detection import ( from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( + rust_control_plane_router, +) from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) @@ -666,7 +673,21 @@ _description = ( def cleanup_router_config_variables(): - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_key_update, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, prisma_client + global \ + master_key, \ + user_config_file_path, \ + otel_logging, \ + user_custom_auth, \ + user_custom_auth_path, \ + user_custom_key_generate, \ + user_custom_key_update, \ + user_custom_sso, \ + user_custom_ui_sso_sign_in_handler, \ + use_background_health_checks, \ + use_shared_health_check, \ + health_check_interval, \ + health_check_concurrency, \ + prisma_client # Set all variables to None master_key = None @@ -686,7 +707,12 @@ def cleanup_router_config_variables(): async def proxy_shutdown_event(): - global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update + global \ + prisma_client, \ + master_key, \ + user_custom_auth, \ + user_custom_key_generate, \ + user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") if prisma_client: verbose_proxy_logger.debug("Disconnecting from Prisma") @@ -756,7 +782,22 @@ async def _initialize_shared_aiohttp_session(): @asynccontextmanager async def proxy_startup_event(app: FastAPI): - global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session + global \ + prisma_client, \ + master_key, \ + use_background_health_checks, \ + llm_router, \ + llm_model_list, \ + general_settings, \ + proxy_budget_rescheduler_min_time, \ + proxy_budget_rescheduler_max_time, \ + litellm_proxy_admin_name, \ + db_writer_client, \ + store_model_in_db, \ + premium_user, \ + _license_check, \ + proxy_batch_polling_interval, \ + shared_aiohttp_session import json init_verbose_loggers() @@ -1910,9 +1951,9 @@ redis_usage_cache: Optional[RedisCache] = ( None # redis cache used for tracking spend, tpm/rpm limits ) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[str] = ( - [] -) # Models that should use native provider background mode instead of polling +native_background_mode: List[ + str +] = [] # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -3219,13 +3260,9 @@ def _write_health_state_to_router_cache( exception_status = getattr(original_exception, "status_code", 500) - if ( - llm_router.health_check_ignore_transient_errors - and exception_status - in ( - 429, - 408, - ) + if llm_router.health_check_ignore_transient_errors and exception_status in ( + 429, + 408, ): continue @@ -4054,7 +4091,9 @@ class ProxyConfig: # Cast to SearchToolTypedDict for type safety try: - search_tool_typed: SearchToolTypedDict = SearchToolTypedDict(**search_tool) # type: ignore + search_tool_typed: SearchToolTypedDict = SearchToolTypedDict( + **search_tool + ) # type: ignore search_tools_parsed.append(search_tool_typed) except Exception as e: verbose_proxy_logger.error( @@ -4129,7 +4168,35 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_key_update, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, proxy_batch_polling_interval, config_passthrough_endpoints + global \ + master_key, \ + user_config_file_path, \ + otel_logging, \ + user_custom_auth, \ + user_custom_auth_path, \ + user_custom_key_generate, \ + user_custom_key_update, \ + user_custom_sso, \ + user_custom_ui_sso_sign_in_handler, \ + use_background_health_checks, \ + use_shared_health_check, \ + health_check_interval, \ + health_check_concurrency, \ + use_queue, \ + proxy_budget_rescheduler_max_time, \ + proxy_budget_rescheduler_min_time, \ + ui_access_mode, \ + litellm_master_key_hash, \ + proxy_batch_write_at, \ + disable_spend_logs, \ + prompt_injection_detection_obj, \ + redis_usage_cache, \ + store_model_in_db, \ + premium_user, \ + open_telemetry_logger, \ + health_check_details, \ + proxy_batch_polling_interval, \ + config_passthrough_endpoints config: dict = await self.get_config(config_file_path=config_file_path) @@ -4406,7 +4473,10 @@ class ProxyConfig: pass elif key == "responses": # Initialize global polling via cache settings - global polling_via_cache_enabled, native_background_mode, polling_cache_ttl + global \ + polling_via_cache_enabled, \ + native_background_mode, \ + polling_cache_ttl background_mode = value.get("background_mode", {}) polling_via_cache_enabled = background_mode.get( "polling_via_cache", False @@ -4459,6 +4529,8 @@ class ProxyConfig: f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}" ) setattr(litellm, key, value) + if key == "request_timeout": + litellm.request_timeout_explicitly_set = True if key in {"s3_audit_callback_params", "s3_callback_params"}: from litellm.integrations.s3_v2 import S3Logger as S3V2Logger from litellm.litellm_core_utils.litellm_logging import ( @@ -5039,8 +5111,7 @@ class ProxyConfig: ### LOAD FROM GOOGLE KMS ### load_google_kms(use_google_kms=True) elif ( - key_management_system - == KeyManagementSystem.AWS_SECRET_MANAGER.value # noqa: F405 + key_management_system == KeyManagementSystem.AWS_SECRET_MANAGER.value # noqa: F405 ): from litellm.secret_managers.aws_secret_manager_v2 import ( AWSSecretsManagerV2, @@ -6575,10 +6646,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[Guardrail] = ( - await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client - ) + guardrails_in_db: List[ + Guardrail + ] = await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -6890,7 +6961,23 @@ async def initialize( use_queue=False, config=None, ): - global user_model, user_api_base, user_debug, user_detailed_debug, user_user_max_tokens, user_request_timeout, user_temperature, user_telemetry, user_headers, experimental, llm_model_list, llm_router, general_settings, master_key, user_custom_auth, prisma_client + global \ + user_model, \ + user_api_base, \ + user_debug, \ + user_detailed_debug, \ + user_user_max_tokens, \ + user_request_timeout, \ + user_temperature, \ + user_telemetry, \ + user_headers, \ + experimental, \ + llm_model_list, \ + llm_router, \ + general_settings, \ + master_key, \ + user_custom_auth, \ + prisma_client from litellm.proxy.common_utils.banner import show_banner show_banner() @@ -7085,57 +7172,122 @@ def _get_client_requested_model_for_streaming(request_data: dict) -> str: return requested_model if isinstance(requested_model, str) else "" +def _is_positive_int_like(value: Any) -> bool: + try: + return int(value) > 0 + except (TypeError, ValueError): + return False + + +def _should_include_fallback_errors(request_data: dict[str, object]) -> bool: + if not general_settings.get("expose_fallback_errors_to_caller"): + return False + return request_data.get("include_fallback_errors") is True + + +def _get_streaming_fallback_metadata( + response_obj: object, +) -> tuple[bool, str | None, list[dict[str, object]]]: + additional_headers = get_hidden_params_dict(response_obj).get("additional_headers") + if not isinstance(additional_headers, dict): + return False, None, [] + + if not _is_positive_int_like( + additional_headers.get("x-litellm-attempted-fallbacks") + ): + return False, None, [] + + fallback_model = additional_headers.get("x-litellm-model-group") + fallback_errors = get_fallback_errors_from_headers(additional_headers) + if isinstance(fallback_model, str) and fallback_model: + return True, fallback_model, fallback_errors + return True, None, fallback_errors + + +def _format_fallback_metadata_sse_event( + *, + fallback_model: str | None, + fallback_errors: list[dict[str, object]], +) -> str: + import time + + payload = { + "id": "litellm-fallback-metadata", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": fallback_model or "", + "choices": [], + "litellm_fallback": { + "fallback_model": fallback_model, + "errors": fallback_errors, + }, + } + return f"data: {json.dumps(payload)}\n\n" + + def _restamp_streaming_chunk_model( *, chunk: Any, requested_model_from_client: str, request_data: dict, model_mismatch_logged: bool, -) -> Tuple[Any, bool]: + fallback_was_attempted: bool = False, + fallback_model_from_metadata: str | None = None, +) -> tuple[Any, bool]: + target_model = ( + fallback_model_from_metadata + if fallback_was_attempted + else requested_model_from_client + ) # Always return the client-requested model name (not provider-prefixed internal identifiers) # on streaming chunks. + # On fallback, use the public OpenAI-compatible model name. This keeps + # provider-prefixed internal identifiers from leaking into the public API. # # Note: This warning is intentionally verbose. A mismatch is a useful signal that an # internal provider/deployment identifier is leaking into the public API, and helps # maintainers/operators catch regressions while preserving OpenAI-compatible output. - if not requested_model_from_client or not isinstance(chunk, (BaseModel, dict)): + if not target_model or not isinstance(chunk, (BaseModel, dict)): return chunk, model_mismatch_logged # For Azure Model Router, preserve the actual model used in each chunk - if _is_azure_model_router_request(requested_model_from_client): + if not fallback_was_attempted and _is_azure_model_router_request( + requested_model_from_client + ): return chunk, model_mismatch_logged # For fastest_response batch completions, preserve the winning model's name # instead of stamping the comma-separated list the client sent. - if request_data.get("fastest_response", False): + if not fallback_was_attempted and request_data.get("fastest_response", False): return chunk, model_mismatch_logged downstream_model = ( chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) ) - if downstream_model == requested_model_from_client: + if downstream_model == target_model: return chunk, model_mismatch_logged - if not model_mismatch_logged and downstream_model != requested_model_from_client: + if not model_mismatch_logged and downstream_model != target_model: verbose_proxy_logger.debug( - "litellm_call_id=%s: streaming chunk model mismatch - requested=%r downstream=%r. Overriding model to requested.", + "litellm_call_id=%s: streaming chunk model mismatch - target=%r downstream=%r fallback_was_attempted=%s. Overriding chunk model to target.", request_data.get("litellm_call_id"), - requested_model_from_client, + target_model, downstream_model, + fallback_was_attempted, ) model_mismatch_logged = True if isinstance(chunk, dict): - chunk["model"] = requested_model_from_client + chunk["model"] = target_model return chunk, model_mismatch_logged try: - setattr(chunk, "model", requested_model_from_client) + chunk.model = target_model except Exception as e: verbose_proxy_logger.error( "litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s", request_data.get("litellm_call_id"), - requested_model_from_client, + target_model, type(chunk), str(e), exc_info=True, @@ -7294,7 +7446,14 @@ async def async_data_generator( requested_model_from_client = _get_client_requested_model_for_streaming( request_data=request_data ) + ( + fallback_was_attempted, + fallback_model_from_metadata, + fallback_errors, + ) = _get_streaming_fallback_metadata(response) model_mismatch_logged = False + fallback_metadata_event_sent = False + include_fallback_errors = _should_include_fallback_errors(request_data) # Use a running string instead of list + join to avoid O(n^2) overhead. # Previously "".join(str_so_far_parts) was called every chunk, re-joining # the entire accumulated response. String += is O(n) amortized total. @@ -7332,13 +7491,37 @@ async def async_data_generator( str_so_far=_str_so_far, ) + # Mid-stream fallbacks surface metadata on individual chunks rather than + # the response wrapper. Keep scanning chunks until a fallback model is + # resolved, then latch it for the rest of the stream. + if fallback_model_from_metadata is None: + ( + chunk_fallback_was_attempted, + chunk_fallback_model, + chunk_fallback_errors, + ) = _get_streaming_fallback_metadata(chunk) + if chunk_fallback_was_attempted: + fallback_was_attempted = True + fallback_model_from_metadata = chunk_fallback_model + fallback_errors = fallback_errors or chunk_fallback_errors + + pending_fallback_event = ( + include_fallback_errors + and fallback_was_attempted + and fallback_errors + and not fallback_metadata_event_sent + ) + chunk, model_mismatch_logged = _restamp_streaming_chunk_model( chunk=chunk, requested_model_from_client=requested_model_from_client, request_data=request_data, model_mismatch_logged=model_mismatch_logged, + fallback_was_attempted=fallback_was_attempted, + fallback_model_from_metadata=fallback_model_from_metadata, ) + raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) elif isinstance(chunk, bytes): @@ -7354,14 +7537,14 @@ async def async_data_generator( raise ValueError( "Raw SSE stream exceeded maximum buffered size without a frame delimiter" ) - continue - if chunk.startswith(("data:", "event:", ":")): + raw_passthrough = True + elif chunk.startswith(("data:", "event:", ":")): yield ( chunk if chunk.endswith(_SSE_FRAME_DELIMITERS) else chunk + "\n\n" ) - continue + raw_passthrough = True elif isinstance(chunk, str) and is_raw_sse_stream: raw_sse_buffer += chunk while True: @@ -7373,15 +7556,23 @@ async def async_data_generator( raise ValueError( "Raw SSE stream exceeded maximum buffered size without a frame delimiter" ) - continue + raw_passthrough = True elif isinstance(chunk, str) and chunk.startswith("data: "): error_message = chunk break - try: - yield _format_streaming_sse_chunk(chunk=chunk) - except Exception as e: - yield f"data: {str(e)}\n\n" + if not raw_passthrough: + try: + yield _format_streaming_sse_chunk(chunk=chunk) + except Exception as e: + yield f"data: {str(e)}\n\n" + + if pending_fallback_event: + yield _format_fallback_metadata_sse_event( + fallback_model=fallback_model_from_metadata, + fallback_errors=fallback_errors, + ) + fallback_metadata_event_sent = True stream_completed = True if not needs_iterator_wrap: @@ -7764,7 +7955,8 @@ class ProxyStartupEvent: teams_pydantic_obj = [NewUserRequestTeam(**team) for team in _teams] await update_default_team_member_budget( - teams=teams_pydantic_obj, user_api_key_dict=UserAPIKeyAuth(token=hash_token(master_key)) # type: ignore + teams=teams_pydantic_obj, + user_api_key_dict=UserAPIKeyAuth(token=hash_token(master_key)), # type: ignore ) @classmethod @@ -8597,7 +8789,13 @@ async def model_list( Hiding is presentation-only: a hidden model can still be called directly. """ - global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + global \ + llm_model_list, \ + general_settings, \ + llm_router, \ + prisma_client, \ + user_api_key_cache, \ + proxy_logging_obj settings = cast(dict[str, object], general_settings) # any-ok: legacy settings @@ -8774,7 +8972,13 @@ async def model_info( scoping, health filtering, paused deployments) drives both endpoints; the listing's public id must resolve to the same internal deployment here. """ - global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + global \ + llm_model_list, \ + general_settings, \ + llm_router, \ + prisma_client, \ + user_api_key_cache, \ + proxy_logging_obj settings = cast(dict[str, object], general_settings) # any-ok: legacy settings @@ -8925,9 +9129,9 @@ async def chat_completion( hasattr(user_api_key_dict, "organization_alias") and user_api_key_dict.organization_alias is not None ): - data["metadata"][ - "user_api_key_org_alias" - ] = user_api_key_dict.organization_alias + data["metadata"]["user_api_key_org_alias"] = ( + user_api_key_dict.organization_alias + ) if ( hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None @@ -9109,9 +9313,9 @@ async def completion( hasattr(user_api_key_dict, "organization_alias") and user_api_key_dict.organization_alias is not None ): - data["metadata"][ - "user_api_key_org_alias" - ] = user_api_key_dict.organization_alias + data["metadata"]["user_api_key_org_alias"] = ( + user_api_key_dict.organization_alias + ) if ( hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None @@ -9325,9 +9529,12 @@ async def embeddings( litellm_params = deployment.get("litellm_params", {}) or {} litellm_model = litellm_params.get("model", "") # Check if this provider supports token arrays - supports_token_arrays = litellm_model in litellm.open_ai_embedding_models or any( - litellm_model.startswith(provider) - for provider in LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS + supports_token_arrays = ( + litellm_model in litellm.open_ai_embedding_models + or any( + litellm_model.startswith(provider) + for provider in LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS + ) ) if not supports_token_arrays: # non-openai/azure embedding model called with token input - decode tokens @@ -9360,9 +9567,9 @@ async def embeddings( hasattr(user_api_key_dict, "organization_alias") and user_api_key_dict.organization_alias is not None ): - data["metadata"][ - "user_api_key_org_alias" - ] = user_api_key_dict.organization_alias + data["metadata"]["user_api_key_org_alias"] = ( + user_api_key_dict.organization_alias + ) if ( hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None @@ -11390,18 +11597,20 @@ def get_direct_access_models( llm_router: Router, ) -> List[str]: """ - Get all models that user has direct access to - """ + Get all models that user has direct access to. - direct_access_models: List[str] = [] - for model in user_db_object.models: - deployments = llm_router.get_model_list(model_name=model) - if deployments is not None: - for deployment in deployments: - model_id = deployment.get("model_info", {}).get("id", None) - if model_id is not None: - direct_access_models.append(model_id) - return direct_access_models + The 'all-proxy-models' sentinel grants direct access to every non-team + deployment, mirroring how get_key_models expands it for the key/team path. + """ + if SpecialModelNames.all_proxy_models.value in user_db_object.models: + return llm_router.get_model_ids(exclude_team_models=True) + + return [ + model_id + for model in user_db_object.models + for deployment in (llm_router.get_model_list(model_name=model) or []) + if (model_id := deployment.get("model_info", {}).get("id", None)) is not None + ] def _filter_models_to_user_accessible(all_models: List[Dict]) -> List[Dict]: @@ -12318,7 +12527,12 @@ async def model_info_v2( } ``` """ - global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router + global \ + llm_model_list, \ + general_settings, \ + user_config_file_path, \ + proxy_config, \ + llm_router # Return empty data array when no models are configured (graceful handling for fresh installs) if llm_router is None or not llm_router.model_list: @@ -13074,7 +13288,13 @@ async def model_info_v1( ``` """ - global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router, user_model + global \ + llm_model_list, \ + general_settings, \ + user_config_file_path, \ + proxy_config, \ + llm_router, \ + user_model # Unit tests call this handler directly; FastAPI normally resolves Query defaults. if not isinstance(include_team_models, bool): @@ -13416,7 +13636,12 @@ async def model_group_info( } ``` """ - global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router + global \ + llm_model_list, \ + general_settings, \ + user_config_file_path, \ + proxy_config, \ + llm_router # Return empty data array when no models are configured (graceful handling for fresh installs) if llm_model_list is None or llm_router is None or not llm_model_list: @@ -14808,7 +15033,14 @@ async def update_config( untouched — this endpoint never persists pre-existing YAML values to DB as a side effect of an unrelated update. """ - global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client + global \ + llm_router, \ + llm_model_list, \ + general_settings, \ + proxy_config, \ + proxy_logging_obj, \ + master_key, \ + prisma_client try: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( @@ -15061,7 +15293,10 @@ async def update_config_general_settings( response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ - "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore + "create": { + "param_name": "general_settings", + "param_value": json.dumps(general_settings), + }, # type: ignore "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) @@ -15073,6 +15308,68 @@ async def update_config_general_settings( return response +# Secret-bearing general_settings fields the segment masker does not match by +# name: database_url and database_extra_connection_params embed DB credentials, +# pass_through_endpoints carry upstream Authorization headers, and +# alert_to_webhook_url is itself a webhook secret +_EXTRA_SECRET_GENERAL_SETTINGS_FIELDS = frozenset( + { + "database_url", + "database_extra_connection_params", + "pass_through_endpoints", + "alert_to_webhook_url", + } +) + + +def _is_secret_general_setting_field(field_name: str) -> bool: + return ( + field_name in _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS + or SENSITIVE_DATA_MASKER.is_sensitive_key(field_name) + ) + + +# Matches the cap on _redact_sensitive_litellm_params (the closest analog in the +# proxy). Past this depth we fail closed by returning "REDACTED" for the whole +# subtree rather than recursing further — better to over-redact a pathological +# config than to silently return a deeply-nested credential verbatim +_REDACT_SECRET_MAX_DEPTH = 10 + + +def _redact_secret_values_in_obj(value: JsonValue, depth: int = 0) -> JsonValue: + """Recursively redact secret leaves inside a structured field so a nested + credential (e.g. aws_web_identity_token under database_args) is never + returned to a non-admin, while non-secret siblings stay visible. At + _REDACT_SECRET_MAX_DEPTH the whole subtree is replaced with "REDACTED" + so depth-overrun fails closed.""" + if depth >= _REDACT_SECRET_MAX_DEPTH: + return "REDACTED" + if isinstance(value, dict): + return { + key: ( + "REDACTED" + if _is_secret_general_setting_field(key) + else _redact_secret_values_in_obj(sub, depth + 1) + ) + for key, sub in value.items() + } + if isinstance(value, list): + return [_redact_secret_values_in_obj(item, depth + 1) for item in value] + return value + + +def _redact_general_setting_value( + field_name: str, value: JsonValue, is_full_admin: bool +) -> JsonValue: + if is_full_admin: + return value + if _is_secret_general_setting_field(field_name): + return "REDACTED" + if isinstance(value, (dict, list)): + return _redact_secret_values_in_obj(value) + return value + + @router.get( "/config/field/info", tags=["config.yaml"], @@ -15125,9 +15422,11 @@ async def get_config_general_settings( general_settings = dict(db_general_settings.param_value) if field_name in general_settings: - field_value = general_settings[field_name] - # Redact plugin_key from plugin configs so the shared credential - # is never returned even to admin-viewer callers. + field_value = _redact_general_setting_value( + field_name, + general_settings[field_name], + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + ) if field_name == "plugins" and isinstance(field_value, list): field_value = [ ( @@ -15183,6 +15482,8 @@ async def get_config_list( }, ) + is_full_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ## get general settings from db db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} @@ -15204,6 +15505,7 @@ async def get_config_list( "maximum_spend_logs_retention_period": {"type": "String"}, "mcp_internal_ip_ranges": {"type": "List"}, "mcp_trusted_proxy_ranges": {"type": "List"}, + "mcp_xff_num_trusted_hops": {"type": "Integer"}, "always_include_stream_usage": {"type": "Boolean"}, "forward_client_headers_to_llm_api": {"type": "Boolean"}, "mcp_required_fields": {"type": "List"}, @@ -15231,7 +15533,11 @@ async def get_config_list( field_name=sub_field, field_type=sub_field_type.__name__, field_description="", # Add custom logic if descriptions are available - field_default_value=general_settings.get(sub_field, None), + field_default_value=_redact_general_setting_value( + sub_field, + general_settings.get(sub_field, None), + is_full_admin, + ), stored_in_db=None, ) for sub_field, sub_field_type in pydantic_class.__annotations__.items() @@ -15246,9 +15552,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[idx].field_description = ( - sub_field_info.description - ) + nested_fields[ + idx + ].field_description = sub_field_info.description idx += 1 _stored_in_db = None @@ -15261,7 +15567,11 @@ async def get_config_list( field_name=field_name, field_type=allowed_args[field_name]["type"], field_description=field_info.description or "", - field_value=general_settings.get(field_name, None), + field_value=_redact_general_setting_value( + field_name, + general_settings.get(field_name, None), + is_full_admin, + ), stored_in_db=_stored_in_db, field_default_value=field_info.default, nested_fields=nested_fields, @@ -15285,7 +15595,9 @@ async def get_config_list( field_name=field_name, field_type=allowed_args[field_name]["type"], field_description=field_info.description or "", - field_value=_field_value, + field_value=_redact_general_setting_value( + field_name, _field_value, is_full_admin + ), stored_in_db=_stored_in_db, field_default_value=field_info.default, nested_fields=nested_fields, @@ -15359,7 +15671,10 @@ async def delete_config_general_settings( response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ - "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore + "create": { + "param_name": "general_settings", + "param_value": json.dumps(general_settings), + }, # type: ignore "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) @@ -15427,9 +15742,9 @@ async def delete_callback( # Remove callback from success_callback list success_callbacks.remove(callback_name) - config.setdefault("litellm_settings", {})[ - "success_callback" - ] = success_callbacks + config.setdefault("litellm_settings", {})["success_callback"] = ( + success_callbacks + ) # Save the updated configuration await proxy_config.save_config(new_config=config) @@ -15473,10 +15788,14 @@ async def get_config(): # return the callbacks and the env variables for the callback """ - global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key + global \ + llm_router, \ + llm_model_list, \ + general_settings, \ + proxy_config, \ + proxy_logging_obj, \ + master_key try: - import base64 - all_available_callbacks = AllCallbacks() config_data = await proxy_config.get_config() @@ -15542,18 +15861,14 @@ async def get_config(): _slack_vars = [ "SLACK_WEBHOOK_URL", ] - _slack_env_vars = {} - for _var in _slack_vars: - env_variable = environment_variables.get(_var, None) - if env_variable is None: - _value = os.getenv("SLACK_WEBHOOK_URL", None) - _slack_env_vars[_var] = _value - else: - # decode + decrypt the value - _decrypted_value = decrypt_value_helper( - value=env_variable, key=_var - ) - _slack_env_vars[_var] = _decrypted_value + _slack_env_vars = { + _var: ( + value + if (value := environment_variables.get(_var)) is not None + else os.getenv(_var) + ) + for _var in _slack_vars + } _slack_env_vars = mask_sensitive_keys( _slack_env_vars, _ALERTING_SENSITIVE_VARS ) @@ -15584,15 +15899,9 @@ async def get_config(): "EMAIL_LOGO_URL", "EMAIL_SUPPORT_CONTACT", ] - _email_env_vars = {} - for _var in _email_vars: - env_variable = environment_variables.get(_var, None) - if env_variable is None: - _email_env_vars[_var] = None - else: - # decode + decrypt the value - _decrypted_value = decrypt_value_helper(value=env_variable, key=_var) - _email_env_vars[_var] = _decrypted_value + _email_env_vars = { + _var: environment_variables.get(_var) for _var in _email_vars + } _email_env_vars = mask_sensitive_keys(_email_env_vars, _ALERTING_SENSITIVE_VARS) alerting_data.append( @@ -16450,6 +16759,7 @@ app.include_router(caching_router) app.include_router(analytics_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) +app.include_router(rust_control_plane_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) app.include_router(team_callback_router) diff --git a/litellm/proxy/read_model_list.py b/litellm/proxy/read_model_list.py new file mode 100644 index 00000000000..2dff8eaf698 --- /dev/null +++ b/litellm/proxy/read_model_list.py @@ -0,0 +1,28 @@ +"""Resolve a proxy config's ``model_list`` for the Rust AI gateway. + +The Rust gateway calls this once at load time (via an embedded interpreter) and +builds its own (Rust) router from the returned ``model_list``. We do NOT call +``ProxyConfig.load_config`` here: that returns a *Python* ``litellm.Router`` (not +usable from Rust) and boots the whole proxy (callbacks, cache, DB, auth) as side +effects. + +Instead we reuse ``ProxyConfig.get_config`` — the actual config reader — so the +gateway inherits the same heavy lifting the proxy does: ``include:`` merging, +``os.environ/`` + secret-manager resolution, and DB-stored models (when a DB is +configured). It has no proxy-setup side effects. Returns the resolved +``model_list``; the Rust side deserializes each entry into its ``Deployment``. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + + +def read_model_list(config_path: str) -> list[dict[str, Any]]: + """Load ``config_path`` via the proxy's own reader and return its + resolved ``model_list``.""" + from litellm.proxy.proxy_server import ProxyConfig + + config = asyncio.run(ProxyConfig().get_config(config_file_path=config_path)) + return config.get("model_list") or [] diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index a69e6734d71..aedfb26326a 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -92,9 +92,7 @@ async def background_streaming_task( # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming output_items: dict[str, dict[str, Any]] = {} # Track output items by ID - accumulated_text = ( - {} - ) # Track accumulated text deltas by (item_id, content_index) + accumulated_text = {} # Track accumulated text deltas by (item_id, content_index) # ResponsesAPIResponse fields to extract from response.completed usage_data = None diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0ba77dcd2f0..29fc6d7c30f 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -536,9 +536,7 @@ async def get_global_activity_model( if db_response is None: return [] - model_ui_data: dict = ( - {} - ) # {"gpt-4": {"daily_data": [], "sum_api_requests": 0, "sum_total_tokens": 0}} + model_ui_data: dict = {} # {"gpt-4": {"daily_data": [], "sum_api_requests": 0, "sum_total_tokens": 0}} for row in db_response: _model = row["model_group"] @@ -690,9 +688,7 @@ async def get_global_activity_exceptions_per_deployment( if db_response is None: return [] - model_ui_data: dict = ( - {} - ) # {"gpt-4": {"daily_data": [], "sum_api_requests": 0, "sum_total_tokens": 0}} + model_ui_data: dict = {} # {"gpt-4": {"daily_data": [], "sum_api_requests": 0, "sum_total_tokens": 0}} for row in db_response: _model = row["api_base"] @@ -2218,9 +2214,7 @@ async def ui_view_request_response_for_request_id( request_id=request_id, ) - custom_loggers = ( - litellm.logging_callback_manager.get_active_additional_logging_utils_from_custom_logger() - ) + custom_loggers = litellm.logging_callback_manager.get_active_additional_logging_utils_from_custom_logger() start_date_obj: Optional[datetime] = None end_date_obj: Optional[datetime] = None if start_date is not None: @@ -2418,7 +2412,9 @@ async def view_spend_logs( ): result: dict = {} for record in response: - dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ") # type: ignore + dt_object = datetime.strptime( + str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ" + ) # type: ignore date = dt_object.date() if date not in result: result[date] = {"users": {}, "models": {}} diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index cefe349aade..674f7efd837 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -621,7 +621,8 @@ async def update_internal_user_settings( isinstance(team, NewUserRequestTeam) for team in settings.teams ): await update_default_team_member_budget( - settings.teams, user_api_key_dict=user_api_key_dict # type: ignore + settings.teams, + user_api_key_dict=user_api_key_dict, # type: ignore ) return await _update_litellm_setting( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 755602cbcc0..781f0e9f301 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -523,7 +523,9 @@ class ProxyLogging: or "outage_alerts" in self.alert_types or "region_outage_alerts" in self.alert_types ): - litellm.logging_callback_manager.add_litellm_callback(self.slack_alerting_instance) # type: ignore + litellm.logging_callback_manager.add_litellm_callback( + self.slack_alerting_instance + ) # type: ignore litellm.logging_callback_manager.add_litellm_success_callback( self.slack_alerting_instance.response_taking_too_long_callback ) @@ -1691,10 +1693,8 @@ class ProxyLogging: for callback in callbacks: if isinstance(callback, str): - resolved: Any = ( - litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) + resolved: Any = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( + cast(_custom_logger_compatible_callbacks_literal, callback) ) else: resolved = callback @@ -3627,7 +3627,9 @@ class PrismaClient: return response elif table_name == "user_notification": if query_type == "find_unique": - response = await UserNotificationsRepository(self).table.find_unique( # type: ignore + response = await UserNotificationsRepository( + self + ).table.find_unique( # type: ignore where={"user_id": user_id} # type: ignore ) elif query_type == "find_all": @@ -3831,7 +3833,9 @@ class PrismaClient: print_verbose( "PrismaClient: Before upsert into litellm_verificationtoken" ) - new_verification_token = await VerificationTokenRepository(self).table.upsert( # type: ignore + new_verification_token = await VerificationTokenRepository( + self + ).table.upsert( # type: ignore where={ "token": hashed_token, }, @@ -5717,19 +5721,19 @@ async def update_daily_tag_spend( """ n_retry_times = 3 try: - if ( - proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis() - ): + if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis(): await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, ) else: - await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db( - prisma_client=prisma_client, - n_retry_times=n_retry_times, - proxy_logging_obj=proxy_logging_obj, + await ( + proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + ) ) except Exception as e: # NOTE: keep this as a plain ``error`` (no traceback) to match the diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 9c2d3050346..9c2d297bfa7 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -42,9 +42,9 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( Raises: HTTPException: If user doesn't have access to the vector store """ - vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = ( - await get_litellm_managed_vector_store(vector_store_id=vector_store_id) - ) + vector_store_to_run: Optional[ + LiteLLM_ManagedVectorStore + ] = await get_litellm_managed_vector_store(vector_store_id=vector_store_id) if vector_store_to_run is not None: if user_api_key_dict is not None: await assert_user_can_access_vector_store( diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 1de68e2ac94..b49c86b3c8b 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -184,7 +184,9 @@ class BaseRAGIngestion(ABC): # Extract text from pages if hasattr(ocr_response, "pages") and ocr_response.pages: # type: ignore return "\n\n".join( - page.markdown for page in ocr_response.pages if hasattr(page, "markdown") # type: ignore + page.markdown + for page in ocr_response.pages + if hasattr(page, "markdown") # type: ignore ) return None diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index 4c79cd26150..b7bd87d1f6c 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -317,7 +317,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Construct upload URL using vertex base URL helper base_url = get_vertex_base_url(self.location) - url = f"{base_url}/upload/v1beta1/" f"{rag_corpus_id}/ragFiles:upload" + url = f"{base_url}/upload/v1beta1/{rag_corpus_id}/ragFiles:upload" # Build metadata for the file with snake_case keys (as per upload API docs) metadata: Dict[str, Any] = { @@ -423,7 +423,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Construct import URL using vertex base URL helper base_url = get_vertex_base_url(self.location) - url = f"{base_url}/v1beta1/" f"{rag_corpus_id}/ragFiles:import" + url = f"{base_url}/v1beta1/{rag_corpus_id}/ragFiles:import" # Build request body with camelCase keys (Vertex AI API format) request_body: Dict[str, Any] = { @@ -447,9 +447,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): "max_embedding_requests_per_min" ) if max_embedding_qpm: - request_body["importRagFilesConfig"][ - "maxEmbeddingRequestsPerMin" - ] = max_embedding_qpm + request_body["importRagFilesConfig"]["maxEmbeddingRequestsPerMin"] = ( + max_embedding_qpm + ) verbose_logger.debug(f"Importing files from GCS: {url}") verbose_logger.debug(f"Request body: {json.dumps(request_body, indent=2)}") diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 03a2f339bea..a29f0a01b7a 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -38,16 +38,14 @@ class LiteLLMCompletionTransformationHandler: Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] ], ]: - litellm_completion_request: dict = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( - model=model, - input=input, - responses_api_request=responses_api_request, - custom_llm_provider=custom_llm_provider, - stream=stream, - extra_headers=extra_headers, - **kwargs, - ) + litellm_completion_request: dict = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=input, + responses_api_request=responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + **kwargs, ) if _is_async: @@ -69,12 +67,10 @@ class LiteLLMCompletionTransformationHandler: ) if isinstance(litellm_completion_response, ModelResponse): - responses_api_response: ResponsesAPIResponse = ( - LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - chat_completion_response=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, - ) + responses_api_response: ResponsesAPIResponse = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + chat_completion_response=litellm_completion_response, + request_input=input, + responses_api_request=responses_api_request, ) return responses_api_response @@ -119,12 +115,10 @@ class LiteLLMCompletionTransformationHandler: ) if isinstance(litellm_completion_response, ModelResponse): - responses_api_response: ResponsesAPIResponse = ( - LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - chat_completion_response=litellm_completion_response, - request_input=request_input, - responses_api_request=responses_api_request, - ) + responses_api_response: ResponsesAPIResponse = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + chat_completion_response=litellm_completion_response, + request_input=request_input, + responses_api_request=responses_api_request, ) return responses_api_response diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 71ff2eb7acf..45ab16b0d4a 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -43,10 +43,10 @@ class ResponsesSessionHandler: verbose_proxy_logger.debug( "inside get_chat_completion_message_history_for_previous_response_id" ) - all_spend_logs: List[SpendLogsPayload] = ( - await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( - previous_response_id - ) + all_spend_logs: List[ + SpendLogsPayload + ] = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + previous_response_id ) verbose_proxy_logger.debug( "found %s spend logs for this response id", len(all_spend_logs) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 767281d43ab..582144e6cb8 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -615,8 +615,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{str(uuid.uuid4())}" text = getattr(litellm_complete_object.choices[0].message, "content", "") or "" # type: ignore - reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" # type: ignore - annotations = getattr(litellm_complete_object.choices[0].message, "annotations", None) # type: ignore + reasoning_content = ( + getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") + or "" + ) # type: ignore + annotations = getattr( + litellm_complete_object.choices[0].message, "annotations", None + ) # type: ignore part: Optional[PART_UNION_TYPES] = None if reasoning_content: @@ -651,7 +656,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{str(uuid.uuid4())}" text = self.litellm_model_response.choices[0].message.content or "" # type: ignore - annotations = getattr(self.litellm_model_response.choices[0].message, "annotations", None) # type: ignore + annotations = getattr( + self.litellm_model_response.choices[0].message, "annotations", None + ) # type: ignore response_annotations = LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( annotations=annotations diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5b5ff122c50..0bd78e59819 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1409,7 +1409,9 @@ class LiteLLMCompletionResponsesConfig: if tool.get("defer_loading"): chat_completion_tool["defer_loading"] = tool.get("defer_loading") # type: ignore if tool.get("allowed_callers"): - chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") # type: ignore + chat_completion_tool["allowed_callers"] = tool.get( + "allowed_callers" + ) # type: ignore if tool.get("input_examples"): chat_completion_tool["input_examples"] = tool.get("input_examples") # type: ignore chat_completion_tools.append( @@ -1522,7 +1524,11 @@ class LiteLLMCompletionResponsesConfig: # Pass through provider_specific_fields as-is if present if provider_specific_fields: - setattr(output_tool_call, "provider_specific_fields", provider_specific_fields) # type: ignore + setattr( + output_tool_call, + "provider_specific_fields", + provider_specific_fields, + ) # type: ignore responses_tools.append(output_tool_call) return responses_tools diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 34c9cdd3d1c..3edcbd430f1 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -58,7 +58,10 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams -from litellm.utils import ProviderConfigManager, client +from litellm.utils import ( + ProviderConfigManager, + client, +) if TYPE_CHECKING: from mcp.types import Tool as MCPTool @@ -304,9 +307,7 @@ async def aresponses_api_with_mcp( # Auto-Execute Tools Handling # If auto-execute tools is True, then we need to execute the tool calls ######################################################### - if should_auto_execute and isinstance( - response, ResponsesAPIResponse - ): # type: ignore + if should_auto_execute and isinstance(response, ResponsesAPIResponse): # type: ignore tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response( response=response ) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 42c46dff47c..839e0232b04 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -632,7 +632,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): try: # Extract tool calls from the response if self.collected_response is not None: - tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response(self.collected_response) # type: ignore[arg-type] + tool_calls = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response( + self.collected_response + ) + ) # type: ignore[arg-type] else: tool_calls = [] if not tool_calls: diff --git a/litellm/router.py b/litellm/router.py index e54eadfb872..0008fcaa30b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -40,7 +40,6 @@ import anyio import httpx import openai from openai import AsyncOpenAI -from pydantic import BaseModel from typing_extensions import overload import litellm @@ -62,6 +61,9 @@ from litellm.constants import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -81,8 +83,10 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_strategy.simple_shuffle import simple_shuffle from litellm.router_strategy.tag_based_routing import get_deployments_for_tag from litellm.router_utils.add_retry_fallback_headers import ( + _HiddenParamsHost, add_fallback_headers_to_response, add_retry_headers_to_response, + get_hidden_params_dict, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -427,9 +431,7 @@ class Router: self.assistants_config = assistants_config self.search_tools = search_tools or [] self.guardrail_list = guardrail_list or [] - self.deployment_names: List = ( - [] - ) # names of models under litellm_params. ex. azure/chatgpt-v-2 + self.deployment_names: List = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = ( @@ -482,9 +484,9 @@ class Router: self.default_max_parallel_requests = default_max_parallel_requests self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() - self.team_pattern_routers: Dict[str, PatternMatchRouter] = ( - {} - ) # {"TEAM_ID": PatternMatchRouter} + self.team_pattern_routers: Dict[ + str, PatternMatchRouter + ] = {} # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} self.complexity_routers: Dict[str, "ComplexityRouter"] = {} self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {} @@ -522,9 +524,7 @@ class Router: if "model" in m["litellm_params"]: self.deployment_latency_map[m["litellm_params"]["model"]] = 0 else: - self.model_list: List = ( - [] - ) # initialize an empty list - to allow _add_deployment and delete_deployment to work + self.model_list: List = [] # initialize an empty list - to allow _add_deployment and delete_deployment to work if allowed_fails is not None: self.allowed_fails = allowed_fails @@ -544,9 +544,7 @@ class Router: self.health_state_cache = DeploymentHealthCache( cache=self.cache, staleness_threshold=float(_staleness) ) - self.failed_calls = ( - InMemoryCache() - ) # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown + self.failed_calls = InMemoryCache() # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown if num_retries is not None: self.num_retries = num_retries @@ -564,6 +562,12 @@ class Router: self._explicit_timeout = timeout # None when user did not pass timeout self.timeout = timeout or litellm.request_timeout + # Per-attempt request_timeout, independent of router_settings.timeout. + # Only stored when a router timeout is also set, since otherwise + # request_timeout already flows through self.timeout above. + self.request_timeout = ( + get_configured_request_timeout() if timeout is not None else None + ) self.stream_timeout = stream_timeout self.retry_after = retry_after @@ -605,9 +609,7 @@ class Router: self.success_calls: defaultdict = defaultdict( int ) # dict to store success_calls made to each model - self.previous_models: List = ( - [] - ) # list to store failed calls (passed in as metadata to next call) + self.previous_models: List = [] # list to store failed calls (passed in as metadata to next call) # make Router.chat.completions.create compatible for openai.chat.completions.create default_litellm_params = default_litellm_params or {} @@ -2165,6 +2167,36 @@ class Router: ) setattr(fallback_item, "usage", combined_usage) + @staticmethod + def _prepare_fallback_hidden_params( + fallback_response: object, + ) -> tuple[dict[str, object], dict[str, object]]: + fallback_hidden_params = get_hidden_params_dict(fallback_response) + fallback_headers = fallback_hidden_params.get("additional_headers") + if not isinstance(fallback_headers, dict): + return fallback_hidden_params, {} + return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + + @staticmethod + def _apply_fallback_hidden_params_to_item( + fallback_item: object, + prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], + ) -> None: + if fallback_item is None or not hasattr(fallback_item, "_hidden_params"): + return + + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + item_hidden_params = get_hidden_params_dict(fallback_item) + item_headers = item_hidden_params.get("additional_headers") + if not isinstance(item_headers, dict): + item_headers = {} + + cast(_HiddenParamsHost, fallback_item)._hidden_params = { + **item_hidden_params, + **fallback_hidden_params, + "additional_headers": {**item_headers, **fallback_headers}, + } + async def _acompletion_streaming_iterator( self, model_response: CustomStreamWrapper, @@ -2257,12 +2289,22 @@ class Router: model_group=model_group, args=(), kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get( + "include_fallback_errors", False + ) + is True, ) ) # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): + prepared_fallback_hidden_params = ( + Router._prepare_fallback_hidden_params(fallback_response) + ) async for fallback_item in fallback_response: # type: ignore + Router._apply_fallback_hidden_params_to_item( + fallback_item, prepared_fallback_hidden_params + ) if ( fallback_item and isinstance(fallback_item, ModelResponseStream) @@ -2686,11 +2728,21 @@ class Router: model_group=model_group, args=(), kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get( + "include_fallback_errors", False + ) + is True, ) ) if hasattr(fallback_response, "__aiter__"): + prepared_fallback_hidden_params = ( + Router._prepare_fallback_hidden_params(fallback_response) + ) async for fallback_item in fallback_response: # type: ignore + Router._apply_fallback_hidden_params_to_item( + fallback_item, prepared_fallback_hidden_params + ) if partial_usage is not None: Router._combine_responses_fallback_usage( fallback_item, partial_usage @@ -2815,7 +2867,13 @@ class Router: ) if hasattr(fallback_response, "__iter__"): + prepared_fallback_hidden_params = ( + Router._prepare_fallback_hidden_params(fallback_response) + ) for fallback_item in fallback_response: + Router._apply_fallback_hidden_params_to_item( + fallback_item, prepared_fallback_hidden_params + ) if ( fallback_item and isinstance(fallback_item, ModelResponseStream) @@ -2901,9 +2959,7 @@ class Router: """ model_name = None deployment = None - _timeout_debug_deployment_dict = ( - {} - ) # this is a temporary dict to debug timeout issues + _timeout_debug_deployment_dict = {} # this is a temporary dict to debug timeout issues try: input_kwargs_for_streaming_fallback = kwargs.copy() input_kwargs_for_streaming_fallback["model"] = model @@ -2972,6 +3028,7 @@ class Router: **kwargs, } input_kwargs.pop("silent_model", None) + input_kwargs.pop("include_fallback_errors", None) _response = litellm.acompletion(**input_kwargs) @@ -3076,7 +3133,18 @@ class Router: - litellm_trace_id - metadata """ - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) + # Normalise an explicit num_retries=None to the router default here (dict.get() + # only falls back when the key is absent, not when its value is None), then to 0 + # if the router default is itself None - mirroring the guard in + # async_function_with_retries, which remains the safety net for paths that bypass + # this setter. + _req_num_retries = kwargs.get("num_retries") + if _req_num_retries is not None: + kwargs["num_retries"] = _req_num_retries + else: + kwargs["num_retries"] = ( + self.num_retries if self.num_retries is not None else 0 + ) kwargs.setdefault("litellm_trace_id", str(uuid.uuid4())) model_group_alias: Optional[str] = None if self._get_model_from_alias(model=model): @@ -3322,6 +3390,7 @@ class Router: "stream_timeout", None ) # timeout set on litellm_params for this deployment or self.stream_timeout # timeout set on router + or self.request_timeout # litellm_settings.request_timeout (per-attempt) or self.default_litellm_params.get("stream_timeout", None) ) @@ -3338,7 +3407,8 @@ class Router: or data.get( "request_timeout", None ) # timeout set on litellm_params for this deployment - or self.timeout # timeout set on router + or self.request_timeout # litellm_settings.request_timeout (per-attempt) + or self.timeout # timeout set on router (router_settings.timeout) or self.default_litellm_params.get("timeout", None) ) return timeout @@ -3424,7 +3494,11 @@ class Router: _tasks = [] for model in models: # add each task but if the task fails - _tasks.append(_async_completion_no_exceptions(model=model, messages=messages, **kwargs)) # type: ignore + _tasks.append( + _async_completion_no_exceptions( + model=model, messages=messages, **kwargs + ) + ) # type: ignore response = await asyncio.gather(*_tasks) return response elif isinstance(messages, list) and all(isinstance(m, list) for m in messages): @@ -3533,7 +3607,9 @@ class Router: Wrapper around self.acompletion that catches exceptions and returns them as a result """ try: - result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) # type: ignore + result = await self.acompletion( + model=model, messages=messages, stream=stream, **kwargs + ) # type: ignore return result except asyncio.CancelledError: verbose_router_logger.debug( @@ -4334,7 +4410,9 @@ class Router: kwargs[k].update(v) # call via litellm.completion() - return litellm.text_completion(**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs}) # type: ignore + return litellm.text_completion( + **{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs} + ) # type: ignore except Exception as e: raise e @@ -6478,6 +6556,7 @@ class Router: model_group: Optional[str], args: tuple, kwargs: dict, + include_fallback_errors: bool = False, ): """ Common utilities for async_function_with_fallbacks @@ -6501,6 +6580,8 @@ class Router: input_kwargs["max_fallbacks"] = self.max_fallbacks if "fallback_depth" not in input_kwargs: input_kwargs["fallback_depth"] = 0 + if include_fallback_errors: + input_kwargs["include_fallback_errors"] = True # ORDER-BASED FALLBACKS: prepend higher order levels to the fallback list # Skip for error types that have their own dedicated fallback handlers @@ -6640,9 +6721,7 @@ class Router: ) verbose_router_logger.info( msg="Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ - to fallbacks, if available.{}".format( - error_message - ) + to fallbacks, if available.{}".format(error_message) ) if litellm.expose_router_debug_in_errors: @@ -6676,9 +6755,7 @@ class Router: ) verbose_router_logger.info( msg="Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ - to fallbacks, if available.{}".format( - error_message - ) + to fallbacks, if available.{}".format(error_message) ) if litellm.expose_router_debug_in_errors: @@ -6739,9 +6816,11 @@ class Router: and litellm.expose_router_debug_in_errors ): # add the available fallbacks to the exception - original_exception.message += ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore - model_group, - fallback_model_group, + original_exception.message += ( + ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore + model_group, + fallback_model_group, + ) ) if len(fallback_failure_exception_str) > 0: original_exception.message += ( # type: ignore @@ -6759,6 +6838,7 @@ class Router: If it fails after num_retries, fall back to another model group """ model_group: Optional[str] = kwargs.get("model") + include_fallback_errors = kwargs.get("include_fallback_errors", False) is True disable_fallbacks: Optional[bool] = kwargs.pop("disable_fallbacks", False) fallbacks: Optional[List] = kwargs.get("fallbacks", self.fallbacks) context_window_fallbacks: Optional[List] = kwargs.get( @@ -6802,6 +6882,7 @@ class Router: model_group, args, kwargs, + include_fallback_errors=include_fallback_errors, ) def _handle_mock_testing_fallbacks( @@ -6868,7 +6949,11 @@ class Router: "model_group_retry_policy", self.model_group_retry_policy ) model_group: Optional[str] = kwargs.get("model") - num_retries = kwargs.pop("num_retries") + num_retries = kwargs.pop("num_retries", None) + if num_retries is None: + # Fall back to the router setting (then 0) so the comparisons below never + # hit `None > int`, which would mask the real upstream error with a TypeError. + num_retries = self.num_retries if self.num_retries is not None else 0 ## ADD MODEL GROUP SIZE TO METADATA - used for model_group_rate_limit_error tracking _metadata: dict = kwargs.get("litellm_metadata", kwargs.get("metadata")) or {} @@ -7520,7 +7605,11 @@ class Router: """ Update RPM usage for a deployment """ - deployment_name = kwargs["litellm_params"]["metadata"].get( + deployment_name = kwargs[ + "litellm_params" + ][ + "metadata" + ].get( "deployment", None ) # handles wildcard routes - by giving the original name sent to `litellm.completion` model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) @@ -7580,9 +7669,7 @@ class Router: for ( k, v, - ) in ( - kwargs.items() - ): # log everything in kwargs except the old previous_models value - prevent nesting + ) in kwargs.items(): # log everything in kwargs except the old previous_models value - prevent nesting if k not in [_metadata_var, "messages", "original_function"]: previous_model[k] = v elif k == _metadata_var and isinstance(v, dict): @@ -9505,7 +9592,8 @@ class Router: ): model_group_info.supports_parallel_function_calling = True if ( - model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True # type: ignore + model_info.get("supports_vision", None) is not None + and model_info["supports_vision"] is True # type: ignore ): model_group_info.supports_vision = True if ( @@ -9525,7 +9613,8 @@ class Router: model_group_info.supports_url_context = True if ( - model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True # type: ignore + model_info.get("supports_reasoning", None) is not None + and model_info["supports_reasoning"] is True # type: ignore ): model_group_info.supports_reasoning = True if ( @@ -9725,17 +9814,19 @@ class Router: # - if healthy_deployments > 1, return model group rate limit headers # - else return the model's rate limit headers """ - if ( - isinstance(response, BaseModel) - and hasattr(response, "_hidden_params") - and isinstance(response._hidden_params, dict) # type: ignore - ): - response._hidden_params.setdefault("additional_headers", {}) # type: ignore - response._hidden_params["additional_headers"][ # type: ignore - "x-litellm-model-group" - ] = model_group + if response is not None and hasattr(response, "_hidden_params"): + hidden_params = getattr(response, "_hidden_params", {}) or {} + if hasattr(hidden_params, "model_dump"): + hidden_params = hidden_params.model_dump() + if not isinstance(hidden_params, dict): + return response + response._hidden_params = hidden_params - additional_headers = response._hidden_params["additional_headers"] # type: ignore + additional_headers = hidden_params.get("additional_headers") + if not isinstance(additional_headers, dict): + additional_headers = {} + hidden_params["additional_headers"] = additional_headers + additional_headers["x-litellm-model-group"] = model_group # Lift QualityRouter routing decision into response headers for # transparency. The decision is stashed in request_kwargs.metadata diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 885798d706c..74451729c73 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -211,9 +211,7 @@ class BaseRoutingStrategy(ABC): return # 2. Fetch all current provider spend from Redis to update in-memory cache - cache_keys = ( - self.get_in_memory_keys_to_update() - ) # if no pattern OR redis cache does not support scan_iter, use in-memory keys + cache_keys = self.get_in_memory_keys_to_update() # if no pattern OR redis cache does not support scan_iter, use in-memory keys cache_keys_list = list(cache_keys) diff --git a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py index 6c9318e83bb..7c1c8f2907d 100644 --- a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py +++ b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py @@ -304,7 +304,7 @@ def run_eval() -> Tuple[int, int, List[dict]]: # Summary print("=" * 70) - print(f"RESULTS: {passed}/{total} passed ({100*passed/total:.1f}%)") + print(f"RESULTS: {passed}/{total} passed ({100 * passed / total:.1f}%)") print("=" * 70) if failures: diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 3adb8d43920..3f2db97e2bf 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -548,9 +548,9 @@ class LowestLatencyLoggingHandler(CustomLogger): deployment = random_valid_deployment[0] metadata_field = self._select_metadata_field(request_kwargs) if request_kwargs is not None and metadata_field in request_kwargs: - request_kwargs[metadata_field][ - "_latency_per_deployment" - ] = _latency_per_deployment + request_kwargs[metadata_field]["_latency_per_deployment"] = ( + _latency_per_deployment + ) return deployment async def async_get_available_deployments( diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 22664dcb704..09c2084c095 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -106,7 +106,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): model_id, deployment.get("model_name", ""), ), - request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="tpm_rpm_limits", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), ) else: @@ -129,7 +132,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): deployment_rpm, result, ), - request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="tpm_rpm_limits", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), ) return deployment @@ -190,7 +196,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): local_result, ), headers={"retry-after": str(60)}, # type: ignore - request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="tpm_rpm_limits", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), num_retries=deployment.get("num_retries"), ) @@ -214,7 +223,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): result, ), headers={"retry-after": str(60)}, # type: ignore - request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="tpm_rpm_limits", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), num_retries=deployment.get("num_retries"), ) @@ -558,7 +570,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): status_code=429, content="", headers={"retry-after": str(60)}, # type: ignore - request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request( + method="tpm_rpm_limits", + url="https://github.com/BerriAI/litellm", + ), # type: ignore ), ) diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 6b921a0db8a..0b927714ca9 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -1,44 +1,99 @@ -from typing import Any, Optional, Union +import json +from typing import Protocol, TypedDict, cast from pydantic import BaseModel -from litellm.types.utils import HiddenParams + +class FallbackErrorInfo(TypedDict): + message: str + type: str + param: str | None + code: str | None -def _add_headers_to_response(response: Any, headers: dict) -> Any: +class _HiddenParamsHost(Protocol): + _hidden_params: dict[str, object] + + +def get_hidden_params_dict(response: object) -> dict[str, object]: + hidden_params: object = cast(object, getattr(response, "_hidden_params", None)) + if isinstance(hidden_params, BaseModel): + return cast("dict[str, object]", hidden_params.model_dump()) + if isinstance(hidden_params, dict): + return cast("dict[str, object]", hidden_params) + return {} + + +def _ensure_additional_headers_dict( + hidden_params: dict[str, object], +) -> dict[str, object]: + additional_headers = hidden_params.get("additional_headers") + if isinstance(additional_headers, dict): + return cast("dict[str, object]", additional_headers) + return {} + + +def get_fallback_error_info(error: Exception) -> FallbackErrorInfo: + message = cast(object, getattr(error, "message", str(error))) + error_type = cast(object, getattr(error, "type", error.__class__.__name__)) + param = cast(object, getattr(error, "param", None)) + code = cast(object, getattr(error, "status_code", getattr(error, "code", None))) + return FallbackErrorInfo( + message=str(message), + type=str(error_type), + param=str(param) if param is not None else None, + code=str(code) if code is not None else None, + ) + + +def _coerce_error_dicts(items: list[object]) -> list[dict[str, object]]: + return [cast("dict[str, object]", item) for item in items if isinstance(item, dict)] + + +def get_fallback_errors_from_headers( + additional_headers: dict[str, object], +) -> list[dict[str, object]]: + existing_errors = additional_headers.get("x-litellm-fallback-errors") + if isinstance(existing_errors, list): + return _coerce_error_dicts(cast("list[object]", existing_errors)) + if isinstance(existing_errors, str): + try: + parsed_errors: object = cast(object, json.loads(existing_errors)) + except json.JSONDecodeError: + return [] + if isinstance(parsed_errors, list): + return _coerce_error_dicts(cast("list[object]", parsed_errors)) + return [] + + +def _add_headers_to_response(response: object, headers: dict[str, object]) -> object: """ Helper function to add headers to a response's hidden params """ - if response is None or not isinstance(response, BaseModel): + if response is None: return response - hidden_params: Optional[Union[dict, HiddenParams]] = getattr( - response, "_hidden_params", {} - ) + if not isinstance(response, BaseModel) and not hasattr(response, "_hidden_params"): + return response - if hidden_params is None: - hidden_params_dict = {} - elif isinstance(hidden_params, HiddenParams): - hidden_params_dict = hidden_params.model_dump() - else: - hidden_params_dict = hidden_params + hidden_params = get_hidden_params_dict(response) + additional_headers = _ensure_additional_headers_dict(hidden_params) + additional_headers.update(headers) + hidden_params["additional_headers"] = additional_headers - hidden_params_dict.setdefault("additional_headers", {}) - hidden_params_dict["additional_headers"].update(headers) - - setattr(response, "_hidden_params", hidden_params_dict) + cast(_HiddenParamsHost, response)._hidden_params = hidden_params return response def add_retry_headers_to_response( - response: Any, + response: object, attempted_retries: int, - max_retries: Optional[int] = None, -) -> Any: + max_retries: int | None = None, +) -> object: """ Add retry headers to the request """ - retry_headers = { + retry_headers: dict[str, object] = { "x-litellm-attempted-retries": attempted_retries, } if max_retries is not None: @@ -48,9 +103,10 @@ def add_retry_headers_to_response( def add_fallback_headers_to_response( - response: Any, + response: object, attempted_fallbacks: int, -) -> Any: + fallback_errors: list[FallbackErrorInfo] | None = None, +) -> object: """ Add fallback headers to the response @@ -64,7 +120,19 @@ def add_fallback_headers_to_response( Note: It's intentional that we don't add max_fallbacks in response headers Want to avoid bloat in the response headers for performance. """ - fallback_headers = { + fallback_headers: dict[str, object] = { "x-litellm-attempted-fallbacks": attempted_fallbacks, } - return _add_headers_to_response(response, fallback_headers) + response = _add_headers_to_response(response, fallback_headers) + if fallback_errors is None or response is None: + return response + + hidden_params = get_hidden_params_dict(response) + additional_headers = _ensure_additional_headers_dict(hidden_params) + merged_errors = get_fallback_errors_from_headers(additional_headers) + [ + cast("dict[str, object]", error) for error in fallback_errors + ] + additional_headers["x-litellm-fallback-errors"] = json.dumps(merged_errors) + hidden_params["additional_headers"] = additional_headers + cast(_HiddenParamsHost, response)._hidden_params = hidden_params + return response diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 5e58479825b..ddec753d362 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -82,43 +82,83 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File if isinstance(file_content, PathLike): return file_content - # Decode the bytes to a string and split into lines - # If file_content is a file-like object, read the bytes - if hasattr(file_content, "read"): - file_content_bytes = file_content.read() # type: ignore - elif isinstance(file_content, tuple): - file_content_bytes = file_content[1] - else: - file_content_bytes = file_content - - # Decode the bytes to a string and split into lines - if isinstance(file_content_bytes, bytes): - file_content_str = file_content_bytes.decode("utf-8") - elif isinstance(file_content_bytes, str): - file_content_str = file_content_bytes + # Iterate the source line-by-line WITHOUT reading it all into memory. A + # spooled upload handle (managed batches stream from it) is read straight + # off its backing; bytes/str are wrapped so they iterate line-by-line. + source = file_content[1] if isinstance(file_content, tuple) else file_content + if hasattr(source, "read"): + if hasattr(source, "seek"): + try: + source.seek(0) # type: ignore[attr-defined] + except (OSError, ValueError): + pass + line_iter: object = source + elif isinstance(source, (bytes, bytearray)): + line_iter = io.BytesIO(bytes(source)) + elif isinstance(source, str): + line_iter = io.StringIO(source) else: return file_content - # Parse JSONL properly, handling potential multiline JSON objects - json_objects = parse_jsonl_with_embedded_newlines(file_content_str) + # Rewrite one row at a time, writing straight into the output buffer + # instead of holding every parsed row in a list. Peak memory stays at + # ~one row plus the output rather than several full copies of the file, + # which the managed-files path depends on (it re-runs this rewrite once + # per target model). Lines are accumulated so JSON objects that span + # multiple physical lines still parse. Streaming the handle also means + # the model rewrite is actually applied to tuple-wrapped upload handles; + # otherwise a restricted body.model would survive and bypass the batch + # model allowlist (which validates the upload target alias). + output = InMemoryFile( + b"", name="modified_file.jsonl", content_type="application/jsonl" + ) + wrote_any = False + buffer = "" + for raw_line in line_iter: # type: ignore[attr-defined] + buffer += ( + raw_line.decode("utf-8") + if isinstance(raw_line, (bytes, bytearray)) + else raw_line + ) + stripped = buffer.strip() + if not stripped: + buffer = "" + continue + try: + json_object = json.loads(stripped) + except json.JSONDecodeError: + continue # object not complete yet; keep accumulating + if isinstance(json_object, dict) and isinstance( + json_object.get("body"), dict + ): + json_object["body"]["model"] = new_model_name + output.write( + (("\n" if wrote_any else "") + json.dumps(json_object)).encode("utf-8") + ) + wrote_any = True + buffer = "" + + if buffer.strip(): + # A row never parsed (truncated/malformed, or it swallowed the rows + # that followed it). Returning the partial `output` would silently + # drop those rows; return the unchanged original so the provider + # rejects the batch loudly instead of accepting a truncated one. + verbose_logger.error( + f"error parsing trailing batch content: {buffer[:100]}..." + ) + if hasattr(source, "seek"): + try: + source.seek(0) # type: ignore[attr-defined] + except (OSError, ValueError): + pass + return file_content # If no valid JSON objects were found, return the original content - if len(json_objects) == 0: + if not wrote_any: return file_content - modified_lines = [] - for json_object in json_objects: - # Replace the model name if it exists - if "body" in json_object: - json_object["body"]["model"] = new_model_name - - # Convert the modified JSON object back to a string - modified_lines.append(json.dumps(json_object)) - - # Reassemble the modified lines and return as bytes - modified_file_content = "\n".join(modified_lines).encode("utf-8") - - return InMemoryFile(modified_file_content, name="modified_file.jsonl", content_type="application/jsonl") # type: ignore + output.seek(0) + return output # type: ignore except (json.JSONDecodeError, UnicodeDecodeError, TypeError): # return the original file content if there is an error replacing the model name diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index b210ea44596..dcfa44381c1 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -38,6 +38,7 @@ class CooldownCache: visible_prefix=50, # Show first 50 characters visible_suffix=0, # Show last 0 characters mask_char="*", # Use * for masking + mask_short_values=False, # Truncate long messages only; keep short ones readable ) def _common_add_cooldown_logic( diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index eb756e3cf8b..f0edc7fc9db 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -6,6 +6,7 @@ from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, + get_fallback_error_info, ) from litellm.types.router import LiteLLMParamsTypedDict @@ -90,6 +91,7 @@ async def run_async_fallback( original_exception: Exception, max_fallbacks: int, fallback_depth: int, + include_fallback_errors: bool = False, **kwargs, ) -> Any: """ @@ -118,6 +120,7 @@ async def run_async_fallback( raise original_exception error_from_fallbacks = original_exception + fallback_errors = (get_fallback_error_info(original_exception),) for mg in fallback_model_group: if mg == original_model_group: @@ -136,6 +139,8 @@ async def run_async_fallback( fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks + if include_fallback_errors: + kwargs["include_fallback_errors"] = include_fallback_errors response = await litellm_router.async_function_with_fallbacks( *args, **kwargs ) @@ -143,6 +148,9 @@ async def run_async_fallback( response = add_fallback_headers_to_response( response=response, attempted_fallbacks=fallback_depth, + fallback_errors=( + list(fallback_errors) if include_fallback_errors else None + ), ) # callback for successfull_fallback_event(): await log_success_fallback_event( @@ -153,6 +161,7 @@ async def run_async_fallback( return response except Exception as e: error_from_fallbacks = e + fallback_errors = fallback_errors + (get_fallback_error_info(e),) await log_failure_fallback_event( original_model_group=original_model_group, kwargs=kwargs, diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 5fd2be9c6dd..cab2041964f 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -252,9 +252,9 @@ class EncryptedContentAffinityCheck(CustomLogger): # _get_metadata_variable_name_from_kwargs would pick "litellm_metadata" # over "metadata" where tags are actually stored. if "litellm_metadata" in request_kwargs: - request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] = True + request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = ( + True + ) request_input = request_kwargs.get("input") model_id = self._extract_model_id_from_input(request_input) diff --git a/litellm/sandbox/main.py b/litellm/sandbox/main.py index 45d3bffb4f9..76d9994c683 100644 --- a/litellm/sandbox/main.py +++ b/litellm/sandbox/main.py @@ -68,7 +68,7 @@ async def acreate_sandbox( provider: str, template: str | None = None, timeout: int | None = None, - allow_internet_access: bool = True, + allow_internet_access: bool | None = None, api_key: str | None = None, api_base: str | None = None, **kwargs, diff --git a/litellm/sandbox/sandbox_tools.py b/litellm/sandbox/sandbox_tools.py index f4a6678f629..509b316ee21 100644 --- a/litellm/sandbox/sandbox_tools.py +++ b/litellm/sandbox/sandbox_tools.py @@ -40,11 +40,14 @@ def _iter_valid_tools(tools: list[dict]) -> Iterator[tuple[str, dict]]: "sandbox_tools: skipping entry missing 'sandbox_provider': %r", tool ) continue - yield name, { - "sandbox_provider": provider, - "api_key": _resolve_secret_value(params.get("api_key")), - "api_base": _resolve_secret_value(params.get("api_base")), - } + yield ( + name, + { + "sandbox_provider": provider, + "api_key": _resolve_secret_value(params.get("api_key")), + "api_base": _resolve_secret_value(params.get("api_base")), + }, + ) def register_sandbox_tools(tools: list[dict]) -> None: diff --git a/litellm/types/completion.py b/litellm/types/completion.py index cb263914be8..a91f6234fad 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -1,8 +1,28 @@ -from typing import Iterable, List, Optional, Union +from __future__ import annotations + +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Coroutine, + Iterable, + List, + Optional, + Union, +) from pydantic import BaseModel, ConfigDict from typing_extensions import Literal, Required, TypedDict +if TYPE_CHECKING: + import httpx + from aiohttp import ClientSession + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm import BaseConfig + from litellm.utils import CustomStreamWrapper, ModelResponse + class ChatCompletionSystemMessageParam(TypedDict, total=False): content: Required[str] @@ -191,3 +211,44 @@ class CompletionRequest(BaseModel): model_list: Optional[List[str]] = None model_config = ConfigDict(protected_namespaces=(), extra="allow") + + +@dataclass(frozen=True, slots=True) +class _CompletionDispatchContext: + _azure_detection_model: str + acompletion: bool + api_base: Optional[str] + api_key: Optional[str] + api_version: Optional[str] + client: Any + custom_llm_provider: str + custom_prompt_dict: dict + extra_headers: Optional[dict] + headers: dict + hf_model_name: Optional[str] + kwargs: dict + litellm_params: dict + logger_fn: Optional[Callable] + logging: LiteLLMLoggingObj + max_retries: Optional[int] + max_tokens: Optional[int] + messages: list + metadata: Optional[dict] + model: str + model_response: ModelResponse + optional_params: dict + organization: Optional[str] + provider_config: Optional[BaseConfig] + shared_session: Optional[ClientSession] + stream: Optional[bool] + temperature: Optional[float] + text_completion: bool + timeout: Optional[Union[float, str, httpx.Timeout]] + top_p: Optional[float] + + +_CompletionDispatchResult = Union[ + Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], + "ModelResponse", + "CustomStreamWrapper", +] diff --git a/litellm/types/files.py b/litellm/types/files.py index bf56894329c..1b2d7e30f1f 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -321,3 +321,21 @@ class TwoStepFileUploadConfig(TypedDict, total=False): upload_request: Required[TwoStepFileUploadRequest] upload_url_location: Required[Literal["headers", "body"]] upload_url_key: str + + +class ResumableChunkedUploadConfig(TypedDict, total=False): + """Drives a memory-bounded resumable upload (GCS JSON API). + + The handler POSTs to the upload URL to open a session, reads the session URI + from ``session_url_header``, then PUTs ``body_stream`` to that URI in + ``chunk_size``-byte chunks (a 256 KiB multiple) using Content-Range, so the + payload is never buffered in full and the transfer is resumable. + + ``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to + avoid importing the llms layer into types. + """ + + body_stream: Required[Any] + chunk_size: int + session_url_header: str + initiate_headers: Dict[str, str] diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index b2e1fb3d46b..5c76736ab90 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -23,7 +23,9 @@ if TYPE_CHECKING: generationConfig: Optional[Any] tools: Optional[ToolConfigDict] # type: ignore[assignment, valid-type] - class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc, valid-type] + class GenerateContentResponse( + GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject + ): # type: ignore[misc, valid-type] _hidden_params: dict = {} pass diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index b5726a11ca0..26a0be36ef4 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -2,6 +2,25 @@ from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field +CHAT_COMPLETION_AGENTIC_SURFACE = "chat_completions" +CODE_INTERPRETER_INTERCEPTION_PREFIX = "_code_interpreter_interception" +NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES = frozenset( + ("_websearch_interception", "_compression_interception") +) +INTERCEPTION_INTERNAL_PREFIXES = frozenset( + ( + *NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + CODE_INTERPRETER_INTERCEPTION_PREFIX, + ) +) + + +def is_interception_internal_key( + key: str, + prefixes: frozenset[str] = INTERCEPTION_INTERNAL_PREFIXES, +) -> bool: + return any(key.startswith(prefix) for prefix in prefixes) + class StandardCustomLoggerInitParams(BaseModel): """ diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 5b1d32cd93c..d203fe73d02 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -698,9 +698,9 @@ class PrometheusMetricLabels: litellm_managed_batch_created_total = _batch_user_labels - litellm_managed_file_size_bytes: List[str] = ( - [] - ) # labels: purpose, file_type, model, api_provider, user (custom) + litellm_managed_file_size_bytes: List[ + str + ] = [] # labels: purpose, file_type, model, api_provider, user (custom) litellm_managed_batch_duration_seconds = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -709,9 +709,9 @@ class PrometheusMetricLabels: litellm_managed_file_created_total = _batch_user_labels - litellm_managed_file_deleted_total: List[str] = ( - [] - ) # only "result" label, added at metric creation + litellm_managed_file_deleted_total: List[ + str + ] = [] # only "result" label, added at metric creation litellm_check_batch_cost_jobs_polled: List[str] = [] diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index b38cd8f58b9..a07642073af 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -954,9 +954,6 @@ class Interaction(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - role: Optional[str] = Field( - None, description="Output only. The role of the interaction." - ) outputs: Optional[List[Content]] = Field( None, description="Output only. Responses from the model." ) @@ -1031,9 +1028,6 @@ class CreateModelInteractionParams(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - role: Optional[str] = Field( - None, description="Output only. The role of the interaction." - ) outputs: Optional[List[Content]] = Field( None, description="Output only. Responses from the model." ) @@ -1101,9 +1095,6 @@ class CreateAgentInteractionParams(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - role: Optional[str] = Field( - None, description="Output only. The role of the interaction." - ) outputs: Optional[List[Content]] = Field( None, description="Output only. Responses from the model." ) @@ -1323,7 +1314,6 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): status: Optional[str] = None created: Optional[str] = None updated: Optional[str] = None - role: Optional[str] = None # Legacy schema field (Api-Revision: 2026-05-07). Remove after June 8, 2026. outputs: Optional[List[Dict[str, Any]]] = None # New schema field (Api-Revision: 2026-05-20). @@ -1356,7 +1346,6 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): status: Optional[str] = None created: Optional[str] = None updated: Optional[str] = None - role: Optional[str] = None # Legacy schema field (Api-Revision: 2026-05-07). Remove after June 8, 2026. outputs: Optional[List[Dict[str, Any]]] = None # New schema field (Api-Revision: 2026-05-20). diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index cbb316eec75..36b6bac90ea 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1100,7 +1100,7 @@ OpenAIImageGenerationOptionalParams = Literal[ OpenAIImageEditOptionalParams = Literal[ "background", "n", - "mask" "output_compression", + "maskoutput_compression", "output_format", "quality", "partial_images", diff --git a/litellm/types/proxy/callback_logs_endpoints.py b/litellm/types/proxy/callback_logs_endpoints.py new file mode 100644 index 00000000000..ef148274ca7 --- /dev/null +++ b/litellm/types/proxy/callback_logs_endpoints.py @@ -0,0 +1,44 @@ +""" +Types for the callback-logs ingest endpoint (POST /v1/callbacks/logs). + +External producers (e.g. the litellm-rust gateway) POST finished logging +payloads here; the proxy replays them through the standard callback fan-out. +""" + +from typing import Any, Literal, Optional + +from pydantic import BaseModel, Field + +from litellm.constants import MAX_CALLBACK_LOG_RECORDS + + +class CallbackLogRecord(BaseModel): + """A single finished logging event to replay through the callbacks.""" + + status: Literal["success", "failure"] + standard_logging_payload: dict[str, Any] + error: Optional[str] = None + + +class CallbackLogsRequest(BaseModel): + """A batch of logging events posted by an external producer.""" + + # Bounded so one POST can't trigger an unbounded callback/DB fan-out (each + # record fires every registered integration). Over the cap → 422. + records: list[CallbackLogRecord] = Field(..., max_length=MAX_CALLBACK_LOG_RECORDS) + + +class CallbackLogFailure(BaseModel): + """A record that failed to replay, identified by its index in the batch.""" + + index: int + error: str + + +class CallbackLogsResponse(BaseModel): + """Per-batch result: counts plus per-record failure detail so the caller can + distinguish a transient callback error from a structurally bad payload.""" + + processed: int + failed: int + failures: list[CallbackLogFailure] = Field(default_factory=list) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py b/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py index af199eed55e..697e794a7f1 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/xecguard.py @@ -33,9 +33,7 @@ class XecGuardConfigModel(GuardrailConfigModel): ) xecguard_model: Optional[str] = Field( default=None, - description=( - "XecGuard scanning model identifier. " "Defaults to 'xecguard_v2'." - ), + description=("XecGuard scanning model identifier. Defaults to 'xecguard_v2'."), ) policy_names: Optional[List[str]] = Field( default=None, diff --git a/litellm/types/router.py b/litellm/types/router.py index 607bfd584fd..a1c571ed7f7 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -110,9 +110,7 @@ class ModelInfo(BaseModel): id: Optional[ str ] # Allow id to be optional on input, but it will always be present as a str in the model instance - db_model: bool = ( - False # used for proxy - to separate models which are stored in the db vs. config. - ) + db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: Optional[datetime.datetime] = None updated_by: Optional[str] = None @@ -180,6 +178,9 @@ class CredentialLiteLLMParams(BaseModel): ## UNIFIED PROJECT/REGION ## region_name: Optional[str] = None + ## OBJECT STORAGE (files / batches) ## + gcs_bucket_name: Optional[str] = None + ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None @@ -436,9 +437,7 @@ class Deployment(BaseModel): elif isinstance(model_info, dict): model_info = ModelInfo(**model_info) - for ( - key - ) in ( + for key in ( SPECIAL_MODEL_INFO_PARAMS ): # ensures custom pricing info is consistently in 'model_info' field = getattr(litellm_params, key, None) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 16e693c7d7c..24d6e84fba7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3163,8 +3163,26 @@ class CustomPricingLiteLLMParams(BaseModel): regional_processing_uplift_multiplier_us: Optional[float] = None +# Server-controlled fields that bound or drive an interceptor's agentic loop +# (depth, cycle fingerprints, ceiling, code-interpreter sandbox state). Listed +# in all_litellm_params so they are treated as LiteLLM-level and excluded from +# get_non_default_completion_params; otherwise the OpenAI param builder sweeps +# any unrecognized top-level key into extra_body and leaks them to the provider. +# This is what lets the loop carry state across rerun calls without a provider +# scrubber. +agentic_loop_internal_litellm_params = [ + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_agentic_loop_api_surface", + "max_agentic_loops", + "_code_interpreter_interception_active", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_converted_stream", +] + all_litellm_params = ( - [ + agentic_loop_internal_litellm_params + + [ "metadata", "litellm_metadata", "litellm_trace_id", @@ -3463,6 +3481,7 @@ class LlmProviders(str, Enum): TENSORMESH = "tensormesh" LIBERTAI = "libertai" PINSTRIPES = "pinstripes" + DARKBLOOM = "darkbloom" LITELLM_AGENT = "litellm_agent" CURSOR = "cursor" BEDROCK_MANTLE = "bedrock_mantle" @@ -3520,6 +3539,7 @@ class SandboxProviders(str, Enum): """ E2B = "e2b" + OPENSANDBOX = "opensandbox" class LiteLLMLoggingBaseClass: diff --git a/litellm/utils.py b/litellm/utils.py index 29f703104da..e0aa8575473 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -799,7 +799,9 @@ def function_setup( # check if callback is a string - e.g. "lago", "openmeter" if isinstance(callback, str): callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore - callback, internal_usage_cache=None, llm_router=None # type: ignore + callback, + internal_usage_cache=None, + llm_router=None, # type: ignore ) if callback is None or any( isinstance(cb, type(callback)) @@ -809,13 +811,21 @@ def function_setup( if callback not in litellm.input_callback: litellm.input_callback.append(callback) # type: ignore if callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback(callback) # type: ignore + litellm.logging_callback_manager.add_litellm_success_callback( + callback + ) # type: ignore if callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback(callback) # type: ignore + litellm.logging_callback_manager.add_litellm_failure_callback( + callback + ) # type: ignore if callback not in litellm._async_success_callback: - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) # type: ignore + litellm.logging_callback_manager.add_litellm_async_success_callback( + callback + ) # type: ignore if callback not in litellm._async_failure_callback: - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) # type: ignore + litellm.logging_callback_manager.add_litellm_async_failure_callback( + callback + ) # type: ignore print_verbose( f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}" ) @@ -1486,9 +1496,9 @@ def client(original_function): ) # Type assertion: logging_obj is guaranteed to be non-None after function_setup - assert ( - logging_obj is not None - ), "logging_obj should not be None after function_setup" + assert logging_obj is not None, ( + "logging_obj should not be None after function_setup" + ) ## LOAD CREDENTIALS load_credentials_from_list(kwargs) @@ -1810,9 +1820,9 @@ def client(original_function): ) # Type assertion: logging_obj is guaranteed to be non-None after function_setup - assert ( - logging_obj is not None - ), "logging_obj should not be None after function_setup" + assert logging_obj is not None, ( + "logging_obj should not be None after function_setup" + ) modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type) if modified_kwargs is not None: @@ -2311,7 +2321,9 @@ def create_pretrained_tokenizer( try: tokenizer = Tokenizer.from_pretrained( - identifier, revision=revision, auth_token=auth_token # type: ignore + identifier, + revision=revision, + auth_token=auth_token, # type: ignore ) except Exception as e: verbose_logger.error( @@ -3163,8 +3175,9 @@ def get_optional_params_transcription( keys = list(non_default_params.keys()) for k in keys: if ( - drop_params is True or litellm.drop_params is True - ) and k not in supported_params: # drop the unsupported non-default values + (drop_params is True or litellm.drop_params is True) + and k not in supported_params + ): # drop the unsupported non-default values non_default_params.pop(k, None) elif k not in supported_params: raise UnsupportedParamsError( @@ -3191,7 +3204,7 @@ def get_optional_params_transcription( model=model, drop_params=drop_params if drop_params is not None else False, ) - elif provider_config is not None: # handles fireworks ai, and any future providers + elif provider_config is not None: # custom audio transcription config supported_params = provider_config.get_supported_openai_params(model=model) _check_valid_arg(supported_params=supported_params) optional_params = provider_config.map_openai_params( @@ -3294,8 +3307,9 @@ def get_optional_params_image_gen( keys = list(non_default_params.keys()) for k in keys: if ( - litellm.drop_params is True or drop_params is True - ) and k not in supported_params: # drop the unsupported non-default values + (litellm.drop_params is True or drop_params is True) + and k not in supported_params + ): # drop the unsupported non-default values non_default_params.pop(k, None) passed_params.pop(k, None) elif k not in supported_params: @@ -3971,11 +3985,7 @@ def pre_process_non_default_params( non_default_params, list ): # fixes https://github.com/BerriAI/litellm/issues/4933 tools = non_default_params["tools"] - for ( - tool - ) in ( - tools - ): # clean out 'additionalProperties = False'. Causes vertexai/gemini OpenAI API Schema errors - https://github.com/langchain-ai/langchainjs/issues/5240 + for tool in tools: # clean out 'additionalProperties = False'. Causes vertexai/gemini OpenAI API Schema errors - https://github.com/langchain-ai/langchainjs/issues/5240 tool_function = tool.get("function", {}) parameters = tool_function.get("parameters", None) if parameters is not None: @@ -5740,8 +5750,8 @@ def _get_potential_model_names( model=model, custom_llm_provider=custom_llm_provider ) combined_stripped_model_name = stripped_model_name - elif custom_llm_provider and model.startswith( - custom_llm_provider + "/" + elif ( + custom_llm_provider and model.startswith(custom_llm_provider + "/") ): # handle case where custom_llm_provider is provided and model starts with custom_llm_provider split_model = model.split("/", 1)[1] combined_model_name = model @@ -7938,7 +7948,9 @@ class ModelResponseIterator: def __init__(self, model_response: ModelResponse, convert_to_delta: bool = False): if convert_to_delta is True: _stream_response = ModelResponseStream() - _stream_response.choices[0].delta.content = model_response.choices[0].message.content # type: ignore + _stream_response.choices[0].delta.content = model_response.choices[ + 0 + ].message.content # type: ignore self.model_response: Union[ModelResponse, ModelResponseStream] = ( _stream_response ) @@ -8915,8 +8927,6 @@ class ProviderConfigManager: ) return AzureSpeechAudioTranscriptionConfig() - if litellm.LlmProviders.FIREWORKS_AI == provider: - return litellm.FireworksAIAudioTranscriptionConfig() elif litellm.LlmProviders.DEEPGRAM == provider: return litellm.DeepgramAudioTranscriptionConfig() elif litellm.LlmProviders.ELEVENLABS == provider: @@ -9733,9 +9743,14 @@ class ProviderConfigManager: Get sandbox (code execution) configuration for a given provider. """ from litellm.llms.e2b.sandbox.transformation import E2BSandboxConfig + from litellm.llms.opensandbox.sandbox.transformation import ( + OpenSandboxSandboxConfig, + ) if provider == SandboxProviders.E2B: return E2BSandboxConfig() + if provider == SandboxProviders.OPENSANDBOX: + return OpenSandboxSandboxConfig() return None @staticmethod diff --git a/migrations/Dockerfile b/migrations/Dockerfile index a78a4e2225a..caca280cbfc 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3c50dde9277..5f3f2294147 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -571,7 +571,7 @@ "output_vector_size": 1536 }, "amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 2e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "max_tokens": 8192, @@ -10443,7 +10443,8 @@ "fast": 6.0 }, "supports_output_config": true, - "supports_max_reasoning_effort": true + "supports_max_reasoning_effort": true, + "supports_speed": true }, "claude-opus-4-6-20260205": { "cache_creation_input_token_cost": 6.25e-06, @@ -10476,7 +10477,8 @@ "fast": 6.0 }, "supports_max_reasoning_effort": true, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -10511,7 +10513,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-opus-4-7-20260416": { "cache_creation_input_token_cost": 6.25e-06, @@ -10546,7 +10549,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, @@ -10615,7 +10619,8 @@ "us": 1.1, "fast": 2.0 }, - "supports_output_config": true + "supports_output_config": true, + "supports_speed": true }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-05-14", @@ -10684,6 +10689,268 @@ "mode": "chat", "output_cost_per_token": 1.923e-06 }, + "cloudflare/@cf/openai/gpt-oss-120b": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-2b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.2-3b-instruct": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 3.35e-07 + }, + "cloudflare/@cf/meta/llama-guard-3-8b": { + "input_cost_per_token": 4.84e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3e-08 + }, + "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 15000, + "max_output_tokens": 15000, + "max_tokens": 15000, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/deepseek-ai/deepseek-r1-distill-qwen-32b": { + "input_cost_per_token": 4.97e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 4.881e-06, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { + "input_cost_per_token": 1.52e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "output_cost_per_token": 2.87e-07 + }, + "cloudflare/@cf/meta/llama-3.2-1b-instruct": { + "input_cost_per_token": 2.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 60000, + "max_output_tokens": 60000, + "max_tokens": 60000, + "mode": "chat", + "output_cost_per_token": 2.01e-07 + }, + "cloudflare/@cf/moonshotai/kimi-k2.6": { + "cache_read_input_token_cost": 1.6e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/zai-org/glm-4.7-flash": { + "input_cost_per_token": 6.05e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta-llama/llama-2-7b-chat-hf-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { + "input_cost_per_token": 2.93e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 2.253e-06, + "supports_function_calling": true + }, + "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { + "input_cost_per_token": 1.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.12e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-06 + }, + "cloudflare/@cf/zai-org/glm-5.2": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/nvidia/nemotron-3-120b-a12b": { + "input_cost_per_token": 5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/aisingapore/gemma-sea-lion-v4-27b-it": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07 + }, + "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.35e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-7b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 3500, + "max_output_tokens": 3500, + "max_tokens": 3500, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/mistralai/mistral-small-3.1-24b-instruct": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { + "input_cost_per_token": 4.85e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.76e-07, + "supports_vision": true + }, + "cloudflare/@cf/openai/gpt-oss-20b": { + "input_cost_per_token": 2e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-4-scout-17b-16e-instruct": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 8.5e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwq-32b": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 1e-06, + "supports_reasoning": true + }, "codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "codestral", @@ -20096,8 +20363,6 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20171,8 +20436,6 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20246,8 +20509,6 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20319,8 +20580,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20362,8 +20621,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20385,8 +20642,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20675,8 +20930,6 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_priority": 1e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -21380,8 +21633,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21775,6 +22026,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21823,6 +22076,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21867,6 +22122,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21911,6 +22168,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21959,6 +22218,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22006,6 +22267,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22046,6 +22309,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22089,6 +22354,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22134,6 +22401,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22180,6 +22449,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22223,6 +22494,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22266,6 +22539,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22304,8 +22579,6 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -22712,8 +22985,6 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22795,8 +23066,6 @@ "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, @@ -39946,24 +40215,6 @@ "litellm_provider": "fireworks_ai", "mode": "chat" }, - "fireworks_ai/accounts/fireworks/models/whisper-v3": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, - "fireworks_ai/accounts/fireworks/models/whisper-v3-turbo": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, "fireworks_ai/accounts/fireworks/models/yi-34b": { "max_tokens": 4096, "max_input_tokens": 4096, @@ -43543,5 +43794,39 @@ "supports_assistant_prefill": true, "supports_reasoning": false, "source": "https://pinstripes.io/pricing" + }, + "darkbloom/gemma-4-26b": { + "input_cost_per_token": 3e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.65e-07, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "darkbloom/gpt-oss-20b": { + "input_cost_per_token": 1.45e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7e-08, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true } } diff --git a/osv-scanner.toml b/osv-scanner.toml index f0f5f045f1a..7ab450945f5 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -2,13 +2,3 @@ id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 reason = "diskcache has no fixed release published; remove this entry once one exists" - -[[IgnoredVulns]] -id = "GHSA-hg6j-4rv6-33pg" -ignoreUntil = 2026-08-15 -reason = "aiohttp held at 3.13.5: vcrpy releases <= 8.1.1 cannot import aiohttp >= 3.14 and the merged upstream fix (vcrpy PR 996) is unreleased; bump aiohttp and drop this entry when a newer vcrpy ships" - -[[IgnoredVulns]] -id = "GHSA-jg22-mg44-37j8" -ignoreUntil = 2026-08-15 -reason = "aiohttp held at 3.13.5: vcrpy releases <= 8.1.1 cannot import aiohttp >= 3.14 and the merged upstream fix (vcrpy PR 996) is unreleased; bump aiohttp and drop this entry when a newer vcrpy ships" diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 05d8aaf3887..847dfdde516 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1833,6 +1833,23 @@ "text_completion": true } }, + "opensandbox": { + "display_name": "OpenSandbox (`opensandbox`)", + "url": "https://open-sandbox.ai/api/", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "sandbox": true + } + }, "openai_like": { "display_name": "OpenAI-like (`openai_like`)", "url": "https://docs.litellm.ai/docs/providers/openai_compatible", @@ -2008,6 +2025,23 @@ "interactions": true } }, + "darkbloom": { + "display_name": "Darkbloom (`darkbloom`)", + "url": "https://docs.litellm.ai/docs/providers/darkbloom", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "predibase": { "display_name": "Predibase (`predibase`)", "url": "https://docs.litellm.ai/docs/providers/predibase", diff --git a/pyproject.toml b/pyproject.toml index b728bfcd514..ae612ae2184 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.90.0" +version = "1.91.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -55,7 +55,7 @@ proxy = [ "fastapi-sso>=0.19.0,<1.0", "PyJWT>=2.13.0,<3.0", "python-multipart>=0.0.27,<1.0", - "cryptography>=46.0.7,<47.0", + "cryptography>=48.0.1,<49.0", "pynacl>=1.6.2,<2.0", "websockets>=15.0.1,<16.0", "boto3>=1.43.1,<2.0", @@ -148,7 +148,6 @@ litellm-proxy = "litellm.proxy.client.cli:cli" dev = [ "diff-cover==9.7.2", "flake8==7.3.0", - "black==26.3.1", "basedpyright==1.39.7", "pytest==9.0.3", "pytest-mock==3.15.1", @@ -182,7 +181,7 @@ dev = [ "parameterized==0.9.0", "openapi-core==0.22.0; python_version < '3.14'", "pytest-timeout==2.4.0", - "vcrpy==8.1.1", + "vcrpy==8.2.1", "pytest-recording==0.13.4", ] proxy-dev = [ @@ -208,7 +207,7 @@ ci = [ "pytest-codspeed==4.3.0", "pytest-retry==1.7.0", "pyarrow==23.0.1", - "langchain==1.2.10", + "langchain==1.3.9", "lunary==1.4.36; python_version == '3.10'", "lunary==1.4.37; python_version >= '3.11'", "logfire==4.6.0", @@ -225,11 +224,8 @@ ci = [ "pylint==4.0.5", "langchain-mcp-adapters==0.2.1", "langchain-openai==1.1.14", - "langgraph==1.0.10", - # langgraph-prebuilt 1.0.9 imports ExecutionInfo/ServerInfo from - # langgraph.runtime, which is not exported until langgraph 1.1.0. - # Pin to 1.0.8 so it pairs correctly with langgraph==1.0.10. - "langgraph-prebuilt==1.0.8", + "langgraph>=1.2.4,<1.3.0", + "langgraph-prebuilt>=1.1.0,<1.3.0", "claude-agent-sdk==0.1.44", ] healthcheck = [ @@ -244,7 +240,7 @@ build-backend = "uv_build" [tool.uv] constraint-dependencies = [ "tornado>=6.5.6", - "aiohttp>=3.13.5,<3.14", + "aiohttp>=3.14.1,<4.0", ] default-groups = ["dev"] required-version = ">=0.10.9" @@ -273,7 +269,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.90.0" +version = "1.91.0" version_files = [ "pyproject.toml:^version", ] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 62ebdb559fc..ae46f020de1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -300,7 +300,7 @@ "slack": 3 }, "RET504": { - "baseline": 709, + "baseline": 702, "slack": 20 }, "RUF010": { diff --git a/ruff.toml b/ruff.toml index 2db4122a30e..082a8f83a0c 100644 --- a/ruff.toml +++ b/ruff.toml @@ -11,7 +11,16 @@ lint.external = [ "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] line-length = 120 -exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_config_yaml/*", "tests/*"] + +# `ruff format` (replacing Black) must wrap at 88, the width Black used and the whole +# history is formatted to. The global line-length stays 120 because E501 and the import +# sorter (I001, strict gate) are tuned to it and ruff has no per-formatter line-length, +# so 88 is passed at the `ruff format --line-length 88` call sites (Makefile + CI). +format.exclude = ["**/enterprise/**"] + +# Was the top-level `exclude`. Scoped to lint so `ruff format` still formats these paths +# (Black did) while `ruff check` keeps skipping them. +lint.exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_config_yaml/*", "tests/*"] [lint.per-file-ignores] diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index 0f9a44703f9..2ef332d91ea 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -1,21 +1,22 @@ #!/usr/bin/env python3 -"""Per-rule count gate for basedpyright. +"""Delta-vs-base per-rule gate for basedpyright. basedpyright's ``--outputjson`` is reduced to a count of errors per *rule* (``reportAny``, ``reportArgumentType``, ...) and checked against a committed budget of the form ``{rule: {baseline, slack}}``, the same shape as -``ruff-strict-budget.json``. A rule fails when its codebase-wide total exceeds -``baseline + slack``. Counts ignore file, line, and column, so a violation -moving anywhere in the tree is invisible; only the per-rule total moves the -needle. +``ruff-strict-budget.json``. A rule fails only when its codebase-wide total is +both over its ceiling (``baseline + slack``) *and* higher than the count on the +base it merges into, so a change is blamed for the errors it adds, never for +drift that already sits in the base. That ``> base`` guard is what stops an +unrelated PR from inheriting a red once two PRs each land near the ceiling and +their sum crosses it: the bystander's count equals its base, so it is spared, +while any PR that actually grows the rule past the cap still fails. -Unlike ``ruff_strict_gate.py`` this does *not* re-run the tool on the merge base -to compute a delta: a second basedpyright pass is minutes and gigabytes, whereas -ruff is milliseconds. The committed budget is the baseline instead -- exactly -how the previous per-file gate worked -- so keep it fresh with ``--update`` -(ratchet), which re-captures every rule's count from the current tree while -preserving each rule's slack. Tool output is read from stdin, so the caller -decides how to invoke basedpyright (and from which cwd). +Head counts are read from stdin (the caller runs basedpyright once and pipes +``--outputjson`` in); the base count is a second basedpyright pass over a +detached worktree at the merge-base, run under the same environment so import +resolution matches. ``--update`` re-captures the absolute per-rule baselines for +the ratchet, preserving each rule's slack. ``--outputjson`` is used rather than text diagnostics because the latter wrap across lines, leaving the ``(reportRule)`` on a continuation line away from the @@ -24,13 +25,21 @@ carries an unambiguous ``rule`` field. """ import argparse +import contextlib import json +import shutil +import subprocess import sys +import tempfile from collections import Counter +from collections.abc import Iterator, Mapping from pathlib import Path -from typing import Mapping, NamedTuple +from typing import NamedTuple REPO_ROOT = Path(__file__).resolve().parent.parent +BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json" +PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json" +DEFAULT_BASE = "origin/litellm_internal_staging" # Bucket for a basedpyright diagnostic with no `rule`. Counted so it's gated. UNCODED = "" @@ -45,6 +54,7 @@ class Breach(NamedTuple): code: str total: int cap: int + added: int def _seed_slack(baseline: int) -> int: @@ -54,18 +64,19 @@ def _seed_slack(baseline: int) -> int: return 10 if baseline >= 50 else 3 -def _to_repo_relative(raw: str) -> str | None: +def _to_relative(raw: str, root: Path) -> str | None: path = Path(raw) - absolute = path if path.is_absolute() else Path.cwd() / path + absolute = path if path.is_absolute() else root / path try: - return absolute.resolve().relative_to(REPO_ROOT).as_posix() + return absolute.resolve().relative_to(root).as_posix() except ValueError: return None -def count_basedpyright(payload: str) -> dict[str, int]: - """Count in-repo basedpyright errors per rule from `--outputjson`. Warnings - and information are ignored; only `severity == "error"` is gated.""" +def count_basedpyright(payload: str, root: Path = REPO_ROOT) -> dict[str, int]: + """Count in-tree basedpyright errors per rule from `--outputjson`. Warnings + and information are ignored; only `severity == "error"` is gated. Files + outside `root` (the venv's site-packages, say) are dropped.""" try: data = json.loads(payload or "{}") except json.JSONDecodeError as exc: @@ -79,21 +90,62 @@ def count_basedpyright(payload: str) -> dict[str, int]: for diag in data.get("generalDiagnostics", []): if diag.get("severity") != "error": continue - if _to_repo_relative(diag.get("file", "")) is None: + if _to_relative(diag.get("file", ""), root) is None: continue counts[diag.get("rule") or UNCODED] += 1 return dict(counts) +def _run(cmd: list[str], cwd: Path = REPO_ROOT) -> str: + proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) + if proc.returncode not in (0, 1): + sys.stderr.write(proc.stderr) + raise SystemExit(f"{cmd[0]} exited {proc.returncode}") + return proc.stdout + + +@contextlib.contextmanager +def _temp_worktree(ref: str) -> Iterator[Path]: + parent = Path(tempfile.mkdtemp(prefix="bpr_base_")) + worktree = parent / "wt" + try: + _run(["git", "worktree", "add", "--detach", str(worktree), ref]) + yield worktree + finally: + subprocess.run( + ["git", "worktree", "remove", "--force", str(worktree)], + cwd=REPO_ROOT, + capture_output=True, + text=True, + ) + shutil.rmtree(parent, ignore_errors=True) + + +def base_counts(ref: str) -> dict[str, int]: + """basedpyright error counts per rule for the merge-base tree. The head + config is copied in so the base is judged by today's rules, and the run uses + the head environment's basedpyright (on PATH) so imports resolve the same.""" + exe = shutil.which("basedpyright") or "basedpyright" + with _temp_worktree(ref) as worktree: + shutil.copy(PYRIGHT_CONFIG, worktree / "pyrightconfig.json") + proc = subprocess.run( + [exe, "--outputjson"], cwd=worktree, capture_output=True, text=True + ) + return count_basedpyright(proc.stdout, root=worktree) + + def evaluate( - counts: Mapping[str, int], budget: Mapping[str, Mapping[str, int]] + head: Mapping[str, int], + base: Mapping[str, int], + budget: Mapping[str, Mapping[str, int]], ) -> list[Breach]: breaches = [] - for code, total in counts.items(): + for code, total in head.items(): spec = budget.get(code) cap = spec["baseline"] + spec["slack"] if spec else DEFAULT_SLACK - if total > cap: - breaches.append(Breach(code, total, cap)) + prior = base.get(code, 0) + if total > cap and total > prior: + breaches.append(Breach(code, total, cap, total - prior)) return sorted(breaches) @@ -107,9 +159,6 @@ def is_vacuous_run( return not counts and any(spec["baseline"] for spec in budget.values()) -BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json" - - def cmd_update(counts: Mapping[str, int]) -> None: existing = json.loads(BUDGET_PATH.read_text()) if BUDGET_PATH.exists() else {} budget = { @@ -127,9 +176,10 @@ def cmd_update(counts: Mapping[str, int]) -> None: ) -def cmd_check(counts: Mapping[str, int]) -> None: +def cmd_check(base_ref: str) -> None: budget = json.loads(BUDGET_PATH.read_text()) - if is_vacuous_run(counts, budget): + head = count_basedpyright(sys.stdin.read()) + if is_vacuous_run(head, budget): expected = sum(spec["baseline"] for spec in budget.values()) print( f"FAIL: basedpyright produced no errors, but {BUDGET_PATH.name} expects " @@ -137,27 +187,44 @@ def cmd_check(counts: Mapping[str, int]) -> None: f"nothing; refusing to certify a vacuous run." ) raise SystemExit(1) - breaches = evaluate(counts, budget) + base_point = _run(["git", "merge-base", base_ref, "HEAD"]).strip() or base_ref + base = base_counts(base_point) + if is_vacuous_run(base, budget): + print( + f"FAIL: basedpyright produced no errors for the base tree at " + f"{base_point[:12]}, so every rule would look freshly added. The base " + f"pass almost certainly crashed; refusing to blame this change for it." + ) + raise SystemExit(1) + breaches = evaluate(head, base, budget) if not breaches: print( - f"OK: every rule is within its basedpyright ceiling ({sum(counts.values())} errors total)" + f"OK: every rule is within its basedpyright ceiling or no higher than base ({sum(head.values())} errors total)" ) return print("FAIL: basedpyright errors exceed the per-rule ceiling:") for breach in breaches: - print(f" {breach.code}: {breach.total} errors over cap {breach.cap}") + print( + f" {breach.code}: total {breach.total} over cap {breach.cap} (this change added {breach.added})" + ) print( - "Resolve the new errors, or run 'make lint-basedpyright-budget-update' if the ceiling should move." + "Reduce the new errors or remove an equal number elsewhere; the ceiling is " + "baseline + slack in basedpyright-code-budget.json." ) + summary = "; ".join(f"{b.code} {b.total}/{b.cap} (+{b.added})" for b in breaches) + print(f"BREACHED RULES: {summary}") raise SystemExit(1) def main() -> None: parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", default=DEFAULT_BASE) parser.add_argument("--update", action="store_true") args = parser.parse_args() - counts = count_basedpyright(sys.stdin.read()) - cmd_update(counts) if args.update else cmd_check(counts) + if args.update: + cmd_update(count_basedpyright(sys.stdin.read())) + else: + cmd_check(args.base) if __name__ == "__main__": diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index cb2ca385ffc..3dc1d116e8d 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -159,12 +159,12 @@ def test_batch_cost_calculator_applies_data_residency_uplift( base_prompt, base_completion = batch_cost_calculator( usage=usage, - model="gpt-5", + model="gpt-5.4", custom_llm_provider="openai", ) regional_prompt, regional_completion = batch_cost_calculator( usage=usage, - model="gpt-5", + model="gpt-5.4", custom_llm_provider="openai", data_residency=data_residency, ) diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index bccb5eaaacb..8a2d5f33805 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -27,7 +27,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import StandardLoggingPayload import socket import httpx -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock def _can_resolve_openai(): @@ -513,10 +513,26 @@ async def test_avertex_batch_prediction(monkeypatch): mock_response.status_code = 200 return mock_response - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - side_effect=mock_side_effect, - ) as mock_global_post: + # Batch jsonl file creation now streams to a GCS resumable session via + # _aresumable_chunked_upload (httpx send), not AsyncHTTPHandler.post, so mock + # that entry point to return the GCS object response. The resumable protocol + # itself is covered in test_vertex_ai_files_streaming.py. + mock_upload_response = httpx.Response( + 200, + json=mock_file_response, + request=httpx.Request("PUT", "https://storage.googleapis.com/upload"), + ) + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + side_effect=mock_side_effect, + ) as mock_global_post, + patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler._aresumable_chunked_upload", + new_callable=AsyncMock, + return_value=mock_upload_response, + ), + ): litellm.set_verbose = True litellm._turn_on_debug() file_name = "vertex_batch_completions.jsonl" diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 254d700ee5a..1d11d676207 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -47,6 +47,7 @@ IGNORE_FUNCTIONS = [ "_read_image_bytes", # max depth set. "_get_masked_values", # max depth set (default 20) to prevent infinite recursion while masking nested sensitive config dicts. "_redact_sensitive_litellm_params", # max depth set (default 10). + "_redact_secret_values_in_obj", # max depth set (default 10, _REDACT_SECRET_MAX_DEPTH); fails closed by returning "REDACTED" at the cap. "_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard. "resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input). "sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth. diff --git a/tests/e2e/budgets/BUDGET_CODE_MATRIX.md b/tests/e2e/budgets/BUDGET_CODE_MATRIX.md new file mode 100644 index 00000000000..63cb41228d2 --- /dev/null +++ b/tests/e2e/budgets/BUDGET_CODE_MATRIX.md @@ -0,0 +1,93 @@ +# Budget Code Matrix + +What LiteLLM actually implements for budgets: every entity that can carry a dollar +budget, how the limit is enforced, and where in the code it happens. This is the +"what we support" reference; the companion `BUDGET_TEST_COVERAGE_MATRIX.md` maps +each row to its tests and the e2e gaps. + +Over-budget surfaces as a `budget_exceeded` error (the live suite +`tests/otel_tests/test_e2e_budgeting.py` asserts `type == "budget_exceeded"`, +`code == "429"`); the underlying `BudgetExceededError` is defined in +`litellm/exceptions.py` (`status_code=400`). Enforcement runs in `common_checks()` +/ `auth_checks.py` at auth time, plus pre-call reservation in +`budget_reservation.py`. + +Legend for "Enforced": **block** = request rejected; **filter** = router skips the +deployment; **alert** = notify only, request proceeds. + +--- + +## 1. Per-entity dollar budgets + +| Entity | Budget stored | Hard `max_budget` | Soft budget | Per-window | Model budget | Reset by `budget_duration` | +|--------|---------------|-------------------|-------------|------------|--------------|----------------------------| +| API key | `LiteLLM_VerificationToken` (direct cols + `budget_id` FK) | block (`_virtual_key_max_budget_check`) | alert (`_virtual_key_soft_budget_check`) + 80% alert | block (`_virtual_key_multi_budget_check`) | block (`model_max_budget_limiter.is_key_within_model_budget`) | keys reset job | +| Internal user | `LiteLLM_UserTable` (direct cols) | block (`common_checks`, only when not on a team) | - | - | via `model_max_budget` json | users reset job | +| Team | `LiteLLM_TeamTable` (direct cols) | block (`_team_max_budget_check`) | alert (`_team_soft_budget_check`) | block (`_team_multi_budget_check`) | via `model_max_budget` | teams reset job | +| Team member | `LiteLLM_TeamMembership` -> `LiteLLM_BudgetTable` | block (`_check_team_member_budget`) | - | - | - | budget-table reset job | +| End-user / customer | `LiteLLM_EndUserTable` -> `LiteLLM_BudgetTable` | block (`_check_end_user_budget`) | - | - | block (`is_end_user_within_model_budget`) | budget-table reset job | +| Organization | `LiteLLM_OrganizationTable` -> `LiteLLM_BudgetTable` | block (`_organization_max_budget_check`) | - | - | via budget-table | budget-table reset job | +| Tag | `LiteLLM_TagTable` -> `LiteLLM_BudgetTable` | block (`_tag_max_budget_check`) | - | - | via budget-table | budget-table reset job | +| Project | `LiteLLM_ProjectTable` -> `LiteLLM_BudgetTable` | block (`_project_max_budget_check`) | alert (`_project_soft_budget_check`) | - | - | budget-table reset job | +| Provider (router) | config `provider_budget_config` (in-memory) | filter (`router_strategy/budget_limiter`) | - | yes (time window) | - | window TTL | +| Global proxy | `litellm.max_budget` (config) | block (`_global_proxy_budget_check`) | - | - | - | - | + +Notes / flags from the code: +- **User budget only enforced off-team**: `common_checks` skips the personal-user + budget when the key belongs to a team (team budget governs instead). +- **Comparison operators are inconsistent**: key/user use `>=`, team/end-user main + budget use `>`. Spend exactly at `max_budget` blocks a key but not a team. +- **Provider budgets are filter-only**: an over-budget provider is removed from + routing; if all are over budget the router raises + `no_deployments_with_provider_budget_routing` (not a per-entity block). +- **Enforcement timing differs by entity**: key / user / org / team-member / tag / + model enforce off real-time reservation counters (block within ~2 calls); + **end-user** enforcement reads `EndUserTable.spend`, which only updates on the + `proxy_batch_write_at` flush, so it lags by that interval (verified live). + +## 2. Budget mechanisms + +| Mechanism | What it does | Code | +|-----------|--------------|------| +| Pre-call reservation | Estimates max request cost, atomically reserves against redis spend counters for key/team/user/end_user/tag/team_member/org before the call; blocks if a counter would exceed | `spend_tracking/budget_reservation.py` | +| Post-call reconciliation | Adjusts the reservation to the actual cost once known | `reconcile_budget_reservation` | +| Read-time enforcement | Auth-time check of current spend vs `max_budget` | `auth_checks.common_checks` + per-entity `_*_max_budget_check` | +| Soft budget / alerts | At `soft_budget` (or 80% of max) fire Slack/email alert, do not block | `_virtual_key_soft_budget_check`, `_team_soft_budget_check`, `budget_alerts` | +| Multi-window budgets | `budget_limits` list of `{budget_duration, max_budget}`; each window enforced + reset independently | `_virtual_key_multi_budget_check`, `reset_budget_windows` | +| Model-level budgets | `model_max_budget` dict (per model: `budget_limit` + `time_period`) on key/user/end_user | `hooks/model_max_budget_limiter.py` | +| Reset by duration | Job zeros `spend`, recomputes `budget_reset_at = now + duration_in_seconds(budget_duration)`, invalidates redis counters | `common_utils/reset_budget_job.py`, `duration_parser.duration_in_seconds` | +| Zero-cost bypass | Models with no configured price bypass budget reservation | `budget_reservation` zero-cost path | + +## 3. Budget management surface (endpoints) + +| Action | Endpoint | Handler | +|--------|----------|---------| +| Create budget | `POST /budget/new` | `new_budget` | +| Update budget | `POST /budget/update` | `update_budget` | +| Budget info | `POST /budget/info` (`{"budgets": [id]}`) | `info_budget` | +| Budget settings | `GET /budget/settings` | `budget_settings` | +| List budgets | `GET /budget/list` | `list_budget` | +| Delete budget | `POST /budget/delete` (`{"id": id}`) | `delete_budget` | +| Set on key | `POST /key/generate`, `/key/update` (`max_budget`, `soft_budget`, `budget_duration`, `model_max_budget`, `budget_id`) | key mgmt | +| Set on user | `POST /user/new` (`max_budget`, `budget_duration`) | internal user | +| Set on team | `POST /team/new` (`max_budget`, `soft_budget`, `team_member_budget`) | team | +| Set on team member | `POST /team/member_add` (`max_budget_in_team`) | team | +| Set on org | `POST /organization/new` (`max_budget`, `soft_budget`, `model_max_budget`) | org | +| Set on customer | `POST /customer/new`, `/customer/update` (`max_budget`, `budget_id`) | customer | +| Set on tag | `POST /tag/new`, `/tag/update` (`max_budget`) | tag mgmt | +| Read budget+spend | `/key/info`, `/user/info`, `/team/info`, `/organization/info`, `/customer/info`, `/budget/info` | per-entity info | + +Endpoint method/shape gotchas verified live: `/organization/delete` is **DELETE** +with `{"organization_ids": [id]}`; `/budget/info` takes `{"budgets": [id]}`; +`model_max_budget` entries use `{"budget_limit", "time_period"}`. + +## 4. Config knobs + +| Setting | Effect | +|---------|--------| +| `litellm.max_budget` | proxy-wide hard cap (global proxy budget) | +| `max_internal_user_budget` / `default_max_internal_user_budget` | default `max_budget` for internal users | +| `internal_user_budget_duration` | default reset duration for internal users | +| `max_end_user_budget` / `max_end_user_budget_id` | default budget for end-users | +| `default_team_params` | default `max_budget` / `budget_duration` / limits for teams | +| `provider_budget_config` (router) | per-provider spend caps + windows | diff --git a/tests/e2e/budgets/BUDGET_TEST_COVERAGE_MATRIX.md b/tests/e2e/budgets/BUDGET_TEST_COVERAGE_MATRIX.md new file mode 100644 index 00000000000..62bfc1fdd41 --- /dev/null +++ b/tests/e2e/budgets/BUDGET_TEST_COVERAGE_MATRIX.md @@ -0,0 +1,78 @@ +# Budget Test Coverage Matrix + +Maps every row of `BUDGET_CODE_MATRIX.md` (what LiteLLM implements) to its tests +and level, then marks the live e2e coverage this suite adds. + +Levels: `unit` mocked (`AsyncMock` on `get_current_spend`/prisma); `router` live +router with fake deployments; `live-e2e` real proxy, real key/team, real requests +until blocked. Status: `covered` / `partial` / `gap`. + +Pre-existing live coverage outside this suite: +- `tests/otel_tests/test_e2e_budgeting.py` - key + team enforcement, budget update. +- `tests/local_testing/test_router_budget_limiter.py` - provider / tag / deployment + budgets at the router. + +This suite (`tests/e2e/budgets/`) adds the missing live coverage and runs +on the shared lifecycle (every entity it creates is deleted on teardown). + +--- + +## Per-entity enforcement + +| Entity | Unit | Pre-existing live | This suite (live) | Status | +|--------|------|-------------------|-------------------|--------| +| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** | +| Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** | +| Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** | +| Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** | +| End-user / customer | `test_custom_auth_end_user_budget.py` | - | `test_end_user_budget_blocks` | **covered (new)** | +| Organization | `test_organization_budget_enforcement.py` (flagged weak) | - | `test_organization_budget_blocks` | **covered (new)** | +| Tag (proxy-level) | - | router only | `test_tag_budget_e2e::test_tag_budget_blocks_tagged_requests` | **covered (new)** | +| Model-level (`model_max_budget`) | `test_unit_test_max_model_budget_limiter.py` | - | `test_model_max_budget_e2e::test_model_max_budget_isolates_per_model` | **covered (new)** | +| Provider (router) | `test_budget_limiter_hotpath.py` | `test_router_budget_limiter.py` | - | **covered** (router) | +| Global proxy (`litellm.max_budget`) | unit | - | - | **gap** (needs a config-level cap; not key-settable) | + +## Budget mechanisms + +| Mechanism | Unit | This suite (live) | Status | +|-----------|------|-------------------|--------| +| Pre-call reservation | `test_budget_reservation.py` | exercised by every enforcement test | **partial** | +| Soft budget / alerts | `SlackAlerting/test_budget_alert_types.py` | `test_soft_budget_e2e::test_soft_budget_does_not_block` | **covered (new)** (block-vs-alert; the alert side-effect itself stays unit) | +| Budget CRUD | `test_budget_endpoints.py` | `test_budget_crud_e2e` (roundtrip + delete) | **covered (new)** | +| Reset scheduling | `test_proxy_budget_reset.py` | `test_budget_crud_e2e::test_budget_duration_schedules_reset_on_key` | **covered (new)** (scheduling; actual zeroing is time-dependent -> unit) | +| Multi-window budgets | `test_multi_budget_windows.py` | - | **gap** (window setup is fiddly; left to unit for now) | +| Read budget+spend | `test_spend_management_endpoints.py` | `/key/info` asserted in CRUD + enforcement | **partial** | + +## Remaining gaps (intentionally not live-tested) + +- **Global proxy budget** (`litellm.max_budget`): set via proxy config, not a + per-key API, so it needs a dedicated proxy boot with that config rather than a + runtime-created entity. Out of scope for the per-entity suite. +- **Multi-window budgets**: the `budget_limits` list shape and per-window reset are + covered by `test_multi_budget_windows.py` (unit); a live version would need to + wait out a short window to see the reset, which is time-dependent. +- **Soft-budget alert delivery**: whether the Slack/email actually fires is not + observable from the proxy API; unit tests own that. The live test pins the + load-bearing behavior (soft does not block). +- **Reset zeroing after the window elapses**: time-dependent; unit tests own the + reset-job logic. The live test pins that `budget_reset_at` is scheduled. + +## This suite's files + +| File | Covers | +|------|--------| +| `test_budget_enforcement_e2e.py` | key / internal-user / end-user / organization / team-member hard enforcement | +| `test_model_max_budget_e2e.py` | per-model caps isolate by model | +| `test_soft_budget_e2e.py` | soft budget alerts but does not block | +| `test_tag_budget_e2e.py` | proxy-level tag budget blocks tagged requests, spares others | +| `test_budget_crud_e2e.py` | `/budget/*` CRUD roundtrip + delete + `budget_reset_at` scheduling | + +## Pattern + timing + +Create the entity with a tiny `max_budget`, drive spend until a `budget_exceeded` +block. The enforcement helper is two-phase: a fast warmup (key/user/org/member/tag/ +model block within ~2 calls off real-time counters), then a poll across the ~60s +batch-write window (end-user enforcement reads table spend that lags). Skip on a +non-budget error (provider down / key missing); fail if the budget is never +enforced. Chat tests use `gpt-5.5` (the model with a working key on the reference +proxy); swap the literal if your proxy differs. diff --git a/tests/e2e/budgets/budget_client.py b/tests/e2e/budgets/budget_client.py new file mode 100644 index 00000000000..af8021f9b93 --- /dev/null +++ b/tests/e2e/budgets/budget_client.py @@ -0,0 +1,418 @@ +"""Client for budget e2e tests: the shared Gateway plus budget-bearing entity +management (user / team / team-member / org / customer / tag / budget-table) and +info reads. + +Over-budget surfaces as a ``budget_exceeded`` error; ``is_budget_block`` detects it +on a chat outcome. Create methods return the new id and raise on failure; tests +register the matching delete with ``resources.defer(...)`` for cleanup. The request +and response models are co-located here because only this suite uses them. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from pydantic import AliasPath, BaseModel, Field, RootModel + +from e2e_gateway import Gateway, build_gateway +from e2e_http import NoBody, StreamingResponse, Success, unwrap +from models import ( + BudgetWindow, + ChatBody, + ChatMessage, + ChatMetadata, + KeyGenerateBody, + ModelBudgetEntry, +) + + +class UserNewBody(BaseModel): + max_budget: float + + +class UserNewResponse(BaseModel): + user_id: str + + +class UserDeleteBody(BaseModel): + user_ids: list[str] + + +class CustomerNewBody(BaseModel): + user_id: str + max_budget: float + + +class OrgNewBody(BaseModel): + organization_alias: str + max_budget: float + + +class OrgNewResponse(BaseModel): + organization_id: str + + +class OrgDeleteBody(BaseModel): + organization_ids: list[str] + + +class TeamMember(BaseModel): + role: str + user_id: str + + +class TeamNewBody(BaseModel): + team_alias: str + max_budget: float | None = None + organization_id: str | None = None + budget_limits: list[BudgetWindow] | None = None + + +class TeamNewResponse(BaseModel): + team_id: str + + +class TeamDeleteBody(BaseModel): + team_ids: list[str] + + +class TeamMemberAddBody(BaseModel): + team_id: str + member: TeamMember + max_budget_in_team: float | None = None + + +class TeamMemberUpdateBody(BaseModel): + team_id: str + user_id: str + max_budget_in_team: float | None = None + budget_duration: str | None = None + + +class TeamMembershipRow(BaseModel): + user_id: str | None = None + budget_reset_at: str | None = Field( + default=None, + validation_alias=AliasPath("litellm_budget_table", "budget_reset_at"), + ) + + +class TeamInfoParams(BaseModel): + team_id: str + + +class TeamInfoResponse(BaseModel): + team_memberships: list[TeamMembershipRow] = [] + + +class TagNewBody(BaseModel): + name: str + max_budget: float + + +class TagDeleteBody(BaseModel): + name: str + + +class BudgetNewBody(BaseModel): + max_budget: float + soft_budget: float | None = None + budget_duration: str | None = None + + +class BudgetNewResponse(BaseModel): + budget_id: str + + +class BudgetDeleteBody(BaseModel): + id: str + + +class BudgetInfoBody(BaseModel): + budgets: list[str] + + +class BudgetRow(BaseModel): + budget_id: str | None = None + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: str | None = None + + +class BudgetInfoResponse(RootModel[list[BudgetRow]]): + pass + + +def is_budget_block(result: StreamingResponse) -> bool: + """True if the call was rejected for being over budget (vs a provider error).""" + return not result.ok and "budget_exceeded" in result.body + + +def model_budget(model: str, limit: float, period: str = "30d") -> dict[str, ModelBudgetEntry]: + """A model_max_budget entry: per-model cap with a reset window.""" + return {model: ModelBudgetEntry(budget_limit=limit, time_period=period)} + + +@dataclass(frozen=True, slots=True) +class BudgetClient: + gateway: Gateway + + # ---- generic key ops (delegate to the shared Gateway) --------------- + + def generate_key( + self, + *, + models: list[str] | None = None, + max_budget: float | None = None, + soft_budget: float | None = None, + budget_duration: str | None = None, + budget_id: str | None = None, + user_id: str | None = None, + team_id: str | None = None, + model_max_budget: dict[str, ModelBudgetEntry] | None = None, + budget_limits: list[BudgetWindow] | None = None, + ) -> str: + return self.gateway.generate_key( + KeyGenerateBody( + models=models or [], + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + budget_id=budget_id, + user_id=user_id, + team_id=team_id, + model_max_budget=model_max_budget, + budget_limits=budget_limits, + ) + ) + + def delete_key(self, key: str) -> None: + self.gateway.delete_key(key) + + def delete_customers(self, user_ids: list[str]) -> None: + self.gateway.delete_customers(user_ids) + + # ---- chat (raw HTTP outcome: a budget block surfaces as a non-2xx) -- + + def chat( + self, + key: str, + model: str, + content: str, + *, + max_tokens: int | None = None, + user: str | None = None, + tags: list[str] | None = None, + ) -> StreamingResponse: + return self.gateway.transport.send( + "/chat/completions", + headers=self.gateway.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + user=user, + metadata=ChatMetadata(tags=tags) if tags else None, + ), + ) + + # ---- internal user -------------------------------------------------- + + def create_user(self, *, max_budget: float) -> str: + return unwrap( + self.gateway.transport.post( + "/user/new", + headers=self.gateway.transport.master, + json=UserNewBody(max_budget=max_budget), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = self.gateway.transport.post( + "/user/delete", + headers=self.gateway.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=NoBody, + ) + + # ---- customer / end-user ------------------------------------------- + + def create_customer(self, customer_id: str, *, max_budget: float) -> str: + resp = self.gateway.transport.send( + "/customer/new", + headers=self.gateway.transport.master, + json=CustomerNewBody(user_id=customer_id, max_budget=max_budget), + ) + assert resp.ok, resp.body + return customer_id + + # ---- organization --------------------------------------------------- + + def create_org(self, *, max_budget: float, alias: str) -> str: + return unwrap( + self.gateway.transport.post( + "/organization/new", + headers=self.gateway.transport.master, + json=OrgNewBody(organization_alias=alias, max_budget=max_budget), + response_type=OrgNewResponse, + ) + ).organization_id + + def delete_org(self, org_id: str) -> None: + _ = self.gateway.transport.delete( + "/organization/delete", + headers=self.gateway.transport.master, + json=OrgDeleteBody(organization_ids=[org_id]), + response_type=NoBody, + ) + + # ---- team ----------------------------------------------------------- + + def create_team( + self, + *, + alias: str, + max_budget: float | None = None, + organization_id: str | None = None, + budget_limits: list[BudgetWindow] | None = None, + ) -> str: + return unwrap( + self.gateway.transport.post( + "/team/new", + headers=self.gateway.transport.master, + json=TeamNewBody( + team_alias=alias, + max_budget=max_budget, + organization_id=organization_id, + budget_limits=budget_limits, + ), + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + _ = self.gateway.transport.post( + "/team/delete", + headers=self.gateway.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + + def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None: + resp = self.gateway.transport.send( + "/team/member_add", + headers=self.gateway.transport.master, + json=TeamMemberAddBody( + team_id=team_id, + member=TeamMember(role="user", user_id=user_id), + max_budget_in_team=max_budget_in_team, + ), + ) + assert resp.ok, resp.body + + def update_team_member( + self, + team_id: str, + user_id: str, + *, + max_budget_in_team: float | None = None, + budget_duration: str | None = None, + ) -> None: + resp = self.gateway.transport.send( + "/team/member_update", + headers=self.gateway.transport.master, + json=TeamMemberUpdateBody( + team_id=team_id, + user_id=user_id, + max_budget_in_team=max_budget_in_team, + budget_duration=budget_duration, + ), + ) + assert resp.ok, resp.body + + def member_budget_reset_at(self, team_id: str, user_id: str) -> str | None: + """The member's per-team budget_reset_at as /team/info reports it, or None if + no reset is scheduled. The reset job advances this each time the window + elapses; a job that skips the row leaves it pinned forever.""" + result = self.gateway.transport.get( + "/team/info", + headers=self.gateway.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoResponse, + ) + match result: + case Success(data=data): + return next( + (row.budget_reset_at for row in data.team_memberships if row.user_id == user_id), + None, + ) + case _: + return None + + # ---- tag ------------------------------------------------------------ + + def create_tag(self, name: str, *, max_budget: float) -> str: + resp = self.gateway.transport.send( + "/tag/new", + headers=self.gateway.transport.master, + json=TagNewBody(name=name, max_budget=max_budget), + ) + assert resp.ok, resp.body + return name + + def delete_tag(self, name: str) -> None: + _ = self.gateway.transport.post( + "/tag/delete", + headers=self.gateway.transport.master, + json=TagDeleteBody(name=name), + response_type=NoBody, + ) + + # ---- budget table --------------------------------------------------- + + def create_budget( + self, + *, + max_budget: float, + soft_budget: float | None = None, + budget_duration: str | None = None, + ) -> str: + return unwrap( + self.gateway.transport.post( + "/budget/new", + headers=self.gateway.transport.master, + json=BudgetNewBody( + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + ), + response_type=BudgetNewResponse, + ) + ).budget_id + + def delete_budget(self, budget_id: str) -> None: + _ = self.gateway.transport.post( + "/budget/delete", + headers=self.gateway.transport.master, + json=BudgetDeleteBody(id=budget_id), + response_type=NoBody, + ) + + def budget_info(self, budget_id: str) -> tuple[BudgetRow, ...]: + result = self.gateway.transport.post( + "/budget/info", + headers=self.gateway.transport.master, + json=BudgetInfoBody(budgets=[budget_id]), + response_type=BudgetInfoResponse, + ) + match result: + case Success(data=data): + return tuple(data.root) + case _: + return () + + +def build_client() -> BudgetClient: + return BudgetClient(gateway=build_gateway()) diff --git a/tests/e2e/budgets/conftest.py b/tests/e2e/budgets/conftest.py new file mode 100644 index 00000000000..236822f4309 --- /dev/null +++ b/tests/e2e/budgets/conftest.py @@ -0,0 +1,16 @@ +"""Budgets suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. BudgetClient holds the shared Gateway, +so the `resources` fixture cleans up keys through it; tests register entity deletes +via `resources.defer(...)`. +""" + +import pytest + +from budget_client import BudgetClient, build_client + + +@pytest.fixture(scope="session") +def client() -> BudgetClient: + return build_client() diff --git a/tests/e2e/budgets/test_budget_crud_e2e.py b/tests/e2e/budgets/test_budget_crud_e2e.py new file mode 100644 index 00000000000..e697eca0051 --- /dev/null +++ b/tests/e2e/budgets/test_budget_crud_e2e.py @@ -0,0 +1,60 @@ +"""Live e2e for the budget management surface (no LLM calls, fast). + +Covers the budget-table CRUD round-trip and that `budget_duration` schedules a +`budget_reset_at`. The actual zeroing after the window is time-dependent, so we +assert the reset is *scheduled* (now + duration), not waited out. +""" + +from datetime import datetime, timezone + +import pytest + +from budget_client import BudgetClient +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + + +def test_budget_crud_roundtrip(client: BudgetClient, resources: ResourceManager) -> None: + budget_id = client.create_budget(max_budget=12.5, soft_budget=10.0, budget_duration="30d") + resources.defer(lambda: client.delete_budget(budget_id)) + + rows = client.budget_info(budget_id) + assert rows, f"/budget/info returned nothing for {budget_id}" + row = rows[0] + assert row.max_budget == 12.5 + assert row.soft_budget == 10.0 + assert row.budget_reset_at, "budget_duration did not schedule a reset" + + # Attach the budget to a key and confirm the key reflects it. + key = client.generate_key(budget_id=budget_id) + resources.defer(lambda: client.delete_key(key)) + info = client.gateway.key_info(key) + linked = info.litellm_budget_table + assert info.budget_id == budget_id or (linked is not None and linked.max_budget == 12.5), ( + f"key does not reflect attached budget: {info.budget_id}, {linked}" + ) + + +def test_budget_delete_removes_it(client: BudgetClient, resources: ResourceManager) -> None: + budget_id = client.create_budget(max_budget=1.0) + resources.defer(lambda: client.delete_budget(budget_id)) + client.delete_budget(budget_id) + assert not client.budget_info(budget_id), "budget still present after delete" + + +def test_budget_duration_schedules_reset_on_key(client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(max_budget=10.0, budget_duration="30d") + resources.defer(lambda: client.delete_key(key)) + + reset_at = client.gateway.key_info(key).budget_reset_at + assert reset_at, "budget_duration did not set budget_reset_at on the key" + + # budget_duration schedules a FUTURE reset. Don't assume now+30d exactly: the + # proxy may align the reset to a calendar boundary (e.g. start of next month), + # so "30d" can land ~12 days out mid-month. Assert it's scheduled ahead. + + # get current time -> assert budget from days_left - budget_duration == days_left + reset_dt = datetime.fromisoformat(str(reset_at).replace("Z", "+00:00")) + days_out = (reset_dt - datetime.now(timezone.utc)).total_seconds() / 86400 + assert 0 < days_out < 40, f"reset should be scheduled ahead, got {days_out:.1f}d out" diff --git a/tests/e2e/budgets/test_budget_enforcement_e2e.py b/tests/e2e/budgets/test_budget_enforcement_e2e.py new file mode 100644 index 00000000000..d03288b637a --- /dev/null +++ b/tests/e2e/budgets/test_budget_enforcement_e2e.py @@ -0,0 +1,148 @@ +"""Live e2e: a tiny max_budget on an entity actually blocks requests. + +Each entity is an E2ECase (lifecycle.E2ECase) driven by run_case: init() creates +the budgeted entity + a key, run() drives spend until a `budget_exceeded` block, +teardown() deletes everything init() created (always runs, even on failure/skip). +Covers the entities with no prior live coverage - internal user, end-user, +organization, team member. See BUDGET_TEST_COVERAGE_MATRIX.md. + +A non-budget error fails hard (never a skip); if calls never get blocked, budget +enforcement is broken -> fail. +""" + +import time +from dataclasses import dataclass, field +from typing import Callable, List, Type + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import run_case + +pytestmark = pytest.mark.e2e + +def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> None: + """Send paid calls until the entity's budget blocks one. Key/user/org/member + block within a couple calls off real-time reservation counters; the end-user + budget enforces off table spend that lands on the batch write, so it takes a + few more. A non-budget error fails hard (never a skip).""" + for _ in range(40): + result = client.chat( + key, + "claude-haiku-4-5", + f"spend {unique_marker()}", + max_tokens=16, + user=user or None, + ) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail("budget never enforced within the call budget") + + +@dataclass +class _BudgetCase: + """Base E2ECase: a key under some budgeted entity must get blocked. + + Subclasses set up the budgeted entity in init() and register every created id + in `_undo` (run LIFO in teardown so a key is deleted before its team/org). + """ + + client: BudgetClient + key: str = "" + _undo: List[Callable[[], None]] = field( + default_factory=list + ) # mutable-ok: per-case teardown registry + + def init(self) -> None: + raise NotImplementedError + + def run(self) -> None: + _assert_budget_blocks(self.client, self.key) + + def teardown(self) -> None: + for undo in reversed(self._undo): + undo() + + +class KeyBudgetCase(_BudgetCase): + def init(self) -> None: + self.key = self.client.generate_key(max_budget=3e-6) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +class InternalUserBudgetCase(_BudgetCase): + def init(self) -> None: + user_id = self.client.create_user(max_budget=3e-6) + self._undo.append(lambda: self.client.delete_user(user_id)) + # personal key (no team) -> the user budget governs + self.key = self.client.generate_key(user_id=user_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +class EndUserBudgetCase(_BudgetCase): + def init(self) -> None: + customer = f"e2e-budget-cust-{unique_marker()}" + self.client.create_customer(customer, max_budget=3e-6) + self._undo.append(lambda: self.client.delete_customers([customer])) + self.key = self.client.generate_key(models=["claude-haiku-4-5"]) + self._undo.append(lambda: self.client.delete_key(self.key)) + self._customer = customer + + def run(self) -> None: + _assert_budget_blocks(self.client, self.key, user=self._customer) + + +class OrganizationBudgetCase(_BudgetCase): + def init(self) -> None: + # Org carries the tiny budget; the team under it has none, so a block here + # is org-level enforcement (the historically weak link). + org_id = self.client.create_org( + max_budget=3e-6, alias=f"e2e-budget-org-{unique_marker()}" + ) + self._undo.append(lambda: self.client.delete_org(org_id)) + team_id = self.client.create_team( + alias=f"e2e-budget-team-{unique_marker()}", organization_id=org_id + ) + self._undo.append(lambda: self.client.delete_team(team_id)) + self.key = self.client.generate_key(team_id=team_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +class TeamMemberBudgetCase(_BudgetCase): + def init(self) -> None: + # Member's per-team budget is tiny while the team has a large budget, so a + # block proves member-level (not team-level) enforcement. + team_id = self.client.create_team( + alias=f"e2e-budget-team-{unique_marker()}", max_budget=100.0 + ) + self._undo.append(lambda: self.client.delete_team(team_id)) + user_id = self.client.create_user(max_budget=100.0) + self._undo.append(lambda: self.client.delete_user(user_id)) + self.client.add_team_member(team_id, user_id, max_budget_in_team=3e-6) + self.key = self.client.generate_key(team_id=team_id, user_id=user_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +def _case_id(case_cls: Type[_BudgetCase]) -> str: + return case_cls.__name__ + + +@pytest.mark.parametrize( + "case_cls", + [ + KeyBudgetCase, + InternalUserBudgetCase, + EndUserBudgetCase, + OrganizationBudgetCase, + TeamMemberBudgetCase, + ], + ids=_case_id, +) +def test_budget_enforcement( + client: BudgetClient, case_cls: Type[_BudgetCase] +) -> None: + run_case(case_cls(client)) diff --git a/tests/e2e/budgets/test_budget_reset_e2e.py b/tests/e2e/budgets/test_budget_reset_e2e.py new file mode 100644 index 00000000000..dcf776db9a2 --- /dev/null +++ b/tests/e2e/budgets/test_budget_reset_e2e.py @@ -0,0 +1,59 @@ +"""Live e2e: a key budget resets (zeroes spend) after its budget_duration. + +Short budget_duration (30s) + the fast-rescheduled reset job: a key blocked for +exceeding its max_budget starts succeeding again once the duration elapses and the +reset job zeroes key.spend. Closes the reset-zeroing gap in +BUDGET_TEST_COVERAGE_MATRIX.md (reset_budget_for_litellm_keys), which the unit +suite covers but no live test did - distinct from the per-window reset in +test_multi_window_budget_e2e.py. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + + +def _call(client: BudgetClient, key: str): + return client.chat( + key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16 + ) + + +def test_key_budget_resets_after_duration( + client: BudgetClient, resources: ResourceManager +) -> None: + key = client.generate_key(max_budget=3e-6, budget_duration="30s") + resources.defer(lambda: client.delete_key(key)) + + # 1. exceed the budget -> litellm returns budget_exceeded + blocked = False + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, "key budget never enforced" + + # 2. once the 30s duration elapses + the reset job runs, key.spend zeroes and + # calls flow again. The window is wall-clock-aligned, so the reset lands up to + # a window later, then the rescheduler (~15-20s) zeroes the spend; allow + # generous headroom over that. A stuck rescheduler is caught by the wait-loop + # timeout, not this elapsed bound. + start = time.monotonic() + while time.monotonic() < start + 150: + time.sleep(5) + result = _call(client, key) + if result.ok: + assert time.monotonic() - start < 120, "reset too slow for a 30s budget" + return + assert is_budget_block(result), f"non-budget error: {result.body[:200]}" + pytest.fail("key budget never reset within 150s") diff --git a/tests/e2e/budgets/test_model_max_budget_e2e.py b/tests/e2e/budgets/test_model_max_budget_e2e.py new file mode 100644 index 00000000000..44e6a333ef0 --- /dev/null +++ b/tests/e2e/budgets/test_model_max_budget_e2e.py @@ -0,0 +1,57 @@ +"""Live e2e: per-model budgets (`model_max_budget`) isolate by model. + +A key caps one model tiny and leaves another generous. Exhausting the capped +model must block *that* model while the other still works - proving the per-model +cap is enforced independently, not as a key-wide budget. Closes the +model_max_budget gap in BUDGET_TEST_COVERAGE_MATRIX.md. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block, model_budget +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +CAPPED_MODEL = "claude-haiku-4-5" +FREE_MODEL = "gemini-2.5-flash" + + +def _call(client: BudgetClient, key: str, model: str): + result = client.chat(key, model, f"hi {unique_marker()}", max_tokens=16) + if not result.ok and not is_budget_block(result): + require_successful_call(result) + return result + + +def test_model_max_budget_isolates_per_model( + client: BudgetClient, resources: ResourceManager +) -> None: + key = client.generate_key( + model_max_budget={ + **model_budget(CAPPED_MODEL, 1e-6), + **model_budget(FREE_MODEL, 1000.0), + } + ) + resources.defer(lambda: client.delete_key(key)) + + # Exhaust the capped model. + blocked = False + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + if is_budget_block(_call(client, key, CAPPED_MODEL)): + blocked = True + break + time.sleep(1) + assert blocked, f"{CAPPED_MODEL} per-model budget never enforced" + + # The other model shares the key but has its own (large) cap -> still works. + other = _call(client, key, FREE_MODEL) + assert not is_budget_block(other), ( + f"{FREE_MODEL} was blocked by {CAPPED_MODEL}'s budget; per-model caps not isolated" + ) + require_successful_call(other) diff --git a/tests/e2e/budgets/test_multi_window_budget_e2e.py b/tests/e2e/budgets/test_multi_window_budget_e2e.py new file mode 100644 index 00000000000..553ad1ce701 --- /dev/null +++ b/tests/e2e/budgets/test_multi_window_budget_e2e.py @@ -0,0 +1,70 @@ +"""Live e2e: multi-window budgets (budget_limits) enforce AND reset per window. + +Short windows make the time limit reachable inside a test: a tight 30s window and +a roomy 1m window. The 30s window blocks once its tiny cap is exceeded, then - once +its 30s elapses and the reset job runs (rescheduled fast via +PROXY_BUDGET_RESCHEDULER_* in docker-compose) - the window resets and calls flow +again. Closes the multi-window gap (enforcement + per-window reset) in +BUDGET_TEST_COVERAGE_MATRIX.md, which the unit suite covered but no live test did. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager +from models import BudgetWindow + +pytestmark = pytest.mark.e2e + +WINDOW_SECONDS = 30 # the tight window; calls succeed again only after it elapses + + +def _call(client: BudgetClient, key: str): + return client.chat( + key, "claude-haiku-4-5", f"window {unique_marker()}", max_tokens=16 + ) + + +def test_short_window_blocks_then_resets( + client: BudgetClient, resources: ResourceManager +) -> None: + key = client.generate_key( + budget_limits=[ + BudgetWindow(budget_duration=f"{WINDOW_SECONDS}s", max_budget=3e-6), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ] + ) + resources.defer(lambda: client.delete_key(key)) + + # 1. exhaust the tight window -> litellm returns budget_exceeded + start = time.monotonic() + blocked = False + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, f"{WINDOW_SECONDS}s window never enforced" + + # 2. the window resets at the next wall-clock-aligned boundary (up to a window + # after start), then the reset job (~15-20s rescheduler) zeroes the spend. + # Allow generous headroom for that alignment + rescheduler latency; a stuck + # rescheduler is caught by the wait-loop timeout, not this elapsed bound. + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, ( + f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" + ) + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"{WINDOW_SECONDS}s window never reset within 150s") diff --git a/tests/e2e/budgets/test_soft_budget_e2e.py b/tests/e2e/budgets/test_soft_budget_e2e.py new file mode 100644 index 00000000000..407de7ae467 --- /dev/null +++ b/tests/e2e/budgets/test_soft_budget_e2e.py @@ -0,0 +1,35 @@ +"""Live e2e: soft_budget alerts but does NOT block. + +A key with a tiny `soft_budget` well under a large `max_budget`: spend crosses the +soft threshold within a couple calls, but requests keep succeeding (soft budget is +advisory). Closes the soft_budget gap in BUDGET_TEST_COVERAGE_MATRIX.md. The alert +side-effect (Slack/email) is not observable from the proxy API, so we assert the +load-bearing behavior: soft != block. +""" + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + + +def test_soft_budget_does_not_block( + client: BudgetClient, resources: ResourceManager +) -> None: + # soft far below max: spend crosses soft immediately, stays under max. + key = client.generate_key(max_budget=1000.0, soft_budget=1e-9) + resources.defer(lambda: client.delete_key(key)) + + for _ in range(3): + result = client.chat( + key, "claude-haiku-4-5", f"hi {unique_marker()}", max_tokens=16 + ) + assert not is_budget_block(result), ( + "soft_budget blocked a request; it must alert only, not block " + f"(body={result.body[:200]})" + ) + require_successful_call(result) # any other non-2xx (e.g. provider down) is a hard fail diff --git a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py new file mode 100644 index 00000000000..a6860aeef43 --- /dev/null +++ b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py @@ -0,0 +1,144 @@ +"""Live e2e: concurrent cold-counter reseeds keep the spend counter equal to DB spend (#26829). + +Regression for the cross-pod spend-counter multiplication. Real requests build a key's +DB spend through the spend writer; the Redis spend counter then expires (the e2e proxy +sets a short default_redis_ttl) and goes cold. The proxy runs several workers sharing one +Redis, so a concurrent burst makes more than one worker reseed the same cold counter at +once. The fix seeds with SET NX - one worker initializes the counter at the DB spend and +the rest read it back - so the counter still equals the DB spend (plus the burst's own +small cost). The pre-#26829 additive reseed stacked the DB spend once per worker, leaving +the counter at ~N x the real spend. + +The test reads the shared counter straight from Redis and asserts it equals the DB spend, +not a multiple. It also asserts the counter actually went cold before the burst, so a proxy +that never expires the counter (no short TTL) fails loudly instead of passing vacuously. +Skipped when the e2e Redis is not reachable. +""" + +import hashlib +import os +import time +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from typing import TYPE_CHECKING + +import pytest + +from budget_client import BudgetClient +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager + +if TYPE_CHECKING: + import redis + from redis.cluster import RedisCluster + +pytestmark = pytest.mark.e2e + +MODEL = "claude-haiku-4-5" +ACCUMULATE_CALLS = 24 +BURST = 6 +# proxy_batch_write_at (60s) flushes the spend to the DB and default_redis_ttl (20s) +# expires the counter; this waits out both. +COLD_WAIT_SECONDS = 80 + + +def _redis() -> "redis.Redis[str] | RedisCluster[str]": + """The proxy's Redis. The deployed runner sets REDIS_HOST to the serverless + ElastiCache, which is always TLS + cluster-mode; without it, fall back to a + local standalone redis for docker-compose runs.""" + import redis + + host = os.getenv("REDIS_HOST") + if not host: + return redis.Redis(host="localhost", port=6380, decode_responses=True, socket_connect_timeout=2) + + from redis.cluster import RedisCluster + + return RedisCluster( + host=host, + port=int(os.getenv("REDIS_PORT", "6379")), + ssl=True, + decode_responses=True, + socket_connect_timeout=2, + ) + + +def _spend_counter(rds: "redis.Redis[str] | RedisCluster[str]", key: str) -> float | None: + """The shared spend counter for `key`, or None if it is cold. A cluster client + can't run a keyspace SCAN that spans shards, so read the key directly - the stage + gateway sets no cache namespace, so the key is the bare ``spend:key:{sha256(key)}``. + A standalone client matches by suffix, so the local cache namespace (litellm.caching) + need not be hard-coded here.""" + from redis.cluster import RedisCluster + + digest = hashlib.sha256(key.encode()).hexdigest() + suffix = f"spend:key:{digest}" + if isinstance(rds, RedisCluster): + raw = rds.get(suffix) + return float(raw) if raw is not None else None + + matches = list(rds.scan_iter(match=f"*{suffix}")) + if not matches: + return None + raw = rds.get(matches[0]) + return float(raw) if raw is not None else None + + +def _chat(client: BudgetClient, key: str) -> StreamingResponse: + return client.chat(key, MODEL, f"reseed {unique_marker()}", max_tokens=16) + + +def _accumulate(client: BudgetClient, key: str, count: int) -> None: + def one(_: int) -> StreamingResponse: + return _chat(client, key) + + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(one, range(count))) + + +def _burst(client: BudgetClient, key: str, count: int) -> None: + """Fire `count` requests that start together, so multiple workers reseed the cold + counter concurrently rather than one warming it before the others arrive.""" + barrier = Barrier(count) + + def one(_: int) -> StreamingResponse: + barrier.wait() + return _chat(client, key) + + with ThreadPoolExecutor(max_workers=count) as pool: + list(pool.map(one, range(count))) + + +def test_cold_counter_reseed_keeps_counter_equal_to_db_spend( + client: BudgetClient, resources: ResourceManager +) -> None: + try: + rds = _redis() + rds.ping() + except Exception as exc: # noqa: BLE001 - any connect failure means skip + pytest.skip(f"e2e redis not reachable (set REDIS_HOST/REDIS_PORT): {exc}") + + key = client.generate_key(max_budget=1.0, models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + _accumulate(client, key, ACCUMULATE_CALLS) + time.sleep(COLD_WAIT_SECONDS) + + assert _spend_counter(rds, key) is None, ( + "the spend counter never went cold; default_redis_ttl must be short enough for it " + "to expire, otherwise the burst reads a warm counter and the reseed is never exercised" + ) + db_spend = client.gateway.key_info(key).spend or 0.0 + assert db_spend > 0, f"no DB spend accumulated from real calls: {db_spend}" + + _burst(client, key, BURST) + time.sleep(3) + + counter = _spend_counter(rds, key) + assert counter is not None, "the burst did not reseed the cold counter" + assert db_spend * 0.95 <= counter < db_spend * 1.7, ( + f"redis spend counter {counter} does not equal DB spend {db_spend} (expected ~equal " + f"plus the burst's small cost); a near-multiple means the cold-counter reseed stacked " + f"the DB spend once per worker instead of seeding it once (#26829)" + ) diff --git a/tests/e2e/budgets/test_tag_budget_e2e.py b/tests/e2e/budgets/test_tag_budget_e2e.py new file mode 100644 index 00000000000..7cec5bc96c1 --- /dev/null +++ b/tests/e2e/budgets/test_tag_budget_e2e.py @@ -0,0 +1,59 @@ +"""Live e2e: proxy-level tag budgets block tagged requests. + +A tag with a tiny budget: requests carrying that tag get blocked once the tag's +spend is exceeded, while a request with a different tag (no budget) still works. +Closes the proxy-level tag-budget gap in BUDGET_TEST_COVERAGE_MATRIX.md (today +only router-level tag budgets are tested). +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +TINY_BUDGET = 1e-6 + + +def _tagged_call(client: BudgetClient, key: str, tag: str): + result = client.chat( + key, + "claude-haiku-4-5", + f"hi {unique_marker()}", + tags=[tag], + max_tokens=16, + ) + if not result.ok and not is_budget_block(result): + require_successful_call(result) + return result + + +def test_tag_budget_blocks_tagged_requests( + client: BudgetClient, scoped_key: str, resources: ResourceManager +) -> None: + budgeted_tag = f"e2e-budget-tag-{unique_marker()}" + client.create_tag(budgeted_tag, max_budget=TINY_BUDGET) + resources.defer(lambda: client.delete_tag(budgeted_tag)) + + # Requests under the budgeted tag get blocked once its spend is exceeded. + blocked = False + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + if is_budget_block(_tagged_call(client, scoped_key, budgeted_tag)): + blocked = True + break + time.sleep(1) + assert blocked, f"tag budget for {budgeted_tag!r} never enforced" + + # A request with an unbudgeted tag on the same key is unaffected. + free_tag = f"e2e-free-tag-{unique_marker()}" + other = _tagged_call(client, scoped_key, free_tag) + assert not is_budget_block(other), ( + f"unbudgeted tag {free_tag!r} was blocked by {budgeted_tag!r}'s budget" + ) + require_successful_call(other) diff --git a/tests/e2e/budgets/test_team_member_budget_e2e.py b/tests/e2e/budgets/test_team_member_budget_e2e.py new file mode 100644 index 00000000000..301617bfdca --- /dev/null +++ b/tests/e2e/budgets/test_team_member_budget_e2e.py @@ -0,0 +1,107 @@ +"""Live e2e: a team member's per-team budget attributes spend and enforces a cap. + +The team carries a large budget while the one enrolled member is capped at a tiny +per-team budget, so any block is member-level, not team-level. Two scenarios share +that single member: +- attribution: the member's calls land in the spend logs tagged with both the team_id + and the member's user_id, so per-member spend can be billed back +- enforcement: once the member's spend passes the per-team budget, calls are blocked + with budget_exceeded while the team's own budget is nowhere near exhausted + +Per-member budgets enforce off batch-written spend (~60s), so a quick burst all goes +through; the block only lands once that spend flushes. +""" + +import time +from collections.abc import Iterator +from dataclasses import dataclass + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import Success, require_successful_call +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage + +pytestmark = pytest.mark.e2e + +MODEL = "claude-haiku-4-5" +TEAM_BUDGET = 100.0 +MEMBER_BUDGET = 3e-6 +BURST = 6 + + +@dataclass(frozen=True, slots=True) +class _Member: + team_id: str + user_id: str + key: str + + +@pytest.fixture(scope="class") +def member(client: BudgetClient) -> Iterator[_Member]: + """A team with a large budget plus one member capped at a tiny per-team budget, + and that member's key. Shared across the class; torn down when it finishes. + Cleanups register progressively and run LIFO best-effort through ResourceManager, + so a partial-setup failure still releases what came before and one failed delete + never strands the rest on the shared proxy.""" + resources = ResourceManager(client=client.gateway) + try: + marker = unique_marker() + team_id = client.create_team(alias=f"e2e-team-member-{marker}", max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_team(team_id)) + user_id = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(user_id)) + client.add_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET) + key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + yield _Member(team_id=team_id, user_id=user_id, key=key) + finally: + resources.teardown() + + +def _send(client: BudgetClient, key: str) -> str | None: + """One member call; its response id (== the spend-log request_id) if it went + through, else None.""" + match client.gateway.chat( + key, + ChatBody( + model=MODEL, + messages=[ChatMessage(role="user", content=f"hi {unique_marker()}")], + max_tokens=16, + ), + ): + case Success(data=response): + return response.id + case _: + return None + + +class TestTeamMemberBudget: + def test_member_spend_attributed_to_team_and_user(self, client: BudgetClient, member: _Member) -> None: + sent = frozenset(rid for rid in (_send(client, member.key) for _ in range(BURST)) if rid) + assert sent, "no member call went through; cannot check attribution" + + rows = client.gateway.poll_logs_for_key( + member.key, predicate=lambda rs: bool(sent & {r.request_id for r in rs}) + ) + logged = [row for row in rows if row.request_id in sent] + assert logged, f"none of the member's {len(sent)} calls reached the spend logs" + + for row in logged: + assert row.team_id == member.team_id, ( + f"call {row.request_id} logged under team {row.team_id}, not the member's team {member.team_id}" + ) + assert row.user == member.user_id, ( + f"call {row.request_id} logged under user {row.user}, not member {member.user_id}" + ) + + def test_member_spend_over_budget_is_blocked(self, client: BudgetClient, member: _Member) -> None: + for _ in range(40): + result = client.chat(member.key, MODEL, f"spend {unique_marker()}", max_tokens=16) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail("per-member budget never enforced within the call budget") diff --git a/tests/e2e/budgets/test_team_member_budget_reset_e2e.py b/tests/e2e/budgets/test_team_member_budget_reset_e2e.py new file mode 100644 index 00000000000..2749f16a26e --- /dev/null +++ b/tests/e2e/budgets/test_team_member_budget_reset_e2e.py @@ -0,0 +1,47 @@ +import time +from datetime import datetime + +import pytest + +from budget_client import BudgetClient +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MEMBER_BUDGET = 1.0 # default member budget is $50, we're testing with a smaller value + +def _as_datetime(value: str) -> datetime: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +def test_team_member_budget_reset_keeps_advancing(client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-member-reset-{unique_marker()}", max_budget=100.0) + resources.defer(lambda: client.delete_team(team_id)) + user_id = client.create_user(max_budget=100.0) + resources.defer(lambda: client.delete_user(user_id)) + + # add the member, then update them onto a short per-team budget window + client.add_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET) + client.update_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET, budget_duration="30s") + + scheduled = client.member_budget_reset_at(team_id, user_id) + assert scheduled, "updating the member with a budget_duration set no budget_reset_at" + first_reset = _as_datetime(scheduled) + + # the member can spend within the team while the window is live + key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + require_successful_call(client.chat(key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16)) + + # once the window elapses the reset job must move budget_reset_at forward; a job + # that skips the member's budget row (the #25109 regression) leaves it pinned at + # first_reset forever + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + current = client.member_budget_reset_at(team_id, user_id) + if current and _as_datetime(current) > first_reset: + return + pytest.fail(f"member budget_reset_at never advanced past {first_reset.isoformat()} in 150s") diff --git a/tests/e2e/budgets/test_team_multi_window_budget_e2e.py b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py new file mode 100644 index 00000000000..c58e74db965 --- /dev/null +++ b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py @@ -0,0 +1,72 @@ +"""Live e2e: a team's multi-window budgets (budget_limits) enforce AND reset per window. + +The team analog of test_multi_window_budget_e2e.py (which covers keys). A team is +created with a tight 30s window and a roomy 1m window; a key on that team blocks once +the tight window's cap is exceeded, then - once the 30s elapses and the reset job runs +(rescheduled fast via PROXY_BUDGET_RESCHEDULER_* in docker-compose) - the window resets +and calls flow again. This exercises the reset_budget_windows TEAM branch (raw SQL over +LiteLLM_TeamTable.budget_limits, the literal #25109 path), which had no live coverage. + +This also guards the /team/new write path: it must json.dumps the window list into +the Json? column. A raw list there made Prisma reject the create with a 500 (the key +path and /team/update already json.dumps first); a regression would fail team creation +here. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager +from models import BudgetWindow + +pytestmark = pytest.mark.e2e + +WINDOW_SECONDS = 30 + + +def _call(client: BudgetClient, key: str): + return client.chat(key, "claude-haiku-4-5", f"team-window {unique_marker()}", max_tokens=16) + + +def test_team_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-window-{unique_marker()}", + budget_limits=[ + BudgetWindow(budget_duration=f"{WINDOW_SECONDS}s", max_budget=3e-6), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ], + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + + # 1. exhaust the tight window -> litellm returns budget_exceeded + start = time.monotonic() + blocked = False + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, f"team {WINDOW_SECONDS}s window never enforced" + + # 2. the window resets at the next wall-clock-aligned boundary (up to a window + # after start), then the reset job (~15-20s rescheduler) zeroes the spend. + # Allow generous headroom for that alignment + rescheduler latency; a stuck + # rescheduler is caught by the wait-loop timeout, not this elapsed bound. + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 150s") diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py new file mode 100644 index 00000000000..cc95c7538dd --- /dev/null +++ b/tests/e2e/conftest.py @@ -0,0 +1,120 @@ +"""Shared fixtures for all live e2e suites under tests/e2e/. + +Design rule: skip on environment, fail on behavior. Live tests (marked `e2e`) +skip when no proxy answers; once a request reaches the proxy, behavior is +asserted. Pure unit coverage of the harness itself carries no `e2e` marker and +runs regardless of whether a proxy is up. + +Lifecycle: the `resources` fixture maps the init -> run -> teardown contract +(lifecycle.E2ECase) onto pytest - setup is init(), the test body is run(), and +teardown deletes every resource the test created on the long-lived proxy. + +Each suite provides its own `client` fixture (a lifecycle.ResourceClient); these +shared fixtures build on it. +""" + +import functools +import sys +from pathlib import Path +from typing import Iterator + +import pytest +import requests + +from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL +from lifecycle import GatewayProvider, ResourceManager + + +_E2E_TEST_RAN = pytest.StashKey[bool]() + + +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "markers", + "e2e: live test that requires a running proxy and real provider keys", + ) + + +def _liveness_reason(label: str, base_url: str) -> str | None: + """None if `base_url` answers its liveness probe, else a skip reason.""" + try: + resp = requests.get(f"{base_url}/health/liveliness", timeout=5) + except requests.RequestException as exc: + return f"No live {label} at {base_url}: {exc}" + if resp.status_code >= 500: + return f"{label} at {base_url} returned {resp.status_code}" + return None + + +@functools.lru_cache(maxsize=1) +def _proxy_skip_reason() -> str | None: + """Probe the proxy once per session. None if it answers, else a skip reason. In + a split deployment the management/admin control plane is a separate service, so + require it too (when it differs) - else its tests would fail rather than skip.""" + reason = _liveness_reason("proxy", PROXY_BASE_URL) + if reason is not None: + return reason + if CONTROL_PLANE_BASE_URL != PROXY_BASE_URL: + return _liveness_reason("control plane", CONTROL_PLANE_BASE_URL) + return None + + +def pytest_runtest_setup(item: pytest.Item) -> None: + """Skip `e2e`-marked tests unless a proxy answers its liveness probe. Unmarked + tests (unit coverage of the harness) don't touch the proxy, so they run even + when none is up.""" + if item.get_closest_marker("e2e") is None: + return + reason = _proxy_skip_reason() + if reason is not None: + pytest.skip(reason) + + +def pytest_runtest_call(item: pytest.Item) -> None: + """Mark that an e2e test body actually ran (not skipped at setup). Skipped + sessions never reach this hook, so the session-finish cleanup can use it as a + guard before truncating the spend-log DB. Tests under `tests/e2e/` without the + `e2e` marker (pure unit coverage for the harness itself) never hit the proxy, + so they must not arm the destructive DB truncate.""" + if item.get_closest_marker("e2e") is None: + return + item.session.stash[_E2E_TEST_RAN] = True + + +def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: + """Once the whole e2e session is done (all suites), truncate the spend logs so + the DB doesn't accumulate test rows. Skipped sessions (no live proxy, no test + actually executed) leave the DB alone so a `DATABASE_URL` pointing at a shared + instance is never wiped without an e2e run. Best-effort: a cleanup failure (no + DB reachable) must not fail the run. The spend_tracking dir goes on sys.path + only for this import and is removed after, so a broader `pytest tests/` run is + not left with a mutated path.""" + if not session.stash.get(_E2E_TEST_RAN, False): + return + spend_dir = str(Path(__file__).parent / "spend_tracking") + sys.path.insert(0, spend_dir) + try: + from spend_e2e_client import reset_spend_logs # pyright: ignore + + reset_spend_logs() + except Exception as exc: # noqa: BLE001 - cleanup is best-effort + print(f"spend-log cleanup skipped: {exc}") + finally: + if spend_dir in sys.path: + sys.path.remove(spend_dir) + + +@pytest.fixture +def resources(client: GatewayProvider) -> Iterator[ResourceManager]: + """init -> run -> teardown: create a manager, run the test, release resources. + Cleanup goes through the shared Gateway, whatever the suite's client adds.""" + manager = ResourceManager(client=client.gateway) + manager.init() + yield manager + manager.teardown() + + +@pytest.fixture +def scoped_key(resources: ResourceManager) -> str: + """A fresh all-models key per test, auto-deleted by the resources teardown.""" + return resources.key() diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py new file mode 100644 index 00000000000..3865804b08f --- /dev/null +++ b/tests/e2e/e2e_config.py @@ -0,0 +1,34 @@ +"""Generic configuration for live e2e tests against a running LiteLLM proxy. + +Shared by every e2e suite under tests/e2e/. Values come from the +environment so the same tests run against localhost or a deployed proxy. +""" + +import os +import uuid + +PROXY_BASE_URL = os.environ.get("LITELLM_PROXY_URL", "http://localhost:4000").rstrip("/") +MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY", "sk-1234") + +# Control-plane (management/admin) base URL. In a split control-plane/data-plane +# deployment the LLM data plane (PROXY_BASE_URL: /chat, /embeddings, native +# passthrough) and the management API (keys, users, teams, orgs, budgets, spend, +# model info, /openapi.json) are served by *different* services. The suite drives +# both through one Transport that routes by path (see transport.SplitTransport). +# Defaults to PROXY_BASE_URL so a monolithic proxy serving everything on one URL +# behaves exactly as before. +CONTROL_PLANE_BASE_URL = os.environ.get( + "LITELLM_CONTROL_PLANE_URL", PROXY_BASE_URL +).rstrip("/") + +# Writes on the proxy are eventually consistent (e.g. spend rows flush on +# proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once. +POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) +POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) +REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) + + +def unique_marker() -> str: + """A short unique token per call/run, so concurrent runs and the shared + response cache never collide on prompts, tags, or customer ids.""" + return uuid.uuid4().hex[:12] diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py new file mode 100644 index 00000000000..b700145d434 --- /dev/null +++ b/tests/e2e/e2e_gateway.py @@ -0,0 +1,211 @@ +"""Gateway: the shared proxy operations, DI'd into every client (composition). + +A frozen-slots dataclass holding a Transport plus poll config. Clients hold a +Gateway and add their own route methods; the lifecycle ResourceManager uses the +Gateway's key/customer methods for cleanup. Read-backs are eventually consistent +(proxy_batch_write_at ~60s) so they poll to a deadline. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass + +from e2e_http import ( + NoBody, + ProbeResult, + Result, + StreamingResponse, + Success, + unwrap, +) +from models import ( + ChatBody, + ChatResponse, + CustomerDeleteBody, + EmbedBody, + EmbedResponse, + KeyDeleteBody, + KeyGenerateBody, + KeyGenerateResponse, + KeyInfo, + KeyInfoParams, + KeyInfoResponse, + ModelInfoEntry, + ModelInfoResponse, + SpendLogRow, + SpendLogs, + SpendLogsParams, +) +from e2e_config import ( + CONTROL_PLANE_BASE_URL, + MASTER_KEY, + POLL_INTERVAL, + POLL_TIMEOUT, + PROXY_BASE_URL, + REQUEST_TIMEOUT, +) +from transport import HttpTransport, SplitTransport, Transport + +RowsPredicate = Callable[[list[SpendLogRow]], bool] + + +@dataclass(frozen=True, slots=True) +class Gateway: + transport: Transport + poll_timeout: float = 120.0 + poll_interval: float = 5.0 + + # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- + + def generate_key(self, body: KeyGenerateBody) -> str: + return unwrap( + self.transport.post( + "/key/generate", + headers=self.transport.master, + json=body, + response_type=KeyGenerateResponse, + ) + ).key + + def delete_key(self, key: str) -> None: + _ = self.transport.post( + "/key/delete", + headers=self.transport.master, + json=KeyDeleteBody(keys=[key]), + response_type=NoBody, + ) + + def delete_customers(self, user_ids: list[str]) -> None: + if not user_ids: + return + _ = self.transport.post( + "/customer/delete", + headers=self.transport.master, + json=CustomerDeleteBody(user_ids=user_ids), + response_type=NoBody, + ) + + def key_info(self, key: str) -> KeyInfo: + return unwrap( + self.transport.get( + "/key/info", + headers=self.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ).info + + def model_info(self) -> list[ModelInfoEntry]: + """Every configured deployment with the price the proxy resolved for it + (config override merged over cost-map defaults).""" + return unwrap( + self.transport.get( + "/model/info", + headers=self.transport.master, + params=NoBody(), + response_type=ModelInfoResponse, + ) + ).data + + # ---- LLM calls ------------------------------------------------------ + + def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]: + return self.transport.post( + "/chat/completions", + headers=self.transport.bearer(key), + json=body, + response_type=ChatResponse, + ) + + def chat_stream(self, key: str, body: ChatBody) -> StreamingResponse: + return self.transport.stream( + "/chat/completions", headers=self.transport.bearer(key), json=body + ) + + def embed(self, key: str, body: EmbedBody) -> Result[EmbedResponse]: + return self.transport.post( + "/embeddings", + headers=self.transport.bearer(key), + json=body, + response_type=EmbedResponse, + ) + + # ---- spend read-back ------------------------------------------------ + + def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]: + result = self.transport.get( + "/spend/logs", + headers=self.transport.master, + params=params, + response_type=SpendLogs, + ) + match result: + case Success(data=logs): + return logs.root + case _: + return [] + + def poll_logs_for_key( + self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None + ) -> list[SpendLogRow]: + return self._poll( + lambda: self.spend_logs(SpendLogsParams(api_key=key)), min_rows, predicate + ) + + def poll_logs_for_request_id( + self, + request_id: str, + *, + min_rows: int = 1, + predicate: RowsPredicate | None = None, + ) -> list[SpendLogRow]: + return self._poll( + lambda: self.spend_logs(SpendLogsParams(request_id=request_id)), + min_rows, + predicate, + ) + + def _poll( + self, + fetch: Callable[[], list[SpendLogRow]], + min_rows: int, + predicate: RowsPredicate | None, + ) -> list[SpendLogRow]: + deadline = time.monotonic() + self.poll_timeout + rows: list[SpendLogRow] = [] + while time.monotonic() < deadline: + rows = fetch() + if len(rows) >= min_rows and (predicate is None or predicate(rows)): + return rows + time.sleep(self.poll_interval) + return rows + + # ---- route probe ---------------------------------------------------- + + def probe(self, path: str, *, params: NoBody) -> ProbeResult: + return self.transport.probe(path, params=params) + + +def build_gateway() -> Gateway: + """The Gateway every suite's client is built from: a SplitTransport that routes + LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the + control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two + base URLs are the same for a monolithic proxy, so routing is then a no-op.""" + return Gateway( + transport=SplitTransport( + data=HttpTransport( + base_url=PROXY_BASE_URL, + master_key=MASTER_KEY, + request_timeout=REQUEST_TIMEOUT, + ), + control=HttpTransport( + base_url=CONTROL_PLANE_BASE_URL, + master_key=MASTER_KEY, + request_timeout=REQUEST_TIMEOUT, + ), + ), + poll_timeout=POLL_TIMEOUT, + poll_interval=POLL_INTERVAL, + ) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py new file mode 100644 index 00000000000..7458f316852 --- /dev/null +++ b/tests/e2e/e2e_http.py @@ -0,0 +1,306 @@ +"""The ONLY module permitted to call ``requests.*``. + +Enforced by tests/code_coverage_tests/check_e2e_no_raw_requests.py. Every request +body / query / header / response is a pydantic model; outcomes are a tagged union +(``Result[R]``) so callers ``match`` on them instead of catching exceptions. + +Named e2e_http (not http) so it does not shadow the stdlib ``http`` package that +requests itself imports. +""" + +from __future__ import annotations + +from typing import Generic, Iterator, Literal, NewType, TypeVar, cast + +import pytest +import requests +from pydantic import BaseModel, ConfigDict, Field + +URL = NewType("URL", str) + + +class Headers(BaseModel): + """Base for header models. Subclasses may alias to hyphenated header names + (e.g. ``x-litellm-api-key``); serialization uses by_alias.""" + + model_config = ConfigDict(populate_by_name=True) + + +class AuthHeaders(Headers): + # litellm accepts either; set whichever the call needs, leave the other None. + authorization: str | None = None + x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key") + + +class NoBody(BaseModel): + """Empty body/query for routes that take none.""" + + +# ---------- Result types ---------- + +R = TypeVar("R", bound=BaseModel) + + +class Success(BaseModel, Generic[R]): + kind: Literal["success"] = "success" + data: R + + +class NetworkError(BaseModel): + kind: Literal["network"] = "network" + message: str + + +class UnauthorizedError(BaseModel): + kind: Literal["unauthorized"] = "unauthorized" + + +class RateLimitedError(BaseModel): + kind: Literal["rate_limited"] = "rate_limited" + retry_after_seconds: int | None = None + # litellm overloads 429 for budget_exceeded too, so keep the body to tell them apart. + body: str = "" + + +class ValidationError(BaseModel): + kind: Literal["validation"] = "validation" + message: str + + +class UnknownApiError(BaseModel): + kind: Literal["unknown"] = "unknown" + status_code: int + body: str + + +type Result[R: BaseModel] = ( + Success[R] + | NetworkError + | UnauthorizedError + | RateLimitedError + | ValidationError + | UnknownApiError +) + + +class ProbeResult(BaseModel): + """A route's reachability: status + body, no schema validation. Healthy == + route exists (not 404) and the handler did not crash (not 5xx).""" + + status_code: int + body: str + + @property + def healthy(self) -> bool: + return 200 <= self.status_code < 500 and self.status_code != 404 + + +class StreamingResponse(BaseModel): + """Raw outcome for calls whose body is provider-native or streamed: status, the + x-litellm-call-id header (== SpendLogs.request_id), the content-type (which + tells streaming `text/event-stream` from non-streaming `application/json`), and + the body. Used by passthrough and streaming, where one validated JSON model + does not fit.""" + + status_code: int + call_id: str | None = None # x-litellm-call-id header + content_type: str | None = None + body: str + chunks: int = 0 # streamed events (0 for non-streaming) + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 300 + + @property + def is_streaming(self) -> bool: + return "text/event-stream" in (self.content_type or "") + + +def _hdr(resp: requests.Response, name: str) -> str | None: + value = resp.headers.get(name) + return value if isinstance(value, str) else None + + +def unwrap[R: BaseModel](result: Result[R]) -> R: + match result: + case Success(data=data): + return data + case _: + raise AssertionError(result) + + +def is_ok[R: BaseModel](result: Result[R]) -> bool: + match result: + case Success(): + return True + case _: + return False + + +def require_successful_call(result: StreamingResponse) -> None: + """A call that should have succeeded but didn't is a hard failure, never a skip: + if the proxy can't make a call it's expected to, the test must fail.""" + if result.ok: + return + pytest.fail( + f"upstream call failed (status {result.status_code}); body={result.body[:300]}" + ) + + +def _headers(headers: BaseModel) -> dict[str, str]: + dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) + return {key: str(value) for key, value in dumped.items()} + + +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()} + + +def _classify[R: BaseModel]( + resp: requests.Response, response_type: type[R] +) -> Result[R]: + if resp.status_code == 401: + return UnauthorizedError() + if resp.status_code == 429: + return RateLimitedError(body=resp.text) + if not resp.ok: + return UnknownApiError(status_code=resp.status_code, body=resp.text) + try: + return Success(data=response_type.model_validate(resp.json())) + except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value + return ValidationError(message=str(exc)) + + +def post[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.post( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + +def get[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.get( + str(url), + headers=_headers(headers), + params=params.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + +def delete[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.delete( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + +def probe( + url: URL, *, headers: BaseModel, params: BaseModel, timeout: float = 30.0 +) -> ProbeResult: + try: + resp = requests.get( + str(url), + headers=_headers(headers), + params=params.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return ProbeResult(status_code=-1, body=str(exc)) + return ProbeResult(status_code=resp.status_code, body=resp.text) + + +def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingResponse: + call_id = _hdr(resp, "x-litellm-call-id") + content_type = _hdr(resp, "content-type") + if not stream or not (200 <= resp.status_code < 300): + return StreamingResponse( + status_code=resp.status_code, + call_id=call_id, + content_type=content_type, + body=resp.text, + ) + lines = cast("Iterator[bytes]", resp.iter_lines()) + chunks = sum(1 for line in lines if line) + return StreamingResponse( + status_code=resp.status_code, + call_id=call_id, + content_type=content_type, + body="", + chunks=chunks, + ) + + +def send( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + timeout: float = 60.0, +) -> StreamingResponse: + """Raw POST returning the unparsed HTTP outcome: status, full body, and the + x-litellm-call-id header. For native/passthrough bodies and for calls judged by + status rather than a typed JSON model (e.g. a budget block is a non-2xx). With + ``stream=True`` the SSE body is consumed and its events counted instead.""" + try: + resp = requests.post( + str(url), + headers=_headers(headers), + params=_params(params), + json=json.model_dump(by_alias=True, exclude_none=True), + stream=stream, + timeout=timeout, + ) + except requests.RequestException as exc: + return StreamingResponse(status_code=-1, body=str(exc)) + return _streaming_outcome(resp, stream) + + +def stream( + url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0 +) -> StreamingResponse: + """Streaming (SSE) call: consumes the stream counting events, and captures the + x-litellm-call-id + content-type headers. Body is elided.""" + return send(url, headers=headers, json=json, stream=True, timeout=timeout) diff --git a/tests/e2e/gateway/litellm-config.yml b/tests/e2e/gateway/litellm-config.yml new file mode 100644 index 00000000000..f4ca48cfee0 --- /dev/null +++ b/tests/e2e/gateway/litellm-config.yml @@ -0,0 +1,170 @@ +# This default config file aims to support most popular model providers out of the box + +#In general, the model name used by the client will be the same as the ones from the provider (For example, you will use "anthropic.claude-3-5-sonnet-20240620-v1:0" when you're calling LiteLLM just like you would when calling Amazon Bedrock directly) +#In the case where there are model name conflicts, a prefix will be used (For example, the Azure and the openAI model names conflict, so when you are using Azure, you will use "azure/gpt-4o-realtime-preview-2024-10-01") + +#Some model providers require additional user-specific configuration (such as Azure which requires you to specify your own api_base with your resource name, and your api_version). +#In this case, the provider is commented out, and you should uncomment it and provide your specific info + +#For more detailed information about each provider, refer to the docs: https://docs.litellm.ai/docs/providers + +#If you are not interested in a particular provider, just remove it from your config.yaml, and redeploy, and it will no longer show up in your LiteLLM deployment + +#If a particular provider is not working, double check your .env file, and make sure you have provided a valid api key for that provider, and then redeploy + +#Full details on guardrails here: https://docs.litellm.ai/docs/proxy/guardrails/bedrock +general_settings: + store_prompts_in_spend_logs: true + master_key: os.environ/LITELLM_MASTER_KEY + proxy_batch_write_at: 60 + database_connection_pool_limit: 10 + # disable_error_logs: True + forward_client_headers_to_llm_api: false + maximum_spend_logs_retention_period: "60d" # GSE-13389: Cleanup logs older than 60 days + maximum_spend_logs_cleanup_cron: "0 1 * * *" # 01:00 UTC daily = 18:00 PDT + database_url: os.environ/DATABASE_URL + control_plane_url: os.environ/CONTROL_PLANE_URL + alerts: ["email"] + proxy_budget_rescheduler_min_time: 15 + proxy_budget_rescheduler_max_time: 20 + +# fallbacks: [{"gpt-4": ["anthropic.claude-3-5-sonnet-20240620-v1:0"]}] #Configure fallbacks for context window exeeded errors (In this example, we will fall back to Claude Sonnet if over 8000 tokens, which is gpt-4's limit) + # default_fallbacks: ["anthropic.claude-3-haiku-20240307-v1:0"] #Configure fallbacks for any error for every model (the above fallback configurations override this one) +# environment_variables: +# STORE_MODEL_IN_DB: 'True' +# LITELLM_LOG: "DEBUG" +litellm_settings: + drop_params: True + # Spend counters inherit this as their Redis TTL, so an idle counter goes cold and + # the next request reseeds it from the DB; kept short to exercise the cross-pod + # reseed path in test_spend_counter_reseed_e2e. Response-cache writes pass their own + # ttl and are unaffected. + default_redis_ttl: 20 + request_timeout: 600 + num_retries: 3 + json_logs: true + store_audit_logs: True + cache: true + cache_params: + type: redis + host: redis + port: 6379 + password: os.environ/REDIS_PASSWORD + namespace: litellm.caching + ttl: 16600 + # max_budget: 1000000000.0 # (float) sets max budget in dollars across the entire proxy across all API keys. Note, the budget does not apply to the master key. That is the only exception. + # budget_duration: 1mo # (str) frequency of budget reset - You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + # max_internal_user_budget: 1000000000.0 # (float) sets default budget in dollars for each internal user. (Doesn't apply to Admins. Doesn't apply to Teams. Doesn't apply to master key) + # internal_user_budget_duration: "1mo" # (str) frequency of budget reset - You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + # success_callback: ["s3_v2"] + # failure_callback: ["s3_v2"] + # service_callback: ["datadog"] + callbacks: ["arize_phoenix", "datadog", "smtp_email", "prometheus", "otel"] + require_auth_for_metrics_endpoint: false + #type: redis-semantic + #similarity_threshold: 0.8 # similarity threshold for semantic cache + #redis_semantic_cache_embedding_model: text-embedding-ada-002 # only works with text-embedding-ada-002 for now... https://github.com/BerriAI/litellm/issues/4001 + +router_settings: + routing_strategy: simple-shuffle + num_retries: 3 + allowed_fails: 5 + cooldown_time: 30 + # When gemini deployments are exhausted (provider 429 / auth), cross over to + # working models. Exercised by tests/e2e/router/test_rate_limiter.py. + fallbacks: + - gemini-2.5-flash: ["gpt-5.5", "claude-haiku-4-5"] + +#ttl: Optional[float] +#default_in_memory_ttl: Optional[float] +#default_in_redis_ttl: Optional[float] + +model_list: + - model_name: gpt-5.5 + litellm_params: + model: openai/gpt-5.5 + api_key: os.environ/OPENAI_API_KEY + + - model_name: claude-haiku-4-5 + litellm_params: + model: anthropic/claude-haiku-4-5 + api_key: os.environ/ANTHROPIC_API_KEY + + # Same underlying model via Vertex AI — distinct routing/auth path + # # (service-account JSON), so it gets its own model_name. + - model_name: gemini-2.5-flash-vertex + litellm_params: + model: vertex_ai/gemini-2.5-flash + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + vertex_credentials: os.environ/VERTEXAI_CREDENTIALS + + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + + # load balancing to a different deployment, if gemini gets rate limited. + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + + # Custom per-token pricing exercised by llm_translation/test_custom_pricing_e2e.py. + # Rates deliberately exceed canonical gemini-2.5-flash (input 3e-7 / output 2.5e-6) + # so an override that is ignored or under-applied reports spend at the base rate + # and fails that test. The test reads these same rates back from this file. + - model_name: custom-priced-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + input_cost_per_token: 0.00005 + output_cost_per_token: 0.0001 + + # embedding models + - model_name: openai-text-embedding-3-small + litellm_params: + model: openai/text-embedding-3-small + api_key: os.environ/OPENAI_API_KEY + + - model_name: gemini-2-embedding + litellm_params: + model: gemini/gemini-2-embedding + api_key: os.environ/GEMINI_API_KEY + + # realtime models + - model_name: openai-realtime + litellm_params: + model: openai/realtime-2 + api_key: os.environ/OPENAI_API_KEY + model_info: + mode: realtime + + +mcp_servers: + deepwiki_mcp: + url: "https://mcp.deepwiki.com/mcp" + auth_type: none + description: "just a test" + + atlassian: + url: "https://mcp.atlassian.com/v1/mcp" + auth_type: oauth2 + authorization_url: https://auth.atlassian.com/authorize + + +guardrails: + - guardrail_name: "presidio-pii" + litellm_params: + guardrail: presidio + mode: pre_call + presidio_analyzer_api_base: os.environ/PRESIDIO_ANALYZER_API_BASE + presidio_anonymizer_api_base: os.environ/PRESIDIO_ANONYMIZER_API_BASE + default_on: false + pii_entities_config: + EMAIL_ADDRESS: BLOCK + CREDIT_CARD: BLOCK + US_SSN: BLOCK + PHONE_NUMBER: BLOCK + + diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py new file mode 100644 index 00000000000..fdf2137584e --- /dev/null +++ b/tests/e2e/lifecycle.py @@ -0,0 +1,117 @@ +"""Lifecycle contract and resource cleanup for stateful e2e tests. + +Shared by every e2e suite under tests/e2e/. The proxy under test is +long-lived and never reset between tests, so anything a test creates (keys, +customers, teams, orgs, users, guardrails, budgets, ...) persists unless +explicitly deleted. Every check follows an init -> run -> teardown lifecycle; +teardown releases each resource init() created, even when run() raises. + +In pytest terms (see conftest.py): the `resources` fixture's setup is init(), +the test body is run(), and the fixture's teardown is teardown(). +""" + +from dataclasses import dataclass, field +from typing import Callable, List, Protocol, runtime_checkable + +from e2e_gateway import Gateway +from models import KeyGenerateBody + + +@runtime_checkable +class E2ECase(Protocol): + """A stateful e2e check run against a long-lived proxy. + + init() acquires resources, run() exercises behaviour and asserts, teardown() + releases everything init() created. teardown() must run even if init() fails + partway or run() raises. + """ + + def init(self) -> None: ... + + def run(self) -> None: ... + + def teardown(self) -> None: ... + + +def run_case(case: E2ECase) -> None: + """Drive a case through its lifecycle: init -> run -> teardown. + + teardown always runs - even when init() fails partway or run() raises (or + skips) - so resources the case already registered on the long-lived proxy are + released. init() is inside the try because cases register cleanups + progressively (e.g. create team, then user, then key), and a failure after + the first creation must still release what came before. + """ + try: + case.init() + case.run() + finally: + case.teardown() + + +@runtime_checkable +class ResourceClient(Protocol): + """Proxy operations the convenience creators use. Resource types without a + creator here are handled generically via ResourceManager.defer(). The Gateway + satisfies this.""" + + def generate_key(self, body: KeyGenerateBody) -> str: ... + + def delete_key(self, key: str) -> None: ... + + def delete_customers(self, user_ids: List[str]) -> None: ... + + +@runtime_checkable +class GatewayProvider(Protocol): + """Every suite's client exposes the shared Gateway, which the resources fixture + uses for cleanup. The client adds its own route methods on top.""" + + @property + def gateway(self) -> Gateway: ... + + +@dataclass +class ResourceManager: + """Registry of teardown actions for resources a test creates on the stateful + proxy. + + Not limited to any resource type: register a cleanup with ``defer()`` for a + key, customer, team, org, user, guardrail, budget, MCP server - anything with + a delete. The two most common resources have sugar (``key``, ``customer``); + everything else is ``resources.defer(lambda: client.delete_team(team_id))``. + + Cleanups run LIFO (so a resource is removed before whatever it depends on) and + best-effort (one failing cleanup never blocks the rest). + """ + + client: ResourceClient + _cleanups: List[Callable[[], None]] = field( + default_factory=list + ) # mutable-ok: append-only teardown registry + + def init(self) -> None: + """No global setup needed today; present for lifecycle symmetry.""" + return None + + def defer(self, cleanup: Callable[[], None]) -> None: + """Register a teardown action for any resource the test just created.""" + self._cleanups.append(cleanup) + + def key(self) -> str: + """Create an all-models virtual key; delete it on teardown.""" + key = self.client.generate_key(KeyGenerateBody(models=[])) + self.defer(lambda: self.client.delete_key(key)) + return key + + def customer(self, customer_id: str) -> str: + """Track an end-user id (from the `user` param); delete it on teardown.""" + self.defer(lambda: self.client.delete_customers([customer_id])) + return customer_id + + def teardown(self) -> None: + for cleanup in reversed(self._cleanups): + try: + cleanup() + except Exception: + pass # best-effort: a failed cleanup must not block the rest diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md new file mode 100644 index 00000000000..5e4a448857f --- /dev/null +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -0,0 +1,84 @@ +# LLM Translation Test Coverage Matrix + +Scope: the proxy's two translation surfaces, end to end against a live proxy. + +1. **Passthrough** - the client speaks the provider's NATIVE API (Gemini + `generateContent`, Anthropic `/v1/messages`); the proxy forwards it and still + logs a costed `SpendLogs` row (`call_type="pass_through_endpoint"`). Routes: + `/gemini`, `/anthropic`, `/vertex_ai`, `/openai`, `/bedrock`, `/cohere`, + `/mistral`, `/vllm`. +2. **Non-passthrough** - the client speaks OpenAI format + (`/chat/completions`, `/embeddings`); litellm translates to/from the provider. + +The two axes that must work in production for each: **passthrough vs +non-passthrough** and **streaming vs non-streaming**, with **cost logged** and +**tool calls** working in every cell. + +Companion: live suite `test_passthrough_e2e.py` (this directory). The +non-passthrough chat/embedding cells are exercised by `../spend_tracking/`. + +Levels: `live` real provider + proxy + SpendLogs row; `unit` mocked. +Status: `covered` / `partial` / `gap`. + +--- + +## Passthrough endpoints (native provider format) + +| Provider | Non-streaming | Streaming | Tool calls | Cost logged | Status | +|----------|---------------|-----------|------------|-------------|--------| +| Gemini (`/gemini/v1beta/models/{m}:generateContent` / `:streamGenerateContent`) | live | live | live | live | **covered** | +| Anthropic (`/anthropic/v1/messages`) | live | live | live | live | **covered** | +| Vertex AI (`/vertex_ai/...`) | - | - | - | - | gap (gcloud auth) | +| OpenAI / Bedrock / Cohere / Mistral / VLLM | - | - | - | - | gap | + +Each covered cell asserts: `call_type == "pass_through_endpoint"`, `spend > 0`, +`status == "success"`, correct `custom_llm_provider`/`model`, row correlated by the +`x-litellm-call-id` header. Gemini non-streaming also pins `request_tags` +propagation; streaming pins `chunks > 0` then a costed row; tool tests assert the +provider emitted a tool call (`functionCall` / `tool_use`) and it was costed. + +Cost on passthrough is computed in the success handler by transforming the native +response to a `ModelResponse` and calling `litellm.completion_cost()`; for +streaming, chunks are buffered and costed after the stream ends. This is the path +most likely to silently break and the one a mock can't prove works. + +## Non-passthrough endpoints (OpenAI-compatible translation) + +| Modality | Non-streaming | Streaming | Tool calls | Cost logged | Status | +|----------|---------------|-----------|------------|-------------|--------| +| Chat | live (spend suite) | live (spend suite) | gap | live | partial | +| Embeddings | live (spend suite) | n/a | n/a | live | covered | +| Responses / image / audio / rerank / realtime | - | - | - | - | gap | + +## This suite's files + +| Test | Cell | +|------|------| +| `test_gemini_passthrough_nonstreaming_logs_cost` | gemini native, non-stream, cost + tags | +| `test_gemini_passthrough_streaming_logs_cost` | gemini native, stream, cost | +| `test_gemini_passthrough_tool_call_logs_cost` | gemini native, tool call, cost | +| `test_anthropic_passthrough_nonstreaming_logs_cost` | anthropic native, non-stream, cost | +| `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | +| `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | + +## Gaps + +- Vertex / OpenAI / Bedrock / Cohere passthrough (same shape; add once the + provider credential is configured; Vertex is closest - route exists, auth stale). +- Non-passthrough tool calls over `/chat/completions` end to end with cost. +- Image / audio / rerank / responses / realtime translation + cost. +- Streaming cost-injection (`include_cost_in_streaming_usage`); passthrough on + client disconnect (partial-usage logging). + +## Adding a provider/modality + +Extend `PassthroughClient` with the native call (it inherits keys, cleanup, and +SpendLogs polling from `ProxyClient`), then add a test that calls it, +`require_successful_call(result)`, and `_costed_row(...)`. + +## Timing + +Passthrough spend is logged asynchronously after the response and lands on the +`proxy_batch_write_at` (~60s) cycle, so cost assertions poll +`/spend/logs?request_id=` to a deadline. Streaming cost is only +known after the stream is fully consumed. diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py new file mode 100644 index 00000000000..fbf008cf085 --- /dev/null +++ b/tests/e2e/llm_translation/conftest.py @@ -0,0 +1,15 @@ +"""LLM-translation suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared +Gateway, so the `resources` fixture cleans up keys this suite creates. +""" + +import pytest + +from passthrough_client import PassthroughClient, build_client + + +@pytest.fixture(scope="session") +def client() -> PassthroughClient: + return build_client() diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py new file mode 100644 index 00000000000..fff4064a328 --- /dev/null +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -0,0 +1,163 @@ +"""Client for LLM-translation e2e tests over the proxy's passthrough endpoints. + +A passthrough request is sent in the PROVIDER's native format (Gemini +generateContent, Anthropic /v1/messages) to the proxy, which forwards it to the +provider and still logs a SpendLogs row (call_type="pass_through_endpoint"). The +litellm virtual key is passed as the provider key; the proxy swaps in the real env +credential. SpendLogs.request_id == the x-litellm-call-id response header. The +native request models are co-located here because only this suite uses them. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from pydantic import BaseModel, Field + +from e2e_gateway import Gateway, build_gateway +from e2e_http import Headers, StreamingResponse +from models import ChatMessage + + +class JsonSchemaProperty(BaseModel): + type: str + + +class JsonSchema(BaseModel): + type: str + properties: dict[str, JsonSchemaProperty] + required: list[str] + + +class GeminiHeaders(Headers): + x_goog_api_key: str = Field(serialization_alias="x-goog-api-key") + content_type: str = Field( + default="application/json", serialization_alias="Content-Type" + ) + tags: str | None = None + + +class AnthropicHeaders(Headers): + x_api_key: str = Field(serialization_alias="x-api-key") + anthropic_version: str = Field( + default="2023-06-01", serialization_alias="anthropic-version" + ) + content_type: str = Field( + default="application/json", serialization_alias="Content-Type" + ) + tags: str | None = None + + +class AltSseParams(BaseModel): + alt: str = "sse" + + +class GeminiPart(BaseModel): + text: str + + +class GeminiContent(BaseModel): + role: str = "user" + parts: list[GeminiPart] + + +class GeminiFunctionDeclaration(BaseModel): + name: str + description: str + parameters: JsonSchema + + +class GeminiTool(BaseModel): + function_declarations: list[GeminiFunctionDeclaration] = Field( + serialization_alias="functionDeclarations" + ) + + +class GeminiGenerateBody(BaseModel): + contents: list[GeminiContent] + tools: list[GeminiTool] | None = None + + +class AnthropicTool(BaseModel): + name: str + description: str + input_schema: JsonSchema + + +class AnthropicMessageBody(BaseModel): + model: str + max_tokens: int + messages: list[ChatMessage] + tools: list[AnthropicTool] | None = None + stream: bool = False + + +def _tags_header(tags: list[str] | None) -> str | None: + return ",".join(tags) if tags else None + + +@dataclass(frozen=True, slots=True) +class PassthroughClient: + gateway: Gateway + + # ---- Gemini native passthrough (/gemini/v1beta/...) ----------------- + + def gemini_generate( + self, + key: str, + model: str, + text: str, + *, + tools: list[GeminiTool] | None = None, + tags: list[str] | None = None, + ) -> StreamingResponse: + return self.gateway.transport.send( + f"/gemini/v1beta/models/{model}:generateContent", + headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=text)])], tools=tools + ), + ) + + def gemini_stream( + self, key: str, model: str, text: str, *, tags: list[str] | None = None + ) -> StreamingResponse: + return self.gateway.transport.send( + f"/gemini/v1beta/models/{model}:streamGenerateContent", + headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=text)])] + ), + params=AltSseParams(), + stream=True, + ) + + # ---- Anthropic native passthrough (/anthropic/v1/messages) ---------- + + def anthropic_message( + self, + key: str, + model: str, + text: str, + *, + max_tokens: int = 64, + tools: list[AnthropicTool] | None = None, + stream: bool = False, + tags: list[str] | None = None, + ) -> StreamingResponse: + return self.gateway.transport.send( + "/anthropic/v1/messages", + headers=AnthropicHeaders(x_api_key=key, tags=_tags_header(tags)), + json=AnthropicMessageBody( + model=model, + max_tokens=max_tokens, + messages=[ChatMessage(role="user", content=text)], + tools=tools, + stream=stream, + ), + stream=stream, + ) + + +def build_client() -> PassthroughClient: + return PassthroughClient(gateway=build_gateway()) diff --git a/tests/e2e/llm_translation/test_custom_pricing_e2e.py b/tests/e2e/llm_translation/test_custom_pricing_e2e.py new file mode 100644 index 00000000000..58faab61aac --- /dev/null +++ b/tests/e2e/llm_translation/test_custom_pricing_e2e.py @@ -0,0 +1,212 @@ +"""Live e2e: a model's custom per-token pricing is loaded, billed, and isolated. + +The gateway config declares ``custom-priced-flash`` (gemini-2.5-flash underneath) +with input/output rates deliberately far above the canonical gemini price, read +back here from the same config file. Three behaviors are checked independently: + +- billing: a real call's logged cost breakdown charges input and output tokens at + the custom rates, each component checked separately (a base-rate bill lands + ~100x lower; a swapped input/output rate passes a total-only check but not this) +- reporting: /model/info surfaces those rates for the model +- isolation: gemini-2.5-flash shares the same underlying gemini/gemini-2.5-flash + but sets no override, so it must keep its own price; an override that leaks into + the shared cost map misprices it. A regression that reintroduces that leak makes + the sibling's rate match the custom one and fails the isolation check. +""" + +import time +from dataclasses import dataclass +from pathlib import Path + +import pytest +import yaml +from pydantic import BaseModel, RootModel + +from e2e_config import unique_marker +from e2e_http import Success, unwrap +from models import ChatBody, ChatMessage, CustomPricing, ModelInfoEntry, SpendLogsParams +from passthrough_client import PassthroughClient + +pytestmark = pytest.mark.e2e + +CUSTOM_MODEL = "custom-priced-flash" +BASE_MODEL = "gemini-2.5-flash" +CONFIG_PATH = Path(__file__).resolve().parents[1] / "gateway" / "litellm-config.yml" + + +@dataclass(frozen=True, slots=True) +class _Rates: + input_per_token: float + output_per_token: float + + +class _ConfiguredModel(BaseModel): + model_name: str + litellm_params: CustomPricing + + +class _GatewayConfig(BaseModel): + model_list: list[_ConfiguredModel] + + +class _CostBreakdown(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + + +class _RowMetadata(BaseModel): + cost_breakdown: _CostBreakdown | None = None + + +class _SpendRow(BaseModel): + request_id: str | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + metadata: _RowMetadata | None = None + + +class _SpendRows(RootModel[list[_SpendRow]]): + pass + + +def _approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +def _configured_pricing(model_name: str) -> _Rates: + """The custom rates declared for `model_name` in the gateway config the proxy + runs with - the source of truth the billed and reported prices are checked + against.""" + config = _GatewayConfig.model_validate(yaml.safe_load(CONFIG_PATH.read_text())) + for entry in config.model_list: + if entry.model_name == model_name: + pricing = entry.litellm_params + assert pricing.input_cost_per_token and pricing.output_cost_per_token, ( + f"{model_name} declares no custom per-token rates in {CONFIG_PATH.name}" + ) + return _Rates(pricing.input_cost_per_token, pricing.output_cost_per_token) + pytest.fail(f"{model_name} not found in {CONFIG_PATH.name}") + + +def _model_info_entry( + entries: list[ModelInfoEntry], model_name: str +) -> ModelInfoEntry: + for entry in entries: + if entry.model_name == model_name: + return entry + pytest.fail(f"{model_name} absent from /model/info; the override did not load") + + +def _poll_breakdown_row( + client: PassthroughClient, key: str, response_id: str | None +) -> _SpendRow: + """Poll /spend/logs until the call's row lands with a cost breakdown (rows + flush ~60s behind the call via proxy_batch_write_at).""" + deadline = time.monotonic() + client.gateway.poll_timeout + while time.monotonic() < deadline: + result = client.gateway.transport.get( + "/spend/logs", + headers=client.gateway.transport.master, + params=SpendLogsParams(api_key=key), + response_type=_SpendRows, + ) + match result: + case Success(data=data): + rows = data.root + case _: + rows = [] + priced = [ + row + for row in rows + if row.metadata + and row.metadata.cost_breakdown + and row.metadata.cost_breakdown.input_cost is not None + ] + for row in priced: + if response_id and row.request_id == response_id: + return row + if priced and response_id is None: + return priced[0] + time.sleep(client.gateway.poll_interval) + pytest.fail("no spend row with a cost breakdown landed before the deadline") + + +def test_custom_pricing_is_billed_at_configured_rate( + client: PassthroughClient, scoped_key: str +) -> None: + rates = _configured_pricing(CUSTOM_MODEL) + + chat = unwrap( + client.gateway.chat( + scoped_key, + ChatBody( + model=CUSTOM_MODEL, + messages=[ + ChatMessage( + role="user", content=f"reply with one word {unique_marker()}" + ) + ], + max_tokens=16, + ), + ) + ) + + row = _poll_breakdown_row(client, scoped_key, chat.id) + assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll + breakdown = row.metadata.cost_breakdown + + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + assert prompt > 0 and completion > 0, f"call tokens not logged on the row: {row}" + + input_cost = breakdown.input_cost + output_cost = breakdown.output_cost + assert input_cost is not None and output_cost is not None, ( + f"row cost breakdown missing input/output cost: {breakdown}" + ) + assert _approx_equal(input_cost, prompt * rates.input_per_token), ( + f"input_cost {input_cost} != {prompt} tokens * {rates.input_per_token} " + f"= {prompt * rates.input_per_token}" + ) + assert _approx_equal(output_cost, completion * rates.output_per_token), ( + f"output_cost {output_cost} != {completion} tokens * {rates.output_per_token} " + f"= {completion * rates.output_per_token}" + ) + + +def test_model_info_reports_custom_pricing(client: PassthroughClient) -> None: + rates = _configured_pricing(CUSTOM_MODEL) + entry = _model_info_entry(client.gateway.model_info(), CUSTOM_MODEL) + + assert entry.litellm_params.input_cost_per_token == rates.input_per_token, ( + f"/model/info litellm_params input rate " + f"{entry.litellm_params.input_cost_per_token} != configured " + f"{rates.input_per_token}" + ) + assert entry.litellm_params.output_cost_per_token == rates.output_per_token, ( + f"/model/info litellm_params output rate " + f"{entry.litellm_params.output_cost_per_token} != configured " + f"{rates.output_per_token}" + ) + + +def test_custom_pricing_is_isolated_from_sibling_deployment( + client: PassthroughClient, +) -> None: + entries = {entry.model_name: entry for entry in client.gateway.model_info()} + custom = entries.get(CUSTOM_MODEL) + base = entries.get(BASE_MODEL) + assert custom is not None, f"{CUSTOM_MODEL} absent from /model/info" + assert base is not None, f"{BASE_MODEL} absent from /model/info" + + # custom-priced-flash overrides pricing; gemini-2.5-flash shares the same + # underlying gemini/gemini-2.5-flash but sets no override, so it must keep its + # own price. Equal rates mean the override leaked into the shared cost map. + assert ( + base.model_info.input_cost_per_token != custom.model_info.input_cost_per_token + ), ( + f"{BASE_MODEL} input rate {base.model_info.input_cost_per_token} matches " + f"{CUSTOM_MODEL}'s override {custom.model_info.input_cost_per_token}; " + f"per-deployment custom pricing is not isolated" + ) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py new file mode 100644 index 00000000000..37d55c665b3 --- /dev/null +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -0,0 +1,159 @@ +"""Live e2e for LLM-translation passthrough endpoints. + +Each test sends a NATIVE provider request through the proxy's passthrough route +and verifies the proxy still logged a costed SpendLogs row +(call_type="pass_through_endpoint"), correlated by the x-litellm-call-id header. + +Covered: gemini ("gemini-2.5-flash") + anthropic ("claude-haiku-4-5"), streaming + +non-streaming, plus native tool calls. See LLM_TRANSLATION_COVERAGE_MATRIX.md. + +A passthrough call returning non-2xx fails hard (never a skip); once it returns +2xx, a missing or zero-cost SpendLogs row fails too. +""" + +import pytest + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from models import SpendLogRow +from passthrough_client import ( + AnthropicTool, + GeminiFunctionDeclaration, + GeminiTool, + JsonSchema, + JsonSchemaProperty, + PassthroughClient, +) + +pytestmark = pytest.mark.e2e + + +def _fetch_cost_breakdown(client: PassthroughClient, result: StreamingResponse) -> SpendLogRow: + """The passthrough call's logged row, polled until it carries a cost. + + Asserts (not skips) that a 2xx passthrough call produced a costed row - the + whole point of passthrough spend tracking. + """ + assert result.call_id, "passthrough response had no x-litellm-call-id header" + rows = client.gateway.poll_logs_for_request_id( + result.call_id, + predicate=lambda rs: (rs[0].spend or 0) > 0, + ) + assert rows, f"no SpendLogs row for passthrough call_id {result.call_id}" + row = rows[0] + assert row.call_type == "pass_through_endpoint" + assert (row.spend or 0) > 0, f"passthrough call was not costed: {row}" + assert row.status == "success" + return row + + +# ---- Gemini passthrough ------------------------------------------------ + + +def test_gemini_passthrough_nonstreaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + tag = f"e2e-passthrough-{unique_marker()}" + result = client.gemini_generate( + scoped_key, "gemini-2.5-flash", "Say hello in one word", tags=[tag, "gemini"] + ) + require_successful_call(result) + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "gemini" + assert "gemini" in (row.model or "") + assert tag in (row.request_tags or []), f"tags not logged: {row.request_tags}" + + +def test_gemini_passthrough_streaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.gemini_stream(scoped_key, "gemini-2.5-flash", "Count to five") + require_successful_call(result) + assert result.chunks > 0, "streaming passthrough produced no events" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "gemini" + + +def test_gemini_passthrough_tool_call_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.gemini_generate( + scoped_key, + "gemini-2.5-flash", + "What is the weather in Paris? Use the get_weather tool.", + tools=[ + GeminiTool( + function_declarations=[ + GeminiFunctionDeclaration( + name="get_weather", + description="Get the weather for a city", + parameters=JsonSchema( + type="object", + properties={"city": JsonSchemaProperty(type="string")}, + required=["city"], + ), + ) + ] + ) + ], + ) + require_successful_call(result) + assert "functionCall" in result.body, "gemini did not emit a tool call" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "gemini" + + +# ---- Anthropic passthrough --------------------------------------------- + + +def test_anthropic_passthrough_nonstreaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.anthropic_message(scoped_key, "claude-haiku-4-5", "Say hello") + require_successful_call(result) + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "anthropic" + assert "claude" in (row.model or "") + + +def test_anthropic_passthrough_streaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.anthropic_message( + scoped_key, "claude-haiku-4-5", "Count to five", stream=True + ) + require_successful_call(result) + assert result.chunks > 0, "streaming passthrough produced no events" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "anthropic" + + +def test_anthropic_passthrough_tool_call_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.anthropic_message( + scoped_key, + "claude-haiku-4-5", + "What is the weather in Paris? Use the get_weather tool.", + tools=[ + AnthropicTool( + name="get_weather", + description="Get the weather for a city", + input_schema=JsonSchema( + type="object", + properties={"city": JsonSchemaProperty(type="string")}, + required=["city"], + ), + ) + ], + ) + require_successful_call(result) + assert "tool_use" in result.body, "anthropic did not emit a tool call" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "anthropic" diff --git a/tests/e2e/models.py b/tests/e2e/models.py new file mode 100644 index 00000000000..fbeb3d44fa5 --- /dev/null +++ b/tests/e2e/models.py @@ -0,0 +1,240 @@ +"""Shared pydantic request/response models for the e2e gateway. + +Only the fields the tests read are modelled; pydantic ignores the rest, so a +response validates without mirroring every proxy field. No untyped dicts. +""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, RootModel + +# ---------- keys ---------- + + +class ModelBudgetEntry(BaseModel): + budget_limit: float + time_period: str + + +class BudgetWindow(BaseModel): + budget_duration: str + max_budget: float + + +class KeyGenerateBody(BaseModel): + models: list[str] = [] + duration: str | None = None + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + user_id: str | None = None + team_id: str | None = None + budget_id: str | None = None + model_max_budget: dict[str, ModelBudgetEntry] | None = None + budget_limits: list[BudgetWindow] | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + + +class KeyGenerateResponse(BaseModel): + key: str + + +class KeyDeleteBody(BaseModel): + keys: list[str] + + +class KeyInfoParams(BaseModel): + key: str + + +class LiteLLMBudgetTable(BaseModel): + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: str | None = None + + +class KeyInfo(BaseModel): + spend: float | None = None + max_budget: float | None = None + budget_reset_at: str | None = None + budget_id: str | None = None + litellm_budget_table: LiteLLMBudgetTable | None = None + + +class KeyInfoResponse(BaseModel): + info: KeyInfo + + +# ---------- customers ---------- + + +class CustomerDeleteBody(BaseModel): + user_ids: list[str] + + +# ---------- chat / embeddings ---------- + + +class ChatMetadata(BaseModel): + tags: list[str] | None = None + + +class ChatMessage(BaseModel): + role: str + content: str + + +class ChatBody(BaseModel): + model: str + messages: list[ChatMessage] + stream: bool = False + max_tokens: int | None = None + user: str | None = None + metadata: ChatMetadata | None = None + + +class OutMessage(BaseModel): + content: str | None = None + + +class ChatChoice(BaseModel): + message: OutMessage | None = None + + +class Usage(BaseModel): + prompt_tokens: int | None = None + completion_tokens: int | None = None + total_tokens: int | None = None + + +class ChatResponse(BaseModel): + id: str | None = None + model: str | None = None + choices: list[ChatChoice] = [] + usage: Usage | None = None + + +class EmbedBody(BaseModel): + model: str + input: str + + +class EmbedResponse(BaseModel): + model: str | None = None + + +# ---------- spend logs ---------- + + +class SpendLogRow(BaseModel): + request_id: str | None = None + model: str | None = None + spend: float | None = None + status: str | None = None + cache_hit: str | None = None + call_type: str | None = None + custom_llm_provider: str | None = None + team_id: str | None = None + user: str | None = None + end_user: str | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + total_tokens: int | None = None + request_tags: list[str] | None = None + + +class SpendLogs(RootModel[list[SpendLogRow]]): + pass + + +class SpendLogsParams(BaseModel): + request_id: str | None = None + api_key: str | None = None + + +# ---------- spend calculate ---------- + + +class SpendCalculateBody(BaseModel): + model: str + messages: list[ChatMessage] + + +class SpendCalculateResponse(BaseModel): + cost: float + + +# ---------- route probing ---------- + + +class DateRangeParams(BaseModel): + start_date: str + end_date: str + + +class RouteSpec(RootModel[dict[str, object]]): + """One /openapi.json path entry: a map of HTTP method -> operation. Only the + method names are read, so the operation specs stay opaque.""" + + @property + def methods(self) -> frozenset[str]: + return frozenset(method.lower() for method in self.root) + + +class OpenAPISchema(BaseModel): + paths: dict[str, RouteSpec] = {} + + +# ---------- model info / custom pricing ---------- + + +class CustomPricing(BaseModel): + """The per-token custom-pricing fields a deployment can override in + litellm_params - the token-cost subset of litellm's CustomPricingLiteLLMParams + the proxy applies to chat spend. All optional: a config sets only what it + overrides, and /model/info echoes the rates the proxy resolved.""" + + model_config = ConfigDict(extra="ignore") + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + + def overrides(self) -> dict[str, float]: + """The rates actually declared (non-null) - e.g. those a config.yml sets.""" + declared = { + "input_cost_per_token": self.input_cost_per_token, + "output_cost_per_token": self.output_cost_per_token, + "cache_read_input_token_cost": self.cache_read_input_token_cost, + "cache_creation_input_token_cost": self.cache_creation_input_token_cost, + } + return {field: rate for field, rate in declared.items() if rate is not None} + + def token_cost(self, prompt_tokens: int, completion_tokens: int) -> float: + """Spend for a fresh (uncached) call under these rates: the proxy's + custom-pricing formula (prompt * input + completion * output).""" + assert ( + self.input_cost_per_token is not None + and self.output_cost_per_token is not None + ), "custom pricing has no per-token rates" + return ( + prompt_tokens * self.input_cost_per_token + + completion_tokens * self.output_cost_per_token + ) + + +class ModelInfoEntry(BaseModel): + """One /model/info row. `litellm_params` is the configured deployment (carries + any custom-pricing override); `model_info` is the price the proxy resolved for + it - the override merged over the cost-map defaults.""" + + model_config = ConfigDict(protected_namespaces=()) + model_name: str + litellm_params: CustomPricing = CustomPricing() + model_info: CustomPricing = CustomPricing() + + +class ModelInfoResponse(BaseModel): + data: list[ModelInfoEntry] = [] diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini new file mode 100644 index 00000000000..7799f6b16a2 --- /dev/null +++ b/tests/e2e/pytest.ini @@ -0,0 +1,7 @@ +[pytest] +# Config when any e2e suite under tests/e2e/ is run directly, e.g. +# uv run pytest tests/e2e/spend_tracking/ -v +# The e2e marker is also registered in conftest.py for runs rooted elsewhere. +addopts = --strict-markers --strict-config +markers = + e2e: live test that requires a running proxy and real provider keys diff --git a/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md b/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md new file mode 100644 index 00000000000..53c4d4ace83 --- /dev/null +++ b/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md @@ -0,0 +1,78 @@ +# Spend Tracking Test Coverage Matrix + +Scope: every distinct spend-tracking code path, mapped to the test that exercises +it and the level it runs at. Highlights where a live e2e check is the only thing +that would catch a regression. + +Companion: live suite `test_spend_tracking_e2e.py` + route breadth +`test_spend_routes.py` (this directory). Offline regression suite: +`tests/test_litellm/proxy/spend_tracking/`. Reference PR: BerriAI/litellm#29956. + +Levels: `unit` mocked; `integration` real DB/cost-map; `live` real provider + +proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. + +--- + +## SpendLogs row construction (`spend_tracking_utils.get_logging_payload`) + +| Path | Existing | Level | Status | Live e2e | +|------|----------|-------|--------|----------| +| `_get_status_for_spend_log` | `test_spend_tracking_utils.py` | unit | covered | yes (status read off the row) | +| cache-hit `request_id` suffix | `test_spend_tracking_utils.py` | unit | covered | yes (`test_cache_hit_is_zero_cost_and_suffixed`) | +| failure status + zero spend | `test_spend_tracking_utils.py` | unit | covered | no (live failure logging is non-deterministic across providers) | +| per-model / per-provider attribution | `test_spend_tracking_utils.py` | unit | covered | yes (`test_each_model_on_a_shared_key_gets_its_own_row`) | +| field population (model/tokens/api_key/team/org) | `test_spend_tracking_utils.py` | unit | partial | yes (asserts real values) | +| `request_tags` propagation | `test_db_spend_update_writer.py` | unit | partial | yes (`test_request_tags_round_trip`) | +| `end_user` attribution | unit | unit | partial | yes (`test_end_user_spend_attributed_on_row`) | + +## Cost calculation by modality + +| Modality | Existing | Status | Live e2e | +|----------|----------|--------|----------| +| Chat (non-stream) | `test_cost_calculator.py`, `local_testing/test_completion_cost.py` | covered | yes (`test_chat_completion_writes_nonzero_spend_row`) | +| Chat (streaming) | `test_streaming_interrupt_spend_tracking.py` | partial | yes (`test_streaming_chat_completion_tracks_spend`) | +| Embedding | `test_cost_calculator.py` (#29956) | partial | yes (`test_embedding_writes_nonzero_spend_row`) | +| Pass-through (gemini/anthropic) | `pass_through_tests/*.test.js` + `llm_translation/` suite | covered | yes (llm_translation suite) | +| Image / audio / rerank / responses / realtime | per-provider unit cost tests | partial/gap | gap | + +## Entity spend aggregation + +| Entity | Existing | Status | Live e2e | +|--------|----------|--------|----------| +| API key | `test_db_spend_update_writer.py`, `test_spend_counters.py` | covered | yes (`test_key_spend_equals_sum_of_logs`) | +| Tag | `test_update_daily_tag_spend.py` | partial | yes (`test_request_tags_round_trip`, propagation only) | +| End-user | `test_proxy_update_spend.py` | covered | yes | +| Spend == sum(logs) consistency | none | gap | yes (key aggregate == sum of rows) | + +## Spend read endpoints (verification surface) + +| Endpoint | Existing | Status | Live e2e | +|----------|----------|--------|----------| +| `/spend/logs` (request_id / api_key) | `test_spend_management_endpoints.py` | covered | yes (primary read path; `test_spend_logs_endpoint_returns_spend` asserts 200 + spend, never 5xx) | +| `/spend/calculate` | `local_testing/test_spend_calculate_endpoint.py` | covered | yes (`test_spend_calculate_returns_nonzero_cost`) | +| `/spend/tags` | `test_spend_management_endpoints.py` | partial | yes (`test_spend_routes.py` route probe) | +| whole spend GET surface (22 routes) | unit per-handler | partial | yes (`test_spend_routes.py` probes each for 404/5xx) | + +## What this suite pins + +| Test | Invariant | +|------|-----------| +| `test_chat_completion_writes_nonzero_spend_row` | nonzero cost, token arithmetic, status, row findable by `response.id` | +| `test_streaming_chat_completion_tracks_spend` | streamed responses still costed | +| `test_embedding_writes_nonzero_spend_row` | embedding cost, `completion_tokens == 0` | +| `test_cache_hit_is_zero_cost_and_suffixed` | cache hits not double-charged; `_cache_hit` suffix | +| `test_key_spend_equals_sum_of_logs` | key aggregate == sum of rows | +| `test_request_tags_round_trip` | tags persist onto the row | +| `test_end_user_spend_attributed_on_row` | `end_user` attributed + costed | +| `test_each_model_on_a_shared_key_gets_its_own_row` | per-model/provider rows, correct model + cost, distinct request_ids matching response id | +| `test_spend_calculate_returns_nonzero_cost` | cost-map smoke (no batch wait) | +| `test_spend_logs_endpoint_returns_spend` | `/spend/logs` returns 200 + the key's spend, never a 5xx (intermittent-500 regression) | +| `test_spend_routes.py` (23) | no spend route 404s or 5xxs | + +## Design + timing + +`proxy_batch_write_at` (~60s) means rows land late; every read polls to a deadline. +Fresh scoped key per test (isolation, xdist-safe, cleaned up). Assert invariants +(`spend > 0`, `total == prompt + completion`, aggregate == sum), not literal +$/token values, so pricing drift is not a failure. Skip on environment (no proxy / +no provider key), fail on behavior (a real 2xx call with a wrong/missing row). diff --git a/tests/e2e/spend_tracking/conftest.py b/tests/e2e/spend_tracking/conftest.py new file mode 100644 index 00000000000..1d01ab3d17a --- /dev/null +++ b/tests/e2e/spend_tracking/conftest.py @@ -0,0 +1,16 @@ +"""Spend-tracking suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. SpendClient exposes the shared Gateway +(GatewayProvider), so the `resources` fixture cleans up keys and customers this +suite creates. +""" + +import pytest + +from spend_e2e_client import SpendClient, build_client + + +@pytest.fixture(scope="session") +def client() -> SpendClient: + return build_client() diff --git a/tests/e2e/spend_tracking/spend_e2e_client.py b/tests/e2e/spend_tracking/spend_e2e_client.py new file mode 100644 index 00000000000..d749d69f1a4 --- /dev/null +++ b/tests/e2e/spend_tracking/spend_e2e_client.py @@ -0,0 +1,167 @@ +"""Spend-tracking e2e client: a Gateway plus the spend-specific read endpoints. + +Generic proxy operations (keys, customers, chat/embed, route probing, SpendLogs +polling) come from the shared Gateway, DI'd in (composition, not inheritance). +This client adds only the spend surface: /spend/calculate, key-spend +polling, and the route probes the breadth test uses. + +Re-exports unwrap / is_ok / unique_marker / SpendLogRow so the tests import their +helpers from one place. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable +from dataclasses import dataclass + +from e2e_config import unique_marker +from e2e_http import ( + NoBody, + ProbeResult, + Result, + StreamingResponse, + is_ok, + unwrap, +) +from e2e_gateway import Gateway, build_gateway +from models import ( + ChatBody, + ChatMessage, + ChatMetadata, + ChatResponse, + DateRangeParams, + EmbedBody, + EmbedResponse, + OpenAPISchema, + SpendCalculateBody, + SpendCalculateResponse, + SpendLogRow, +) + +__all__ = [ + "SpendClient", + "build_client", + "reset_spend_logs", + "unique_marker", + "unwrap", + "is_ok", + "SpendLogRow", + "ProbeResult", +] + + +def reset_spend_logs() -> None: + """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes + spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses + DATABASE_URL (default: the local docker postgres on its mapped host port; note + the in-container `@db` host isn't resolvable from the host, so default to + localhost). + """ + import psycopg + + url = os.environ.get( + "DATABASE_URL", + "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm", + ) + with psycopg.connect(url) as conn: + _ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"') + + +def _chat_body( + model: str, + content: str, + *, + max_tokens: int | None = None, + tags: list[str] | None = None, + user: str | None = None, + stream: bool = False, +) -> ChatBody: + return ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + stream=stream, + user=user, + metadata=ChatMetadata(tags=tags) if tags else None, + ) + + +@dataclass(frozen=True, slots=True) +class SpendClient: + gateway: Gateway + + def chat( + self, + key: str, + model: str, + content: str, + *, + max_tokens: int | None = None, + tags: list[str] | None = None, + user: str | None = None, + ) -> Result[ChatResponse]: + return self.gateway.chat( + key, _chat_body(model, content, max_tokens=max_tokens, tags=tags, user=user) + ) + + def chat_stream( + self, key: str, model: str, content: str, *, max_tokens: int | None = None + ) -> StreamingResponse: + return self.gateway.chat_stream( + key, _chat_body(model, content, max_tokens=max_tokens, stream=True) + ) + + def embed(self, key: str, model: str, content: str) -> Result[EmbedResponse]: + return self.gateway.embed(key, EmbedBody(model=model, input=content)) + + def poll_logs_for_key( + self, + key: str, + *, + min_rows: int = 1, + predicate: Callable[[list[SpendLogRow]], bool] | None = None, + ) -> list[SpendLogRow]: + return self.gateway.poll_logs_for_key( + key, min_rows=min_rows, predicate=predicate + ) + + def calculate_spend(self, model: str, content: str) -> float: + return unwrap( + self.gateway.transport.post( + "/spend/calculate", + headers=self.gateway.transport.master, + json=SpendCalculateBody( + model=model, messages=[ChatMessage(role="user", content=content)] + ), + response_type=SpendCalculateResponse, + ) + ).cost + + def poll_key_spend(self, key: str, *, minimum: float = 0.0) -> float: + deadline = time.monotonic() + self.gateway.poll_timeout + spend = 0.0 + while time.monotonic() < deadline: + spend = self.gateway.key_info(key).spend or 0.0 + if spend > minimum: + return spend + time.sleep(self.gateway.poll_interval) + return spend + + def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: + return self.gateway.transport.probe(path, params=params) + + def openapi(self) -> OpenAPISchema: + return unwrap( + self.gateway.transport.get( + "/openapi.json", + headers=self.gateway.transport.master, + params=NoBody(), + response_type=OpenAPISchema, + ) + ) + + +def build_client() -> SpendClient: + return SpendClient(gateway=build_gateway()) diff --git a/tests/e2e/spend_tracking/test_spend_routes.py b/tests/e2e/spend_tracking/test_spend_routes.py new file mode 100644 index 00000000000..e3c96a4d578 --- /dev/null +++ b/tests/e2e/spend_tracking/test_spend_routes.py @@ -0,0 +1,96 @@ +"""Breadth check: query every route on the spend read surface and show what it +returns. + +Spend tracking sprawls across many routes (model-cost / key / user / team / org / +customer aggregation, tags, and activity reports). Most are served with +`include_in_schema=False`, so they do NOT appear in `/openapi.json` - discovery +from the schema alone misses ~70% of the surface. So we probe a curated, verified +list directly, plus any spend route the schema does list (to auto-catch new ones). + +Each probe captures status AND body, so a failure shows the proxy's actual error +(a 500 traceback, a 404 meaning the route was removed) rather than a bare code. +Run with `-rA` (or `-s`) to print every route's response, not just failures. + +Healthy == route exists (not 404) and handler did not crash (not 5xx). A 4xx +(missing params / auth nuance) still means the route is wired and ran. Cheap and +fast: no batch-write wait, no provider calls. +""" + +from datetime import datetime, timedelta, timezone + +import pytest + +from models import DateRangeParams +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +# Verified present and responsive on a live proxy. One per row of the spend +# surface: key / user / team / org / customer aggregation, model-cost, tags, +# activity. +SPEND_ROUTES = ( + "/spend/keys", + "/spend/users", + "/spend/tags", + "/spend/logs", + "/spend/logs/ui", + "/global/spend", + "/global/spend/keys", + "/global/spend/teams", + "/global/spend/models", + "/global/spend/provider", + "/global/spend/report", + "/global/spend/tags", + "/global/spend/logs", + "/global/spend/all_tag_names", + "/global/activity", + "/global/activity/model", + "/global/activity/exceptions", + "/key/list", + "/user/list", + "/team/list", + "/organization/list", + "/customer/list", +) + +_SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity") + + +def _date_range() -> DateRangeParams: + # Satisfies date-required endpoints (report/activity/provider); ignored elsewhere. + end = datetime.now(timezone.utc).date() + start = end - timedelta(days=1) + return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) + + +@pytest.mark.parametrize("route", SPEND_ROUTES) +def test_spend_route_responsive(client: SpendClient, route: str) -> None: + result = client.probe(route, params=_date_range()) + print(f"{route} -> {result.status_code}\n{result.body[:600]}") + assert result.healthy, f"{route} -> {result.status_code}\n{result.body[:600]}" + + +def test_schema_listed_spend_routes_are_responsive(client: SpendClient) -> None: + """Probe any spend GET route the schema lists that isn't in SPEND_ROUTES.""" + schema = client.openapi() + assert schema.paths, "/openapi.json had no paths" + + discovered = [ + path + for path, spec in schema.paths.items() + if "get" in spec.methods + and "{" not in path + and any(path.startswith(prefix) for prefix in _SPEND_PREFIXES) + ] + extras = [path for path in discovered if path not in SPEND_ROUTES] + + params = _date_range() + results = [(path, client.probe(path, params=params)) for path in extras] + for path, result in results: + print(f"{path} -> {result.status_code}") + offenders = [ + f"{path} -> {result.status_code}\n{result.body[:600]}" + for path, result in results + if not result.healthy + ] + assert not offenders, "non-responsive schema spend routes:\n" + "\n".join(offenders) diff --git a/tests/e2e/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py new file mode 100644 index 00000000000..8c9e913b10f --- /dev/null +++ b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py @@ -0,0 +1,328 @@ +"""Live end-to-end spend-tracking tests against a running proxy. + +Run against a proxy started with the gateway config. Coverage rationale: +SPEND_TRACKING_COVERAGE_MATRIX.md. + +Model names are literals from that config: chat tests hit "gemini-2.5-flash", +embedding tests hit "openai-text-embedding-3-small". + +Every test: fresh scoped key (isolation) -> real provider call -> unwrap (hard +fail if the proxy couldn't make a call it should) -> poll /spend/logs to a +deadline (rows land ~60s later via proxy_batch_write_at) -> assert invariants on +the real row (spend, token arithmetic, status, cache). + +Assertions target invariants, not literals: a regression in the spend pipeline +fails the test; a pricing or token-count drift does not. +""" + +import time +from collections.abc import Callable + +import pytest + +from e2e_http import Success +from lifecycle import ResourceManager +from models import SpendLogs, SpendLogsParams +from spend_e2e_client import SpendClient, SpendLogRow, unique_marker, unwrap + +pytestmark = pytest.mark.e2e + + +def _approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +def _summarize(rows: list[SpendLogRow]) -> list[dict[str, object]]: + fields = { + "request_id", + "model", + "spend", + "status", + "cache_hit", + "prompt_tokens", + "completion_tokens", + "total_tokens", + } + return [row.model_dump(include=fields) for row in rows] + + +def _require_row( + rows: list[SpendLogRow], predicate: Callable[[SpendLogRow], bool], what: str +) -> SpendLogRow: + matches = [r for r in rows if predicate(r)] + assert matches, ( + f"no SpendLogs row {what} after polling; saw {len(rows)} row(s): " + f"{_summarize(rows)}" + ) + return matches[0] + + +def test_chat_completion_writes_nonzero_spend_row( + client: SpendClient, scoped_key: str +) -> None: + chat = unwrap( + client.chat( + scoped_key, + "gemini-2.5-flash", + f"reply with one word {unique_marker()}", + max_tokens=16, + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.status == "success" for r in rs) + ) + row = _require_row(rows, lambda r: r.status == "success", "for the chat call") + + assert (row.spend or 0) > 0, f"chat row should cost > 0: {_summarize(rows)}" + assert row.status == "success" + assert row.cache_hit != "True", "fresh call must not be a cache hit" + assert "gemini-2.5-flash" in (row.model or "") + + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + total = row.total_tokens or 0 + assert prompt > 0 and completion > 0 + assert total == prompt + completion, f"token arithmetic broken: {_summarize(rows)}" + + if chat.id: + assert any( + r.request_id == chat.id for r in rows + ), f"row request_id != client response.id ({chat.id})" + + +def test_streaming_chat_completion_tracks_spend( + client: SpendClient, scoped_key: str +) -> None: + result = client.chat_stream( + scoped_key, + "gemini-2.5-flash", + f"count to three {unique_marker()}", + max_tokens=64, + ) + assert ( + result.ok + ), f"stream failed (status {result.status_code}): {result.body[:300]}" + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + row = _require_row( + rows, lambda r: (r.spend or 0) > 0, "with nonzero spend for the stream" + ) + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + assert ( + prompt > 0 and completion > 0 + ), f"streaming tokens not tracked: {_summarize(rows)}" + assert (row.total_tokens or 0) == prompt + completion + + +def test_embedding_writes_nonzero_spend_row( + client: SpendClient, scoped_key: str +) -> None: + _ = unwrap( + client.embed( + scoped_key, + "openai-text-embedding-3-small", + f"vectorize this sentence {unique_marker()}", + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + row = _require_row( + rows, lambda r: (r.spend or 0) > 0, "with nonzero spend for the embedding" + ) + assert (row.prompt_tokens or 0) > 0 + assert (row.completion_tokens or 0) == 0, "embeddings have no completion tokens" + assert "text-embedding-3-small" in (row.model or "") + + +def test_cache_hit_is_zero_cost_and_suffixed( + client: SpendClient, scoped_key: str +) -> None: + # Unique marker shared by both calls: call 1 is a guaranteed cache MISS (fresh + # content, paid), call 2 repeats the identical request and HITS the cache just + # populated. The marker keeps each run isolated - a fixed prompt would persist + # in the shared response cache across runs and make both calls hit (flaky). + prompt = f"What is the capital of France? Answer in one word. {unique_marker()}" + _ = unwrap(client.chat(scoped_key, "gemini-2.5-flash", prompt, max_tokens=16)) + _ = unwrap(client.chat(scoped_key, "gemini-2.5-flash", prompt, max_tokens=16)) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.cache_hit == "True" for r in rs) + ) + cache_rows = [r for r in rows if r.cache_hit == "True"] + if not cache_rows: + pytest.skip( + "no cache-hit row observed; caching may be disabled on this proxy. " + f"rows seen: {_summarize(rows)}" + ) + + cache_row = cache_rows[0] + assert ( + cache_row.spend or 0 + ) == 0.0, f"cache hit was charged (double-charge regression): {_summarize(rows)}" + assert "_cache_hit" in (cache_row.request_id or ""), ( + "cache-hit row missing the _cache_hit request_id suffix; " + "duplicate-key collisions will silently drop rows" + ) + paid_rows = [r for r in rows if r.cache_hit != "True"] + assert any( + (r.spend or 0) > 0 for r in paid_rows + ), f"the non-cached call should still be charged: {_summarize(rows)}" + + +def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> None: + for _ in range(2): + _ = unwrap( + client.chat( + scoped_key, + "gemini-2.5-flash", + f"say hi {unique_marker()}", + max_tokens=16, + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, + min_rows=2, + predicate=lambda rs: sum((r.spend or 0) for r in rs) > 0, + ) + assert len(rows) >= 2, f"expected >=2 rows for the key, saw {_summarize(rows)}" + logs_total = sum((r.spend or 0) for r in rows) + assert logs_total > 0 + + key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999) + assert _approx_equal( + key_spend, logs_total + ), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}" + + +def test_request_tags_round_trip(client: SpendClient, scoped_key: str) -> None: + tag = f"e2e-spend-{unique_marker()}" + _ = unwrap( + client.chat( + scoped_key, "gemini-2.5-flash", "tagged request", tags=[tag], max_tokens=16 + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(tag in (r.request_tags or []) for r in rs) + ) + _require_row( + rows, lambda r: tag in (r.request_tags or []), f"carrying request tag {tag!r}" + ) + + +def test_end_user_spend_attributed_on_row( + client: SpendClient, scoped_key: str, resources: ResourceManager +) -> None: + customer = resources.customer(f"e2e-cust-{unique_marker()}") + _ = unwrap( + client.chat(scoped_key, "gemini-2.5-flash", "hi", user=customer, max_tokens=16) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.end_user == customer for r in rs) + ) + row = _require_row( + rows, lambda r: r.end_user == customer, f"attributed to end_user {customer!r}" + ) + assert (row.spend or 0) > 0, f"end-user row should cost > 0: {_summarize(rows)}" + + +def test_each_model_on_a_shared_key_gets_its_own_row( + client: SpendClient, scoped_key: str +) -> None: + """One key calling two different models, on two providers, gets one spend row per + call - each carrying its own model and a nonzero cost, under distinct request_ids + that match the call's response id. Pins per-model/per-provider attribution: a + regression that stamps the wrong model on the row, bills a call's cost to the + sibling deployment, or collapses both calls onto one request_id fails here.""" + gemini = unwrap( + client.chat( + scoped_key, "gemini-2.5-flash", f"one word {unique_marker()}", max_tokens=16 + ) + ) + claude = unwrap( + client.chat( + scoped_key, "claude-haiku-4-5", f"one word {unique_marker()}", max_tokens=16 + ) + ) + + def both_models_costed(rows: list[SpendLogRow]) -> bool: + costed = [r.model or "" for r in rows if (r.spend or 0) > 0] + return any("gemini-2.5-flash" in m for m in costed) and any( + "claude-haiku-4-5" in m for m in costed + ) + + rows = client.poll_logs_for_key(scoped_key, min_rows=2, predicate=both_models_costed) + gemini_row = _require_row( + rows, lambda r: "gemini-2.5-flash" in (r.model or ""), "for the gemini call" + ) + claude_row = _require_row( + rows, lambda r: "claude-haiku-4-5" in (r.model or ""), "for the claude call" + ) + + assert (gemini_row.spend or 0) > 0, f"gemini row should cost > 0: {_summarize(rows)}" + assert (claude_row.spend or 0) > 0, f"claude row should cost > 0: {_summarize(rows)}" + assert ( + gemini_row.request_id != claude_row.request_id + ), f"two distinct calls collapsed onto one request_id: {_summarize(rows)}" + if gemini.id: + assert ( + gemini_row.request_id == gemini.id + ), f"gemini row request_id {gemini_row.request_id} != response id {gemini.id}" + if claude.id: + assert ( + claude_row.request_id == claude.id + ), f"claude row request_id {claude_row.request_id} != response id {claude.id}" + + +def test_spend_calculate_returns_nonzero_cost(client: SpendClient) -> None: + cost = client.calculate_spend( + "gemini-2.5-flash", "estimate the cost of this request" + ) + assert cost > 0, ( + "/spend/calculate returned 0 for gemini-2.5-flash; " + "cost map may be missing this model" + ) + + +def test_spend_logs_endpoint_returns_spend( + client: SpendClient, scoped_key: str +) -> None: + """The /spend/logs read endpoint returns a 200 carrying the key's spend, never a + 5xx. Regression for intermittent 500s (DB query / serialization errors under load) + on this endpoint: every poll asserts a success response, not just a truthy row + list, so a 500 fails loudly instead of being swallowed as 'no rows yet'; the + call's nonzero spend must surface before the deadline.""" + unwrap( + client.chat( + scoped_key, "gemini-2.5-flash", f"spend logs {unique_marker()}", max_tokens=16 + ) + ) + + gateway = client.gateway + deadline = time.monotonic() + gateway.poll_timeout + while True: + result = gateway.transport.get( + "/spend/logs", + headers=gateway.transport.master, + params=SpendLogsParams(api_key=scoped_key), + response_type=SpendLogs, + ) + assert isinstance(result, Success), f"/spend/logs did not return 200 OK: {result}" + rows = result.data.root + if sum((r.spend or 0) for r in rows) > 0: + return + if time.monotonic() >= deadline: + pytest.fail( + f"/spend/logs never surfaced the key's spend before the deadline; " + f"saw {_summarize(rows)}" + ) + time.sleep(gateway.poll_interval) diff --git a/tests/e2e/test_lifecycle.py b/tests/e2e/test_lifecycle.py new file mode 100644 index 00000000000..d3c559dd2ed --- /dev/null +++ b/tests/e2e/test_lifecycle.py @@ -0,0 +1,46 @@ +"""Unit coverage for the lifecycle harness (lifecycle.run_case). + +Cases register cleanups progressively during init() (create team, then user, then +key), so a failure partway through init() must still release whatever was already +created on the long-lived shared proxy. This guards that contract. +""" + +from dataclasses import dataclass, field +from typing import Callable, List + +import pytest + +from lifecycle import run_case + + +@dataclass +class _PartialInitCase: + """init() registers a cleanup, then raises before finishing - mirroring a real + case that creates a resource, registers its delete, then fails on the next + step.""" + + released: List[str] = field(default_factory=list) + _undo: List[Callable[[], None]] = field(default_factory=list) + + def init(self) -> None: + self._undo.append(lambda: self.released.append("first")) + raise RuntimeError("init failed after registering the first resource") + + def run(self) -> None: + raise AssertionError("run() must not execute when init() failed") + + def teardown(self) -> None: + for undo in reversed(self._undo): + undo() + + +def test_run_case_releases_resources_when_init_fails_partway() -> None: + case = _PartialInitCase() + + with pytest.raises(RuntimeError, match="init failed"): + run_case(case) + + assert case.released == ["first"], ( + "a resource registered before init() failed must still be released, or it " + "leaks on the long-lived shared proxy" + ) diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py new file mode 100644 index 00000000000..37412fc0cf5 --- /dev/null +++ b/tests/e2e/transport.py @@ -0,0 +1,244 @@ +"""Transport: the typed request primitives clients use, behind a Protocol. + +`Transport` is what each client depends on (composition + DI); `HttpTransport` is +the concrete frozen-slots dataclass that fulfils it via the e2e_http wrapper. No +client touches requests.* or builds raw dicts; they pass pydantic models here. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from pydantic import BaseModel + +import e2e_http +from e2e_http import URL, AuthHeaders, ProbeResult, Result, StreamingResponse + + +class Transport(Protocol): + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: ... + + def stream( + self, path: str, *, headers: BaseModel, json: BaseModel + ) -> StreamingResponse: ... + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: ... + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + ) -> Result[R]: ... + + def delete[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: ... + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ... + + def bearer(self, key: str) -> AuthHeaders: ... + + @property + def master(self) -> AuthHeaders: ... + + +@dataclass(frozen=True, slots=True) +class HttpTransport: + base_url: str + master_key: str + request_timeout: float = 60.0 + + def _url(self, path: str) -> URL: + return URL(f"{self.base_url.rstrip('/')}{path}") + + def bearer(self, key: str) -> AuthHeaders: + return AuthHeaders(authorization=f"Bearer {key}") + + @property + def master(self) -> AuthHeaders: + return self.bearer(self.master_key) + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return e2e_http.post( + self._url(path), + headers=headers, + json=json, + response_type=response_type, + timeout=self.request_timeout, + ) + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + ) -> Result[R]: + return e2e_http.get( + self._url(path), + headers=headers, + params=params, + response_type=response_type, + timeout=self.request_timeout, + ) + + def delete[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return e2e_http.delete( + self._url(path), + headers=headers, + json=json, + response_type=response_type, + timeout=self.request_timeout, + ) + + def stream( + self, path: str, *, headers: BaseModel, json: BaseModel + ) -> StreamingResponse: + return e2e_http.stream( + self._url(path), headers=headers, json=json, timeout=self.request_timeout + ) + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + return e2e_http.send( + self._url(path), + headers=headers, + json=json, + params=params, + stream=stream, + timeout=self.request_timeout, + ) + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + return e2e_http.probe( + self._url(path), + headers=self.master, + params=params, + timeout=self.request_timeout, + ) + + +# Top-level management/admin route groups. In a split deployment these are served +# by the control plane (a different service from the LLM data plane). LLM routes +# (/chat, /embeddings, and native passthrough like /gemini, /anthropic) are NOT +# here and fall through to the data plane. Matched as path prefixes. +CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( + "/key", + "/user", + "/team", + "/organization", + "/customer", + "/tag", + "/budget", + "/model/info", + "/spend", + "/global", + "/openapi.json", +) + + +def is_control_plane_path(path: str) -> bool: + """True if `path` is a management/admin route (served by the control plane in a + split deployment), false for LLM data-plane routes.""" + return path.startswith(CONTROL_PLANE_PREFIXES) + + +@dataclass(frozen=True, slots=True) +class SplitTransport: + """A Transport that dispatches each call by path to one of two backends: the + management/admin control plane or the LLM data plane. + + Litellm can run as a split control-plane/data-plane deployment where the two + surfaces live on different services. Clients here stay plane-agnostic — they + keep calling ``transport.post("/budget/new", ...)`` or + ``transport.send("/chat/completions", ...)`` — and routing happens in one place + by path (see ``CONTROL_PLANE_PREFIXES``). When ``control`` and ``data`` share a + base URL (the monolithic default), routing is a no-op. ``bearer``/``master`` + are plane-agnostic (same master key both planes), so they come from ``data``. + """ + + data: HttpTransport + control: HttpTransport + + def _route(self, path: str) -> HttpTransport: + return self.control if is_control_plane_path(path) else self.data + + def bearer(self, key: str) -> AuthHeaders: + return self.data.bearer(key) + + @property + def master(self) -> AuthHeaders: + return self.data.master + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return self._route(path).post( + path, headers=headers, json=json, response_type=response_type + ) + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + ) -> Result[R]: + return self._route(path).get( + path, headers=headers, params=params, response_type=response_type + ) + + def delete[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return self._route(path).delete( + path, headers=headers, json=json, response_type=response_type + ) + + def stream( + self, path: str, *, headers: BaseModel, json: BaseModel + ) -> StreamingResponse: + return self._route(path).stream(path, headers=headers, json=json) + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + return self._route(path).send( + path, headers=headers, json=json, params=params, stream=stream + ) + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + return self._route(path).probe(path, params=params) diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index fa22ff6b392..f4c307e9c8a 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2502,19 +2502,34 @@ async def test_bedrock_image_url_sync_client(): mock_post.assert_called_once() -def test_bedrock_error_handling_streaming(): +@pytest.mark.parametrize( + "exception_type, expected_status_code", + [ + ("internalServerException", 500), + ("serviceUnavailableException", 503), + ("modelTimeoutException", 408), + ("modelStreamErrorException", 424), + ("validationException", 400), + ], +) +def test_bedrock_error_handling_streaming(exception_type, expected_status_code): + """Bedrock event-stream error events arrive with botocore's hard-coded + status_code=400; the decoder must surface the modeled HTTP status instead + (e.g. internalServerException -> 500). For 5xx this is what makes the error + retryable downstream; for all types it replaces the misleading 400 with the + true code. Regression for #24608.""" from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, BedrockError, ) - from unittest.mock import patch, Mock + from unittest.mock import Mock event = Mock() event.to_response_dict = Mock( return_value={ "status_code": 400, "headers": { - ":exception-type": "serviceUnavailableException", + ":exception-type": exception_type, ":content-type": "application/json", ":message-type": "exception", }, @@ -2525,11 +2540,10 @@ def test_bedrock_error_handling_streaming(): decoder = AWSEventStreamDecoder( model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0" ) - with pytest.raises(Exception) as e: + with pytest.raises(BedrockError) as e: decoder._parse_message_from_event(event) - assert isinstance(e.value, BedrockError) assert "Bedrock is unable to process your request." in e.value.message - assert e.value.status_code == 400 + assert e.value.status_code == expected_status_code @pytest.mark.parametrize( diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py new file mode 100644 index 00000000000..099d73fed87 --- /dev/null +++ b/tests/llm_translation/test_bedrock_embedding_pricing.py @@ -0,0 +1,34 @@ +""" +Tests for AWS Bedrock embedding model pricing in the model cost map. + +Regression test for the Amazon Titan Text Embeddings V2 commercial price, +which was previously set 10x too high (2e-07 instead of 2e-08). +AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens +(= $0.00002 per 1K tokens = 2e-08 per token). +""" + +import importlib + + +class TestBedrockEmbeddingPricing: + """Test suite for Bedrock embedding model pricing in the cost map.""" + + def test_titan_embed_v2_commercial_input_cost(self, monkeypatch): + """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08).""" + # Scope the local-cost-map flag to this test only, so it does not leak + # into sibling tests. monkeypatch restores the environment on teardown. + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + import litellm.litellm_core_utils.get_model_cost_map + import litellm + + # Reload so the cost map is re-read from the local file with the flag set. + importlib.reload(litellm.litellm_core_utils.get_model_cost_map) + importlib.reload(litellm) + + model = litellm.model_cost["amazon.titan-embed-text-v2:0"] + + assert model["input_cost_per_token"] == 2e-08 + assert model["output_cost_per_token"] == 0.0 + assert model["litellm_provider"] == "bedrock" + assert model["mode"] == "embedding" diff --git a/tests/llm_translation/test_cloudflare.py b/tests/llm_translation/test_cloudflare.py index 5a6a0008398..54c5d9e4e07 100644 --- a/tests/llm_translation/test_cloudflare.py +++ b/tests/llm_translation/test_cloudflare.py @@ -9,9 +9,7 @@ import pytest from litellm import acompletion, completion from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -FAKE_API_BASE = ( - "https://fake-cloudflare.example.com/client/v4/accounts/fake-acct/ai/run/" -) +FAKE_API_BASE = "https://fake-cloudflare.example.com/client/v4/accounts/fake-acct/ai/v1" FAKE_API_KEY = "fake-cf-api-key" @@ -26,28 +24,78 @@ def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: def _chat_response() -> Dict[str, Any]: return { - "result": { - "response": "I am a large language model created to assist you.", - }, - "success": True, - "errors": [], - "messages": [], + "id": "chatcmpl-cf", + "object": "chat.completion", + "created": 1234567890, + "model": "@cf/meta/llama-2-7b-chat-int8", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "I am a large language model created to assist you.", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 8, "completion_tokens": 11, "total_tokens": 19}, + } + + +def _tool_call_response() -> Dict[str, Any]: + return { + "id": "chatcmpl-cf-tools", + "object": "chat.completion", + "created": 1234567890, + "model": "@cf/meta/llama-2-7b-chat-int8", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "New York"}', + }, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 20, "completion_tokens": 9, "total_tokens": 29}, } def _streaming_chunks() -> list[str]: + base = { + "id": "chatcmpl-cf", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "@cf/meta/llama-2-7b-chat-int8", + } return [ - json.dumps({"response": "I am"}), - json.dumps({"response": " a language"}), - json.dumps({"response": " model."}), - ] - - -def _streaming_chunks_response_text() -> list[str]: - return [ - json.dumps({"response_text": "I am"}), - json.dumps({"response_text": " a language"}), - json.dumps({"response_text": " model."}), + json.dumps({**base, "choices": [{"index": 0, "delta": {"content": "I am"}}]}), + json.dumps( + {**base, "choices": [{"index": 0, "delta": {"content": " a language"}}]} + ), + json.dumps( + { + **base, + "choices": [ + { + "index": 0, + "delta": {"content": " model."}, + "finish_reason": "stop", + } + ], + } + ), ] @@ -85,6 +133,48 @@ def test_completion_cloudflare(sync_mode): assert response.choices[0].message.content is not None assert "language model" in response.choices[0].message.content.lower() + called_url = mock_post.call_args.kwargs.get("url") or mock_post.call_args.args[0] + assert called_url.endswith("/ai/v1/chat/completions") + assert "/ai/run/" not in called_url + + +def test_completion_cloudflare_tool_calls_sent_to_openai_endpoint(): + messages = [{"role": "user", "content": "weather in New York?"}] + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + mock_resp = _make_mock_response(_tool_call_response()) + + with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: + response = completion( + model="cloudflare/@cf/meta/llama-2-7b-chat-int8", + messages=messages, + tools=tools, + tool_choice="auto", + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + ) + mock_post.assert_called_once() + + sent_body = json.loads(mock_post.call_args.kwargs["data"]) + assert sent_body["tools"] == tools + assert sent_body["tool_choice"] == "auto" + + assert response.choices[0].finish_reason == "tool_calls" + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0].function.name == "get_weather" + @pytest.mark.parametrize("sync_mode", [True, False]) def test_completion_cloudflare_stream(sync_mode): @@ -153,76 +243,3 @@ def test_completion_cloudflare_stream(sync_mode): if c.choices[0].delta.content ) assert "language" in content.lower() - - -@pytest.mark.parametrize("sync_mode", [True, False]) -def test_completion_cloudflare_stream_response_text(sync_mode): - """Newer Cloudflare Workers AI models (e.g. Nemotron) emit `response_text` - instead of `response` in streamed chunks. The iterator must surface that - text so streaming output is not silently empty. - """ - messages = [{"role": "user", "content": "what llm are you"}] - raw_chunks = _streaming_chunks_response_text() - - if sync_mode: - - def _iter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: - response = completion( - model="cloudflare/@cf/nvidia/nemotron-mini-4b-instruct", - messages=messages, - max_tokens=15, - stream=True, - api_base=FAKE_API_BASE, - api_key=FAKE_API_KEY, - ) - chunks_received = list(response) - mock_post.assert_called_once() - else: - - async def _aiter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - - async def _run(): - with patch.object( - AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp - ) as mock_post: - resp = await acompletion( - model="cloudflare/@cf/nvidia/nemotron-mini-4b-instruct", - messages=messages, - max_tokens=15, - stream=True, - api_base=FAKE_API_BASE, - api_key=FAKE_API_KEY, - ) - received = [] - async for chunk in resp: - received.append(chunk) - mock_post.assert_called_once() - return received - - chunks_received = asyncio.run(_run()) - - assert len(chunks_received) > 0 - content = "".join( - c.choices[0].delta.content - for c in chunks_received - if c.choices[0].delta.content - ) - assert "language" in content.lower() diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py index 204f4d9e31b..27059581e4d 100644 --- a/tests/llm_translation/test_fireworks_ai_translation.py +++ b/tests/llm_translation/test_fireworks_ai_translation.py @@ -7,9 +7,10 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm -from litellm import transcription +from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, +) from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig -from base_audio_transcription_unit_tests import BaseLLMAudioTranscriptionTest fireworks = FireworksAIConfig() @@ -69,74 +70,16 @@ def test_map_response_format(): assert result == {"response_format": response_format} -_AUDIO_FILE_PATH = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "gettysburg.wav" -) - - -class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest): - def get_base_audio_transcription_call_args(self) -> dict: - return { - "model": "fireworks_ai/whisper-v3", - "api_base": "https://audio-prod.api.fireworks.ai/v1", - } - - def get_custom_llm_provider(self) -> litellm.LlmProviders: - return litellm.LlmProviders.FIREWORKS_AI - - def test_audio_transcription(self): - from unittest.mock import MagicMock - - from openai.types.audio import Transcription - - audio_file = open(_AUDIO_FILE_PATH, "rb") - mock_client = MagicMock() - mock_client.audio.transcriptions.create.return_value = Transcription( - text="four score and seven years ago" - ) - - transcript = transcription( - **self.get_base_audio_transcription_call_args(), - file=audio_file, - api_key="fw-test-key", - client=mock_client, - ) - - assert transcript.text == "four score and seven years ago" - sent = mock_client.audio.transcriptions.create.call_args.kwargs - assert sent["model"] == "whisper-v3" - assert sent["file"] is audio_file - - @pytest.mark.asyncio - async def test_audio_transcription_async(self): - from unittest.mock import AsyncMock, MagicMock - - from openai.types.audio import Transcription - - audio_file = open(_AUDIO_FILE_PATH, "rb") - raw_response = MagicMock() - raw_response.headers = {} - raw_response.parse.return_value = Transcription( - text="four score and seven years ago" - ) - mock_client = MagicMock() - mock_client.audio.transcriptions.with_raw_response.create = AsyncMock( - return_value=raw_response - ) - - transcript = await litellm.atranscription( - **self.get_base_audio_transcription_call_args(), - file=audio_file, - api_key="fw-test-key", - client=mock_client, - ) - - assert transcript.text == "four score and seven years ago" - sent = ( - mock_client.audio.transcriptions.with_raw_response.create.call_args.kwargs - ) - assert sent["model"] == "whisper-v3" - assert sent["file"] is audio_file +def test_get_supported_openai_params_transcription_returns_none(): + # Fireworks AI deprecated audio transcription on 2026-06-10; the endpoint + # is decommissioned. Returning None (not chat-completion params) signals + # to callers that transcription is unsupported for this provider. + result = get_supported_openai_params( + model="fireworks_ai/accounts/fireworks/models/whisper-v3", + custom_llm_provider="fireworks_ai", + request_type="transcription", + ) + assert result is None @pytest.mark.parametrize( diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 9a69f513069..4d67f1e426a 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1988,11 +1988,15 @@ class TestConvertToStreamingResponseAsync: return chunks chunks = asyncio.run(run()) - assert len(chunks) == 1 - assert chunks[0].id == "msg_async_1" - assert chunks[0].model == "claude-3" - assert chunks[0].choices[0].delta.content == "Hi there" - assert chunks[0].usage.prompt_tokens == 3 + # Cached replay is sliced into word-shaped chunks to preserve + # streaming cadence; joining the slices reconstructs the content. + assert len(chunks) == 2 + assert all(c.id == "msg_async_1" for c in chunks) + assert all(c.model == "claude-3" for c in chunks) + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "Hi there" + assert chunks[0].choices[0].finish_reason is None + assert chunks[-1].choices[0].finish_reason == "stop" + assert chunks[-1].usage.prompt_tokens == 3 class TestHandleInvalidParallelToolCalls: diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 36e47e3c2f4..05a58a135d2 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -605,8 +605,33 @@ def test_no_messages_yields_user_text(): assert contents == expected_output -def test_convert_url(): - convert_url_to_base64("https://picsum.photos/id/237/200/300") +def test_convert_url(monkeypatch): + import base64 + from unittest.mock import MagicMock + + import httpx + + from litellm.litellm_core_utils.prompt_templates.image_handling import ( + in_memory_cache, + ) + + url = "https://picsum.photos/id/237/200/300" + image_bytes = b"\x89PNG\r\n\x1a\nfake-png-bytes" + + mock_client = MagicMock() + mock_client.get.return_value = httpx.Response( + 200, content=image_bytes, headers={"Content-Type": "image/png"} + ) + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + monkeypatch.setattr(litellm, "module_level_client", mock_client, raising=False) + in_memory_cache.flush_cache() + + result = convert_url_to_base64(url) + + expected = "data:image/png;base64," + base64.b64encode(image_bytes).decode("utf-8") + assert result == expected + mock_client.get.assert_called_once() def test_azure_tool_call_invoke_helper(): diff --git a/tests/mcp_tests/test_mcp_auth_priority.py b/tests/mcp_tests/test_mcp_auth_priority.py index ad6e9438edd..7ae0f59afe5 100644 --- a/tests/mcp_tests/test_mcp_auth_priority.py +++ b/tests/mcp_tests/test_mcp_auth_priority.py @@ -45,7 +45,18 @@ async def test_mcp_server_works_without_config_auth_value(): @pytest.mark.parametrize("token_key", ["authentication_token", "auth_value"]) async def test_mcp_server_config_auth_value_header_used(token_key): - """Ensure auth header is sent when auth token configured in config""" + """Ensure the configured auth token is emitted as the upstream Authorization header. + + The token is resolved through the v2 credential resolver and rides on the client's + httpx.Auth, so assert the header it writes onto the request rather than the (now + credential-free) _get_auth_headers() dict. + """ + import httpx + + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + config = { "test_server": { "url": "https://api.example.com/mcp", @@ -60,7 +71,8 @@ async def test_mcp_server_config_auth_value_header_used(token_key): server = next(iter(manager.config_mcp_servers.values())) client = await manager._create_mcp_client(server) - headers = client._get_auth_headers() - assert headers["Authorization"] == "Bearer example_token" + assert isinstance(client._resolved_auth, StaticHeaderAuth) + emitted = next(client._resolved_auth.auth_flow(httpx.Request("POST", server.url))) + assert emitted.headers["Authorization"] == "Bearer example_token" assert client.auth_type == MCPAuth.bearer_token diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index eea2f2721ab..5f8fcbf835e 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1089,7 +1089,9 @@ async def test_list_tools_only_returns_allowed_servers(monkeypatch): mock_client_constructor, ): # Call list_tools - tools = await test_manager.list_tools(user_api_key_auth=MagicMock()) + from litellm.proxy._types import UserAPIKeyAuth + + tools = await test_manager.list_tools(user_api_key_auth=UserAPIKeyAuth()) # Should only return tools from server_a assert len(tools) == 1 # The server should use the server_name as prefix since no alias is provided diff --git a/tests/pyrightconfig.json b/tests/pyrightconfig.json new file mode 100644 index 00000000000..5757c97f812 --- /dev/null +++ b/tests/pyrightconfig.json @@ -0,0 +1,11 @@ +{ + "include": ["e2e"], + "exclude": ["**/node_modules", "**/__pycache__"], + "pythonVersion": "3.12", + "typeCheckingMode": "strict", + "enableTypeIgnoreComments": false, + "reportMissingImports": false, + "reportPrivateImportUsage": false, + "reportExplicitAny": "error", + "reportAny": "error" +} \ No newline at end of file diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index 1b8f713a437..b8760906645 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -1,17 +1,10 @@ import sys import os -import traceback -from dotenv import load_dotenv -from fastapi import Request -from datetime import datetime sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from litellm import Router import pytest -import litellm -from unittest.mock import patch, MagicMock, AsyncMock import json from io import BytesIO @@ -76,6 +69,29 @@ def test_tuple_input(sample_jsonl_bytes): assert result.content_type == "application/jsonl" +def test_tuple_with_file_handle_rewrites_model(sample_jsonl_bytes): + """Security regression: when the tuple's content element is a file handle + (batch uploads stream from the spooled upload handle), the model must still + be rewritten. Otherwise a restricted body.model survives unmodified and + bypasses the batch model allowlist, which only checks the upload target.""" + new_model = "approved-target-model" + handle = BytesIO(sample_jsonl_bytes) + test_tuple = ("test.jsonl", handle, "application/json") + + result = replace_model_in_jsonl(test_tuple, new_model) + + assert isinstance(result, InMemoryFile) + rows = [ + json.loads(line) + for line in result.getvalue().decode("utf-8").splitlines() + if line.strip() + ] + assert rows, "rewrite must produce rows" + # every row now carries the rewritten target, not the original (restricted) model + assert all(row["body"]["model"] == new_model for row in rows) + assert all(row["body"]["model"] != "gpt-5.5" for row in rows) + + def test_file_like_object(sample_file_like): """Test with file-like object input""" new_model = "claude-3" @@ -129,9 +145,9 @@ def test_should_replace_model_in_jsonl(): """Test that should_replace_model_in_jsonl returns the correct value""" from litellm.router_utils.batch_utils import should_replace_model_in_jsonl - assert should_replace_model_in_jsonl(purpose="batch") == True - assert should_replace_model_in_jsonl(purpose="test") == False - assert should_replace_model_in_jsonl(purpose="user_data") == False + assert should_replace_model_in_jsonl(purpose="batch") is True + assert should_replace_model_in_jsonl(purpose="test") is False + assert should_replace_model_in_jsonl(purpose="user_data") is False def test_parse_jsonl_with_embedded_newlines_simple(): @@ -217,6 +233,63 @@ def test_parse_jsonl_with_embedded_newlines_whitespace_only(): assert len(result) == 0 +def test_replace_model_in_jsonl_malformed_middle_row_returns_original(): + """Regression: a malformed/truncated middle row must not silently drop the + rows that follow it. The streaming rewrite accumulates physical lines into a + buffer; a row that never parses poisons the buffer so every later valid row + is concatenated into it and dropped. Returning that partial rewrite would + ship a truncated batch with no error to the caller. Instead the original + content is returned unchanged so the provider rejects the bad batch loudly.""" + content = ( + b'{"custom_id":"a","body":{"model":"x"}}\n' + b'{"custom_id":"b","body":{"model":\n' # truncated, never completes + b'{"custom_id":"c","body":{"model":"x"}}\n' + ) + + result = replace_model_in_jsonl(content, "new-model") + + assert ( + result == content + ), "must return the original unchanged, not a partial rewrite" + + +def test_replace_model_in_jsonl_malformed_row_seekable_handle_rewound(): + """When the source is a seekable handle that gets consumed during the failed + rewrite, it must be rewound to 0 so the caller can re-read the full original.""" + content = ( + b'{"custom_id":"a","body":{"model":"x"}}\n' + b'{"custom_id":"b","body":{"model":\n' + b'{"custom_id":"c","body":{"model":"x"}}\n' + ) + handle = BytesIO(content) + + result = replace_model_in_jsonl(handle, "new-model") + + assert result is handle + assert handle.read() == content, "handle must be rewound for the caller to re-read" + + +def test_replace_model_in_jsonl_multi_row_rewrites_every_model(): + """Happy path: a well-formed multi-row file gets every row's model rewritten + and no row is dropped.""" + content = ( + b'{"custom_id":"a","body":{"model":"old1"}}\n' + b'{"custom_id":"b","body":{"model":"old2"}}\n' + b'{"custom_id":"c","body":{"model":"old3"}}\n' + ) + + result = replace_model_in_jsonl(content, "new-model") + + assert isinstance(result, InMemoryFile) + rows = [ + json.loads(line) + for line in result.getvalue().decode("utf-8").splitlines() + if line.strip() + ] + assert [row["custom_id"] for row in rows] == ["a", "b", "c"] + assert all(row["body"]["model"] == "new-model" for row in rows) + + def test_replace_model_in_jsonl_with_embedded_newlines(): """Test that replace_model_in_jsonl works correctly with embedded newlines in content""" # Create a JSONL with embedded newlines in the message content diff --git a/tests/search_tests/test_searchapi_search.py b/tests/search_tests/test_searchapi_search.py index 5ef9d922b89..d16868502a4 100644 --- a/tests/search_tests/test_searchapi_search.py +++ b/tests/search_tests/test_searchapi_search.py @@ -46,10 +46,9 @@ class TestSearchAPIConfig: assert result["Content-Type"] == "application/json" - @patch("litellm.llms.searchapi.search.transformation.get_secret_str") - def test_validate_environment_without_api_key(self, mock_get_secret): + def test_validate_environment_without_api_key(self, monkeypatch): """Test environment validation without API key raises error.""" - mock_get_secret.return_value = None + monkeypatch.delenv("SEARCHAPI_API_KEY", raising=False) config = SearchAPIConfig() headers = {} diff --git a/tests/search_tests/test_searxng_search.py b/tests/search_tests/test_searxng_search.py index 45b0f3214d9..c12d44183b0 100644 --- a/tests/search_tests/test_searxng_search.py +++ b/tests/search_tests/test_searxng_search.py @@ -318,13 +318,11 @@ class TestSearXNGSearchHeaders: assert headers["Content-Type"] == "application/json" assert headers["Authorization"] == "Bearer test-key-123" - def test_headers_with_env_api_key(self): + def test_headers_with_env_api_key(self, monkeypatch): """Test that headers use SEARXNG_API_KEY from env.""" - with patch( - "litellm.llms.searxng.search.transformation.get_secret_str", - return_value="env-key-456", - ): - headers = self.config.validate_environment(headers={}) + monkeypatch.setenv("SEARXNG_API_KEY", "env-key-456") + + headers = self.config.validate_environment(headers={}) assert headers["Authorization"] == "Bearer env-key-456" 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 c9e500b4a5b..704dd7f92f1 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,6 +1,5 @@ import asyncio import os -import ssl import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -543,5 +542,45 @@ class TestExecuteSessionOperationSurfacesTransportError: assert result == "done" +class TestMCPClientResolvedAuth: + """A pre-resolved httpx.Auth is attached to the upstream client's auth= slot.""" + + @pytest.mark.asyncio + async def test_resolved_auth_feeds_the_auth_slot(self): + resolved = httpx.Auth() + client = MCPClient( + server_url="https://upstream.example.com", resolved_auth=resolved + ) + http_client = client._create_httpx_client_factory()() + try: + assert http_client.auth is resolved + finally: + await http_client.aclose() + + @pytest.mark.asyncio + async def test_resolved_auth_takes_precedence_over_aws_auth(self): + resolved = httpx.Auth() + client = MCPClient( + server_url="https://upstream.example.com", + resolved_auth=resolved, + aws_auth=httpx.Auth(), + ) + http_client = client._create_httpx_client_factory()() + try: + assert http_client.auth is resolved + finally: + await http_client.aclose() + + @pytest.mark.asyncio + async def test_without_resolved_auth_falls_back_to_aws_auth(self): + aws = httpx.Auth() + client = MCPClient(server_url="https://upstream.example.com", aws_auth=aws) + http_client = client._create_httpx_client_factory()() + try: + assert http_client.auth is aws + finally: + await http_client.aclose() + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py index f33814b86df..7ff58ba6324 100644 --- a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py +++ b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py @@ -1,8 +1,8 @@ """ Unit tests for CodeInterpreterInterceptionLogger. -All sandbox dependencies are injected (dependency injection, no monkeypatch): -a FakeSandbox stands in for the real e2b config and records how it is called. +All sandbox dependencies are injected: a FakeSandbox stands in for the real e2b +config and records how it is called. """ import time @@ -12,13 +12,17 @@ import pytest from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, LITELLM_CODE_EXECUTION_TOOL_NAME, + _INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY, + _SANDBOX_KEY, +) +from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + is_interception_internal_key, ) from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult from litellm.types.utils import CallTypes -_ACTIVE_KEY = "_code_interpreter_interception_active" -_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" - class FakeHandle: def __init__(self, sandbox_id="sbx_fake"): @@ -51,6 +55,13 @@ class FakeLogging: def __init__(self, litellm_call_id="k1"): self.litellm_call_id = litellm_call_id self.model_call_details = {} + self.dynamic_success_callbacks = [] + + def pre_call(self, *args, **kwargs): + return None + + def post_call(self, *args, **kwargs): + return None def _function_call_item(call_id="c1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): @@ -62,6 +73,17 @@ def _function_call_item(call_id="c1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): } +def _chat_function_call_item(call_id="call_1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): + return { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": '{"code":"print(40 + 2)"}', + }, + } + + class FakeResponse: def __init__(self, output): self.output = output @@ -74,6 +96,18 @@ def _iter_messages(plan): return patch.messages +def test_interception_internal_key_prefix_sets_preserve_code_interpreter_state(): + assert is_interception_internal_key("_code_interpreter_interception_active") + assert not is_interception_internal_key( + "_code_interpreter_interception_active", + prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + ) + assert is_interception_internal_key( + "_websearch_interception_converted_stream", + prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + ) + + @pytest.mark.asyncio async def test_build_plan_runs_code_and_feeds_output_back(): sandbox = FakeSandbox(stdout="42") @@ -133,6 +167,30 @@ async def test_pre_call_converts_code_interpreter_tool(): assert LITELLM_CODE_EXECUTION_TOOL_NAME in names +@pytest.mark.asyncio +async def test_pre_call_converts_code_interpreter_tool_for_chat_completions(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": {"type": "code_interpreter"}, + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + tool = result["tools"][0] + assert tool["type"] == "function" + assert tool["function"]["name"] == LITELLM_CODE_EXECUTION_TOOL_NAME + assert tool["function"]["parameters"]["required"] == ["code"] + assert result["tool_choice"] == { + "type": "function", + "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}, + } + assert result["litellm_metadata"][_ACTIVE_KEY] is True + assert result["litellm_metadata"][_SANDBOX_KEY] == result[_SANDBOX_KEY] + + @pytest.mark.asyncio @pytest.mark.parametrize( "tool_choice", @@ -184,9 +242,24 @@ async def test_pre_call_noop_on_non_responses(): "custom_llm_provider": "openai", } + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aembedding) + + assert result is None + + +@pytest.mark.asyncio +async def test_pre_call_noop_on_chat_completion_without_code_interpreter_tool(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "web_search"}], + "custom_llm_provider": "openai", + } + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) assert result is None + assert _ACTIVE_KEY not in kwargs + assert _SANDBOX_KEY not in kwargs @pytest.mark.asyncio @@ -524,14 +597,142 @@ async def test_gate_rechecks_provider_scope(): assert should_run is False +@pytest.mark.asyncio +async def test_chat_completion_gate_detects_code_execution_tool_call(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = { + "choices": [ + {"message": {"tool_calls": [_chat_function_call_item(call_id="call_123")]}} + ] + } + + should_run, payload = await logger.async_should_run_agentic_loop( + response=response, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={ + _ACTIVE_KEY: True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert should_run is True + assert payload["tool_calls"][0]["id"] == "call_123" + assert payload["tool_calls"][0]["arguments"] == '{"code":"print(40 + 2)"}' + + +@pytest.mark.asyncio +async def test_chat_completion_gate_refuses_without_server_active_marker(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = {"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]} + + should_run, payload = await logger.async_should_run_agentic_loop( + response=response, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE}, + ) + + assert should_run is False + assert payload == {} + + +@pytest.mark.asyncio +async def test_chat_completion_build_plan_runs_code_and_appends_tool_message(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + native_chat_tool = {"type": "code_interpreter", "container": {"type": "auto"}} + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response={ + "choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}] + }, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "tools": [native_chat_tool], + "tool_choice": {"type": "code_interpreter", "container": {"type": "auto"}}, + "temperature": 0, + }, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={ + "acompletion": True, + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert sandbox.run_calls[0]["code"] == "print(40 + 2)" + patch = plan.request_patch + assert patch is not None + assert patch.tools == [ + { + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": "Execute python code in a sandbox and return stdout.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + assert patch.optional_params == {"temperature": 0} + assert patch.kwargs == { + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + } + assert patch.messages is not None + assert patch.messages[-2]["role"] == "assistant" + assert patch.messages[-2]["tool_calls"][0]["id"] == "call_1" + assert patch.messages[-1] == { + "role": "tool", + "tool_call_id": "call_1", + "content": "42", + } + assert plan.metadata["code_interpreter_calls"][0]["code"] == "print(40 + 2)" + + @pytest.mark.asyncio async def test_pre_call_strips_client_forged_marker_on_initial_request(): - """A client cannot pre-set the active marker on the original request.""" + """A client cannot pre-set the active marker on the original request: with no + native code_interpreter tool, any client-supplied interception markers in + litellm_metadata are scrubbed and the active flag in kwargs is cleared.""" logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) kwargs = { "tools": [{"type": "web_search"}], "custom_llm_provider": "openai", _ACTIVE_KEY: True, + "litellm_metadata": { + _ACTIVE_KEY: True, + _SANDBOX_KEY: "client-forged", + "safe_user_value": "kept", + }, } await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) @@ -540,6 +741,42 @@ async def test_pre_call_strips_client_forged_marker_on_initial_request(): "no native code_interpreter tool was present, so a client-supplied " "active marker must be cleared" ) + assert kwargs["litellm_metadata"] == {"safe_user_value": "kept"} + + +@pytest.mark.asyncio +async def test_pre_call_strips_forged_loop_controls_then_mints_own_markers(): + """On an INITIAL request (no server-set _agentic_loop_depth) a client cannot + smuggle loop-control state: forged _agentic_loop_depth / max_agentic_loops and + interception markers in litellm_metadata are stripped before the interceptor + activates, so the only interception markers that survive are the ones the + server mints for the converted code_interpreter tool.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "litellm_metadata": { + _ACTIVE_KEY: True, + _SANDBOX_KEY: "client-forged", + "_agentic_loop_depth": 99, + "max_agentic_loops": 999, + "safe_user_value": "kept", + }, + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + metadata = result["litellm_metadata"] + assert metadata["safe_user_value"] == "kept" + assert "_agentic_loop_depth" not in metadata, "forged loop depth must be stripped" + assert "max_agentic_loops" not in metadata, "forged loop cap must be stripped" + assert metadata[_ACTIVE_KEY] is True + assert metadata[_SANDBOX_KEY] == result[_SANDBOX_KEY] + assert metadata[_SANDBOX_KEY] != "client-forged", ( + "the surviving sandbox key must be the server-minted one, not the forged " + "value the client supplied" + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/interactions/test_google_interactions_integration.py b/tests/test_litellm/interactions/test_google_interactions_integration.py index cfff26d51ef..9c651cc94f5 100644 --- a/tests/test_litellm/interactions/test_google_interactions_integration.py +++ b/tests/test_litellm/interactions/test_google_interactions_integration.py @@ -299,7 +299,6 @@ class TestGoogleInteractionsResponseStructure: assert hasattr(response, "outputs") assert hasattr(response, "usage") assert hasattr(response, "model") or hasattr(response, "agent") - assert hasattr(response, "role") assert hasattr(response, "created") assert hasattr(response, "updated") diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 44c0b6e5b02..209e99895db 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -162,7 +162,8 @@ class TestResponseCompliance: # Keep this aligned with the live spec. schema = spec_dict["components"]["schemas"]["Interaction"] - # Output fields (readOnly). + # Output fields (readOnly). `role` was removed from the `Interaction` + # schema by Google; it now lives only on `Turn`. output_fields = [ "id", "status", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 9b3152fae07..7f3d5a959a1 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1499,19 +1499,20 @@ def _local_model_cost_map(): @pytest.mark.parametrize("data_residency", ["eu", "us"]) def test_data_residency_applies_uplift(data_residency, _local_model_cost_map): - """gpt-5 should apply the regional processing uplift multiplier when - data_residency is set.""" + """gpt-5.4 should apply the regional processing uplift multiplier when + data_residency is set. gpt-5.4+ (released 2026-03-05) carry the 10% uplift; + gpt-5 and older models do not.""" from litellm.types.utils import Usage usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) base = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", ) regional = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", data_residency=data_residency, @@ -1526,6 +1527,23 @@ def test_data_residency_applies_uplift(data_residency, _local_model_cost_map): assert regional[1] == pytest.approx(base[1] * 1.10, rel=1e-9) +@pytest.mark.parametrize("model", ["gpt-5", "gpt-5-mini", "gpt-5-nano", "gpt-5-pro", "gpt-4o", "gpt-4.1"]) +def test_data_residency_no_uplift_for_pre_march_2026_models(model, _local_model_cost_map): + """Models released before 2026-03-05 must not have the regional uplift.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + regional = generic_cost_per_token( + model=model, usage=usage, custom_llm_provider="openai", data_residency="eu" + ) + + assert base == regional, ( + f"{model} should not have a regional uplift, but cost changed with data_residency" + ) + + def test_data_residency_no_uplift_for_unmarked_model(_local_model_cost_map): """A model without a regional_processing_uplift_multiplier_* entry should fall back to base pricing, not error.""" @@ -1555,12 +1573,12 @@ def test_data_residency_none_no_uplift(_local_model_cost_map): usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) base = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", ) explicit_none = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", data_residency=None, @@ -1576,13 +1594,13 @@ def test_data_residency_composes_with_service_tier(_local_model_cost_map): usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) priority_base = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", service_tier="priority", ) priority_eu = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", service_tier="priority", diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py new file mode 100644 index 00000000000..2c3dae21058 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -0,0 +1,282 @@ +""" +Tests for the cache-hit replay generators in +``litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response``. + +These generators are used by ``LLMCachingHandler._convert_cached_stream_response`` +to replay a cached non-streaming ``ModelResponse`` as a stream when the +incoming request has ``stream=True``. The fix in this test file ensures the +replay yields multiple word-shaped chunks instead of a single one-shot +content frame, restoring per-token cadence on cache hits. +""" + +import pytest + +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _split_assembled_content_for_replay, + convert_to_streaming_response, + convert_to_streaming_response_async, +) +from litellm.types.utils import ModelResponseStream + + +def _async_payload(content="Hello world! How are you?"): + return { + "id": "chatcmpl-test-async", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "system_fingerprint": "fp_test", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + }, + } + + +async def _collect_async(payload): + return [chunk async for chunk in convert_to_streaming_response_async(payload)] + + +# ---------- helper ---------- + + +@pytest.mark.parametrize( + "text", + [ + "Hello world!", + "Sure! Here's a list of 25 fruits:\n\n1. Apple\n2. Banana\n3. Orange\n", + " leading whitespace matters", + "Hi", + "你好世界", + ], +) +def test_split_is_lossless(text): + assert "".join(_split_assembled_content_for_replay(text)) == text + + +def test_split_returns_empty_for_none_and_empty(): + assert _split_assembled_content_for_replay(None) == [] + assert _split_assembled_content_for_replay("") == [] + + +def test_split_returns_empty_for_whitespace_only(): + # Must short-circuit before the regex: findall backtracks quadratically + # on all-whitespace input. + assert _split_assembled_content_for_replay(" ") == [] + assert _split_assembled_content_for_replay(" \n\t" * 10000) == [] + + +# ---------- async generator ---------- + + +@pytest.mark.asyncio +async def test_async_yields_multiple_content_chunks_with_lossless_join(): + text = "Sure! Here's a list of fruits: apple, banana, orange." + chunks = await _collect_async(_async_payload(content=text)) + assert len(chunks) > 1 + assert all(isinstance(c, ModelResponseStream) for c in chunks) + reassembled = "".join((c.choices[0].delta.content or "") for c in chunks) + assert reassembled == text + + +@pytest.mark.asyncio +async def test_async_finish_reason_only_on_last_chunk(): + chunks = await _collect_async(_async_payload()) + finish_reasons = [c.choices[0].finish_reason for c in chunks] + assert finish_reasons[-1] == "stop" + assert all(fr is None for fr in finish_reasons[:-1]) + + +@pytest.mark.asyncio +async def test_async_role_only_on_first_chunk(): + chunks = await _collect_async(_async_payload()) + assert chunks[0].choices[0].delta.role == "assistant" + for c in chunks[1:]: + assert c.choices[0].delta.role is None + + +@pytest.mark.asyncio +async def test_async_usage_attached_to_last_chunk_only(): + chunks = await _collect_async(_async_payload()) + usage_frames = [c for c in chunks if getattr(c, "usage", None) is not None] + assert len(usage_frames) == 1 + assert usage_frames[0] is chunks[-1] + assert usage_frames[0].usage.completion_tokens == 7 + assert usage_frames[0].usage.prompt_tokens == 5 + assert usage_frames[0].usage.total_tokens == 12 + assert usage_frames[0].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_async_empty_content_yields_single_finish_frame(): + # tool-calls-only-style response: short-circuit to the single-yield path. + payload = _async_payload(content=None) + payload["usage"] = None + chunks = await _collect_async(payload) + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content is None + assert chunks[0].choices[0].finish_reason == "stop" + assert chunks[0].choices[0].delta.role == "assistant" + assert getattr(chunks[0], "usage", None) is None + + +@pytest.mark.asyncio +async def test_async_tool_calls_and_function_call_only_on_first_chunk(): + # A cached response combining multi-word content with tool_calls must not + # repeat the tool_calls on every slice — downstream handlers accumulate + # tool-call deltas and would collect them N times. + payload = _async_payload(content="I'll look that up for you") + payload["choices"][0]["message"]["tool_calls"] = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "index": 0, + } + ] + chunks = await _collect_async(payload) + assert len(chunks) > 1 + first_tool_calls = chunks[0].choices[0].delta.tool_calls + assert first_tool_calls is not None and len(first_tool_calls) == 1 + assert first_tool_calls[0].function.name == "lookup" + for c in chunks[1:]: + assert c.choices[0].delta.tool_calls is None + assert c.choices[0].delta.function_call is None + + +@pytest.mark.asyncio +async def test_async_logprobs_only_on_first_chunk(): + payload = _async_payload(content="cached token logprobs") + payload["choices"][0]["logprobs"] = {"content": []} + chunks = await _collect_async(payload) + assert len(chunks) > 1 + assert chunks[0].choices[0].logprobs is not None + for c in chunks[1:]: + assert c.choices[0].logprobs is None + + +@pytest.mark.asyncio +async def test_async_metadata_propagated_to_every_chunk(): + chunks = await _collect_async(_async_payload()) + for c in chunks: + assert c.id == "chatcmpl-test-async" + assert c.model == "gpt-4o-mini" + assert c.system_fingerprint == "fp_test" + assert c.created == 1700000000 + + +# ---------- sync generator (parity smoke test) ---------- + + +def test_sync_multi_chunk_and_lossless_join(): + text = "Sure! Here's a list of fruits: apple, banana, orange." + payload = { + "id": "chatcmpl-test-sync", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "system_fingerprint": "fp_test", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + } + chunks = list(convert_to_streaming_response(payload)) + assert len(chunks) > 1 + reassembled = "".join((c.choices[0].delta.content or "") for c in chunks) + assert reassembled == text + # Sync path must also honor the "usage on last chunk" invariant. + assert chunks[-1].choices[0].finish_reason == "stop" + assert getattr(chunks[-1], "usage", None) is not None + assert chunks[-1].usage.total_tokens == 12 + + +def test_sync_delta_and_choice_metadata_only_on_first_chunk(): + thinking_blocks = [ + {"type": "thinking", "thinking": "cached thinking", "signature": "sig"} + ] + payload = { + "id": "chatcmpl-test-sync", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "cached content slices", + "reasoning_content": "cached reasoning", + "thinking_blocks": thinking_blocks, + }, + "finish_reason": "stop", + "logprobs": {"content": []}, + "enhancements": {"source": "cache"}, + } + ], + } + chunks = list(convert_to_streaming_response(payload)) + assert len(chunks) > 1 + first_choice = chunks[0].choices[0] + assert getattr(first_choice.delta, "reasoning_content", None) == "cached reasoning" + assert getattr(first_choice.delta, "thinking_blocks", None) == thinking_blocks + assert first_choice.logprobs is not None + assert first_choice.enhancements == {"source": "cache"} + for c in chunks[1:]: + choice = c.choices[0] + assert getattr(choice.delta, "reasoning_content", None) is None + assert getattr(choice.delta, "thinking_blocks", None) is None + assert choice.logprobs is None + assert getattr(choice, "enhancements", None) is None + + +def test_sync_non_enumerated_delta_fields_only_on_first_chunk(): + # annotations (and any other Delta field beyond the role/tool_call/reasoning + # set) must also stay on the first slice. Rebuilding later slices as a bare + # content delta drops the whole class, so this holds without enumerating + # every field by hand. + annotations = [ + { + "type": "url_citation", + "url_citation": { + "url": "https://example.com", + "title": "Example", + "start_index": 0, + "end_index": 1, + }, + } + ] + payload = { + "id": "chatcmpl-test-sync", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "cached content slices here", + "annotations": annotations, + }, + "finish_reason": "stop", + } + ], + } + chunks = list(convert_to_streaming_response(payload)) + assert len(chunks) > 1 + assert getattr(chunks[0].choices[0].delta, "annotations", None) == annotations + for c in chunks[1:]: + assert getattr(c.choices[0].delta, "annotations", None) is None diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py new file mode 100644 index 00000000000..f1196ab4692 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -0,0 +1,422 @@ +""" +Tests for the provider-agnostic chat completion agentic loop dispatcher +(`litellm/litellm_core_utils/chat_completion_agentic_loop.py`) and the +code-interpreter interception integration that drives it. + +The load-bearing regression here protects a reviewer requirement: the internal +agentic/interception control fields must NEVER reach the outbound provider HTTP +request body. The relevant fields are: + + _agentic_loop_depth + _agentic_loop_fingerprints + _agentic_loop_api_surface + max_agentic_loops + _code_interpreter_interception_active + _code_interpreter_interception_sandbox_key + _code_interpreter_interception_converted_stream + +A scrubber in gpt_transformation.py used to strip these. That scrubber was +removed, so `test_internal_control_fields_never_leak_into_provider_body` proves +they stay out of the body even without it. +""" + +import os +import sys +from typing import Any, Dict, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, +) +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import ( + Choices, + Function, + ChatCompletionMessageToolCall, + Message, + ModelResponse, +) + +# The internal control fields that must never reach a provider request body. +_INTERNAL_CONTROL_FIELDS = ( + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_agentic_loop_api_surface", + "max_agentic_loops", + "_code_interpreter_interception_active", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_converted_stream", + "litellm_metadata", +) + + +@pytest.fixture +def restore_callbacks(): + """Save/restore litellm.callbacks so a registered fake logger never pollutes + other tests in the suite.""" + saved = list(litellm.callbacks) + try: + yield + finally: + litellm.callbacks = saved + + +class _SandboxResult: + def __init__(self, stdout: str) -> None: + self.stdout = stdout + self.error = None + + +class FakeSandboxConfig: + """Injected sandbox so the interception loop runs no real network / E2B.""" + + def __init__(self) -> None: + self.created = 0 + self.deleted = 0 + self.run_codes: List[str] = [] + + async def acreate_sandbox(self) -> Any: + self.created += 1 + return MagicMock(id="sandbox-123") + + async def arun_code(self, container: Any, code: str) -> _SandboxResult: + self.run_codes.append(code) + return _SandboxResult(stdout="42\n") + + async def adelete_sandbox(self, container: Any) -> None: + self.deleted += 1 + + +def _tool_call_model_response() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_abc", + type="function", + function=Function( + name="litellm_code_execution", + arguments='{"code": "print(6*7)"}', + ), + ) + ], + ), + ) + ] + ) + + +def _plain_model_response(content: str = "The answer is 42") -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="stop", + message=Message(role="assistant", content=content), + ) + ] + ) + + +def _raw_response_for(model_response: ModelResponse) -> MagicMock: + """Wrap a ModelResponse as the OpenAI `with_raw_response.create` return value + (an object exposing `.headers` and `.parse()` -> something with model_dump).""" + parsed = MagicMock() + parsed.model_dump.return_value = model_response.model_dump() + raw = MagicMock() + raw.headers = {} + raw.parse.return_value = parsed + return raw + + +# --------------------------------------------------------------------------- +# A) PROVIDER-PAYLOAD REGRESSION +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_internal_control_fields_never_leak_into_provider_body(restore_callbacks): + """Drive a real acompletion with a native code_interpreter tool through the + interception logger + agentic loop, capturing every outbound OpenAI request + body. None of the internal control fields may appear at top-level or inside + extra_body on ANY of the captured calls.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandboxConfig()) + litellm.callbacks = [logger] + + # First create -> model emits a code_execution tool call (triggers the loop). + # Second create -> model returns a plain answer (loop terminates). + create = AsyncMock( + side_effect=[ + _raw_response_for(_tool_call_model_response()), + _raw_response_for(_plain_model_response()), + ] + ) + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create = create + + response = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + tools=[{"type": "code_interpreter"}], + tool_choice={"type": "code_interpreter"}, + api_key="sk-test", + client=mock_client, + ) + + # The loop must have actually fired (sanity: two provider calls). + assert create.await_count == 2, ( + "expected the agentic loop to issue a follow-up provider call; " + f"got {create.await_count} call(s)" + ) + + for idx, call in enumerate(create.await_args_list): + body = call.kwargs + extra_body = body.get("extra_body") or {} + for field in _INTERNAL_CONTROL_FIELDS: + assert field not in body, ( + f"provider call #{idx}: internal field {field!r} leaked into " + f"top-level request body: {sorted(body.keys())}" + ) + assert field not in extra_body, ( + f"provider call #{idx}: internal field {field!r} leaked into " + f"extra_body: {sorted(extra_body.keys())}" + ) + # The native code_interpreter tool must have been swapped for the + # function tool, never sent raw to OpenAI as a chat-completions request. + for tool in body.get("tools") or []: + assert tool.get("type") != "code_interpreter" + + # The final response is the post-loop answer, not the tool-call turn. + assert response.choices[0].message.content == "The answer is 42" + + +# --------------------------------------------------------------------------- +# B) DISPATCHER UNIT TESTS +# --------------------------------------------------------------------------- + + +class _LoggingStub: + """Minimal logging_obj: dispatcher only reads dynamic_success_callbacks and + litellm_call_id off it.""" + + litellm_call_id = "call-test" + dynamic_success_callbacks: List[Any] = [] + + +class _GateOnlyLogger(CustomLogger): + """Overrides the gate to fire, but builds a plan from request_patch.""" + + def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None: + super().__init__() + self._plan = plan + self._tool_calls = tool_calls + self.cleanup_calls = 0 + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict[str, Any], + ) -> Tuple[bool, Dict[str, Any]]: + return True, self._tool_calls + + async def async_build_agentic_loop_plan( + self, + tools: Dict[str, Any], + model: str, + messages: List[Dict[str, Any]], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict[str, Any], + logging_obj: Any, + stream: bool, + kwargs: Dict[str, Any], + ) -> AgenticLoopPlan: + return self._plan + + async def async_agentic_loop_cleanup_hook( + self, plan: AgenticLoopPlan, kwargs: Dict[str, Any] + ) -> None: + self.cleanup_calls += 1 + + +def _patched_messages() -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "what is 6*7?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "litellm_code_execution", + "arguments": '{"code": "print(6*7)"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc", "content": "42\n"}, + ] + + +@pytest.mark.asyncio +async def test_dispatcher_returns_none_when_no_callback_gates(restore_callbacks): + """No callback overrides the gate -> dispatcher returns None so the caller + keeps the original response untouched.""" + litellm.callbacks = [] + + result = await maybe_run_chat_completion_agentic_loop( + response=_plain_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_dispatcher_runs_followup_with_incremented_depth_and_patched_messages( + restore_callbacks, +): + """A gating logger with a request_patch -> the dispatcher calls + litellm.acompletion exactly once with _agentic_loop_depth == 1 and the + patched messages. Loop-control state rides as litellm-level kwargs and is + mirrored into litellm_metadata; the provider-surface transient + _agentic_loop_api_surface is never forwarded. (Provider-body stripping of + these litellm-level kwargs is asserted separately in test A.)""" + followup = _plain_model_response("done") + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages()), + ) + logger = _GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]}) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock(return_value=followup) + with patch.object(litellm, "acompletion", acompletion_mock): + result = await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={"temperature": 0.1}, + kwargs={"_code_interpreter_interception_active": True}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + assert result is followup + acompletion_mock.assert_awaited_once() + call_kwargs = acompletion_mock.await_args.kwargs + + assert call_kwargs["_agentic_loop_depth"] == 1 + assert call_kwargs["messages"] == _patched_messages() + # Preserved non-internal optional param survives the rerun. + assert call_kwargs["temperature"] == 0.1 + # Loop-control state is carried at the litellm level for the follow-up. + assert call_kwargs["max_agentic_loops"] >= 1 + assert "_agentic_loop_fingerprints" in call_kwargs + # Interception markers are mirrored into litellm_metadata for the follow-up. + assert ( + call_kwargs["litellm_metadata"]["_code_interpreter_interception_active"] is True + ) + # The transient surface marker is NOT forwarded to the follow-up call. + assert "_agentic_loop_api_surface" not in call_kwargs + # Cleanup hook always runs. + assert logger.cleanup_calls == 1 + + +@pytest.mark.asyncio +async def test_dispatcher_raises_when_depth_reaches_max_agentic_loops( + restore_callbacks, +): + """depth >= max_agentic_loops -> ValueError mentioning max_agentic_loops, + before any follow-up call is attempted.""" + logger = _GateOnlyLogger( + plan=AgenticLoopPlan(run_agentic_loop=True), + tool_calls={"tool_calls": [{"id": "call_abc"}]}, + ) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + with pytest.raises(ValueError, match="max_agentic_loops"): + await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + acompletion_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callbacks): + """A tool_calls fingerprint already present in _agentic_loop_fingerprints -> + ValueError about the repeated fingerprint (cycle guard), with no follow-up + call.""" + import json + + # The dispatcher fingerprints the whole value the gate returns as its second + # tuple element, so the seeded fingerprint must mirror that dict exactly. + gate_tool_calls = { + "tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}] + } + fingerprint = json.dumps(gate_tool_calls, sort_keys=True, default=str) + + logger = _GateOnlyLogger( + plan=AgenticLoopPlan(run_agentic_loop=True), + tool_calls=gate_tool_calls, + ) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + with pytest.raises(ValueError, match="fingerprint"): + await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={ + "_agentic_loop_depth": 0, + "max_agentic_loops": 3, + "_agentic_loop_fingerprints": [fingerprint], + }, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + acompletion_mock.assert_not_awaited() diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py index 3c280c6ba92..84900e3f2ed 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py @@ -132,3 +132,17 @@ def test_azure_base_model_detection_preserved(): assert params is not None assert "reasoning_effort" in params assert "tools" in params + + +def test_sambanova_embeddings_request_returns_list_not_none(): + """The sambanova embeddings branch resolved the config but dropped the result, + so embedding requests got ``None`` instead of the supported-params list while the + chat branch returned correctly. A list (the sambanova embeddings config exposes no + extra params, hence ``[]``) must reach the caller.""" + embedding_params = get_supported_openai_params( + model="E5-Mistral-7B-Instruct", + custom_llm_provider="sambanova", + request_type="embeddings", + ) + + assert embedding_params == [] diff --git a/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py b/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py new file mode 100644 index 00000000000..4e016622f1c --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py @@ -0,0 +1,58 @@ +"""Unit tests for litellm.litellm_core_utils.request_timeout_resolver. + +The resolver decides whether ``litellm.request_timeout`` was *explicitly configured* +(env REQUEST_TIMEOUT / litellm_settings, or a non-default runtime value) versus left +at the package default. This is what lets request_timeout act as an independent +per-attempt timeout instead of being indistinguishable from "nobody set it". +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))) + +import litellm +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) + + +@pytest.fixture +def restore_request_timeout(): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + try: + yield + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + +def test_default_value_without_flag_is_unset(restore_request_timeout): + litellm.request_timeout = DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + assert get_configured_request_timeout() is None + + +def test_explicit_flag_returns_value(restore_request_timeout): + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + assert get_configured_request_timeout() == 300.0 + + +def test_explicit_flag_preserves_value_equal_to_default(restore_request_timeout): + # The case the bare ``!= default`` heuristic gets wrong: a user who explicitly + # configures the default value still means it explicitly. + litellm.request_timeout = DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = True + assert get_configured_request_timeout() == float(DEFAULT_REQUEST_TIMEOUT_SECONDS) + + +def test_non_default_runtime_value_treated_as_explicit(restore_request_timeout): + # SDK users assigning litellm.request_timeout directly (no flag) must keep working. + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = False + assert get_configured_request_timeout() == 300.0 diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index 6808c4821c1..7239636fd48 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -126,6 +126,49 @@ def test_lists_with_sensitive_keys_are_masked(): assert masked["tags"] == ["prod", "test"] +def test_short_secrets_are_fully_masked(): + """ + Regression test: secrets at or below the reveal threshold (visible_prefix + + visible_suffix, 8 by default) were returned verbatim instead of masked. + An exactly-8-char value hit masked_length == 0 and round-tripped unchanged; + anything shorter hit the early return. Both leaked short credentials (e.g. an + 8-char redis password) in plaintext through mask_dict. + """ + masker = SensitiveDataMasker() + + # Boundary: exactly 8 chars previously returned verbatim. + assert masker._mask_value("abcd1234") == "********" + # Below threshold previously hit the early return and leaked verbatim. + assert masker._mask_value("sk-12") == "*****" + # Values above the threshold must still partially reveal, not over-mask. + assert masker._mask_value("abcd12345") == "abcd*2345" + + masked = masker.mask_dict({"redis_password": "pass1234", "api_key": "sk-7a"}) + assert masked["redis_password"] == "********" + assert masked["api_key"] == "*****" + + +def test_mask_short_values_false_keeps_short_values_readable(): + """ + mask_short_values=False opts out of full masking so short values are returned + as-is. This preserves the truncation use (e.g. CooldownCache shows the first 50 + chars of an exception and only masks longer tails), while longer values are still + partially masked. + """ + masker = SensitiveDataMasker( + visible_prefix=50, visible_suffix=0, mask_short_values=False + ) + + short = "Test exception for structure validation" + assert masker._mask_value(short) == short + + long_value = "x" * 60 + masked = masker._mask_value(long_value) + assert masked.startswith("x" * 50) + assert masked.endswith("*" * 10) + assert len(masked) == 60 + + def test_cost_per_token_fields_not_masked(): """ Regression test: cost fields like input_cost_per_token contain "token" in their name diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 734e49161c4..81af0ad3e6f 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -878,6 +878,114 @@ def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging): assert "invalid maxOutputTokens" in str(excinfo.value) +def _bedrock_error_event(exception_type: str): + """A mocked botocore event-stream error event: status_code is botocore's + hard-coded 400, with the real type in the :exception-type header.""" + event = Mock() + event.to_response_dict = Mock( + return_value={ + "status_code": 400, + "headers": { + ":exception-type": exception_type, + ":content-type": "application/json", + ":message-type": "exception", + }, + "body": b'{"message":"Bedrock had an internal error."}', + } + ) + return event + + +@pytest.mark.asyncio +async def test_bedrock_midstream_internal_server_error_wraps_for_fallback( + logging_obj: Logging, +): + """End-to-end regression for https://github.com/BerriAI/litellm/issues/24608: + a Bedrock mid-stream internalServerException event (botocore stamps it 400) + must flow through the real decoder, gain its modeled 500 status, and wrap + into MidStreamFallbackError so the Router can run streaming fallback. + + Calls the real AWSEventStreamDecoder, so reverting the decoder status fix + makes the decoder raise BedrockError(400) and the gate raises BadRequestError + directly -> this test fails without the fix.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + decoder = AWSEventStreamDecoder(model="anthropic.claude-3-sonnet-20240229-v1:0") + + async def _bedrock_stream(): + decoder._parse_message_from_event( + _bedrock_error_event("internalServerException") + ) + yield # unreachable; the line above raises + + async def _make_call(**kwargs): + return _bedrock_stream() + + response = CustomStreamWrapper( + completion_stream=None, + model="anthropic.claude-3-sonnet-20240229-v1:0", + logging_obj=logging_obj, + custom_llm_provider="bedrock", + make_call=_make_call, + ) + + with pytest.raises(MidStreamFallbackError): + await response.__anext__() + + +@pytest.mark.asyncio +async def test_bedrock_5xx_wraps_for_midstream_fallback(logging_obj: Logging): + """Gate contract: a Bedrock 5xx (here 503 serviceUnavailableException) wraps + into MidStreamFallbackError so the Router can run streaming fallback.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.bedrock.chat.invoke_handler import BedrockError + + async def _raise_503(**kwargs): + raise BedrockError( + status_code=503, + message="serviceUnavailableException Bedrock is unavailable.", + ) + + response = CustomStreamWrapper( + completion_stream=None, + model="anthropic.claude-3-sonnet-20240229-v1:0", + logging_obj=logging_obj, + custom_llm_provider="bedrock", + make_call=_raise_503, + ) + + with pytest.raises(MidStreamFallbackError): + await response.__anext__() + + +@pytest.mark.asyncio +async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): + """Gate contract: a Bedrock validationException (400) is a client error and + must surface directly, never wrapped into MidStreamFallbackError.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.bedrock.chat.invoke_handler import BedrockError + + async def _raise_400(**kwargs): + raise BedrockError( + status_code=400, + message="validationException malformed input.", + ) + + response = CustomStreamWrapper( + completion_stream=None, + model="anthropic.claude-3-sonnet-20240229-v1:0", + logging_obj=logging_obj, + custom_llm_provider="bedrock", + make_call=_raise_400, + ) + + with pytest.raises(Exception) as excinfo: + await response.__anext__() + assert not isinstance(excinfo.value, MidStreamFallbackError) + assert getattr(excinfo.value, "status_code", None) == 400 + + @pytest.mark.asyncio async def test_async_streaming_read_timeout_triggers_midstream_fallback( logging_obj: Logging, diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 2876b56f516..b111b65e3af 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1886,6 +1886,89 @@ def test_anthropic_model_supports_effort_param_rejects_non_supporting_models(mod assert AnthropicConfig._model_supports_effort_param(model) is False +@pytest.mark.parametrize( + "model", + [ + "claude-opus-4-6", + "claude-opus-4-7", + "claude-opus-4-8", + "claude-opus-4-6-20260205", + "claude-opus-4-7-20260416", + ], +) +def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model): + assert AnthropicConfig._model_supports_speed_param(model) is True + + +@pytest.mark.parametrize( + "model", + [ + "claude-sonnet-4-6", + "claude-fable-5", + "claude-3-haiku-20240307", + "vertex_ai/claude-opus-4-8", + "azure_ai/claude-opus-4-8", + "anthropic.claude-opus-4-8", + ], +) +def test_anthropic_model_supports_speed_param_rejects_non_supporting_models(model): + assert AnthropicConfig._model_supports_speed_param(model) is False + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "azure_ai", "bedrock"]) +def test_anthropic_model_supports_speed_param_rejects_non_anthropic_providers( + custom_llm_provider, +): + """Fast mode is direct-Anthropic-only. Vertex/Azure/Bedrock strip their prefix + before the shared transform runs, so the bare Opus id must still be rejected.""" + assert ( + AnthropicConfig._model_supports_speed_param( + "claude-opus-4-8", custom_llm_provider + ) + is False + ) + assert ( + AnthropicConfig._model_supports_speed_param("claude-opus-4-8", "anthropic") + is True + ) + + +def test_vertex_anthropic_drops_speed_for_opus_with_drop_params(monkeypatch): + """Regression: vertex_ai Opus must drop ``speed`` even though the prefix-stripped + ``claude-opus-4-8`` maps to a fast-mode-capable direct-Anthropic entry.""" + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( + VertexAIAnthropicConfig, + ) + + monkeypatch.setattr(litellm, "drop_params", True) + result = VertexAIAnthropicConfig().transform_request( + model="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"speed": "fast", "max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + + assert "speed" not in result + + +def test_vertex_anthropic_raises_on_speed_without_drop_params(monkeypatch): + """Regression: vertex_ai Opus raises rather than forwarding an unsupported + ``speed`` when neither global nor per-request drop_params is set.""" + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( + VertexAIAnthropicConfig, + ) + + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + VertexAIAnthropicConfig().map_openai_params( + non_default_params={"speed": "fast"}, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + + def test_translate_system_message_skips_empty_string_content(): """ Test that translate_system_message skips system messages with empty string content. @@ -3766,6 +3849,61 @@ def test_fast_mode_parameter_mapping(): assert result["speed"] == "fast" +def test_anthropic_drop_params_strips_speed_for_unsupported_models(): + """``drop_params=True`` strips unsupported ``speed`` for non-Opus models.""" + config = AnthropicConfig() + messages = [{"role": "user", "content": "Hello"}] + + original = litellm.drop_params + litellm.drop_params = True + try: + result = config.transform_request( + model="claude-sonnet-4-6", + messages=messages, + optional_params={"speed": "fast", "max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + finally: + litellm.drop_params = original + + assert "speed" not in result + + +def test_anthropic_drop_params_keeps_speed_for_supporting_models(): + """``drop_params=True`` must not strip ``speed`` on Opus fast-mode models.""" + config = AnthropicConfig() + messages = [{"role": "user", "content": "Hello"}] + + original = litellm.drop_params + litellm.drop_params = True + try: + result = config.transform_request( + model="claude-opus-4-6", + messages=messages, + optional_params={"speed": "fast", "max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + finally: + litellm.drop_params = original + + assert result.get("speed") == "fast" + + +def test_speed_raises_clean_error_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params={"speed": "fast"}, + optional_params={}, + model="claude-sonnet-4-6", + drop_params=False, + ) + + def test_map_openai_params_max_tokens_normalized_to_int(): """ Test that map_openai_params normalizes max_tokens to an integer (e.g. 0.7 -> 1). 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 0300b6f3f51..f93f2404fa4 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 @@ -508,6 +508,35 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c assert result[0]["input"] == {"location": "Boston"} +def test_translate_openai_content_to_anthropic_sanitizes_colon_dot_tool_call_ids(): + """Cross-provider ids like ``functions.Bash:0`` must be normalized for Anthropic replay.""" + openai_choices = [ + Choices( + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionAssistantToolCall( + id="functions.Bash:0", + type="function", + function=Function( + name="Bash", + arguments='{"command": "ls"}', + ), + ) + ], + ) + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter._translate_openai_content_to_anthropic(choices=openai_choices) + + assert len(result) == 1 + assert result[0]["type"] == "tool_use" + assert result[0]["id"] == "functions_Bash_0" + + def test_translate_openai_response_to_anthropic_text_and_tool_calls(): """`translate_openai_response_to_anthropic` should surface assistant text even when tools fire.""" openai_response = ModelResponse( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index b1e1d789d74..a58873f5387 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -105,6 +105,48 @@ async def test_anthropic_messages_sanitizes_empty_text_blocks_before_dispatch(): assert len(msgs[0]["content"]) == 2 # caller untouched +@pytest.mark.asyncio +async def test_anthropic_messages_sanitizes_tool_use_ids_before_dispatch(): + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "functions.Bash:0", + "name": "Bash", + "input": {}, + } + ], + } + ] + captured = {} + + def fake_handler(*args, **kwargs): + captured["messages"] = kwargs.get("messages") + return "stub" + + fake_loop = MagicMock() + fake_loop.run_in_executor = lambda _e, func: _async_return(func()) + + with ( + patch.object(handler, "anthropic_messages_handler", side_effect=fake_handler), + patch("asyncio.get_event_loop", return_value=fake_loop), + ): + await handler.anthropic_messages( + max_tokens=100, + messages=msgs, + model="anthropic/claude-sonnet-4-5-20250929", + custom_llm_provider="anthropic", + api_key="k", + ) + + assert captured["messages"][0]["content"][0]["id"] == "functions_Bash_0" + assert msgs[0]["content"][0]["id"] == "functions.Bash:0" + + async def _async_return(value): return value diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py new file mode 100644 index 00000000000..6900f1062bf --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py @@ -0,0 +1,117 @@ +import litellm +import pytest +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) +from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + AnthropicMessagesRequestUtils, +) + + +def test_messages_drop_params_strips_speed_for_unsupported_models(): + original = litellm.drop_params + litellm.drop_params = True + try: + optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={ + "max_tokens": 1024, + "speed": "fast", + "messages": [{"role": "user", "content": "Hello"}], + }, + model="claude-sonnet-4-6", + drop_params=False, + ) + config = AnthropicMessagesConfig() + headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + optional_params=dict(optional_params), + litellm_params={}, + ) + result = config.transform_anthropic_messages_request( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=dict(optional_params), + litellm_params={}, + headers=headers, + ) + finally: + litellm.drop_params = original + + assert "speed" not in optional_params + assert "speed" not in result + assert "fast-mode-2026-02-01" not in headers.get("anthropic-beta", "") + + +def test_messages_drop_params_keeps_speed_for_supporting_models(): + original = litellm.drop_params + litellm.drop_params = True + try: + optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"max_tokens": 1024, "speed": "fast"}, + model="claude-opus-4-6", + drop_params=False, + ) + config = AnthropicMessagesConfig() + headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="claude-opus-4-6", + messages=[{"role": "user", "content": "Hello"}], + optional_params=dict(optional_params), + litellm_params={}, + ) + result = config.transform_anthropic_messages_request( + model="claude-opus-4-6", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=dict(optional_params), + litellm_params={}, + headers=headers, + ) + finally: + litellm.drop_params = original + + assert optional_params.get("speed") == "fast" + assert result.get("speed") == "fast" + assert "fast-mode-2026-02-01" in headers.get("anthropic-beta", "") + + +def test_messages_raises_when_speed_unsupported_and_drop_params_false(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"max_tokens": 1024, "speed": "fast"}, + model="claude-sonnet-4-6", + drop_params=False, + ) + + +def test_messages_drops_speed_for_vertex_opus_with_drop_params(monkeypatch): + """Regression: a vertex_ai Opus passthrough must drop ``speed`` even though the + prefix-stripped model id maps to a fast-mode-capable direct-Anthropic entry.""" + monkeypatch.setattr(litellm, "drop_params", True) + optional_params = ( + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"max_tokens": 1024, "speed": "fast"}, + model="claude-opus-4-8", + drop_params=False, + custom_llm_provider="vertex_ai", + ) + ) + + assert "speed" not in optional_params + + +def test_messages_raises_for_vertex_opus_without_drop_params(monkeypatch): + """Regression: vertex_ai Opus passthrough raises rather than forwarding an + unsupported ``speed`` when drop_params is unset.""" + monkeypatch.setattr(litellm, "drop_params", False) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"max_tokens": 1024, "speed": "fast"}, + model="claude-opus-4-8", + drop_params=False, + custom_llm_provider="vertex_ai", + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py index 3ce076640e8..f0252e13336 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py @@ -6,6 +6,7 @@ Regression tests for the /v1/messages request-parse fast paths: while resolving the (static) type hints only once per process. """ +import litellm from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( AnthropicMessagesRequestUtils, _anthropic_messages_optional_param_keys, @@ -54,3 +55,36 @@ def test_empty_params(): ) == {} ) + + +def test_drop_params_strips_speed_for_unsupported_model(): + original = litellm.drop_params + litellm.drop_params = True + try: + result = ( + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"speed": "fast", "temperature": 0.5}, + model="claude-sonnet-4-6", + ) + ) + finally: + litellm.drop_params = original + + assert result == {"temperature": 0.5} + assert "speed" not in result + + +def test_drop_params_keeps_speed_for_supporting_model(): + original = litellm.drop_params + litellm.drop_params = True + try: + result = ( + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"speed": "fast"}, + model="claude-opus-4-6", + ) + ) + finally: + litellm.drop_params = original + + assert result == {"speed": "fast"} diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a09e55d4ed7..03b72c504b8 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1329,6 +1329,52 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_text_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["tool_result"] + def test_sanitize_tool_use_ids_in_anthropic_messages(self): + from litellm.llms.anthropic.common_utils import ( + sanitize_tool_use_ids_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "functions.Bash:0", + "name": "Bash", + "input": {}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "functions.Bash:0", + "content": "ok", + } + ], + }, + ] + out = sanitize_tool_use_ids_in_anthropic_messages(msgs) + assert out[0]["content"][0]["id"] == "functions_Bash_0" + assert out[1]["content"][0]["tool_use_id"] == "functions_Bash_0" + assert msgs[0]["content"][0]["id"] == "functions.Bash:0" + + def test_normalize_anthropic_tool_use_id_strips_thought_signature(self): + from litellm.litellm_core_utils.prompt_templates.factory import ( + THOUGHT_SIGNATURE_SEPARATOR, + ) + from litellm.llms.anthropic.common_utils import normalize_anthropic_tool_use_id + + base = "call_abc123" + sig = "CiIBDDnWx+/a==" + assert ( + normalize_anthropic_tool_use_id(f"{base}{THOUGHT_SIGNATURE_SEPARATOR}{sig}") + == base + ) + def test_anthropic_messages_config_http_retry_helpers(self): import httpx diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py index 32838701949..bc26268ee92 100644 --- a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py +++ b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py @@ -66,9 +66,8 @@ class TestAPISerpentConfig: assert headers["X-API-Key"] == "test-api-key" assert headers["Content-Type"] == "application/json" - @patch("litellm.llms.apiserpent.search.transformation.get_secret_str") - def test_validate_environment_without_api_key(self, mock_get_secret): - mock_get_secret.return_value = None + def test_validate_environment_without_api_key(self, monkeypatch): + monkeypatch.delenv("APISERPENT_API_KEY", raising=False) with pytest.raises(ValueError, match="APISERPENT_API_KEY is not set"): APISerpentSearchConfig().validate_environment({}) diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py new file mode 100644 index 00000000000..a1353d57038 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -0,0 +1,329 @@ +""" +Regression tests for the host-aware server-credential fallback guard in +``BaseSearchConfig``. + +A caller-supplied ``api_base`` is honored when building the request URL, so +falling back to a server-configured secret while the caller controls the host +would send the operator's credential to an attacker. The guard must refuse that +combination for every provider that carries a server-managed secret, while +leaving keyless providers and legitimate operator overrides untouched. +""" + +from typing import Dict, Tuple, Type +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.llms.apiserpent.search.transformation import APISerpentSearchConfig +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + _is_trusted_search_api_base, +) +from litellm.llms.brave.search.transformation import BraveSearchConfig +from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig +from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig +from litellm.llms.fastcrw.search.transformation import FastCRWSearchConfig +from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig +from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig +from litellm.llms.linkup.search.transformation import LinkupSearchConfig +from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig +from litellm.llms.perplexity.search.transformation import PerplexitySearchConfig +from litellm.llms.searchapi.search.transformation import SearchAPIConfig +from litellm.llms.searxng.search.transformation import SearXNGSearchConfig +from litellm.llms.serper.search.transformation import SerperSearchConfig +from litellm.llms.tavily.search.transformation import TavilySearchConfig +from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig +from litellm.llms.you_com.search.transformation import YouComSearchConfig + +ATTACKER_BASE = "https://attacker.example.com" + +# Every *_API_BASE override env var that could otherwise mark the attacker host +# as trusted; cleared before each test so the suite is hermetic. +_BASE_ENV_VARS = ( + "SERPER_API_BASE", + "TAVILY_API_BASE", + "PERPLEXITY_API_BASE", + "APISERPENT_API_BASE", + "EXA_API_BASE", + "BRAVE_API_BASE", + "FIRECRAWL_API_BASE", + "LINKUP_API_BASE", + "SEARCHAPI_API_BASE", + "GOOGLE_PSE_API_BASE", + "PARALLEL_AI_API_BASE", + "YOUCOM_API_BASE", + "SEARXNG_API_BASE", + "DATAFORSEO_API_BASE", + "TINYFISH_API_BASE", + "CRW_API_BASE", +) + + +@pytest.fixture(autouse=True) +def _clear_base_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + for var in _BASE_ENV_VARS: + monkeypatch.delenv(var, raising=False) + + +# (config, {server secret env vars}, caller_api_key honored as-is, extra env for full validate) +ProviderSpec = Tuple[Type[BaseSearchConfig], Dict[str, str], str, Dict[str, str]] + +PROVIDERS: Tuple[ProviderSpec, ...] = ( + (SerperSearchConfig, {"SERPER_API_KEY": "srv"}, "caller-key", {}), + (TavilySearchConfig, {"TAVILY_API_KEY": "srv"}, "caller-key", {}), + (PerplexitySearchConfig, {"PERPLEXITYAI_API_KEY": "srv"}, "caller-key", {}), + (APISerpentSearchConfig, {"APISERPENT_API_KEY": "srv"}, "caller-key", {}), + (ExaAISearchConfig, {"EXA_API_KEY": "srv"}, "caller-key", {}), + (BraveSearchConfig, {"BRAVE_API_KEY": "srv"}, "caller-key", {}), + (FirecrawlSearchConfig, {"FIRECRAWL_API_KEY": "srv"}, "caller-key", {}), + (LinkupSearchConfig, {"LINKUP_API_KEY": "srv"}, "caller-key", {}), + (SearchAPIConfig, {"SEARCHAPI_API_KEY": "srv"}, "caller-key", {}), + ( + GooglePSESearchConfig, + {"GOOGLE_PSE_API_KEY": "srv"}, + "caller-key", + {"GOOGLE_PSE_ENGINE_ID": "engine"}, + ), + (ParallelAISearchConfig, {"PARALLEL_API_KEY": "srv"}, "caller-key", {}), + (YouComSearchConfig, {"YOUCOM_API_KEY": "srv"}, "caller-key", {}), + (SearXNGSearchConfig, {"SEARXNG_API_KEY": "srv"}, "caller-key", {}), + ( + DataForSEOSearchConfig, + {"DATAFORSEO_LOGIN": "srv", "DATAFORSEO_PASSWORD": "pw"}, + "login:password", + {}, + ), + (TinyfishSearchConfig, {"TINYFISH_API_KEY": "srv"}, "caller-key", {}), + (FastCRWSearchConfig, {"CRW_API_KEY": "srv"}, "caller-key", {}), +) + +_IDS = tuple(spec[0].__name__ for spec in PROVIDERS) + + +@pytest.mark.parametrize( + "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS +) +def test_server_secret_refused_for_caller_api_base( + config_cls: Type[BaseSearchConfig], + server_env: Dict[str, str], + caller_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + for key, value in {**server_env, **extra_env}.items(): + monkeypatch.setenv(key, value) + + with pytest.raises(ValueError, match="Refusing to send the server-configured"): + config_cls().validate_environment(headers={}, api_base=ATTACKER_BASE) + + +@pytest.mark.parametrize( + "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS +) +def test_caller_supplied_key_is_honored_for_custom_api_base( + config_cls: Type[BaseSearchConfig], + server_env: Dict[str, str], + caller_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + for key, value in {**server_env, **extra_env}.items(): + monkeypatch.setenv(key, value) + + # An explicit caller key is the caller's own credential, so pointing it at + # the caller's own host must be allowed. + config_cls().validate_environment( + headers={}, api_key=caller_key, api_base=ATTACKER_BASE + ) + + +@pytest.mark.parametrize( + "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS +) +def test_server_secret_used_without_caller_api_base( + config_cls: Type[BaseSearchConfig], + server_env: Dict[str, str], + caller_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + for key, value in {**server_env, **extra_env}.items(): + monkeypatch.setenv(key, value) + + # No caller-supplied api_base -> the request targets the trusted default, so + # the server secret is still used and nothing is refused. + config_cls().validate_environment(headers={}) + + +def test_keyless_provider_allows_caller_api_base( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("SEARXNG_API_KEY", raising=False) + + headers = SearXNGSearchConfig().validate_environment( + headers={}, api_base="https://my-searxng.internal" + ) + + assert "Authorization" not in headers + + +def test_operator_env_base_override_is_trusted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SERPER_API_KEY", "srv") + monkeypatch.setenv("SERPER_API_BASE", "https://serper.internal.corp") + + # Mirrors the second validate_environment call in the search handler, which + # receives the already-resolved operator base as api_base. + headers = SerperSearchConfig().validate_environment( + headers={}, api_base="https://serper.internal.corp/search" + ) + + assert headers["X-API-KEY"] == "srv" + + +class TestResolveServerApiKey: + def test_caller_key_short_circuits(self) -> None: + result = BaseSearchConfig().resolve_server_api_key( + caller_api_key="mine", + caller_api_base=ATTACKER_BASE, + key_env_vars=("SERPER_API_KEY",), + base_env_var="SERPER_API_BASE", + default_api_base="https://google.serper.dev", + ) + assert result == "mine" + + def test_returns_none_when_no_server_secret( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("SEARXNG_API_KEY", raising=False) + result = BaseSearchConfig().resolve_server_api_key( + caller_api_key=None, + caller_api_base=ATTACKER_BASE, + key_env_vars=("SEARXNG_API_KEY",), + base_env_var="SEARXNG_API_BASE", + default_api_base=None, + ) + assert result is None + + def test_first_set_env_var_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PARALLEL_AI_API_KEY", raising=False) + monkeypatch.setenv("PARALLEL_API_KEY", "second") + result = BaseSearchConfig().resolve_server_api_key( + caller_api_key=None, + caller_api_base=None, + key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"), + base_env_var="PARALLEL_AI_API_BASE", + default_api_base="https://api.parallel.ai", + ) + assert result == "second" + + +class TestIsTrustedSearchApiBase: + def test_matches_default_host(self) -> None: + assert _is_trusted_search_api_base( + "https://google.serper.dev/search", "https://google.serper.dev", None + ) + + def test_foreign_host_untrusted(self) -> None: + assert not _is_trusted_search_api_base( + ATTACKER_BASE, "https://google.serper.dev", None + ) + + def test_env_override_host_trusted(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SERPER_API_BASE", "https://serper.internal.corp") + assert _is_trusted_search_api_base( + "https://serper.internal.corp/search", + "https://google.serper.dev", + "SERPER_API_BASE", + ) + + def test_schemeless_candidate_untrusted(self) -> None: + # Without a scheme urlsplit puts the value in the path, leaving an empty + # netloc; an unparseable host must never be treated as trusted. + assert not _is_trusted_search_api_base( + "attacker.example.com", "https://google.serper.dev", None + ) + + +@pytest.mark.asyncio +async def test_asearch_does_not_leak_server_key_to_caller_api_base( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """End-to-end regression on the reported vector: a search call with a foreign + api_base and no caller key must fail without any outbound request carrying the + server-configured key.""" + monkeypatch.setenv("SERPER_API_KEY", "sk-server-secret") + monkeypatch.delenv("SERPER_API_BASE", raising=False) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get, + ): + with pytest.raises(Exception): + await litellm.asearch( + query="secrets", + search_provider="serper", + api_base=ATTACKER_BASE, + ) + + mock_post.assert_not_called() + mock_get.assert_not_called() + + +@pytest.mark.parametrize( + "provider, key_env, server_key, extra_env", + [ + ("searchapi", "SEARCHAPI_API_KEY", "sk-server-searchapi", {}), + ( + "google_pse", + "GOOGLE_PSE_API_KEY", + "sk-server-google", + {"GOOGLE_PSE_ENGINE_ID": "engine-id"}, + ), + ], +) +@pytest.mark.asyncio +async def test_query_param_key_not_leaked_with_dummy_caller_key( + provider: str, + key_env: str, + server_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Providers that send the key as a URL query param resolve it in + transform_search_request, not validate_environment. A caller who passes a + dummy api_key to clear the validate_environment short-circuit must not cause + the server key to be placed in the URL sent to their own api_base.""" + monkeypatch.setenv(key_env, server_key) + for name, value in extra_env.items(): + monkeypatch.setenv(name, value) + + captured: Dict[str, str] = {} + + async def fake_get(self, *args, **kwargs): # type: ignore[no-untyped-def] + captured["url"] = kwargs.get("url") or (args[0] if args else "") + raise RuntimeError("stop after capturing the outbound url") + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + fake_get, + ): + with pytest.raises(Exception): + await litellm.asearch( + query="secrets", + search_provider=provider, + api_key="sk-CALLER-DUMMY", + api_base=ATTACKER_BASE, + ) + + assert captured["url"], "expected an outbound request to be attempted" + assert server_key not in captured["url"] + assert "sk-CALLER-DUMMY" in captured["url"] diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index bbefdd621f0..f7f8f582abc 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -125,6 +125,74 @@ def test_mantle_messages_url_construction(): assert url == "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages" +_VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com" + + +def test_mantle_chat_url_honors_api_base_host(): + config = AmazonMantleConfig() + url = config.get_complete_url( + api_base=_VPC_ENDPOINT, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + +def test_mantle_chat_url_honors_api_base_full_path_without_duplication(): + config = AmazonMantleConfig() + full = f"{_VPC_ENDPOINT}/anthropic/v1/messages" + url = config.get_complete_url( + api_base=full, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == full + + +def test_mantle_messages_url_honors_api_base_host(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=_VPC_ENDPOINT, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + assert "api.aws" not in url + + +def test_mantle_messages_url_honors_api_base_with_trailing_slash(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=f"{_VPC_ENDPOINT}/", + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-gov-west-1"}, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + +def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={ + "aws_region_name": "us-gov-west-1", + "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT, + }, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + def test_mantle_transform_request_strips_prefix_and_adds_model(): config = AmazonMantleConfig() request = config.transform_request( @@ -247,3 +315,35 @@ async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body() assert requests[0]["path"] == "/anthropic/v1/messages" assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" assert "aws_bedrock_project_id" not in requests[0]["body"] + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_routes_to_vpc_api_base(): + import litellm + + urls = [] + + async def mock_post(self, url, data=None, headers=None, **kwargs): + urls.append(str(url)) + return _anthropic_response(str(url)) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + api_base=_VPC_ENDPOINT, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-gov-west-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert len(urls) == 1 + assert urls[0] == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + assert "api.aws" not in urls[0] diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py index cecb6024de1..1a46015ceb9 100644 --- a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py @@ -3,25 +3,190 @@ import pytest from litellm.llms.cloudflare.chat.transformation import CloudflareChatConfig -def test_get_complete_url_encodes_model_path_segment(): +def test_supported_params_include_tools_and_tool_choice(): config = CloudflareChatConfig() - assert ( - config.get_complete_url( - api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", - api_key="cf-key", - model="@cf/meta/llama?x=1#frag", - optional_params={}, - litellm_params={}, - ) - == "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/%40cf/meta/llama%3Fx%3D1%23frag" + params = config.get_supported_openai_params(model="@cf/meta/llama-2-7b-chat-int8") + + assert "tools" in params + assert "tool_choice" in params + assert "stream" in params + assert "max_tokens" in params + + +def test_get_complete_url_defaults_to_openai_compatible_endpoint(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base=None, + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, ) - with pytest.raises(ValueError, match="dot path segment"): + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + assert "/ai/run/" not in url + + +def test_get_complete_url_appends_chat_completions_to_explicit_base(): + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/v1", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + assert "/ai/run/" not in url + + +def test_get_complete_url_is_idempotent_for_full_base(): + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + + +def test_get_complete_url_falls_back_to_account_id_when_base_is_empty(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + + +def test_get_complete_url_raises_when_account_id_and_base_missing(monkeypatch): + monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False) + config = CloudflareChatConfig() + + with pytest.raises(ValueError, match="Missing CLOUDFLARE_ACCOUNT_ID"): config.get_complete_url( - api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_base=None, api_key="cf-key", - model="../../accounts/other", + model="@cf/meta/llama-2-7b-chat-int8", optional_params={}, litellm_params={}, ) + + +def test_get_complete_url_raises_when_account_id_is_empty(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", " ") + config = CloudflareChatConfig() + + with pytest.raises(ValueError, match="Missing CLOUDFLARE_ACCOUNT_ID"): + config.get_complete_url( + api_base=None, + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + +def test_get_complete_url_migrates_legacy_ai_run_base(): + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + assert "/ai/run" not in url + + +def test_transform_request_passes_tools_through_in_openai_format(): + config = CloudflareChatConfig() + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ] + messages = [{"role": "user", "content": "weather in nyc?"}] + + body = config.transform_request( + model="@cf/meta/llama-2-7b-chat-int8", + messages=messages, + optional_params={"tools": tools, "tool_choice": "auto"}, + litellm_params={}, + headers={}, + ) + + assert body["messages"] == messages + assert body["model"] == "@cf/meta/llama-2-7b-chat-int8" + assert body["tools"] == tools + assert body["tool_choice"] == "auto" + + +def test_validate_environment_requires_api_key(): + config = CloudflareChatConfig() + + with pytest.raises(ValueError, match="Missing Cloudflare API Key"): + config.validate_environment( + headers={}, + model="@cf/meta/llama-2-7b-chat-int8", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + +def test_validate_environment_sets_bearer_and_content_type(): + config = CloudflareChatConfig() + + headers = config.validate_environment( + headers={}, + model="@cf/meta/llama-2-7b-chat-int8", + messages=[], + optional_params={}, + litellm_params={}, + api_key="cf-key", + ) + + assert headers["Authorization"] == "Bearer cf-key" + assert headers["Content-Type"] == "application/json" diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index bf835b5d8f9..7bd1d7a6031 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -851,3 +851,56 @@ async def test_async_get_forwards_per_request_timeout(): } finally: await handler.close() + + +class TestDefaultCachedClientTimeoutHonorsRequestTimeout: + """Cached default httpx clients must fall back to an explicit litellm.request_timeout. + + Regression for LIT-2369: get_async_httpx_client / _get_httpx_client hardcoded a + 600s default and never consulted litellm.request_timeout, so provider calls with + no per-model timeout (e.g. Bedrock) hung for 600s. + """ + + @pytest.fixture + def restore_request_timeout(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + try: + yield + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_default_when_request_timeout_unset(self, restore_request_timeout): + from litellm.llms.custom_httpx.http_handler import ( + _DEFAULT_TIMEOUT, + _default_cached_client_timeout, + ) + + litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT + + def test_uses_explicit_request_timeout(self, restore_request_timeout): + from litellm.llms.custom_httpx.http_handler import ( + _default_cached_client_timeout, + ) + + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + resolved = _default_cached_client_timeout() + assert resolved.read == 300.0 + assert resolved.connect == 5.0 + + def test_cached_async_client_built_with_explicit_request_timeout( + self, restore_request_timeout + ): + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.utils import LlmProviders + + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + litellm.in_memory_llm_clients_cache = LLMClientCache() + client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK) + assert client.timeout.read == 300.0 diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index f7d445d0788..64ae30daa70 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -10,13 +10,21 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path import litellm +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + LITELLM_CODE_EXECUTION_TOOL_NAME, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, _google_genai_streaming_hidden_params, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams +_ACTIVE_KEY = "_code_interpreter_interception_active" +_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler @@ -116,6 +124,117 @@ def test_response_api_handler_streams_when_provider_transform_adds_stream(): assert client.post.call_args.kwargs["json"]["stream"] is True +def test_response_api_handler_runs_agentic_hooks_in_sync_path(monkeypatch): + handler = BaseLLMHTTPHandler() + config = Mock() + config.validate_environment.return_value = {} + config.get_complete_url.return_value = "https://chatgpt.example.com/responses" + config.transform_responses_api_request.return_value = { + "model": "gpt-5", + "input": "hi", + } + config.sign_request.return_value = ({}, None) + initial_response = Mock() + final_response = Mock() + config.transform_response_api_response.return_value = initial_response + + client = HTTPHandler(client=httpx.Client()) + client.post = Mock( + return_value=httpx.Response( + 200, + request=httpx.Request("POST", "https://chatgpt.example.com/responses"), + ) + ) + logging_obj = Mock() + + monkeypatch.setattr(handler, "_has_agentic_completion_hook", Mock(return_value=True)) + hook_mock = AsyncMock(return_value=final_response) + monkeypatch.setattr(handler, "_call_agentic_completion_hooks", hook_mock) + + response = handler.response_api_handler( + model="gpt-5", + input="hi", + responses_api_provider_config=config, + response_api_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + client=client, + ) + + assert response is final_response + hook_mock.assert_awaited_once() + assert hook_mock.call_args.kwargs["api_surface"] == "responses" + assert hook_mock.call_args.kwargs["messages"] == [ + {"role": "user", "content": "hi"} + ] + + +def test_response_api_handler_runs_responses_pre_call_hook_before_transform(): + handler = BaseLLMHTTPHandler() + config = Mock() + config.validate_environment.return_value = {} + config.get_complete_url.return_value = "https://api.openai.com/v1/responses" + config.sign_request.return_value = ({}, None) + initial_response = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="gpt-5", + ) + config.transform_response_api_response.return_value = initial_response + + def transform_responses_api_request(**kwargs): + return { + "model": kwargs["model"], + "input": kwargs["input"], + **kwargs["response_api_optional_request_params"], + } + + config.transform_responses_api_request.side_effect = transform_responses_api_request + client = HTTPHandler(client=httpx.Client()) + client.post = Mock( + return_value=httpx.Response( + 200, + request=httpx.Request("POST", "https://api.openai.com/v1/responses"), + ) + ) + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = [] + + old_callbacks = list(litellm.callbacks) + litellm.callbacks = [CodeInterpreterInterceptionLogger()] + try: + response = handler.response_api_handler( + model="gpt-5", + input="use code", + responses_api_provider_config=config, + response_api_optional_request_params={ + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}] + }, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(api_key="sk-test"), + logging_obj=logging_obj, + client=client, + ) + finally: + litellm.callbacks = old_callbacks + + assert response is initial_response + transform_kwargs = config.transform_responses_api_request.call_args.kwargs + tools = transform_kwargs["response_api_optional_request_params"]["tools"] + assert not any(tool.get("type") == "code_interpreter" for tool in tools) + assert any( + tool.get("type") == "function" + and tool.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + for tool in tools + ) + hook_litellm_params = transform_kwargs["litellm_params"] + assert hook_litellm_params.get(_ACTIVE_KEY) is True + assert hook_litellm_params.get(_SANDBOX_KEY) + + @pytest.mark.asyncio async def test_async_response_api_handler_streams_when_provider_transform_adds_stream(): handler = BaseLLMHTTPHandler() diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index 3ee53bb46cd..7a3f372582f 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -719,3 +719,93 @@ class TestMistralFileHandling: # Check that file_ids are modified to match Mistral's expected format assert result[0]["content"][1]["file_id"] == "file-12345" # type: ignore assert result[0]["content"][2]["file_id"] == "file-67890" # type: ignore + + +class TestMistralStripsOutputOnlyFields: + """Mistral rejects unknown input fields with a 422 ``extra_forbidden``. + + LiteLLM attaches ``reasoning_content`` / ``thinking_blocks`` to assistant + responses, so replaying an assistant turn verbatim must not forward them. + Regression for https://github.com/BerriAI/litellm/issues/30835. + """ + + def test_assistant_reasoning_content_is_dropped(self): + messages = cast( + List[AllMessageValues], + [ + {"role": "user", "content": "Question?"}, + { + "role": "assistant", + "content": "Follow-up", + "reasoning_content": "Some internal reasoning text.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "step", "signature": "mistral"} + ], + }, + ], + ) + + result = cast( + List[AllMessageValues], + MistralConfig()._transform_messages( + messages=messages, model="mistral-medium-3-5" + ), + ) + + assistant_message = result[-1] + assert "reasoning_content" not in assistant_message + assert "thinking_blocks" not in assistant_message + assert assistant_message["content"] == "Follow-up" + assert assistant_message["role"] == "assistant" + + def test_non_assistant_messages_are_untouched(self): + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": "Question?", "reasoning_content": "noise"}], + ) + + result = cast( + List[AllMessageValues], + MistralConfig()._transform_messages( + messages=messages, model="mistral-medium-3-5" + ), + ) + + assert result[0].get("reasoning_content") == "noise" + + def test_reasoning_content_dropped_when_image_present(self): + """The image branch returns early, so stripping must run before it.""" + messages = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/cat.png"}, + }, + ], + }, + { + "role": "assistant", + "content": "A cat.", + "reasoning_content": "leaked reasoning", + }, + ], + ) + + with patch.object( + MistralConfig, + "_transform_messages_sync", + side_effect=lambda transformed, model: transformed, + ): + result = cast( + List[AllMessageValues], + MistralConfig()._transform_messages( + messages=messages, model="mistral-medium-3-5", is_async=False + ), + ) + + assert "reasoning_content" not in result[-1] diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index 39a4964f5f4..c8743e1809d 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -2,6 +2,7 @@ Tests for JSON-based provider configuration system. """ +import json import os import sys from unittest.mock import MagicMock, patch @@ -244,6 +245,99 @@ class TestPinstripes: assert result["temperature"] == 0.7 +class TestDarkbloom: + def test_darkbloom_json_config_exists(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + darkbloom = JSONProviderRegistry.get("darkbloom") + assert darkbloom is not None + assert darkbloom.base_url == "https://api.darkbloom.dev/v1" + assert darkbloom.api_key_env == "DARKBLOOM_API_KEY" + assert darkbloom.api_base_env == "DARKBLOOM_API_BASE" + assert darkbloom.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_darkbloom_provider_resolution(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="darkbloom/gemma-4-26b", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gemma-4-26b" + assert provider == "darkbloom" + assert api_key is None + assert api_base == "https://api.darkbloom.dev/v1" + + def test_darkbloom_dynamic_config(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.darkbloom.dev/v1" + + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.darkbloom.dev/v1", "test-key" + ) + assert api_base == "https://custom.darkbloom.dev/v1" + assert api_key == "test-key" + + def test_darkbloom_complete_url_appends_endpoint(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + url = config.get_complete_url( + api_base="https://api.darkbloom.dev/v1", + api_key="test-key", + model="darkbloom/gemma-4-26b", + optional_params={}, + litellm_params={}, + stream=True, + ) + + assert url == "https://api.darkbloom.dev/v1/chat/completions" + + def test_darkbloom_provider_config_manager(self): + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="gemma-4-26b", provider=LlmProviders.DARKBLOOM + ) + + assert config is not None + assert config.custom_llm_provider == "darkbloom" + + def test_darkbloom_model_cost_map(self): + with open( + os.path.join(workspace_path, "model_prices_and_context_window.json") + ) as f: + model_cost = json.load(f) + + expected_models = { + "darkbloom/gemma-4-26b": (3e-08, 1.65e-07), + "darkbloom/gpt-oss-20b": (1.45e-08, 7e-08), + } + for model, (input_cost, output_cost) in expected_models.items(): + assert model in model_cost + assert model_cost[model]["litellm_provider"] == "darkbloom" + assert model_cost[model]["max_output_tokens"] == 32768 + assert model_cost[model]["supports_function_calling"] is True + assert model_cost[model]["supports_tool_choice"] is True + assert model_cost[model]["input_cost_per_token"] == input_cost + assert model_cost[model]["output_cost_per_token"] == output_cost + + class TestPublicAIIntegration: """Integration tests for PublicAI provider""" diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index b5c1a86205b..7be295826e3 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -293,7 +293,10 @@ class TestParallelAISearch: ], ) @pytest.mark.asyncio - async def test_custom_api_base_appends_v1_search(self, api_base): + async def test_custom_api_base_appends_v1_search(self, api_base, monkeypatch): + # Operator points at an internal base via the env override (a trusted + # host), so the server key is still used and the URL is normalized. + monkeypatch.setenv("PARALLEL_AI_API_BASE", api_base) with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=AsyncMock, @@ -303,7 +306,6 @@ class TestParallelAISearch: await litellm.asearch( query="AI developments", search_provider="parallel_ai", - api_base=api_base, ) call_args = mock_post.call_args @@ -312,6 +314,23 @@ class TestParallelAISearch: == "https://proxy.internal.example.com/v1/search" ) + @pytest.mark.asyncio + async def test_caller_api_base_without_key_is_refused(self, monkeypatch): + # A caller-supplied api_base (untrusted host) while relying on the + # server key must be refused without any outbound request. + monkeypatch.setenv("PARALLEL_API_KEY", "server-secret") + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + with pytest.raises(Exception, match="Refusing to send"): + await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + api_base="https://attacker.example.com", + ) + mock_post.assert_not_called() + @pytest.mark.asyncio async def test_missing_api_key_raises(self, monkeypatch): monkeypatch.delenv("PARALLEL_API_KEY", raising=False) diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py index e2d1ab72c5e..46c1e457d7c 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py @@ -9,7 +9,7 @@ import json import math import os import sys -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest @@ -120,10 +120,10 @@ class TestPerplexityCostCalculator: # Expected costs: # Input: 100 tokens * $2e-6 = $0.0002 # Output: 50 tokens * $8e-6 = $0.0004 - # Search: 3 queries * ($0.005 / 1000) = $0.000015 - # Total completion cost: $0.000415 + # Search: 3 queries * $0.005 per request = $0.015 + # Total completion cost: $0.0154 expected_prompt_cost = 100 * 2e-6 - expected_completion_cost = (50 * 8e-6) + (3 / 1000 * 0.005) + expected_completion_cost = (50 * 8e-6) + (3 * 0.005) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-6) @@ -195,10 +195,10 @@ class TestPerplexityCostCalculator: # Total prompt cost = $0.00026 # Output (text): (50 - 15) tokens * $8e-6 = $0.00028 # Reasoning: 15 tokens * $3e-6 = $0.000045 - # Search: 2 queries * ($0.005 / 1000) = $0.00001 - # Total completion cost = $0.000335 + # Search: 2 queries * $0.005 per request = $0.01 + # Total completion cost = $0.010325 expected_prompt_cost = (100 * 2e-6) + (30 * 2e-6) - expected_completion_cost = ((50 - 15) * 8e-6) + (15 * 3e-6) + (2 / 1000 * 0.005) + expected_completion_cost = ((50 - 15) * 8e-6) + (15 * 3e-6) + (2 * 0.005) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-6) @@ -311,7 +311,7 @@ class TestPerplexityCostCalculator: # Calculate expected total cost (reasoning is a subset of completion_tokens) expected_prompt_cost = (100 * 2e-6) + (15 * 2e-6) # Input + citation expected_completion_cost = ( - ((50 - 10) * 8e-6) + (10 * 3e-6) + (1 / 1000 * 0.005) + ((50 - 10) * 8e-6) + (10 * 3e-6) + (1 * 0.005) ) # Output (text) + reasoning + search expected_total = expected_prompt_cost + expected_completion_cost @@ -361,7 +361,7 @@ class TestPerplexityCostCalculator: expected_completion_cost = ( ((50 - reasoning_tokens) * 8e-6) + (reasoning_tokens * 3e-6) - + (search_queries / 1000 * 0.005) + + (search_queries * 0.005) ) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_integration.py b/tests/test_litellm/llms/perplexity/test_perplexity_integration.py index e59fbc9f272..8691e6a1ee5 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_integration.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_integration.py @@ -9,7 +9,6 @@ import json import math import os import sys -from unittest.mock import Mock, patch import pytest @@ -106,8 +105,8 @@ class TestPerplexityIntegration: expected_prompt_cost = (100 * 2e-6) + (citation_tokens * 2e-6) expected_completion_cost = ( - ((50 - 10) * 8e-6) + (10 * 3e-6) + (2 / 1000 * 0.005) - ) + ((50 - 10) * 8e-6) + (10 * 3e-6) + (2 * 0.005) + ) # Output (text) + reasoning + search expected_total = expected_prompt_cost + expected_completion_cost assert math.isclose(total_cost, expected_total, rel_tol=1e-6) @@ -152,8 +151,8 @@ class TestPerplexityIntegration: expected_prompt_cost = (200 * 2e-6) + (40 * 2e-6) expected_completion_cost = ( - ((100 - 25) * 8e-6) + (25 * 3e-6) + (3 / 1000 * 0.005) - ) + ((100 - 25) * 8e-6) + (25 * 3e-6) + (3 * 0.005) + ) # Output (text) + reasoning + search assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6) @@ -262,9 +261,9 @@ class TestPerplexityIntegration: expected_prompt_cost = (50000 * 2e-6) + (5000 * 2e-6) expected_completion_cost = ( - ((25000 - 10000) * 8e-6) + (10000 * 3e-6) + (100 / 1000 * 0.005) - ) - expected_total = expected_prompt_cost + expected_completion_cost + ((25000 - 10000) * 8e-6) + (10000 * 3e-6) + (100 * 0.005) + ) # $0.65 + expected_total = expected_prompt_cost + expected_completion_cost # $0.76 assert math.isclose(total_cost, expected_total, rel_tol=1e-6) assert total_cost > 0.25 @@ -326,7 +325,7 @@ class TestPerplexityIntegration: # Should calculate costs correctly expected_prompt_cost = (100 * 2e-6) + (10 * 2e-6) - expected_completion_cost = (50 * 8e-6) + (1 / 1000 * 0.005) + expected_completion_cost = (50 * 8e-6) + (1 * 0.005) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6) diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py index 5496486765c..9870d30d488 100644 --- a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py +++ b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py @@ -63,23 +63,17 @@ class TestTinyfishSearchConfig: assert headers["X-API-Key"] == "sk-tinyfish-test" assert headers["Accept"] == "application/json" - def test_validate_environment_from_env(self): + def test_validate_environment_from_env(self, monkeypatch): + monkeypatch.setenv("TINYFISH_API_KEY", "sk-from-env") config = TinyfishSearchConfig() - with patch( - "litellm.llms.tinyfish.search.transformation.get_secret_str", - return_value="sk-from-env", - ): - headers = config.validate_environment(headers={}) + headers = config.validate_environment(headers={}) assert headers["X-API-Key"] == "sk-from-env" - def test_validate_environment_missing_key(self): + def test_validate_environment_missing_key(self, monkeypatch): + monkeypatch.delenv("TINYFISH_API_KEY", raising=False) config = TinyfishSearchConfig() - with patch( - "litellm.llms.tinyfish.search.transformation.get_secret_str", - return_value=None, - ): - with pytest.raises(ValueError, match="TINYFISH_API_KEY"): - config.validate_environment(headers={}) + with pytest.raises(ValueError, match="TINYFISH_API_KEY"): + config.validate_environment(headers={}) def test_validate_environment_uses_api_base_kwarg(self): config = TinyfishSearchConfig() diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index d4586134b13..122518d4acb 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -8,8 +8,8 @@ Regression test for: UTF-8 codec error when uploading binary files """ import io +import json import pytest -from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -137,11 +137,11 @@ class TestVertexAIBinaryFileUpload: ), "Binary file data should remain as bytes" @pytest.mark.asyncio - async def test_jsonl_file_upload_returns_string(self): + async def test_jsonl_file_upload_returns_resumable_stream(self): """ - Test that JSONL files (text) are correctly transformed to strings. - - This ensures we handle both binary and text files correctly. + Test that JSONL batch files are transformed into a resumable-upload config + carrying a streaming body (not a buffered bytes payload), so the handler + can stream the upload to GCS in bounded chunks. """ # Create mock JSONL content mock_jsonl_content = ( @@ -164,10 +164,16 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) - # JSONL files should be transformed to string - assert isinstance( - transformed_request, str - ), f"Expected string for JSONL file, got {type(transformed_request)}" + assert ( + isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + ), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}" + + stream = transformed_request["resumable_chunked_upload"]["body_stream"] + decoded = json.loads(b"".join(stream.iter_bytes()).decode("utf-8")) + assert ( + "request" in decoded + ), "JSONL transform must wrap each row in {'request': ...}" @pytest.mark.asyncio async def test_mixed_file_types_in_sequence(self): @@ -208,7 +214,7 @@ class TestVertexAIBinaryFileUpload: optional_params={}, litellm_params={}, ) - assert isinstance(result2, str) + assert isinstance(result2, dict) and "resumable_chunked_upload" in result2 # Test 3: Upload another binary file binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" @@ -251,7 +257,7 @@ class TestVertexAIBinaryFileUpload: }, "text_files": { "input_type": "str or bytes", - "output_type": "str", + "output_type": "bytes", "examples": ["JSONL", "CSV", "TXT"], "http_method": "POST", "encoding": "UTF-8", diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py new file mode 100644 index 00000000000..cd556c48b6b --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -0,0 +1,696 @@ +""" +Tests for the streaming OpenAI -> Vertex JSONL batch transform. + +The transform converts batch uploads entry-by-entry rather than materializing +the payload in full intermediate lists (decoded str, parsed dicts, transformed +dicts, joined output), which keeps peak memory bounded on large uploads. + +These tests lock in the behaviour that would regress if the streaming path were +replaced by a list-based pipeline: + 1. Byte-for-byte output parity with a list pipeline (wire format). + 2. The streaming transform peaks at a clear fraction of a list pipeline on the + same input (relative differential, robust to GC noise). + 3. ``get_object_name`` only parses the first JSONL row, so a payload whose + later rows are not valid JSON does not raise. + 4. A tuple-wrapped file handle uploaded through the real create_file ordering + keeps every row, including entry 0 (no partial upload from a consumed + cursor). +""" + +import gc +import io +import json +import time +import tracemalloc + +import httpx +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.files.transformation import BaseFileUploadStream +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.vertex_ai.files.transformation import ( + VertexAIFilesConfig, + _OpenAIToVertexBatchUploadStream, + _get_litellm_batch_custom_id_from_labels, + _iter_openai_jsonl_entries, + _iter_openai_jsonl_lines, + _openai_batch_jsonl_entry_to_vertex_wrapped_request, +) +from litellm.types.llms.openai import CreateFileRequest + + +def _resumable_stream(transformed) -> BaseFileUploadStream: + """Pull the streaming body out of a resumable-upload transform result.""" + return transformed["resumable_chunked_upload"]["body_stream"] + + +def _join_upload_body(transformed) -> bytes: + """Materialize a transform result's upload body for byte-level assertions.""" + if isinstance(transformed, dict) and "resumable_chunked_upload" in transformed: + return b"".join(_resumable_stream(transformed).iter_bytes()) + if isinstance(transformed, BaseFileUploadStream): + return b"".join(transformed.iter_bytes()) + if isinstance(transformed, str): + return transformed.encode("utf-8") + return transformed + + +def _make_openai_jsonl_bytes(n_rows: int, padding: int = 400) -> bytes: + pad = "x" * padding + rows = [] + for i in range(n_rows): + rows.append( + json.dumps( + { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "user", "content": f"{pad} {i}"}], + "max_tokens": 4, + }, + } + ) + ) + return ("\n".join(rows)).encode("utf-8") + + +def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> str: + """Row-by-row reference output built eagerly from the live single-entry + transform, so the streaming path can be checked against it for parity.""" + entries = [json.loads(line) for line in content.splitlines() if line.strip()] + return "\n".join( + json.dumps( + _openai_batch_jsonl_entry_to_vertex_wrapped_request( + entry, cfg._map_openai_to_vertex_params + ) + ) + for entry in entries + ) + + +class TestStreamingOutputParity: + def test_transform_create_file_request_returns_resumable_stream_parity(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(300) + request: CreateFileRequest = { + "file": ("batch.jsonl", raw, "application/jsonl"), + "purpose": "batch", + } + + out = cfg.transform_create_file_request( + model="", create_file_data=request, optional_params={}, litellm_params={} + ) + + # A batch upload must be a resumable-upload config carrying a streaming + # body, so the handler can chunk it; a buffered bytes/str return would + # defeat the OOM fix. + assert isinstance(out, dict) and "resumable_chunked_upload" in out + assert isinstance(_resumable_stream(out), BaseFileUploadStream) + assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string( + cfg, raw.decode("utf-8") + ) + + +class TestFileLikeInputNotPartiallyConsumed: + """ + In ``llm_http_handler.create_file`` the object-name step + (get_complete_file_url -> get_object_name) runs before + transform_create_file_request, and both read the same create_file_data + source. When the file is a tuple-wrapped open handle, the streaming reader + must still emit every row including entry 0: ``_iter_openai_jsonl_lines`` + rewinds a seekable source (seek(0)) before each pass, so the object-name + step's partial read of the cursor does not consume the upload. A partial + upload missing the first request would be silent and hard to catch, so this + locks the full-payload invariant in. + """ + + def test_filehandle_create_file_keeps_first_entry(self): + cfg = VertexAIFilesConfig() + n_rows = 25 + raw = _make_openai_jsonl_bytes(n_rows) + create_file_data: CreateFileRequest = { + "file": ("batch.jsonl", io.BytesIO(raw), "application/jsonl"), + "purpose": "batch", + } + + # Object-name step first (as the handler does), then the transform, both + # reading the same live BytesIO handle. + cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=create_file_data, + ) + out = cfg.transform_create_file_request( + model="", + create_file_data=create_file_data, + optional_params={}, + litellm_params={}, + ) + + lines = _join_upload_body(out).decode("utf-8").splitlines() + assert len(lines) == n_rows, "no batch row may be dropped from the upload" + first_labels = json.loads(lines[0])["request"]["labels"] + assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" + + +class TestStreamingLineIterator: + def test_skips_blank_and_whitespace_lines(self): + content = b'{"a": 1}\n\n \n{"b": 2}\n' + assert list(_iter_openai_jsonl_lines(content)) == ['{"a": 1}', '{"b": 2}'] + + def test_handles_crlf_and_missing_trailing_newline(self): + content = b'{"a": 1}\r\n{"b": 2}' + assert [json.loads(line) for line in _iter_openai_jsonl_lines(content)] == [ + {"a": 1}, + {"b": 2}, + ] + + def test_accepts_str_bytes_tuple_and_filelike(self): + expected = [{"a": 1}, {"b": 2}] + text = '{"a": 1}\n{"b": 2}\n' + for source in ( + text, + text.encode("utf-8"), + ("name.jsonl", text.encode("utf-8"), "application/jsonl"), + io.BytesIO(text.encode("utf-8")), + ): + assert list(_iter_openai_jsonl_entries(source)) == expected + + def test_str_input_without_trailing_newline(self): + assert list(_iter_openai_jsonl_lines('{"a": 1}\n{"b": 2}')) == [ + '{"a": 1}', + '{"b": 2}', + ] + + def test_pathlike_input_is_read_line_by_line(self, tmp_path): + path = tmp_path / "batch.jsonl" + path.write_bytes(b'{"a": 1}\n{"b": 2}\n') + assert list(_iter_openai_jsonl_entries(path)) == [{"a": 1}, {"b": 2}] + + def test_unsupported_content_type_raises(self): + with pytest.raises(ValueError, match="Unsupported file content type"): + list(_iter_openai_jsonl_lines(12345)) # type: ignore[arg-type] + + def test_non_seekable_handle_raises_instead_of_dropping_first_row(self): + # The handle is read twice (object-name probe, then body). A non-seekable + # handle can't rewind, so it must fail loudly rather than silently resume + # mid-stream and omit the opening batch request. + class _NonSeekable: + def __init__(self, raw: bytes): + self._buf = io.BytesIO(raw) + + def read(self, *args): + return self._buf.read(*args) + + def __iter__(self): + return iter(self._buf) + + def seek(self, *args): + raise io.UnsupportedOperation("not seekable") + + handle = _NonSeekable( + b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n' + ) + with pytest.raises(ValueError, match="seekable"): + list(_iter_openai_jsonl_lines(handle)) + + def test_is_lazy_does_not_parse_past_first_entry(self): + # Second row is invalid JSON; pulling only the first entry must not raise. + content = b'{"custom_id": "first"}\nnot-json-at-all\n' + gen = _iter_openai_jsonl_entries(content) + assert next(gen)["custom_id"] == "first" + with pytest.raises(json.JSONDecodeError): + next(gen) + + +class TestGetObjectNameLazyParse: + def test_only_parses_first_row_for_model(self): + cfg = VertexAIFilesConfig() + # Tail rows are deliberately not valid JSON. Parsing the whole payload + # would raise here; a first-row-only parse must not. + raw = ( + b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\n' + b"garbage line that is not json\n" + ) + object_name = cfg.get_object_name( + ("batch.jsonl", raw, "application/jsonl"), purpose="batch" + ) + assert "gemini-2.5-flash" in object_name + + +class TestStreamingPeakMemory: + """ + Differential guard: the streaming transform must stay well under the peak + that a list pipeline incurs on the same input. If the hot path builds full + intermediate lists, the streaming assertion fails. + + The assertion that matters is the *relative* one: ``streaming_peak`` must be + a clear fraction of ``list_peak`` on the identical input. Absolute + ``tracemalloc`` ratios drift with GC timing and the live set carried in from + earlier tests, so they make poor CI gates; the relative comparison cancels + that shared noise and is exactly what regresses (toward 1.0) when the hot + path builds full intermediate lists. ``gc.collect()`` before each + measurement removes any garbage the previous run left behind. + """ + + def _measure(self, fn): + gc.collect() + tracemalloc.start() + try: + fn() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + def test_streaming_peak_well_below_list_pipeline(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(8000) + content_str = raw.decode("utf-8") + + def drain_stream(): + # Consume the upload body one row at a time, as the chunked uploader + # does, without accumulating it. + for _ in _OpenAIToVertexBatchUploadStream( + raw, cfg._map_openai_to_vertex_params + ).iter_bytes(): + pass + + streaming_peak = self._measure(drain_stream) + list_peak = self._measure( + lambda: _reference_vertex_jsonl_string(cfg, content_str) + ) + + # Core guard: the lazily consumed streaming body peaks well under a list + # pipeline that materializes every transformed row. Building full + # intermediate lists in the hot path pushes this ratio back toward 1.0. + assert streaming_peak < list_peak * 0.6, ( + f"streaming peak {streaming_peak} not a clear win over list pipeline " + f"{list_peak} (ratio {streaming_peak / list_peak:.2f})" + ) + + def test_get_object_name_does_not_scale_with_payload(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(8000) + file_data = ("batch.jsonl", raw, "application/jsonl") + + # The payload bytes already exist before measurement starts, so a lazy + # first-row parse should allocate only a small fraction of the payload; + # parsing every row would blow past this bound. + peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch")) + assert ( + peak / len(raw) < 2.0 + ), "get_object_name should not copy the whole payload" + + +class TestPathSourcedStreaming: + """ + The proxy spools large batch uploads to a temp file and passes a pathlib.Path + as the file content instead of pre-reading bytes, so the transform streams + from disk. These lock in that a Path source yields identical output, keeps + every row, stays memory-bounded, and is re-iterable (multi-model uploads). + """ + + def _write_jsonl(self, tmp_path, n_rows, padding=400): + raw = _make_openai_jsonl_bytes(n_rows, padding=padding) + path = tmp_path / "batch.jsonl" + path.write_bytes(raw) + return path, raw + + def _batch_request(self, path) -> CreateFileRequest: + return {"file": ("batch.jsonl", path, "application/jsonl"), "purpose": "batch"} + + def test_transform_from_path_matches_legacy_and_keeps_all_rows(self, tmp_path): + cfg = VertexAIFilesConfig() + n_rows = 200 + path, raw = self._write_jsonl(tmp_path, n_rows) + data = self._batch_request(path) + + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=data, + ) + assert "uploadType=resumable" in url + + out = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + assert isinstance(out, dict) and "resumable_chunked_upload" in out + body = _join_upload_body(out).decode("utf-8") + assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) + lines = body.splitlines() + assert len(lines) == n_rows, "no batch row may be dropped from a Path source" + first_labels = json.loads(lines[0])["request"]["labels"] + assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" + + def test_path_source_peak_stays_below_payload(self, tmp_path): + cfg = VertexAIFilesConfig() + path, raw = self._write_jsonl(tmp_path, 8000) + data = self._batch_request(path) + + def run(): + cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=data, + ) + out = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + for _ in _resumable_stream(out).iter_bytes(): + pass # drain without accumulating + + gc.collect() + tracemalloc.start() + try: + run() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + # Streaming from disk must not materialize the payload. Reading the whole + # file into bytes (the pre-fix path) would push peak past the file size. + assert peak < len(raw) * 0.3, ( + f"peak {peak} not bounded vs payload {len(raw)} " + f"(ratio {peak / len(raw):.2f})" + ) + + def test_path_source_stream_is_reiterable(self, tmp_path): + cfg = VertexAIFilesConfig() + path, _ = self._write_jsonl(tmp_path, 50) + data = self._batch_request(path) + + out = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + stream = _resumable_stream(out) + first = b"".join(stream.iter_bytes()) + second = b"".join(stream.iter_bytes()) + assert first == second and len(first) > 0 + + +_GCS_OBJECT_JSON = { + "id": "test-bucket/litellm-vertex-files/x/123", + "name": "litellm-vertex-files/x", + "size": "0", + "timeCreated": "2026-01-01T00:00:00.000000Z", + "purpose": "batch", +} + + +class _FixedBytesStream(BaseFileUploadStream): + """Streaming body of exact, controllable bytes for protocol-edge tests.""" + + def __init__(self, data: bytes, piece: int = 64): + self._data = data + self._piece = piece + + def iter_bytes(self): + for i in range(0, len(self._data), self._piece): + yield self._data[i : i + self._piece] + + +def _logging_obj() -> Logging: + return Logging( + model="", + messages=[], + stream=False, + call_type="acreate_file", + start_time=time.time(), + litellm_call_id="test", + function_id="", + ) + + +def _gcs_resumable_mock(session_url: str, final_status: int = 200): + """A fake GCS resumable endpoint: POST opens a session (URI in Location), + each PUT appends and returns 308 until the final chunk returns 200/201.""" + state = {"received": bytearray(), "ranges": [], "methods": [], "urls": []} + + async def handler(request: httpx.Request) -> httpx.Response: + state["methods"].append(request.method) + state["urls"].append(str(request.url)) + if request.method == "POST": + return httpx.Response(200, headers={"location": session_url}) + body = await request.aread() + content_range = request.headers["content-range"] + state["ranges"].append(content_range) + state["received"].extend(body) + if content_range.rsplit("/", 1)[-1] == "*": + return httpx.Response( + 308, headers={"range": f"bytes=0-{len(state['received']) - 1}"} + ) + return httpx.Response(final_status, json=_GCS_OBJECT_JSON) + + return handler, state + + +def _async_handler_with(mock) -> AsyncHTTPHandler: + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock)) + return handler + + +class TestResumableUploadUrl: + def test_batch_jsonl_uses_resumable_upload_type(self): + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"), + "purpose": "batch", + } + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + assert "uploadType=resumable" in url + assert "uploadType=media" not in url + + def test_batch_text_plain_uses_resumable_upload_type(self): + # Clients often label a .jsonl batch upload as text/plain; it must still + # take the streaming/resumable path, not the buffered media path. + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "text/plain"), + "purpose": "batch", + } + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + assert "uploadType=resumable" in url + assert "uploadType=media" not in url + + def test_binary_upload_stays_simple_media(self): + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"), + "purpose": "user_data", + } + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + assert "uploadType=media" in url + assert "uploadType=resumable" not in url + + +class TestResumableStreamBody: + def test_stream_matches_legacy_pipeline(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(120) + stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params) + assert b"".join(stream.iter_bytes()).decode( + "utf-8" + ) == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) + + def test_stream_is_reiterable_for_retries(self): + # A one-shot generator would make a transport retry upload an empty body; + # iter_bytes() must yield the full payload every call. + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(40) + stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params) + first = b"".join(stream.iter_bytes()) + second = b"".join(stream.iter_bytes()) + assert first == second and len(first) > 0 + + def test_stream_is_reiterable_for_seekable_file_like_input(self): + # A seekable handle (BytesIO, temp file) must be rewound between calls; + # otherwise the first iter_bytes() exhausts it and a retry would upload + # an empty body silently. + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(40) + stream = _OpenAIToVertexBatchUploadStream( + io.BytesIO(raw), cfg._map_openai_to_vertex_params + ) + first = b"".join(stream.iter_bytes()) + second = b"".join(stream.iter_bytes()) + assert first == second and len(first) > 0 + + +class TestResumableChunking: + def test_intermediate_chunks_are_exactly_chunk_size(self): + pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4)) + assert pieces == [b"xxxx", b"xxxx", b"xx"] + + def test_exact_multiple_yields_no_trailing_empty(self): + # An exactly chunk-aligned stream yields only full chunks; the upload + # finalizes on the last data chunk instead of an extra empty request. + pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4)) + assert pieces == [b"xxxx", b"xxxx"] + + def test_empty_stream_yields_nothing(self): + # A 0-byte stream yields no chunks; the caller finalizes with one empty + # request (bytes */0). + assert list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([]), 4)) == [] + + def test_default_chunk_size_is_256kib_multiple(self): + assert BaseLLMHTTPHandler._RESUMABLE_CHUNK_SIZE % (256 * 1024) == 0 + + def test_content_range_intermediate_uses_star_total(self): + assert ( + BaseLLMHTTPHandler._resumable_content_range(0, 4096, is_final=False) + == "bytes 0-4095/*" + ) + + def test_content_range_final_uses_real_total(self): + assert ( + BaseLLMHTTPHandler._resumable_content_range(8192, 100, is_final=True) + == "bytes 8192-8291/8292" + ) + + def test_content_range_empty_finalize(self): + assert ( + BaseLLMHTTPHandler._resumable_content_range(8192, 0, is_final=True) + == "bytes */8192" + ) + + +@pytest.mark.asyncio +class TestResumableUploadProtocol: + """End-to-end against a faked GCS resumable endpoint. These are the tests + that fail if the handler buffers the whole body, drops bytes, mislabels a + Content-Range, follows the 308 instead of continuing, or skips finalize.""" + + async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200): + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("batch.jsonl", raw, "application/jsonl"), + "purpose": "batch", + } + api_base = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + transformed = cfg.transform_create_file_request( + model="", create_file_data=request, optional_params={}, litellm_params={} + ) + transformed["resumable_chunked_upload"]["chunk_size"] = chunk_size + expected = _join_upload_body(transformed) + + session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" + mock, state = _gcs_resumable_mock(session_url, final_status=final_status) + response = await BaseLLMHTTPHandler().async_create_file( + transformed_request=transformed, + litellm_params={}, + provider_config=cfg, + headers={"Authorization": "Bearer x"}, + api_base=api_base, + logging_obj=_logging_obj(), + client=_async_handler_with(mock), + timeout=None, + ) + return expected, state, response, session_url, api_base + + async def test_streams_in_chunks_and_reassembles(self): + raw = _make_openai_jsonl_bytes(300) + chunk_size = 4096 + expected, state, response, session_url, api_base = await self._run( + raw, chunk_size + ) + + # One session-open POST, then a sequence of chunk PUTs. + assert state["methods"][0] == "POST" + assert set(state["methods"][1:]) == {"PUT"} + assert state["methods"].count("PUT") >= 2, "payload must span multiple chunks" + + # POST opens a resumable session; every chunk goes to the session URI. + assert "uploadType=resumable" in state["urls"][0] + assert all(u == session_url for u in state["urls"][1:]) + + # Every non-final chunk is exactly chunk_size with an unknown-total range; + # the final chunk carries the real total. + intermediate = state["ranges"][:-1] + for index, content_range in enumerate(intermediate): + assert ( + content_range + == f"bytes {index * chunk_size}-{(index + 1) * chunk_size - 1}/*" + ) + total = len(expected) + last_offset = len(intermediate) * chunk_size + if last_offset == total: # payload landed on a chunk boundary + assert state["ranges"][-1] == f"bytes */{total}" + else: + assert state["ranges"][-1] == f"bytes {last_offset}-{total - 1}/{total}" + + # The bytes GCS received are exactly the transformed batch payload. + assert bytes(state["received"]) == expected + assert response.object == "file" + + async def test_exact_multiple_finalizes_on_last_data_chunk(self): + # A body that is an exact multiple of the chunk size finalizes on its + # last data chunk (bytes (TOTAL-chunk)-(TOTAL-1)/TOTAL), with no extra + # empty finalize request. + chunk_size = 256 + total = chunk_size * 3 + stream = _FixedBytesStream(b"a" * total) + config = {"body_stream": stream, "chunk_size": chunk_size} + session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" + mock, state = _gcs_resumable_mock(session_url) + + response = await BaseLLMHTTPHandler()._aresumable_chunked_upload( + client=_async_handler_with(mock), + initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", + base_headers={"Authorization": "Bearer x"}, + config=config, + timeout=None, + ) + + assert state["ranges"][-1] == f"bytes {total - chunk_size}-{total - 1}/{total}" + assert "*" not in state["ranges"][-1] + assert bytes(state["received"]) == b"a" * total + assert response.status_code == 200 + + async def test_failed_chunk_raises(self): + raw = _make_openai_jsonl_bytes(80) + with pytest.raises(Exception): + await self._run(raw, chunk_size=4096, final_status=403) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 7c063c72607..8c5305ee67b 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -14,8 +14,8 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, - VertexAIJsonlFilesTransformation, _get_litellm_batch_custom_id_from_labels, + _openai_batch_jsonl_entry_to_vertex_wrapped_request, _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent @@ -33,7 +33,7 @@ class TestParseGcsUri: def test_should_parse_standard_gs_uri(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/object.jsonl" bucket, encoded = config._parse_gcs_uri( - file_id, litellm_params={"bucket_name": "my-bucket"} + file_id, litellm_params={"gcs_bucket_name": "my-bucket"} ) assert bucket == "my-bucket" assert encoded == urllib.parse.quote( @@ -43,7 +43,7 @@ class TestParseGcsUri: def test_should_parse_uri_with_nested_publisher_path(self, config): uri = "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" bucket, encoded = config._parse_gcs_uri( - uri, litellm_params={"bucket_name": "litellm-local"} + uri, litellm_params={"gcs_bucket_name": "litellm-local"} ) assert bucket == "litellm-local" expected_path = ( @@ -56,7 +56,7 @@ class TestParseGcsUri: "gs://my-bucket/litellm-vertex-files/some/path", safe="" ) bucket, encoded = config._parse_gcs_uri( - encoded_uri, litellm_params={"bucket_name": "my-bucket"} + encoded_uri, litellm_params={"gcs_bucket_name": "my-bucket"} ) assert bucket == "my-bucket" assert encoded == urllib.parse.quote("litellm-vertex-files/some/path", safe="") @@ -64,21 +64,21 @@ class TestParseGcsUri: def test_should_reject_bucket_only(self, config): with pytest.raises(ValueError, match="object name"): config._parse_gcs_uri( - "gs://my-bucket", litellm_params={"bucket_name": "my-bucket"} + "gs://my-bucket", litellm_params={"gcs_bucket_name": "my-bucket"} ) def test_should_reject_no_gs_prefix(self, config): with pytest.raises(ValueError, match="gs://"): config._parse_gcs_uri( "my-bucket/litellm-vertex-files/object.txt", - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) def test_should_reject_unmanaged_object_path(self, config): with pytest.raises(ValueError, match="LiteLLM-managed"): config._parse_gcs_uri( "gs://my-bucket/private/object.txt", - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) def test_should_reject_request_supplied_legacy_flag(self, config): @@ -86,7 +86,7 @@ class TestParseGcsUri: config._parse_gcs_uri( "gs://my-bucket/private/object.txt", litellm_params={ - "bucket_name": "my-bucket", + "gcs_bucket_name": "my-bucket", "allow_legacy_cloud_file_ids": True, }, ) @@ -96,7 +96,7 @@ class TestParseGcsUri: bucket, encoded = config._parse_gcs_uri( "gs://my-bucket/private/object.txt", litellm_params={ - "bucket_name": "my-bucket", + "gcs_bucket_name": "my-bucket", "_litellm_internal_model_credentials": trusted_credentials, }, ) @@ -109,7 +109,7 @@ class TestParseGcsUri: config._parse_gcs_uri( "gs://my-bucket/private/object.txt", litellm_params={ - "bucket_name": "my-bucket", + "gcs_bucket_name": "my-bucket", "_litellm_internal_model_credentials": { "allow_legacy_cloud_file_ids": True }, @@ -121,7 +121,7 @@ class TestParseGcsUri: bucket, encoded = config._parse_gcs_uri( "gs://my-bucket/team-a/private/object.txt", litellm_params={ - "bucket_name": "my-bucket/team-a", + "gcs_bucket_name": "my-bucket/team-a", "_litellm_internal_model_credentials": trusted_credentials, }, ) @@ -135,7 +135,7 @@ class TestParseGcsUri: config._parse_gcs_uri( "gs://my-bucket/team-b/private/object.txt", litellm_params={ - "bucket_name": "my-bucket/team-a", + "gcs_bucket_name": "my-bucket/team-a", "_litellm_internal_model_credentials": trusted_credentials, }, ) @@ -144,7 +144,7 @@ class TestParseGcsUri: with pytest.raises(ValueError, match="configured storage bucket"): config._parse_gcs_uri( "gs://other-bucket/litellm-vertex-files/object.txt", - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) @@ -156,7 +156,7 @@ class TestCreateFileUrl: model="", optional_params={}, litellm_params={ - "bucket_name": "safe-bucket", + "gcs_bucket_name": "safe-bucket", "litellm_metadata": {"gcs_bucket_name": "attacker-bucket"}, }, data={ @@ -182,7 +182,7 @@ class TestTransformRetrieveFile: url, params = config.transform_retrieve_file_request( file_id=file_id, optional_params={}, - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) expected_encoded = urllib.parse.quote( "litellm-vertex-files/path/to/file.jsonl", safe="" @@ -243,7 +243,7 @@ class TestTransformFileContent: url, params = config.transform_file_content_request( file_content_request={"file_id": file_id}, optional_params={}, - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") assert ( @@ -378,7 +378,7 @@ class TestTransformDeleteFile: url, params = config.transform_delete_file_request( file_id=file_id, optional_params={}, - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") assert ( @@ -854,6 +854,106 @@ class TestVertexBatchOutputTransformation: ) assert transformed_content == invalid_content + def test_binary_content_passthrough(self, config): + """A binary file (PDF/video) whose first bytes are not valid UTF-8 must be + returned unchanged. The row-by-row transform only engages for a JSONL + batch output and must never line-parse or corrupt binary content.""" + binary = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\n" + b"\x00\x01\x02\xff\xfe" * 64 + assert config._try_transform_vertex_batch_output_to_openai(binary) == binary + + def test_streaming_transform_peaks_below_list_pipeline(self, config): + """The output transform must stream row-by-row, not build a list of every + parsed row and a second list of transformed rows. This guards against a + regression to the list pipeline, which peaks at several full copies and + OOMs on large result files. The relative comparison cancels shared noise + (per-row transform cost, GC timing) and only the list overhead differs. + """ + import gc + import tracemalloc + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + def vertex_row(index: int) -> dict: + return { + "status": "", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "labels": {"litellm_custom_id": f"r-{index}"}, + }, + "response": { + "candidates": [ + { + "content": { + "parts": [{"text": "hello " * 20}], + "role": "model", + }, + "finishReason": "STOP", + } + ], + "modelVersion": "gemini-2.0-flash-001", + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 20, + "totalTokenCount": 30, + }, + }, + } + + content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode( + "utf-8" + ) + + def list_pipeline() -> bytes: + gemini_config = VertexGeminiConfig() + logging_obj = Logging( + model="", + messages=[], + stream=False, + call_type="batch_transform", + start_time=0.1, + litellm_call_id="", + function_id="", + ) + logging_obj.optional_params = {} + mock_response = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + request=httpx.Request("POST", "https://example.com"), + ) + rows = content.decode("utf-8").strip().split("\n") + transformed = [ + json.dumps( + config._transform_single_vertex_batch_output_to_openai( + json.loads(row), gemini_config, logging_obj, mock_response + ) + ) + for row in rows + ] + return "\n".join(transformed).encode("utf-8") + + def peak_of(fn) -> int: + gc.collect() + tracemalloc.start() + try: + fn() + return tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + streaming_peak = peak_of( + lambda: config._try_transform_vertex_batch_output_to_openai(content) + ) + list_peak = peak_of(list_pipeline) + + assert streaming_peak < list_peak * 0.75, ( + f"streaming peak {streaming_peak} is not a clear win over the list " + f"pipeline {list_peak} (ratio {streaming_peak / list_peak:.2f})" + ) + class TestTryTransformDoesNotMutateCallerLoggingObj: """Regression tests: _try_transform_vertex_batch_output_to_openai must not mutate @@ -953,12 +1053,23 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: assert transformed["response"]["status_code"] == 200 +def _wrap_entries(openai_jsonl_content): + """Vertex-wrapped requests for a list of OpenAI batch entries, built via the + live single-entry transform that the streaming upload path uses.""" + cfg = VertexAIFilesConfig() + return [ + _openai_batch_jsonl_entry_to_vertex_wrapped_request( + entry, cfg._map_openai_to_vertex_params + ) + for entry in openai_jsonl_content + ] + + class TestVertexBatchCustomIdLabels: """Test custom_id handling in batch transformations""" def test_custom_id_added_to_labels_in_vertex_request(self): """Test that custom_id from OpenAI format is added as a label in Vertex AI format""" - transformation = VertexAIJsonlFilesTransformation() openai_jsonl_content = [ { @@ -973,11 +1084,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) assert len(vertex_jsonl_content) == 1 vertex_request = vertex_jsonl_content[0] @@ -992,7 +1099,6 @@ class TestVertexBatchCustomIdLabels: def test_long_custom_id_round_trips_across_raw_label_chunks(self): """Test that long custom_ids are not truncated in raw labels.""" - transformation = VertexAIJsonlFilesTransformation() custom_id_a = "shared-prefix-that-is-longer-than-thirty-six-bytes-A" custom_id_b = "shared-prefix-that-is-longer-than-thirty-six-bytes-B" @@ -1009,11 +1115,7 @@ class TestVertexBatchCustomIdLabels: for custom_id in (custom_id_a, custom_id_b) ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) labels_a = vertex_jsonl_content[0]["request"]["labels"] labels_b = vertex_jsonl_content[1]["request"]["labels"] @@ -1028,7 +1130,6 @@ class TestVertexBatchCustomIdLabels: def test_multiple_requests_each_get_their_own_label(self): """Test that multiple requests each get their own custom_id label""" - transformation = VertexAIJsonlFilesTransformation() openai_jsonl_content = [ { @@ -1043,11 +1144,7 @@ class TestVertexBatchCustomIdLabels: for i in range(3) ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) assert len(vertex_jsonl_content) == 3 @@ -1063,7 +1160,6 @@ class TestVertexBatchCustomIdLabels: def test_request_without_custom_id_has_no_label(self): """Test that requests without custom_id don't get a label""" - transformation = VertexAIJsonlFilesTransformation() openai_jsonl_content = [ { @@ -1076,11 +1172,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) # Should not have labels if no custom_id was provided assert "labels" not in vertex_jsonl_content[0]["request"] @@ -1090,7 +1182,6 @@ class TestVertexBatchCustomIdLabels: Test the full round trip: OpenAI format -> Vertex AI format -> Vertex AI output -> OpenAI output Verify that custom_id is preserved through the entire flow. """ - transformation = VertexAIJsonlFilesTransformation() config = VertexAIFilesConfig() # Step 1: Transform OpenAI input to Vertex AI format (mixed case exercises raw label) @@ -1106,11 +1197,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_input = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_input - ) - ) + vertex_input = _wrap_entries(openai_input) # Verify both labels are GCP-safe and encoded raw preserves round-trip. assert ( @@ -1154,7 +1241,6 @@ class TestVertexBatchCustomIdLabels: def test_custom_id_label_sanitization(self): """Test that custom_id values are sanitized to meet GCP label constraints""" - transformation = VertexAIJsonlFilesTransformation() # Test sanitization function assert _sanitize_gcp_label_value("MyRequest-1") == "myrequest-1" @@ -1179,11 +1265,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_input = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_input - ) - ) + vertex_input = _wrap_entries(openai_input) # Verify both labels are safe for GCP labels. assert ( @@ -1192,3 +1274,47 @@ class TestVertexBatchCustomIdLabels: raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != "MyRequest-1" assert _sanitize_gcp_label_value(raw_label) == raw_label + + +class TestConfiguredBucketNameResolution: + def test_should_resolve_new_gcs_bucket_name_key(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + assert ( + config._get_configured_bucket_name({"gcs_bucket_name": "my-new-bucket"}) + == "my-new-bucket" + ) + + def test_should_resolve_legacy_bucket_name_key(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + assert ( + config._get_configured_bucket_name({"bucket_name": "my-legacy-bucket"}) + == "my-legacy-bucket" + ) + + def test_should_prefer_new_key_over_legacy(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + assert ( + config._get_configured_bucket_name( + {"gcs_bucket_name": "new", "bucket_name": "legacy"} + ) + == "new" + ) + + def test_should_fall_back_to_env(self, config, monkeypatch): + monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket") + assert config._get_configured_bucket_name({}) == "env-bucket" + + def test_should_raise_when_no_bucket_anywhere(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + with pytest.raises(ValueError, match="GCS bucket_name is required"): + config._get_configured_bucket_name({}) + + def test_legacy_kwarg_survives_get_litellm_params(self): + from litellm.litellm_core_utils.get_litellm_params import ( + OPTIONAL_KWARGS_KEYS, + get_litellm_params, + ) + + assert "bucket_name" in OPTIONAL_KWARGS_KEYS + params = get_litellm_params(bucket_name="my-legacy-bucket") + assert params.get("bucket_name") == "my-legacy-bucket" diff --git a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index 1ebd704be34..1f171496cce 100644 --- a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -346,3 +346,74 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) "Vertex AI Realtime" in record.message and "session.update" in record.message for record in caplog.records ) + + +async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( + monkeypatch, +): + """Regression: forwarding client ?model=/?intent= to the Vertex Live WSS URL causes 1007 errors. + + Exercises ``async_realtime`` end-to-end so that re-adding ``_append_query_params`` + (the reverted bug) would push ``model=``/``intent=`` onto the backend URL and fail here. + """ + import websockets + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + captured = {} + + def fake_connect(url, *args, **kwargs): + captured["url"] = url + raise RuntimeError("stop before establishing the backend connection") + + monkeypatch.setattr(websockets, "connect", fake_connect) + + await BaseLLMHTTPHandler().async_realtime( + model="gemini-live-2.5-flash-native-audio", + websocket=AsyncMock(), + logging_obj=MagicMock(), + provider_config=cfg, + headers={}, + query_params={ + "model": "gemini-live-2.5-flash-native-audio", + "intent": "chat", + }, + ) + + assert "?" not in captured["url"] + assert "model=" not in captured["url"] + assert "intent=" not in captured["url"] + + +def test_vertex_function_call_output_omits_id(): + """Regression: Vertex Live rejects ``id`` on toolResponse.functionResponses (1007).""" + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + cfg._tool_call_id_to_name["call_abc123"] = "terminate_call" + + messages = cfg.transform_realtime_request( + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": "call_abc123", + "output": '{"status": "ok"}', + }, + } + ), + "gemini-live-2.5-flash-native-audio", + session_configuration_request="existing", + ) + + assert len(messages) == 1 + payload = json.loads(messages[0]) + function_response = payload["toolResponse"]["functionResponses"][0] + assert "id" not in function_response + assert function_response["name"] == "terminate_call" + assert function_response["response"] == {"status": "ok"} diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py new file mode 100644 index 00000000000..7e028064e4c --- /dev/null +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -0,0 +1,333 @@ +"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" + +import importlib +import sys +import types + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse + +# `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.ocr.rust_bridge") + +MODEL = "mistral/mistral-ocr-latest" +DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + +FAKE_OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "hello world"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": None, + "usage_info": {"pages_processed": 1}, + "object": "ocr", +} + + +class RecordingBridge: + """A fake ``RustOcr`` callable that records the args it was handed.""" + + def __init__(self): + self.calls = [] + + def __call__( + self, model, document, api_key, api_base, optional_params, timeout_seconds + ): + self.calls.append( + { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "optional_params": optional_params, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_OCR_RESPONSE) + + +class RecordingLogging: + """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" + + def __init__(self): + self.pre_call_kwargs = None + + def pre_call(self, *, input, api_key, additional_args): + self.pre_call_kwargs = { + "input": input, + "api_key": api_key, + "additional_args": additional_args, + } + + +class FakeOCRConfig: + """A stand-in ``BaseOCRConfig`` that echoes the request it would build.""" + + def validate_environment( + self, *, headers, model, api_key, api_base, litellm_params + ): + return {"authorization": f"Bearer {api_key}"} + + def get_complete_url(self, *, api_base, model, optional_params, litellm_params): + return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" + + +@pytest.fixture(autouse=True) +def _reset_rust_flag(): + """Keep the global toggle isolated between tests.""" + rust_bridge.use_litellm_rust(False, ocr=None) + yield + rust_bridge.use_litellm_rust(False, ocr=None) + + +@pytest.fixture +def fake_bridge(): + """Enable the Rust path with an injected recording bridge (no native wheel).""" + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + return bridge + + +def test_use_litellm_rust_toggles_flag(): + assert rust_bridge.rust_ocr_enabled() is False + litellm.use_litellm_rust() + assert rust_bridge.rust_ocr_enabled() is True + litellm.use_litellm_rust(False) + assert rust_bridge.rust_ocr_enabled() is False + + +def test_load_rust_ocr_returns_injected_impl(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + assert rust_bridge.load_rust_ocr() is bridge + + +def test_toggle_without_ocr_arg_preserves_injected_impl(): + """Regression: routine enable/disable calls must not clobber a prior injection. + + Earlier, ``use_litellm_rust()`` unconditionally assigned the keyword default + of ``None`` to ``_rust_ocr_impl``, silently dropping a custom bridge whenever + a caller toggled the flag without re-passing ``ocr=``. + """ + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + litellm.use_litellm_rust(False) + assert rust_bridge.load_rust_ocr() is bridge + litellm.use_litellm_rust(True) + assert rust_bridge.load_rust_ocr() is bridge + + +def test_explicit_ocr_none_clears_injected_impl(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + litellm.use_litellm_rust(True, ocr=None) + assert rust_bridge.load_rust_ocr() is None + + +def test_load_rust_ocr_none_when_extension_absent(): + """With no injected impl and no compiled wheel, the loader returns None so the + caller degrades to the Python path instead of raising ImportError.""" + litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI + assert rust_bridge.load_rust_ocr() is None + + +def test_load_rust_ocr_uses_compiled_extension(monkeypatch): + """With no injected impl but a compiled ``litellm_python_bridge`` importable, + the loader returns the extension's ``ocr`` callable. The native wheel isn't + built in CI, so stand in a fake module via ``sys.modules``.""" + fake_module = types.ModuleType("litellm_python_bridge") + fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) + + litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension + assert rust_bridge.load_rust_ocr() is fake_module.ocr + + +def test_timeout_to_seconds_handles_float_timeout_and_none(): + assert ocr_main._timeout_to_seconds(12.5) == 12.5 + assert ocr_main._timeout_to_seconds(None) is None + assert ocr_main._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 + + +def test_run_rust_ocr_forwards_args_and_wraps_response(): + bridge = RecordingBridge() + logging_obj = RecordingLogging() + + response = ocr_main._run_rust_ocr( + rust_ocr=bridge, + logging_obj=logging_obj, + provider_config=FakeOCRConfig(), + resolve_api_key=lambda _name: None, + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://proxy.internal", + optional_params={"include_image_base64": True}, + litellm_params={}, + timeout_seconds=12.5, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + call = bridge.calls[0] + assert call == { + "model": "mistral-ocr-latest", + "document": DOCUMENT, + "api_key": "sk-test", + "api_base": "https://proxy.internal", + "optional_params": {"include_image_base64": True}, + "timeout_seconds": 12.5, + } + + +def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): + """No explicit api_key: the resolver (get_secret_str in production) supplies it, + so secret-manager backends (AWS/Azure/GCP/Vault) work like the Python path.""" + bridge = RecordingBridge() + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + logging_obj=RecordingLogging(), + provider_config=FakeOCRConfig(), + resolve_api_key=lambda name: ( + "sk-from-vault" if name == "MISTRAL_API_KEY" else None + ), + model="mistral-ocr-latest", + document=DOCUMENT, + api_key=None, + api_base=None, + optional_params={}, + litellm_params={}, + timeout_seconds=None, + ) + + assert bridge.calls[0]["api_key"] == "sk-from-vault" + + +def test_run_rust_ocr_prefers_explicit_key_over_resolver(): + bridge = RecordingBridge() + resolver_calls = [] + + def _resolver(name): + resolver_calls.append(name) + return "sk-from-vault" + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + logging_obj=RecordingLogging(), + provider_config=FakeOCRConfig(), + resolve_api_key=_resolver, + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-explicit", + api_base=None, + optional_params={}, + litellm_params={}, + timeout_seconds=None, + ) + + assert bridge.calls[0]["api_key"] == "sk-explicit" + assert resolver_calls == [] # resolver never consulted when a key is supplied + + +def test_run_rust_ocr_runs_pre_call_logging(): + """The Rust shortcut must run pre_call so callbacks and spend tracking fire.""" + logging_obj = RecordingLogging() + + ocr_main._run_rust_ocr( + rust_ocr=RecordingBridge(), + logging_obj=logging_obj, + provider_config=FakeOCRConfig(), + resolve_api_key=lambda _name: None, + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://api.mistral.ai/v1", + optional_params={"include_image_base64": True}, + litellm_params={}, + timeout_seconds=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 + # The logged request mirrors what Rust sends: resolved URL + headers. + assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" + assert additional_args["headers"] == {"authorization": "Bearer sk-test"} + + +def test_ocr_routes_to_rust_when_enabled(fake_bridge): + response = litellm.ocr( + model=MODEL, + document=DOCUMENT, + api_key="sk-test", + 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] + # Provider prefix is stripped before reaching the bridge. + assert call["model"] == "mistral-ocr-latest" + assert call["document"] == DOCUMENT + assert call["api_key"] == "sk-test" + # Raw OCR params ride along in optional_params; Rust filters to supported keys. + assert call["optional_params"].get("include_image_base64") is True + + +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): + """When no explicit timeout is given, the library default (request_timeout) + must still be forwarded so the Rust path matches the Python path's deadline.""" + from litellm.constants import request_timeout + + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + 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.use_litellm_rust(False, ocr=bridge) + + assert rust_bridge.rust_ocr_enabled() is False + # 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.""" + litellm.use_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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py new file mode 100644 index 00000000000..5481d60a22a --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -0,0 +1,141 @@ +"""Tests for the v1 -> v2 bridge. + +`to_server_spec` maps the migrated modes (none + the static-header family, shared-key) and +defers everything else to v1 by returning None; `to_subject` maps the principal; `raise_public` +maps each CredError onto its HTTP status. These pin the parity-critical mapping before the graft. +""" + +import base64 +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_public, + to_server_spec, + to_subject, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ApiKeyConfig, + CredError, + NoneConfig, + SharedKey, +) +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def _server(**kwargs) -> MCPServer: + return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs) + + +def test_none_maps_to_none_config(): + spec = to_server_spec(_server(auth_type=None)) + assert spec is not None + assert isinstance(spec.config, NoneConfig) + + +def test_api_key_maps_to_x_api_key_shared(): + spec = to_server_spec(_server(auth_type=MCPAuth.api_key, authentication_token="k")) + assert spec is not None and isinstance(spec.config, ApiKeyConfig) + assert spec.config.header_name == "X-API-Key" + assert spec.config.value_prefix == "" + assert isinstance(spec.config.key_source, SharedKey) + assert spec.config.key_source.value.get_secret_value() == "k" + + +@pytest.mark.parametrize( + "auth_type, prefix", + [ + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.token, "token"), + (MCPAuth.authorization, ""), + ], +) +def test_authorization_schemes_map_with_their_prefix(auth_type, prefix): + spec = to_server_spec(_server(auth_type=auth_type, authentication_token="t")) + assert spec is not None and isinstance(spec.config, ApiKeyConfig) + assert spec.config.header_name == "Authorization" + assert spec.config.value_prefix == prefix + assert spec.config.key_source.value.get_secret_value() == "t" + + +def test_basic_scheme_base64_encodes_the_token(): + spec = to_server_spec( + _server(auth_type=MCPAuth.basic, authentication_token="user:pass") + ) + assert spec is not None and isinstance(spec.config, ApiKeyConfig) + assert spec.config.value_prefix == "Basic" + expected = base64.b64encode(b"user:pass").decode() + assert spec.config.key_source.value.get_secret_value() == expected + + +@pytest.mark.parametrize( + "server", + [ + _server(auth_type=MCPAuth.api_key), # no token configured + _server(auth_type=MCPAuth.bearer_token), # no token configured + _server(auth_type=MCPAuth.oauth2), + _server(auth_type=MCPAuth.oauth2_token_exchange), + _server(auth_type=MCPAuth.aws_sigv4), + _server( + auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"] + ), + ], +) +def test_unmigrated_modes_defer_to_v1(server): + # A None spec is the defer signal; the caller falls back to v1. + assert to_server_spec(server) is None + + +@pytest.mark.parametrize( + "server", + [ + _server(auth_type=MCPAuth.api_key, is_byok=True), + # BYOK rides on auth_type, so it must defer for every scheme, not just api_key. A stray + # static token must not route a BYOK server to a v2 shared-key spec with the wrong value. + _server(auth_type=MCPAuth.bearer_token, is_byok=True, authentication_token="x"), + _server(auth_type=MCPAuth.basic, is_byok=True, authentication_token="x"), + _server( + auth_type=MCPAuth.authorization, is_byok=True, authentication_token="x" + ), + _server(auth_type=MCPAuth.token, is_byok=True, authentication_token="x"), + _server(auth_type=None, is_byok=True), + ], +) +def test_byok_defers_regardless_of_auth_type(server): + assert to_server_spec(server) is None + + +def test_to_subject_unauthenticated_is_empty_with_inbound_token(): + subject = to_subject(None, "inbound-jwt") + assert subject.tenant_id == "" + assert subject.subject_id == "" + assert subject.inbound_token is not None + assert subject.inbound_token.get_secret_value() == "inbound-jwt" + + +def test_to_subject_maps_principal_fields(): + principal = SimpleNamespace(org_id="org1", team_id="team1", user_id="user1") + subject = to_subject(principal, None) + assert subject.tenant_id == "org1" + assert subject.subject_id == "user1" + assert subject.inbound_token is None + + +@pytest.mark.parametrize( + "error, status", + [ + (CredError.of_unauthorized("x"), 401), + (CredError.of_misconfigured("x"), 500), + (CredError.of_upstream_unavailable("x"), 503), + (CredError.of_unsupported_mode("x"), 500), + (CredError.of_precondition_required("x"), 412), + (CredError.of_not_implemented("x"), 501), + ], +) +def test_raise_public_maps_each_error_to_its_status(error, status): + with pytest.raises(HTTPException) as exc_info: + raise_public(error) + assert exc_info.value.status_code == status diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py new file mode 100644 index 00000000000..75be6dfc157 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -0,0 +1,104 @@ +"""Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. + +`none` and `api_key` (shared-key source) are implemented; every other arm, plus the `api_key` +BYOK source, returns a typed `not_implemented` error until its mode lands. Parametrizing the +stubs over one config each also guards reachability: a dropped `case` would hit `assert_never` +and raise instead of returning the stub. +""" + +import httpx +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + ApiKeyConfig, + AuthorizationCodeConfig, + AwsSigV4Config, + Byok, + ClientCredentialsConfig, + Error, + NoneConfig, + NoOpAuth, + Ok, + PassthroughConfig, + ServerSpec, + SharedKey, + StaticHeaderAuth, + Subject, + TokenExchangeConfig, + UpstreamCredentialProvider, +) + +_SUBJECT = Subject(tenant_id="", subject_id="") + + +def _spec(config): + return ServerSpec( + server_id="s", resource="https://upstream.example.com", config=config + ) + + +def _emitted(auth: httpx.Auth) -> httpx.Headers: + request = httpx.Request("GET", "https://upstream.example.com/mcp") + flow = auth.auth_flow(request) + next(flow) + flow.close() + return request.headers + + +@pytest.mark.asyncio +async def test_none_mode_yields_a_no_op_auth(): + result = await UpstreamCredentialProvider().resolve_credentials( + _SUBJECT, _spec(NoneConfig()) + ) + assert isinstance(result, Ok) + assert isinstance(result.ok, NoOpAuth) + + +@pytest.mark.asyncio +async def test_api_key_shared_emits_the_configured_header(): + config = ApiKeyConfig( + header_name="X-API-Key", + value_prefix="", + key_source=SharedKey(value=SecretStr("secret-key")), + ) + result = await UpstreamCredentialProvider().resolve_credentials( + _SUBJECT, _spec(config) + ) + assert isinstance(result, Ok) + assert isinstance(result.ok, StaticHeaderAuth) + assert _emitted(result.ok)["X-API-Key"] == "secret-key" + + +@pytest.mark.asyncio +async def test_api_key_shared_honors_authorization_scheme(): + config = ApiKeyConfig( + header_name="Authorization", + value_prefix="Bearer", + key_source=SharedKey(value=SecretStr("tok")), + ) + result = await UpstreamCredentialProvider().resolve_credentials( + _SUBJECT, _spec(config) + ) + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer tok" + + +_STUBBED = [ + ("api_key_byok", ApiKeyConfig(key_source=Byok())), + ("passthrough", PassthroughConfig()), + ("client_credentials", ClientCredentialsConfig()), + ("token_exchange", TokenExchangeConfig()), + ("authorization_code", AuthorizationCodeConfig()), + ("aws_sigv4", AwsSigV4Config(region="us-east-1")), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("label, config", _STUBBED) +async def test_unbuilt_arms_fail_closed_with_not_implemented(label, config): + result = await UpstreamCredentialProvider().resolve_credentials( + _SUBJECT, _spec(config) + ) + assert isinstance(result, Error) + assert result.error.tag == "not_implemented" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index 468bd946ae9..d299239f68e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -41,9 +41,13 @@ class TestMask: def test_empty_returns_none_label(self): assert MCPDebug._mask("") == "(none)" - def test_short_value_unchanged(self): - # visible_prefix=6 + visible_suffix=4 = 10, so <= 10 chars unchanged - assert MCPDebug._mask("sk-1234") == "sk-1234" + def test_short_value_masked(self): + # Short auth values must not be echoed verbatim in debug headers, even though + # visible_prefix + visible_suffix would otherwise reveal the whole value. + masked = MCPDebug._mask("sk-1234") + assert "sk-1234" not in masked + assert set(masked) == {"*"} + assert len(masked) == len("sk-1234") def test_long_value_masked(self): result = MCPDebug._mask("Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index f44552c2943..ca326d197ae 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -3313,6 +3313,7 @@ async def test_list_tools_filters_by_key_team_permissions(): server.server_id = "server1" server.name = "Test Server" server.alias = "test" + server.short_prefix = None server.allowed_tools = None server.disallowed_tools = None server.server_name = "server1" @@ -3423,6 +3424,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): server.server_id = "server1" server.name = "Test Server" server.alias = "test" + server.short_prefix = None server.allowed_tools = None server.disallowed_tools = None server.server_name = "server1" @@ -3519,6 +3521,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): server.server_id = "server1" server.name = "Test Server" server.alias = "test" + server.short_prefix = None server.allowed_tools = None server.disallowed_tools = None server.server_name = "server1" @@ -3620,7 +3623,9 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): server = MagicMock() server.server_id = "gitmcp_server" server.name = "GITMCP" - server.alias = "gitmcp" + server.alias = "GITMCP" + server.short_prefix = None + server.server_name = "GITMCP" server.allowed_tools = None server.disallowed_tools = None 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 8dbee1daa36..35a67391315 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 @@ -4568,5 +4568,173 @@ class TestGetPublicMCPServersLegacyMode: assert sorted(s.server_id for s in result) == ["a", "b"] +class TestCreateMcpClientV2Graft: + """The PR4 v2-resolver graft in ``_create_mcp_client``. + + Migrated HTTP/SSE modes (``none`` plus the static ``api_key`` family) resolve through the + injected ``UpstreamCredentialProvider`` into the ``resolved_auth`` slot; every other mode, + and every stdio server, defers to v1's ``auth_value`` path unchanged. + """ + + def _http_server(self, **overrides: Any) -> MCPServer: + base: Dict[str, Any] = dict( + server_id="http-graft", + name="graft_server", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + ) + base.update(overrides) + return MCPServer(**base) + + async def test_none_mode_resolves_to_noop_auth(self): + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + NoOpAuth, + ) + + client = await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=None) + ) + + assert isinstance(client._resolved_auth, NoOpAuth) + assert client._mcp_auth_value is None + + @pytest.mark.parametrize( + "auth_type, token, expected_name, expected_value", + [ + (MCPAuth.api_key, "k-123", "X-API-Key", "k-123"), + (MCPAuth.bearer_token, "b-123", "Authorization", "Bearer b-123"), + (MCPAuth.token, "t-123", "Authorization", "token t-123"), + (MCPAuth.authorization, "raw-123", "Authorization", "raw-123"), + ], + ) + async def test_static_family_emits_expected_header( + self, auth_type, token, expected_name, expected_value + ): + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + + client = await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=auth_type, authentication_token=token) + ) + + assert isinstance(client._resolved_auth, StaticHeaderAuth) + assert client._resolved_auth.header_name == expected_name + assert client._resolved_auth._header_value.get_secret_value() == expected_value + assert client._mcp_auth_value is None + + async def test_basic_mode_base64_encodes(self): + import base64 + + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + + client = await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=MCPAuth.basic, authentication_token="user:pass") + ) + + encoded = base64.b64encode(b"user:pass").decode() + assert isinstance(client._resolved_auth, StaticHeaderAuth) + assert client._resolved_auth.header_name == "Authorization" + assert ( + client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}" + ) + + async def test_deferred_mode_uses_v1_auth_value(self): + client = await MCPServerManager()._create_mcp_client( + self._http_server( + auth_type=MCPAuth.oauth2, authentication_token="legacy-token" + ) + ) + + assert client._resolved_auth is None + assert client._mcp_auth_value == "legacy-token" + + async def test_static_token_missing_defers_to_v1(self): + client = await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=MCPAuth.api_key, authentication_token=None) + ) + + assert client._resolved_auth is None + + async def test_stdio_migrated_auth_type_still_defers_to_v1(self): + client = await MCPServerManager()._create_mcp_client( + MCPServer( + server_id="stdio-graft", + name="stdio_graft", + transport=MCPTransport.stdio, + command="node", + args=["server.js"], + auth_type=MCPAuth.api_key, + authentication_token="k-stdio", + ) + ) + + assert client.transport_type == MCPTransport.stdio + assert client._resolved_auth is None + assert client._mcp_auth_value == "k-stdio" + + async def test_resolver_error_maps_to_http_exception(self): + from litellm.proxy._experimental.mcp_server.outbound_credentials import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + CredError, + ) + + class _UnauthorizedProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_unauthorized("denied")) + + manager = MCPServerManager(cred_provider=_UnauthorizedProvider()) + + with pytest.raises(HTTPException) as exc: + await manager._create_mcp_client(self._http_server(auth_type=None)) + + assert exc.value.status_code == 401 + + async def test_per_request_override_defers_to_v1(self): + # A per-request override (mcp_auth_header) must win over the shared static token, + # exactly as v1 did, so a migrated static server defers to v1 when one is present. + client = await MCPServerManager()._create_mcp_client( + self._http_server( + auth_type=MCPAuth.bearer_token, authentication_token="shared-tok" + ), + mcp_auth_header="caller-override", + ) + + assert client._resolved_auth is None + assert client._mcp_auth_value == "caller-override" + + async def test_conflicting_extra_header_skips_resolved_auth_on_v2(self): + # An Authorization already supplied via extra_headers (guardrail hook like the JWT + # signer, static_headers, or a forwarded caller header) must win. The server stays on + # the v2 path but skips resolved_auth, so nothing overwrites the inbound header. + client = await MCPServerManager()._create_mcp_client( + self._http_server( + auth_type=MCPAuth.bearer_token, authentication_token="shared-tok" + ), + extra_headers={"Authorization": "Bearer hook-jwt"}, + ) + + assert client._resolved_auth is None + assert client._mcp_auth_value is None + assert client._get_auth_headers()["Authorization"] == "Bearer hook-jwt" + + async def test_none_with_extra_header_stays_v2_without_clobbering(self): + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + NoOpAuth, + ) + + # none resolves to NoOpAuth, which writes no header, so it cannot clobber an inbound + # Authorization; it stays on the v2 path and the inbound header is preserved verbatim. + client = await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=None), + extra_headers={"Authorization": "Bearer hook-jwt"}, + ) + + assert isinstance(client._resolved_auth, NoOpAuth) + assert client._get_auth_headers()["Authorization"] == "Bearer hook-jwt" + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index d80d15e2140..67b72dbb695 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -676,9 +676,11 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): @pytest.mark.asyncio async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_challenge(): """ - OAuth2 server with ``delegate_auth_to_upstream=True`` should let the - upstream MCP server's RFC 9728 challenge reach the client instead of - pre-emptively returning LiteLLM's gateway authorization_uri challenge. + OAuth2 server with ``delegate_auth_to_upstream=True`` where the client + already presents a bearer token the upstream rejects: the request bypasses + the preemptive challenge (a token is present) and reaches the session + manager, whose upstream ``MCPUpstreamAuthError`` is surfaced verbatim so the + client sees the upstream's RFC 9728 challenge rather than the gateway's. """ from fastapi import HTTPException @@ -722,9 +724,7 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha delegated_server.needs_user_oauth_token = True delegated_server.server_id = "delegated-oauth-server" - upstream_challenge = ( - 'Bearer resource_metadata="https://upstream.example.com/.well-known/oauth-protected-resource"' - ) + upstream_challenge = 'Bearer resource_metadata="https://upstream.example.com/.well-known/oauth-protected-resource"' with ( patch( @@ -735,7 +735,7 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha None, ["delegated_oauth_server"], None, - None, + {"Authorization": "Bearer upstream-token-the-upstream-rejects"}, None, ), ), @@ -863,16 +863,24 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): @pytest.mark.asyncio -async def test_handle_streamable_http_mcp_delegated_server_without_token_reaches_session_manager(): +async def test_handle_streamable_http_mcp_delegated_server_without_token_returns_preemptive_resource_metadata_401(): """ OAuth2 server with ``delegate_auth_to_upstream=True`` and no stored token - should not receive LiteLLM's gateway authorization_uri challenge. The - request continues so the upstream MCP server can emit its RFC 9728 challenge. + must fail fast with a 401 carrying a ``resource_metadata=`` challenge that + points at the gateway's proxied oauth-protected-resource well-known, so the + MCP client starts PKCE against the upstream IdP. On ``initialize`` the + gateway answers locally and never probes upstream, so this preemptive + challenge is the only thing that can drive the client into the OAuth flow; + falling through to the session manager (or emitting the gateway + ``authorization_uri=`` challenge) leaves the client with no tools and no + sign-in prompt. """ + from fastapi import HTTPException + try: from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, - session_manager_stateless, + session_manager_stateful, ) except ImportError: pytest.skip("MCP server not available") @@ -880,7 +888,12 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_reaches scope = { "type": "http", "method": "POST", - "path": "/mcp", + "path": "/mcp/delegated_oauth_server", + "_original_path": "/delegated_oauth_server/mcp", + "scheme": "https", + "query_string": b"", + "root_path": "", + "server": ("litellm.example.com", 443), "headers": [ (b"content-type", b"application/json"), (b"host", b"litellm.example.com"), @@ -889,7 +902,7 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_reaches receive = AsyncMock( return_value={ "type": "http.request", - "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + "body": b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', "more_body": False, } ) @@ -900,6 +913,7 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_reaches delegated_server.auth_type = MCPAuth.oauth2 delegated_server.delegate_auth_to_upstream = True delegated_server.needs_user_oauth_token = True + delegated_server.server_id = "delegated-oauth-server" with ( patch( @@ -936,12 +950,20 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_reaches return_value=delegated_server, ), patch.object( - session_manager_stateless, + session_manager_stateful, "handle_request", new_callable=AsyncMock, ) as mock_handle_request, ): - await handle_streamable_http_mcp(scope, receive, send) + with pytest.raises(HTTPException) as exc_info: + await handle_streamable_http_mcp(scope, receive, send) assert mock_get_stored_token.await_count == 1 - assert mock_handle_request.await_count == 1 + assert mock_handle_request.await_count == 0 + assert exc_info.value.status_code == 401 + challenge = exc_info.value.headers["www-authenticate"] + assert "resource_metadata=" in challenge + assert "authorization_uri=" not in challenge + assert ( + "/.well-known/oauth-protected-resource/delegated_oauth_server/mcp" in challenge + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index b41cdfb1576..e7709165119 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -235,6 +235,132 @@ class TestFetchMCPToolsetsAccess: mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1", "ts-2"]) +class TestToolsetPrefixResolution: + """Regression for LIT-3419. + + Toolsets store bare tool names; the live tools come back prefixed with the + server's own prefix. Reconciling them must strip exactly that prefix, not + chop at the first separator, otherwise tools on a server whose prefix + contains the separator (a hyphenated alias, or the UUID server_id used when + a server has no alias) are silently dropped from the toolset. + """ + + # alias, server_name, server_id; the clean-alias row worked before the fix, + # the hyphenated-alias and no-alias (UUID prefix) rows did not. + PREFIX_CASES = [ + ("deepwiki", None, "srv-clean"), + ("deep-wiki", None, "srv-hyphen"), + (None, None, "117c814c-1a2b-3c4d-9e8f"), + ] + + @staticmethod + def _server(alias, server_name, server_id): + from types import SimpleNamespace + + return SimpleNamespace( + alias=alias, + server_name=server_name, + server_id=server_id, + short_prefix=None, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("alias, server_name, server_id", PREFIX_CASES) + async def test_filter_keeps_tools_when_prefix_contains_separator( + self, alias, server_name, server_id + ): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_key_team_permissions, + ) + from litellm.proxy._experimental.mcp_server.utils import ( + add_server_prefix_to_name, + get_server_prefix, + ) + + server = self._server(alias, server_name, server_id) + prefix = get_server_prefix(server) + live_tools = [ + MCPTool( + name=add_server_prefix_to_name(name, prefix), + inputSchema={"type": "object"}, + ) + for name in ("read_wiki_contents", "read_wiki_structure", "not_granted") + ] + # Bare names as stored in the toolset / resolved into the permission dict. + allowed = ["read_wiki_contents", "read_wiki_structure"] + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server." + "MCPRequestHandler.get_allowed_tools_for_server", + new=AsyncMock(return_value=allowed), + ), + patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.get_mcp_server_by_id", + return_value=server, + ), + ): + result = await filter_tools_by_key_team_permissions( + tools=live_tools, + server_id=server_id, + user_api_key_auth=_make_auth(), + ) + + assert sorted(t.name for t in result) == sorted( + add_server_prefix_to_name(name, prefix) + for name in ("read_wiki_contents", "read_wiki_structure") + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("alias, server_name, server_id", PREFIX_CASES) + async def test_resolve_unprefixes_stored_names_with_separator_prefix( + self, alias, server_name, server_id + ): + from types import SimpleNamespace + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.utils import ( + add_server_prefix_to_name, + get_server_prefix, + ) + + server = self._server(alias, server_name, server_id) + # A caller (e.g. the management API) may persist already-prefixed names; + # resolution must reduce them to the true bare name regardless of prefix. + stored = add_server_prefix_to_name( + "read_wiki_contents", get_server_prefix(server) + ) + toolset = SimpleNamespace(tools=[{"server_id": server_id, "tool_name": stored}]) + cache = MagicMock( + async_get_cache=AsyncMock(return_value=None), + async_set_cache=AsyncMock(), + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager." + "global_mcp_server_manager.get_mcp_server_by_id", + return_value=server, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch( + "litellm.proxy._experimental.mcp_server.toolset_db.list_mcp_toolsets", + new=AsyncMock(return_value=[toolset]), + ), + ): + result = await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=["ts-1"] + ) + + assert result == {server_id: ["read_wiki_contents"]} + + class TestMCPActiveToolsetContextVar: """Tests for _mcp_active_toolset_id ContextVar — clients cannot inject it.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 2558df8533b..cebc265a148 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -452,6 +452,430 @@ async def test_semantic_filter_hook_skips_no_tools(): print("✅ Hook correctly skips requests without tools") +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_native_tools(): + """ + Regression test: mixed MCP + native tools. + + Given: 5 MCP tools (registered in _tool_map) + 2 native OpenAI-format + function tools (not in _tool_map) + When: The hook filters tools + Then: The native tools must survive unconditionally, and only MCP + tools go through the semantic filter. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + # --- MCP tools (registered in the semantic router) --- + mcp_tools = [ + MCPTool( + name=f"mcp_tool_{i}", + description=f"MCP tool {i}", + inputSchema={"type": "object"}, + ) + for i in range(5) + ] + filter_instance._build_router(mcp_tools) + + # --- Native OpenAI-format function tools (NOT in _tool_map) --- + native_tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": {"type": "object", "properties": {}}, + }, + }, + { + "type": "function", + "function": { + "name": "search_web", + "description": "Search the web", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ] + + # Combine: MCP tools + native tools + all_tools = list(mcp_tools) + native_tools + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "What is the weather?"}], + "tools": all_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # Native tools must survive + native_in_result = [ + t for t in filtered if isinstance(t, dict) and t.get("type") == "function" + ] + assert ( + len(native_in_result) == 2 + ), f"Both native tools must survive, got {len(native_in_result)}" + + # MCP tools should be filtered (top_k=2) + mcp_in_result = [t for t in filtered if not isinstance(t, dict)] + assert ( + len(mcp_in_result) <= 2 + ), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" + + # Total should be native + filtered MCP + assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}" + + # Filter stats should be emitted (MCP tools were present) + assert "litellm_semantic_filter_stats" in result["metadata"] + + # Stats should report MCP-only counts, not inflated with native tools + stats = result["metadata"]["litellm_semantic_filter_stats"] + mcp_before, mcp_after = stats.split("->") + assert ( + int(mcp_before) == 5 + ), f"Stats 'from' should be MCP count (5), got {mcp_before}" + assert int(mcp_after) == len( + mcp_in_result + ), f"Stats 'to' should match filtered MCP count, got {mcp_after}" + + print( + f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} " + f"({len(native_in_result)} native + {len(mcp_in_result)} MCP), " + f"stats={stats}" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_all_native_tools(): + """ + Regression test: all-native request. + + Given: Only native OpenAI-format function tools (none registered in + the MCP semantic router) + When: The hook processes the request + Then: All tools pass through, and NO spurious semantic filter response + headers are emitted (no litellm_semantic_filter_stats in metadata). + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + mock_router = Mock() + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=3, + similarity_threshold=0.3, + enabled=True, + ) + + # Build router with some MCP tools (so tool_router is not None) + mcp_tools = [ + MCPTool( + name="some_mcp_tool", + description="An MCP tool", + inputSchema={"type": "object"}, + ) + ] + + from litellm.types.utils import Embedding, EmbeddingResponse + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance._build_router(mcp_tools) + + # --- Only native tools in the request --- + native_tools = [ + { + "type": "function", + "function": { + "name": f"native_func_{i}", + "description": f"Native function {i}", + "parameters": {"type": "object", "properties": {}}, + }, + } + for i in range(3) + ] + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "tools": native_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # All native tools must pass through + assert ( + len(filtered) == 3 + ), f"All 3 native tools must pass through, got {len(filtered)}" + + # No spurious semantic filter stats (P2 fix) + assert ( + "litellm_semantic_filter_stats" not in result["metadata"] + ), "Should NOT emit semantic filter stats for all-native-tool requests" + + print( + f"✅ Hook passes through all {len(filtered)} native tools, " + f"no spurious filter headers emitted" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_responses_api_name_collision(): + """ + Regression test: Responses API native tool with MCP-matching name. + + Given: A Responses-API native tool whose top-level ``name`` collides + with an MCP canonical name in ``_tool_map`` + When: The hook classifies tools + Then: The native tool must NOT be sent to the semantic filter, even + though its name matches an MCP canonical. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + # Register an MCP tool with name "github-search" + mcp_tools = [ + MCPTool( + name="github-search", + description="Search GitHub repos", + inputSchema={"type": "object"}, + ) + ] + filter_instance._build_router(mcp_tools) + + # Responses API native tool with SAME name as MCP canonical + responses_api_tool = { + "type": "function", + "name": "github-search", + "description": "Caller-owned search tool", + "parameters": {"type": "object"}, + } + + hook = SemanticToolFilterHook(filter_instance) + + # Verify classification: should be native, not MCP + assert not hook._is_mcp_tool(responses_api_tool), ( + "Responses API tool with type=function + top-level name " + "should be classified as native, not MCP" + ) + + # Full hook test: all-native request should preserve tools + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Search GitHub"}], + "tools": [responses_api_tool], + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + # All tools are native → hook returns data with all tools preserved + filtered = (result or data)["tools"] + assert len(filtered) == 1, f"Native tool must survive, got {len(filtered)}" + assert filtered[0]["name"] == "github-search" + + print("✅ Responses API tool with MCP-matching name correctly classified as native") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_tool_order(): + """ + Regression test: tool ordering preservation. + + Given: An interleaved request [mcp_tool_A, native_tool, mcp_tool_B] + When: The hook filters tools (all MCP tools survive) + Then: The output order must match the original request order, + NOT native-first. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=5, + similarity_threshold=0.3, + enabled=True, + ) + + # Register MCP tools + mcp_tool_a = MCPTool( + name="github-search", + description="Search GitHub", + inputSchema={"type": "object"}, + ) + mcp_tool_b = MCPTool( + name="github-issue", + description="Create GitHub issue", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([mcp_tool_a, mcp_tool_b]) + + # Mock filter_tools to return both MCP tools (deterministic) + filter_instance.filter_tools = AsyncMock( # type: ignore[method-assign] + return_value=[mcp_tool_a, mcp_tool_b] + ) + + # Native tool (interleaved between MCP tools) + native_tool = { + "type": "function", + "function": { + "name": "weather_lookup", + "description": "Look up weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + + # Original order: [mcp_A, native, mcp_B] + original_tools = [mcp_tool_a, native_tool, mcp_tool_b] + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Search GitHub and check weather"}], + "tools": original_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # All tools should survive + assert len(filtered) == 3, f"Expected 3 tools, got {len(filtered)}" + + # Order must be preserved: [mcp_A, native, mcp_B] + assert filtered[0] is mcp_tool_a, "First tool should be mcp_tool_a" + assert filtered[1] is native_tool, "Second tool should be native_tool" + assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b" + + print( + "✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering" + ) + + class TestGetToolsByNames: """ Regression coverage for SemanticMCPToolFilter._get_tools_by_names @@ -489,9 +913,7 @@ class TestGetToolsByNames: {"name": "send_email", "description": "send mail"}, ] - matched = filter_instance._get_tools_by_names( - ["send_email"], available_tools - ) + matched = filter_instance._get_tools_by_names(["send_email"], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "send_email" @@ -503,9 +925,7 @@ class TestGetToolsByNames: client_name = "litellm_" + canonical available_tools = [{"name": client_name, "description": "scrape"}] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 # Must return the incoming tool unchanged so the client-facing @@ -516,13 +936,9 @@ class TestGetToolsByNames: """Some clients use dash as alias separator; accept that too.""" filter_instance = self._make_filter() canonical = "weather_svc-get_weather" - available_tools = [ - {"name": "mcp-" + canonical, "description": "weather"} - ] + available_tools = [{"name": "mcp-" + canonical, "description": "weather"}] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "mcp-" + canonical @@ -552,9 +968,7 @@ class TestGetToolsByNames: {"name": "litellm_" + canonical, "description": "wrapped"}, ] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == canonical @@ -567,9 +981,7 @@ class TestGetToolsByNames: separator-anchored suffixes of ``litellm_api-fs-read_file``. """ filter_instance = self._make_filter() - available_tools = [ - {"name": "litellm_api-fs-read_file", "description": "read"} - ] + available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}] matched = filter_instance._get_tools_by_names( ["fs-read_file", "api-fs-read_file"], available_tools @@ -590,9 +1002,7 @@ class TestGetToolsByNames: {"name": "my_" + canonical, "description": "plain search"}, ] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "my_" + canonical diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index bf90e9ebef6..662ef585c6b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( get_server_prefix, is_short_mcp_tool_prefix_enabled, iter_known_server_prefixes, + strip_known_server_prefix, ) from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -143,6 +144,57 @@ class TestIterKnownServerPrefixes: assert compute_short_server_prefix(server.server_id) in prefixes +# --------------------------------------------------------------------------- +# strip_known_server_prefix — inverse of add_server_prefix_to_name (LIT-3419) +# --------------------------------------------------------------------------- + + +class TestStripKnownServerPrefix: + """Removes a server's exact known prefix, even when the prefix itself + contains the separator (hyphenated alias or UUID server_id fallback) where + a first-separator split would cut inside the prefix and mangle the name.""" + + def test_clean_prefix_round_trips(self): + server = _make_server(alias="deepwiki", server_name="deepwiki") + live = add_server_prefix_to_name( + "read_wiki_contents", get_server_prefix(server) + ) + assert strip_known_server_prefix(live, server) == "read_wiki_contents" + + def test_hyphenated_alias_does_not_mangle(self): + server = _make_server(alias="deep-wiki", server_name="deep-wiki") + assert get_server_prefix(server) == "deep-wiki" + # the first-separator split would yield "wiki-read_wiki_contents" + assert ( + strip_known_server_prefix("deep-wiki-read_wiki_contents", server) + == "read_wiki_contents" + ) + + def test_uuid_fallback_prefix_does_not_mangle(self): + server = MCPServer( + server_id="117c814c-1a2b-3c4d-9e8f", + name="deepwiki", + alias=None, + server_name=None, + transport="http", + ) + live = add_server_prefix_to_name( + "read_wiki_contents", get_server_prefix(server) + ) + assert live.startswith("117c814c-1a2b-3c4d-9e8f-") + assert strip_known_server_prefix(live, server) == "read_wiki_contents" + + def test_unprefixed_name_returned_unchanged(self): + server = _make_server(alias="deep-wiki", server_name="deep-wiki") + assert ( + strip_known_server_prefix("read_wiki_contents", server) + == "read_wiki_contents" + ) + + def test_none_server_falls_back_to_first_separator_split(self): + assert strip_known_server_prefix("svc-tool", None) == "tool" + + # --------------------------------------------------------------------------- # Manager-level behaviour: list + reverse-lookup # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 702bae77339..c8dc0ea5ed6 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -351,6 +351,43 @@ async def test_can_key_call_model_all_team_models_no_team_id_is_denied(): assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied +@pytest.mark.asyncio +async def test_can_team_access_model_all_team_models_expands_router_models(): + from litellm import Router + from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import can_team_access_model + + team_object = LiteLLM_TeamTable( + team_id="team-123", + models=[SpecialModelNames.all_team_models.value], + ) + router = Router( + model_list=[ + { + "model_name": "allowed-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}, + } + ] + ) + + assert ( + await can_team_access_model( + model="allowed-model", + team_object=team_object, + llm_router=router, + ) + is True + ) + with pytest.raises(ProxyException) as exc_info: + await can_team_access_model( + model="blocked-model", + team_object=team_object, + llm_router=router, + ) + + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + + @pytest.mark.asyncio async def test_get_key_object_should_reconnect_once_on_db_connection_error(): mock_prisma_client = MagicMock() diff --git a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py index 9444e4ebd2d..e0585ab04f1 100644 --- a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py +++ b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py @@ -5,11 +5,22 @@ Tests that internal callers see all MCP servers while external callers only see servers with available_on_public_internet=True. """ +import logging from unittest.mock import MagicMock, patch +import pytest from fastapi import Request +from pydantic import ValidationError -from litellm.proxy.auth.ip_address_utils import IPAddressUtils +import litellm.proxy.auth.ip_address_utils as ip_mod +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ConfigGeneralSettings +from litellm.proxy.auth.ip_address_utils import ( + IPAddressUtils, + _HopCount, + _HopCountInvalid, + _HopCountUnset, +) from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -77,6 +88,44 @@ class TestMCPClientIPExtraction: assert result == "" assert IPAddressUtils.is_internal_ip(result) is False + def test_no_trusted_ranges_warning_matches_fail_closed_behavior(self, caplog): + ip_mod._warned_xff_without_trusted_ranges = False + + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = "203.0.113.5" + request.headers = {"x-forwarded-for": "1.2.3.4, 10.0.0.1"} + general_settings = {"use_x_forwarded_for": True} + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + assert ( + IPAddressUtils.is_request_from_trusted_proxy( + request, general_settings=general_settings + ) + is False + ) + + warning = next( + ( + record.getMessage() + for record in caplog.records + if "mcp_trusted_proxy_ranges" in record.getMessage() + ), + None, + ) + assert ( + warning is not None + ), "Expected a warning containing 'mcp_trusted_proxy_ranges' but none was logged" + + assert "fails closed" in warning + assert "treated as external" in warning + assert "client IPs will use the proxy's literal request values" not in warning + + assert ( + IPAddressUtils.get_mcp_client_ip(request, general_settings=general_settings) + == "" + ) + def test_private_proxy_peer_does_not_grant_internal_access(self): # Regression: behind an internal reverse proxy with use_x_forwarded_for # enabled but mcp_trusted_proxy_ranges unset, the direct peer is the @@ -128,6 +177,332 @@ class TestMCPClientIPExtraction: assert result == "203.0.113.5" +def _make_request(client_host, xff): + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = client_host + request.headers = {"x-forwarded-for": xff} + return request + + +class TestExtractClientIpFromXffHops: + def test_single_hop_picks_rightmost(self): + assert ( + IPAddressUtils.extract_client_ip_from_xff_hops( + "10.0.0.99, 203.0.113.9", num_trusted_hops=1 + ) + == "203.0.113.9" + ) + + def test_two_hops_picks_second_from_right(self): + assert ( + IPAddressUtils.extract_client_ip_from_xff_hops( + "10.0.0.99, 203.0.113.9, 172.16.0.1", num_trusted_hops=2 + ) + == "203.0.113.9" + ) + + def test_whitespace_and_empty_entries_are_ignored(self): + assert ( + IPAddressUtils.extract_client_ip_from_xff_hops( + " 1.1.1.1 , , 203.0.113.9 ", num_trusted_hops=1 + ) + == "203.0.113.9" + ) + + def test_chain_shorter_than_hops_returns_none(self): + assert ( + IPAddressUtils.extract_client_ip_from_xff_hops( + "203.0.113.9", num_trusted_hops=2 + ) + is None + ) + + def test_zero_or_negative_hops_returns_none(self): + assert ( + IPAddressUtils.extract_client_ip_from_xff_hops( + "203.0.113.9", num_trusted_hops=0 + ) + is None + ) + + def test_invalid_selected_entry_returns_none(self): + assert ( + IPAddressUtils.extract_client_ip_from_xff_hops( + "not-an-ip, 10.0.0.1", num_trusted_hops=2 + ) + is None + ) + + +class TestResolveNumTrustedHops: + def test_unset_is_unset(self): + assert IPAddressUtils._resolve_num_trusted_hops(None) == _HopCountUnset() + + def test_int_value(self): + assert IPAddressUtils._resolve_num_trusted_hops(2) == _HopCount(2) + + def test_numeric_string_is_coerced(self): + assert IPAddressUtils._resolve_num_trusted_hops("3") == _HopCount(3) + + def test_zero_is_invalid_not_disabled(self): + assert IPAddressUtils._resolve_num_trusted_hops(0) == _HopCountInvalid() + + def test_non_numeric_value_is_invalid(self): + assert IPAddressUtils._resolve_num_trusted_hops("abc") == _HopCountInvalid() + + def test_below_minimum_warns_so_misconfig_is_visible(self): + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger" + ) as mock_logger: + assert IPAddressUtils._resolve_num_trusted_hops(0) == _HopCountInvalid() + assert IPAddressUtils._resolve_num_trusted_hops(-3) == _HopCountInvalid() + assert mock_logger.warning.call_count == 2 + + def test_unset_does_not_warn(self): + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger" + ) as mock_logger: + assert IPAddressUtils._resolve_num_trusted_hops(None) == _HopCountUnset() + mock_logger.warning.assert_not_called() + + def test_valid_value_does_not_warn(self): + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger" + ) as mock_logger: + assert IPAddressUtils._resolve_num_trusted_hops(2) == _HopCount(2) + mock_logger.warning.assert_not_called() + + +class TestConfigGeneralSettingsHopsValidation: + """mcp_xff_num_trusted_hops must reject sub-minimum values at config-parse time.""" + + @pytest.mark.parametrize("bad_value", [0, -1]) + def test_below_minimum_is_rejected(self, bad_value): + with pytest.raises(ValidationError): + ConfigGeneralSettings(mcp_xff_num_trusted_hops=bad_value) + + def test_valid_value_and_unset_are_accepted(self): + assert ( + ConfigGeneralSettings(mcp_xff_num_trusted_hops=1).mcp_xff_num_trusted_hops + == 1 + ) + assert ConfigGeneralSettings().mcp_xff_num_trusted_hops is None + + +class TestXffTrustedHopsAccessControl: + def test_spoofed_internal_leftmost_is_defeated(self): + request = _make_request("10.0.0.5", "10.0.0.99, 203.0.113.9") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + "mcp_xff_num_trusted_hops": 1, + }, + ) + + assert result == "203.0.113.9" + assert IPAddressUtils.is_internal_ip(result) is False + + def test_genuine_internal_client_is_preserved(self): + request = _make_request("10.0.0.5", "10.0.0.50") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + "mcp_xff_num_trusted_hops": 1, + }, + ) + + assert result == "10.0.0.50" + assert IPAddressUtils.is_internal_ip(result) is True + + def test_two_hops_skips_proxy_appended_addresses(self): + request = _make_request("10.0.0.5", "10.0.0.99, 203.0.113.9, 172.16.0.1") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + "mcp_xff_num_trusted_hops": 2, + }, + ) + + assert result == "203.0.113.9" + + def test_short_chain_fails_closed(self): + request = _make_request("10.0.0.5", "203.0.113.9") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + "mcp_xff_num_trusted_hops": 2, + }, + ) + + assert result == "" + assert IPAddressUtils.is_internal_ip(result) is False + + def test_hops_without_trusted_ranges_still_fails_closed(self): + request = _make_request("203.0.113.5", "10.0.0.99, 8.8.8.8") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_xff_num_trusted_hops": 1, + }, + ) + + assert result == "" + + def test_unset_hops_keeps_legacy_leftmost_behavior(self): + request = _make_request("10.0.0.5", "10.0.0.99, 203.0.113.9") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + + assert result == "10.0.0.99, 203.0.113.9" + + @pytest.mark.parametrize("bad_value", [0, -1, "abc", 1.5]) + def test_invalid_hops_config_fails_closed_not_legacy(self, bad_value): + request = _make_request("10.0.0.5", "10.0.0.99, 203.0.113.9") + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + "mcp_xff_num_trusted_hops": bad_value, + }, + ) + + assert result == "" + assert IPAddressUtils.is_internal_ip(result) is False + + +class TestXffPresentButDisabledWarning: + """When an XFF header arrives but use_x_forwarded_for is off, the proxy must + loudly warn (the internal-only check is silently trusting the load balancer's + IP) yet still serve the request, so a crafted header can't DoS a no-LB deploy.""" + + def _reset_warning_flag(self): + from litellm.proxy.auth import ip_address_utils + + ip_address_utils._warned_xff_present_but_disabled = False + + def _request_with_xff(self): + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = "10.0.0.7" + request.headers = {"x-forwarded-for": "8.8.8.8"} + return request + + def test_warns_and_does_not_fail_when_xff_present_but_disabled(self): + self._reset_warning_flag() + request = self._request_with_xff() + + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger.error" + ) as mock_error: + result = IPAddressUtils.get_mcp_client_ip( + request, general_settings={"use_x_forwarded_for": False} + ) + + # Does not hard-fail: falls back to the direct peer (the load balancer). + assert result == "10.0.0.7" + mock_error.assert_called_once() + assert "use_x_forwarded_for" in str(mock_error.call_args) + + def test_warning_is_one_shot(self): + self._reset_warning_flag() + + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger.error" + ) as mock_error: + IPAddressUtils.get_mcp_client_ip( + self._request_with_xff(), + general_settings={"use_x_forwarded_for": False}, + ) + IPAddressUtils.get_mcp_client_ip( + self._request_with_xff(), + general_settings={"use_x_forwarded_for": False}, + ) + + # One-shot so a flood of crafted XFF headers cannot spam the logs. + mock_error.assert_called_once() + + def test_re_arms_after_xff_is_enabled_then_disabled_again(self): + self._reset_warning_flag() + + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger.error" + ) as mock_error: + IPAddressUtils.get_mcp_client_ip( + self._request_with_xff(), + general_settings={"use_x_forwarded_for": False}, + ) + # Operator fixes the config; observing it enabled re-arms the warning. + IPAddressUtils.get_mcp_client_ip( + self._request_with_xff(), + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + # Config rolls back to disabled: the misconfiguration must warn again. + IPAddressUtils.get_mcp_client_ip( + self._request_with_xff(), + general_settings={"use_x_forwarded_for": False}, + ) + + assert mock_error.call_count == 2 + + def test_no_warning_without_xff_header(self): + self._reset_warning_flag() + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = "10.0.0.7" + request.headers = {} + + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger.error" + ) as mock_error: + IPAddressUtils.get_mcp_client_ip( + request, general_settings={"use_x_forwarded_for": False} + ) + + mock_error.assert_not_called() + + def test_no_warning_when_xff_enabled(self): + self._reset_warning_flag() + + with patch( + "litellm.proxy.auth.ip_address_utils.verbose_proxy_logger.error" + ) as mock_error: + IPAddressUtils.get_mcp_client_ip( + self._request_with_xff(), + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + + mock_error.assert_not_called() + + class TestMCPServerIPFiltering: """Tests that external callers only see public MCP servers.""" diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 02b1f698132..261485e8965 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -543,3 +543,77 @@ async def test_get_available_models_for_user_expands_query_team_wildcard( ) assert "openai/gpt-4o-mini" in result + + +def test_get_key_models_all_team_models_recursive_team(): + """GH#30619: when key and team both have all-team-models, + the sentinel should expand to proxy_model_list.""" + from litellm.proxy.auth.model_checks import get_key_models + from litellm.proxy._types import SpecialModelNames + + user_api_key_dict = type( + "obj", (object,), + { + "models": [SpecialModelNames.all_team_models.value], + "team_id": "team-1", + "team_models": [SpecialModelNames.all_team_models.value], + }, + )() + proxy_model_list = ["model-a", "model-b"] + result = get_key_models(user_api_key_dict, proxy_model_list, {}) + assert SpecialModelNames.all_team_models.value not in result + assert set(result) == {"model-a", "model-b"} + + +def test_get_key_models_all_team_models_keeps_mixed_team_entries(): + from litellm.proxy.auth.model_checks import get_key_models + from litellm.proxy._types import SpecialModelNames + + user_api_key_dict = type( + "obj", + (object,), + { + "models": [SpecialModelNames.all_team_models.value], + "team_id": "team-1", + "team_models": [ + SpecialModelNames.all_team_models.value, + "restricted-model", + ], + }, + )() + result = get_key_models(user_api_key_dict, ["model-a", "model-b"], {}) + assert SpecialModelNames.all_team_models.value not in result + assert set(result) == {"model-a", "model-b", "restricted-model"} + + +def test_get_team_models_all_team_models_expands(): + """GH#30619: all-team-models in team_models should expand.""" + from litellm.proxy.auth.model_checks import get_team_models + from litellm.proxy._types import SpecialModelNames + + result = get_team_models( + [SpecialModelNames.all_team_models.value], + ["model-a", "model-b"], + {}, + ) + assert SpecialModelNames.all_team_models.value not in result + assert set(result) == {"model-a", "model-b"} + + +def test_get_team_models_all_team_models_expands_with_access_groups(): + """GH#30619: all-team-models with include_model_access_groups + should include access group keys.""" + from litellm.proxy.auth.model_checks import get_team_models + from litellm.proxy._types import SpecialModelNames + + result = get_team_models( + [SpecialModelNames.all_team_models.value], + ["model-a", "model-b"], + {"group-1": ["g1-model"], "group-2": ["g2-model"]}, + include_model_access_groups=True, + ) + assert SpecialModelNames.all_team_models.value not in result + assert "model-a" in result + assert "model-b" in result + assert "group-1" in result + assert "group-2" in result 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 b2212068a5b..573bd5ae584 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -16,7 +16,11 @@ from unittest.mock import patch import pytest -from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.db_url_settings import ( + DatabaseURLSettings, + unsupported_db_scheme, + unsupported_db_scheme_message, +) def _apply() -> bool: @@ -27,6 +31,7 @@ def _apply() -> bool: _MANAGED_DB_ENV_VARS = ( "IAM_TOKEN_DB_AUTH", "DATABASE_URL", + "DIRECT_URL", "DATABASE_URL_READ_REPLICA", "DATABASE_HOST", "DATABASE_PORT", @@ -287,3 +292,87 @@ def test_password_reader_uses_own_credentials(monkeypatch): os.environ["DATABASE_URL_READ_REPLICA"] == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db" ) + + +@pytest.mark.parametrize( + "url", + [ + "postgresql://u:p@host:5432/db", + "postgres://u:p@host:5432/db", + "POSTGRESQL://u:p@host:5432/db", + "postgresql://host/db?schema=public", + ], +) +def test_unsupported_db_scheme_accepts_postgres(url): + assert unsupported_db_scheme(url) is None + + +@pytest.mark.parametrize( + "url,scheme", + [ + ("sqlite:///data/litellm.db", "sqlite"), + ("sqlite:///./local.db", "sqlite"), + ("mysql://u:p@host:3306/db", "mysql"), + ("mssql://host/db", "mssql"), + ], +) +def test_unsupported_db_scheme_rejects_non_postgres(url, scheme): + assert unsupported_db_scheme(url) == scheme + + +def test_unsupported_db_scheme_does_not_echo_schemeless_credentials(): + """A malformed schemeless DSN must not leak its embedded credentials + through the return value (which callers log).""" + leaky = "litellm:s3cr3t_password@db.internal:5432/litellm" + + result = unsupported_db_scheme(leaky) + + assert result is not None + assert "s3cr3t_password" not in result + assert "db.internal" not in result + + +def test_apply_to_env_rejects_pinned_sqlite_writer(monkeypatch): + """Componentized entrypoints pin DATABASE_URL and call apply_to_env; a + sqlite writer must raise here rather than reach Prisma.""" + monkeypatch.setenv("DATABASE_URL", "sqlite:///data/litellm.db") + + with pytest.raises(RuntimeError, match="sqlite"): + _apply() + + # The bad URL must not have been propagated as a usable connection string. + assert os.environ["DATABASE_URL"] == "sqlite:///data/litellm.db" + + +def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): + """DIRECT_URL reaches Prisma the same way DATABASE_URL does; a non-postgres + direct URL must be rejected in apply_to_env, matching the CLI startup guard.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv("DIRECT_URL", "sqlite:///data/litellm.db") + + with pytest.raises(RuntimeError, match="DIRECT_URL.*sqlite"): + _apply() + + +def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" + ) + + with pytest.raises(RuntimeError, match="DATABASE_URL_READ_REPLICA.*mysql"): + _apply() + + +def test_apply_to_env_accepts_pinned_postgres(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@host:5432/db") + + # Operator-pinned URL: nothing reassembled, no error. + assert _apply() is False + + +def test_unsupported_db_scheme_message_names_var_and_scheme(): + msg = unsupported_db_scheme_message("DIRECT_URL", "sqlite") + assert "DIRECT_URL" in msg + assert "sqlite" in msg + assert "postgresql://" in msg diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index a6f6e651487..1d4d39ec140 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -14,155 +14,131 @@ from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + +def _models(file_content_as_dict): + """Distinct body.model values, mirroring how the rate limiter collects the + models from a streamed batch file before the access check.""" + return [ + entry["body"]["model"] + for entry in file_content_as_dict + if (entry.get("body") or {}).get("model") + ] + + # --------------------------------------------------------------------------- # Token counter — covers all three batch payload shapes # --------------------------------------------------------------------------- def test_token_counter_counts_chat_messages(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], } - ] + } ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_text_completion_prompt(): - """Pre-fix this returned 0 tokens (the function only inspected + """Pre-fix this returned 0 tokens (the counter only inspected `messages`), letting `prompt`-style batches slip past TPM limits.""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} - ] + tokens = _count_entry_tokens( + {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_embedding_input_string(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - {"body": {"model": "text-embedding-3-small", "input": "hello world"}} - ] + tokens = _count_entry_tokens( + {"body": {"model": "text-embedding-3-small", "input": "hello world"}} ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_embedding_input_list(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "text-embedding-3-small", - "input": ["hello", "world"], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "text-embedding-3-small", + "input": ["hello", "world"], } - ] + } ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_text_completion_prompt_list(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-3.5-turbo-instruct", - "prompt": ["alpha", "beta"], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": ["alpha", "beta"], } - ] + } ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_pre_tokenized_prompt_int_list(): """OpenAI's text-completion API accepts a single pre-tokenized prompt as a list of ints. Each int is one token; pre-fix this shape was silently counted as zero, leaving a TPM bypass.""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-3.5-turbo-instruct", - "prompt": [1, 2, 3, 4, 5], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [1, 2, 3, 4, 5], } - ] + } ) - assert usage.prompt_tokens == 5 + assert tokens == 5 def test_token_counter_counts_pre_tokenized_prompt_list_of_int_lists(): """Multiple pre-tokenized prompts (`list[list[int]]`) — the most important bypass shape. A 1000-token batch must report 1000 tokens, not zero.""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-3.5-turbo-instruct", - "prompt": [[1] * 250, [2] * 250, [3] * 500], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [[1] * 250, [2] * 250, [3] * 500], } - ] + } ) - assert usage.prompt_tokens == 1000 + assert tokens == 1000 def test_token_counter_counts_pre_tokenized_input_for_embeddings(): """Same shape applies to embeddings (`input`).""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "text-embedding-3-small", - "input": [[1, 2, 3], [4, 5, 6]], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "text-embedding-3-small", + "input": [[1, 2, 3], [4, 5, 6]], } - ] + } ) - assert usage.prompt_tokens == 6 - - -# --------------------------------------------------------------------------- -# Model extractor -# --------------------------------------------------------------------------- - - -def test_model_extractor_returns_distinct_models(): - from litellm.batches.batch_utils import _get_models_from_batch_input_file_content - - models = _get_models_from_batch_input_file_content( - [ - {"body": {"model": "gpt-4o", "messages": []}}, - {"body": {"model": "gpt-4o", "messages": []}}, # duplicate - {"body": {"model": "gpt-4o-mini", "messages": []}}, - {"body": {}}, # missing model - ] - ) - assert models == ["gpt-4o", "gpt-4o-mini"] + assert tokens == 6 # --------------------------------------------------------------------------- @@ -211,7 +187,7 @@ async def test_pre_call_rejects_unauthorized_model_in_batch_file(): with pytest.raises(HTTPException) as exc: await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) assert exc.value.status_code == 403 @@ -250,7 +226,7 @@ async def test_pre_call_allows_all_team_models_key_when_model_in_team_allowlist( with patch("litellm.proxy.proxy_server.llm_router", None): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) @@ -297,7 +273,7 @@ async def test_pre_call_uses_current_team_allowlist_for_all_team_models_key(): ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) assert exc_info.value.status_code == 403 @@ -358,7 +334,7 @@ async def test_pre_call_allows_all_team_models_key_via_current_team_object(): ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) mock_get_team_object.assert_awaited_once() @@ -421,7 +397,7 @@ async def test_pre_call_denies_all_team_models_key_via_member_scope(): ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) assert exc_info.value.status_code == 403 @@ -479,7 +455,7 @@ async def test_pre_call_fails_closed_when_current_team_fetch_fails_for_all_team_ ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) assert exc_info.value.status_code == expected_status @@ -524,7 +500,7 @@ async def test_pre_call_allows_authorized_model_in_batch_file(): # Should not raise await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) @@ -744,7 +720,7 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), target_model_names=[proxy_alias], ) @@ -837,7 +813,7 @@ async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup( ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), target_model_names=[batch_alias], ) @@ -863,11 +839,11 @@ async def test_pre_call_skips_check_when_no_models_present(): # entirely. await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=[], + models=_models([]), ) await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=[{"body": {}}], + models=_models([{"body": {}}]), ) @@ -1390,3 +1366,272 @@ async def test_count_input_file_usage_raises_on_non_bytes_content(): user_api_key_dict=UserAPIKeyAuth(api_key="sk", models=["*"]), data={}, ) + + +# Streaming input counting — peak memory must not scale with a full dict list +# --------------------------------------------------------------------------- + + +def _make_batch_input_bytes(n_rows: int, padding: int = 200) -> bytes: + import json as _json + + pad = "x" * padding + rows = [] + for i in range(n_rows): + rows.append( + _json.dumps( + { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o" if i % 2 else "gpt-3.5-turbo", + "messages": [{"role": "user", "content": f"{pad} {i}"}], + }, + } + ) + ) + return ("\n".join(rows)).encode("utf-8") + + +def test_iter_batch_input_entries_matches_dict_list(): + from litellm.batches.batch_utils import ( + _get_file_content_as_dictionary, + _iter_batch_input_entries, + ) + + raw = _make_batch_input_bytes(50) + streamed = list(_iter_batch_input_entries(raw)) + assert streamed == _get_file_content_as_dictionary(raw) + assert streamed[0]["custom_id"] == "request-0" + # tolerant of blank lines and a missing trailing newline + assert list(_iter_batch_input_entries(raw + b"\n\n")) == streamed + + +def test_streaming_count_peak_below_dict_list(): + import gc + import tracemalloc + + from litellm.batches.batch_utils import ( + _get_file_content_as_dictionary, + _iter_batch_input_entries, + ) + + raw = _make_batch_input_bytes(8000) + + def _measure(fn): + gc.collect() + tracemalloc.start() + try: + fn() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + def _stream(): + count = 0 + models: set = set() + for entry in _iter_batch_input_entries(raw): + count += 1 + model = (entry.get("body") or {}).get("model") + if model: + models.add(model) + return count + + def _build_list(): + return len(_get_file_content_as_dictionary(raw)) + + stream_peak = _measure(_stream) + list_peak = _measure(_build_list) + assert stream_peak < list_peak * 0.5, ( + f"streaming count peak {stream_peak} is not a clear win over the dict " + f"list {list_peak} (ratio {stream_peak / list_peak:.2f})" + ) + + +@pytest.mark.asyncio +async def test_count_input_file_usage_streams_without_building_list(): + """count_input_file_usage must count requests/tokens in one streaming pass. + Mocks the download; asserts the count is correct and that the dict-list + helper is never called (a revert to the list approach would call it).""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + raw = _make_batch_input_bytes(10) + fake_content = MagicMock() + fake_content.content = raw + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary" + ) as mock_dict_list, + ): + usage = await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=None, + ) + + assert usage.request_count == 10 + assert usage.total_tokens > 0 + mock_dict_list.assert_not_called() + + +def _one_row_batch_bytes(model: str) -> bytes: + import json as _json + + return ( + _json.dumps( + { + "custom_id": "r0", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": "x"}], + }, + } + ) + + "\n" + ).encode("utf-8") + + +@pytest.mark.asyncio +async def test_count_input_file_usage_enforces_models_when_token_counting_fails(): + """Security regression: a row whose content makes token counting raise must + NOT skip the model allowlist check. async_pre_call_hook swallows non-HTTP + exceptions and submits the batch, so a raised counting error would otherwise + fail open. The access check must still run and deny the restricted model.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + fake_content = MagicMock() + fake_content.content = _one_row_batch_bytes("restricted-model") + user = UserAPIKeyAuth( + api_key="sk-x", + user_id="bob", + models=["only-allowed"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + def _boom(*args, **kwargs): + raise ValueError("unsupported content part: input_audio") + + deny = AsyncMock(side_effect=Exception("model not in allowlist")) + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", new=_boom), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=deny), + patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + # The access check ran despite token counting failing, and denied the model. + deny.assert_awaited() + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_count_input_file_usage_estimates_tokens_when_counting_fails_for_allowed_model(): + """A token-counting failure for an allowed model must not hard-block the batch + (the pre-streaming behavior let such batches through), but it also must not + zero the token total, which would let a caller evade the TPM limit by sending + rows the counter cannot measure. The row falls back to a conservative + size-based estimate so the batch proceeds with a non-zero count.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + fake_content = MagicMock() + fake_content.content = _one_row_batch_bytes("allowed-model") + user = UserAPIKeyAuth( + api_key="sk-x", + user_id="bob", + models=["allowed-model"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + def _boom(*args, **kwargs): + raise ValueError("unsupported content part: file") + + allow = AsyncMock(return_value=True) + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", new=_boom), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=allow), + patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])), + ): + usage = await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + allow.assert_awaited() + assert usage.request_count == 1 + # Estimated, not zeroed: a crafted uncountable row can't evade the TPM limit. + assert usage.total_tokens > 0 + + +@pytest.mark.asyncio +async def test_count_input_file_usage_collects_models_after_malformed_line(): + """A malformed JSONL line must not abort model collection. A restricted model + named on a row AFTER a malformed line must still be collected and denied by the + allowlist check, otherwise a caller could hide a restricted model behind a bad + row.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + fake_content = MagicMock() + fake_content.content = ( + _one_row_batch_bytes("only-allowed") + + b"{ this is not valid json\n" + + _one_row_batch_bytes("restricted-model") + ) + user = UserAPIKeyAuth( + api_key="sk-x", + user_id="bob", + models=["only-allowed"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + async def _deny_restricted(model, **kwargs): + if model == "restricted-model": + raise Exception("model not in allowlist") + return True + + deny = AsyncMock(side_effect=_deny_restricted) + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=deny), + patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + assert exc.value.status_code == 403 diff --git a/tests/test_litellm/proxy/logging_endpoints/__init__.py b/tests/test_litellm/proxy/logging_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py new file mode 100644 index 00000000000..40e89329b8d --- /dev/null +++ b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py @@ -0,0 +1,195 @@ +"""Unit tests for POST /v1/callbacks/logs (replay logging payloads → callbacks).""" + +import time + +import pytest +from fastapi import HTTPException + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( + CallbackLogsReplayer, + ingest_callback_logs, +) +from litellm.types.proxy.callback_logs_endpoints import ( + CallbackLogRecord, + CallbackLogsRequest, +) + +REQ_ID = "cb-logs-unit-test-1" + + +def _sample_payload(**overrides): + payload = { + "id": REQ_ID, + "litellm_call_id": REQ_ID, + "call_type": "acompletion", + "stream": False, + "response_cost": 0.0123, + "custom_llm_provider": "openai", + "total_tokens": 42, + "prompt_tokens": 30, + "completion_tokens": 12, + "startTime": time.time() - 2, + "endTime": time.time(), + "model": "gpt-4o-mini", + "metadata": { + "user_api_key_hash": "rust-gateway-test-key", + "user_api_key_user_id": "user-cb-logs-test", + "user_api_key_team_id": "team-cb-logs-test", + }, + "messages": [{"role": "user", "content": "hi"}], + } + payload.update(overrides) + return payload + + +def test_epoch_to_datetime_handles_float_and_fallback(): + dt = CallbackLogsReplayer._epoch_to_datetime(1_700_000_000.5) + assert dt.year == 2023 + # Non-numeric input must not raise — falls back to "now". + assert CallbackLogsReplayer._epoch_to_datetime(None) is not None + + +def test_build_logging_obj_seeds_model_call_details(): + obj = CallbackLogsReplayer._build_logging_obj(_sample_payload()) + details = obj.model_call_details + # Prebuilt payload is set so the handler skips rebuilding it. + assert details["standard_logging_object"]["id"] == REQ_ID + assert details["response_cost"] == 0.0123 + assert details["call_type"] == "acompletion" + # Metadata is mapped to the keys the cost-tracking callback reads. + md = details["litellm_params"]["metadata"] + assert md["user_api_key"] == "rust-gateway-test-key" + assert md["user_api_key_user_id"] == "user-cb-logs-test" + assert md["user_api_key_team_id"] == "team-cb-logs-test" + + +def test_response_obj_carries_usage(): + obj = CallbackLogsReplayer._response_obj_from_payload(_sample_payload()) + assert obj["usage"]["total_tokens"] == 42 + assert obj["usage"]["prompt_tokens"] == 30 + assert obj["usage"]["completion_tokens"] == 12 + + +@pytest.mark.asyncio +async def test_success_record_invokes_success_handler(monkeypatch): + captured = {} + + async def fake_success(self, result=None, start_time=None, end_time=None, **kwargs): + captured["standard_logging_object"] = self.model_call_details.get( + "standard_logging_object" + ) + captured["result"] = result + + monkeypatch.setattr(LiteLLMLogging, "async_success_handler", fake_success) + + body = CallbackLogsRequest( + records=[ + CallbackLogRecord( + status="success", standard_logging_payload=_sample_payload() + ) + ] + ) + resp = await ingest_callback_logs( + body, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + ) + assert resp.processed == 1 and resp.failed == 0 + assert captured["standard_logging_object"]["id"] == REQ_ID + assert captured["result"]["usage"]["total_tokens"] == 42 + + +@pytest.mark.asyncio +async def test_failure_record_invokes_failure_handler(monkeypatch): + captured = {} + + async def fake_failure( + self, exception, traceback_exception, start_time=None, end_time=None + ): + captured["exception"] = str(exception) + + monkeypatch.setattr(LiteLLMLogging, "async_failure_handler", fake_failure) + + body = CallbackLogsRequest( + records=[ + CallbackLogRecord( + status="failure", + standard_logging_payload=_sample_payload(), + error="upstream exploded", + ) + ] + ) + resp = await ingest_callback_logs( + body, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + ) + assert resp.processed == 1 and resp.failed == 0 + assert captured["exception"] == "upstream exploded" + + +@pytest.mark.asyncio +async def test_non_admin_is_rejected(monkeypatch): + async def fake_success(self, **kwargs): + return None + + monkeypatch.setattr(LiteLLMLogging, "async_success_handler", fake_success) + + body = CallbackLogsRequest( + records=[ + CallbackLogRecord( + status="success", standard_logging_payload=_sample_payload() + ) + ] + ) + with pytest.raises(HTTPException) as exc_info: + await ingest_callback_logs( + body, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_one_bad_record_does_not_sink_the_batch(monkeypatch): + calls = {"n": 0} + + async def flaky_success( + self, result=None, start_time=None, end_time=None, **kwargs + ): + calls["n"] += 1 + if calls["n"] == 1: + raise ValueError("boom on first record") + + monkeypatch.setattr(LiteLLMLogging, "async_success_handler", flaky_success) + + body = CallbackLogsRequest( + records=[ + CallbackLogRecord( + status="success", standard_logging_payload=_sample_payload() + ), + CallbackLogRecord( + status="success", standard_logging_payload=_sample_payload() + ), + ] + ) + resp = await ingest_callback_logs( + body, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + ) + assert resp.processed == 1 and resp.failed == 1 + # The failed record is reported back by index + error, not silently dropped. + assert len(resp.failures) == 1 + assert resp.failures[0].index == 0 + assert "boom on first record" in resp.failures[0].error + + +def test_batch_over_limit_is_rejected(): + from litellm.constants import MAX_CALLBACK_LOG_RECORDS + from pydantic import ValidationError + + # One over the cap must fail validation (422 at the API boundary), bounding + # the callback/DB fan-out a single POST can trigger. + too_many = [ + CallbackLogRecord(status="success", standard_logging_payload=_sample_payload()) + for _ in range(MAX_CALLBACK_LOG_RECORDS + 1) + ] + with pytest.raises(ValidationError): + CallbackLogsRequest(records=too_many) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 0b5b5fb6ceb..f40904e234d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1197,6 +1197,136 @@ class TestListMCPServers: # Non-admin viewers get no env var config at all (not even names). assert result.env_vars is None + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_sanitizes_for_view_only_admin(self): + """PROXY_ADMIN_VIEW_ONLY must NOT see credential-bearing fields. + + It previously passed the _user_has_admin_view gate (which also grants view-only + admins) and only had the explicit `credentials` field cleared, leaking secrets + embedded in url/static_headers/env_vars. Only a FULL PROXY_ADMIN may see those. + This test exercises the real role helpers (no patching of the gate).""" + mock_server = LiteLLM_MCPServerTable.model_construct( + server_id="leaky-server", + server_name="Leaky Server", + alias="Leaky Server", + transport=MCPTransport.http, + url="https://leaky.example.com/mcp?api_key=sk-embedded-in-url", + static_headers={"Authorization": "Bearer sk-secret-header"}, + env={"UPSTREAM_TOKEN": "sk-secret-env"}, + env_vars=[ + {"name": "GLOBAL_KEY", "value": "super-secret", "scope": "global"}, + ], + credentials={"auth_value": "sk-explicit-credential"}, + ) + + mock_prisma_client = MagicMock() + + mock_health_result = generate_mock_mcp_server_db_record( + server_id="leaky-server", alias="Leaky Server" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + mock_manager = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.health_check_server = AsyncMock(return_value=mock_health_result) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="leaky-server", + user_api_key_dict=mock_user_auth, + ) + + assert result.server_id == "leaky-server" + assert result.credentials is None + assert result.url is None + assert result.static_headers is None + assert result.env == {} + assert result.env_vars is None + + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_full_admin_still_sees_secrets(self): + """the fix must not over-redact for FULL PROXY_ADMIN, + who needs url/static_headers/env to populate the edit form.""" + mock_server = LiteLLM_MCPServerTable.model_construct( + server_id="admin-server", + server_name="Admin Server", + alias="Admin Server", + transport=MCPTransport.http, + url="https://admin.example.com/mcp", + static_headers={"Authorization": "Bearer sk-secret-header"}, + credentials={"auth_value": "sk-explicit-credential"}, + ) + + mock_prisma_client = MagicMock() + + mock_health_result = generate_mock_mcp_server_db_record( + server_id="admin-server", alias="Admin Server" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + mock_manager = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.health_check_server = AsyncMock(return_value=mock_health_result) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="admin-server", + user_api_key_dict=mock_user_auth, + ) + + # credentials field is always redacted; the rest must survive for full admin. + assert result.credentials is None + assert result.url == "https://admin.example.com/mcp" + assert result.static_headers == {"Authorization": "Bearer sk-secret-header"} + class TestTeamScopedMCPServerAccess: """Tests for cross-team information disclosure and restricted key bypass fixes.""" @@ -3338,6 +3468,10 @@ async def test_store_mcp_oauth_user_credential_returns_status(): return_value=generate_mock_mcp_server_db_record(server_id=server_id) ), ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential", new=AsyncMock(return_value=None), @@ -3693,18 +3827,12 @@ def _server_with_env_vars(server_id: str = "srv-env"): @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_role, expected_global_value", - [ - (LitellmUserRoles.PROXY_ADMIN, "super-secret"), - (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ""), - ], -) -async def test_fetch_single_mcp_server_redacts_global_env_for_view_only_admin( - user_role, expected_global_value -): - """Read-only admins must not receive admin-supplied global env var secrets; - full admins still see them so the edit form can pre-fill.""" +async def test_fetch_single_mcp_server_env_vars_full_admin_vs_view_only(): + """full admins see admin-supplied global env var secrets so the edit + form can pre-fill; read-only admins now go through the non-admin sanitizer, which + drops env_vars entirely (the names alone, e.g. ADMIN_API_KEY, leak what secrets the + admin configured). Previously the view-only case merely blanked the global value + while keeping the names, which still leaked configuration metadata.""" server = _server_with_env_vars() health_result = generate_mock_mcp_server_db_record(server_id=server.server_id) @@ -3712,34 +3840,39 @@ async def test_fetch_single_mcp_server_redacts_global_env_for_view_only_admin( health_result.last_health_check = datetime.now() health_result.health_check_error = None - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", - AsyncMock(return_value=None), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", - AsyncMock(return_value=health_result), - ), - ): - result = await mgmt_endpoints.fetch_mcp_server( - request=_make_mock_request(), - server_id=server.server_id, - user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), - ) + async def _fetch(user_role): + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", + AsyncMock(return_value=health_result), + ), + ): + return await mgmt_endpoints.fetch_mcp_server( + request=_make_mock_request(), + server_id=server.server_id, + user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), + ) - by_name = {ev.name: ev for ev in result.env_vars} - assert by_name["ADMIN_API_KEY"].value == expected_global_value - # Per-user placeholders are always preserved. + full_admin = await _fetch(LitellmUserRoles.PROXY_ADMIN) + by_name = {ev.name: ev for ev in full_admin.env_vars} + assert by_name["ADMIN_API_KEY"].value == "super-secret" assert by_name["USER_TOKEN"].value == "placeholder-hint" + + view_only = await _fetch(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + assert view_only.env_vars is None + # The source record must never be mutated. assert {ev.name: ev.value for ev in server.env_vars}[ "ADMIN_API_KEY" @@ -3747,18 +3880,67 @@ async def test_fetch_single_mcp_server_redacts_global_env_for_view_only_admin( @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_role, expected_global_value", - [ - (LitellmUserRoles.PROXY_ADMIN, "super-secret"), - (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ""), - ], -) -async def test_fetch_all_mcp_servers_redacts_global_env_for_view_only_admin( - user_role, expected_global_value -): +async def test_fetch_all_mcp_servers_env_vars_full_admin_vs_view_only(): + """same posture as the single-server fetch. Full admins + keep the env var values; view-only admins get env_vars dropped via the non-admin + sanitizer rather than only having the global value blanked.""" server = _server_with_env_vars() + async def _fetch_all(user_role): + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", + return_value="view_all", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.get_all_mcp_servers_unfiltered", + AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy.proxy_server.prisma_client", + None, + ), + ): + return await mgmt_endpoints.fetch_all_mcp_servers( + user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), + ) + + full_admin = await _fetch_all(LitellmUserRoles.PROXY_ADMIN) + by_name = {ev.name: ev for ev in full_admin[0].env_vars} + assert by_name["ADMIN_API_KEY"].value == "super-secret" + assert by_name["USER_TOKEN"].value == "placeholder-hint" + + view_only = await _fetch_all(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + assert view_only[0].env_vars is None + + assert {ev.name: ev.value for ev in server.env_vars}[ + "ADMIN_API_KEY" + ] == "super-secret" + + +def _leaky_list_server() -> "LiteLLM_MCPServerTable": + """A server whose url/static_headers/env carry embedded secrets, for the + list-endpoint sanitization tests. ``model_construct`` skips validation so + the raw values survive verbatim.""" + return LiteLLM_MCPServerTable.model_construct( + server_id="leaky-list-server", + server_name="Leaky List Server", + alias="Leaky List Server", + transport=MCPTransport.http, + url="https://leaky.example.com/mcp?api_key=sk-embedded-in-url", + static_headers={"Authorization": "Bearer sk-secret-header"}, + env={"UPSTREAM_TOKEN": "sk-secret-env"}, + env_vars=[ + {"name": "GLOBAL_KEY", "value": "super-secret", "scope": "global"}, + ], + credentials={"auth_value": "sk-explicit-credential"}, + ) + + +async def _fetch_all_via_view_all(user_role: LitellmUserRoles): + """Drive GET /v1/mcp/server in view_all mode for the given role using the + real role helpers (the full-admin gate is never patched).""" + server = _leaky_list_server() with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", @@ -3776,13 +3958,46 @@ async def test_fetch_all_mcp_servers_redacts_global_env_for_view_only_admin( result = await mgmt_endpoints.fetch_all_mcp_servers( user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), ) + return server, result - by_name = {ev.name: ev for ev in result[0].env_vars} - assert by_name["ADMIN_API_KEY"].value == expected_global_value - assert by_name["USER_TOKEN"].value == "placeholder-hint" - assert {ev.name: ev.value for ev in server.env_vars}[ - "ADMIN_API_KEY" - ] == "super-secret" + +@pytest.mark.asyncio +async def test_list_mcp_servers_sanitized_for_view_only_admin(): + """PROXY_ADMIN_VIEW_ONLY listing servers must go through the non-admin + sanitizer: url and static_headers cleared, env emptied, env_vars dropped. + A mutation swapping _user_is_full_admin() back to _user_has_admin_view() + (which also grants view-only admins) would return the raw url/headers and + fail this. The real role helpers are exercised; the gate is not patched.""" + source, result = await _fetch_all_via_view_all( + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + + assert len(result) == 1 + sanitized = result[0] + assert sanitized.server_id == "leaky-list-server" + assert sanitized.url is None + assert sanitized.static_headers is None + assert sanitized.env == {} + assert sanitized.env_vars is None + assert sanitized.credentials is None + + # The source record must never be mutated by sanitization. + assert source.url == "https://leaky.example.com/mcp?api_key=sk-embedded-in-url" + assert source.static_headers == {"Authorization": "Bearer sk-secret-header"} + + +@pytest.mark.asyncio +async def test_list_mcp_servers_full_admin_still_sees_secrets(): + """The view-only redaction must not over-redact for a FULL PROXY_ADMIN, + who needs url/static_headers to populate the edit form. Only the explicit + credentials field is cleared for full admins on the list endpoint.""" + _, result = await _fetch_all_via_view_all(LitellmUserRoles.PROXY_ADMIN) + + assert len(result) == 1 + raw = result[0] + assert raw.url == "https://leaky.example.com/mcp?api_key=sk-embedded-in-url" + assert raw.static_headers == {"Authorization": "Bearer sk-secret-header"} + assert raw.credentials is None def _make_env_var_server( @@ -4347,8 +4562,13 @@ class TestMCPUserEnvVarsAccessControl: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", - AsyncMock(return_value=[_make_env_var_server(server_id="other")]), + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + AsyncMock(return_value=["other"]), ), patch.object(mgmt_endpoints, "get_user_env_vars", get_user_env_vars), ): @@ -4378,7 +4598,12 @@ class TestMCPUserEnvVarsAccessControl: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", AsyncMock(return_value=[]), ), patch.object(mgmt_endpoints, "merge_user_env_vars", merge_mock), @@ -4412,7 +4637,12 @@ class TestMCPUserEnvVarsAccessControl: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", AsyncMock(return_value=[]), ), patch.object(mgmt_endpoints, "delete_user_env_vars", delete_mock), @@ -4444,8 +4674,13 @@ class TestMCPUserEnvVarsAccessControl: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", - AsyncMock(return_value=[server]), + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + AsyncMock(return_value=["srv-1"]), ), patch.object( mgmt_endpoints, @@ -4465,13 +4700,13 @@ class TestMCPUserEnvVarsAccessControl: @pytest.mark.asyncio async def test_admin_bypasses_access_check(self): - """Proxy admins must not be filtered by get_all_mcp_servers_for_user.""" + """Proxy admins must not be filtered by the allowed-server check.""" server = _make_env_var_server( server_id="srv-1", env_vars=_ENV_VARS_MIXED, static_headers=_STATIC_HEADERS_MIXED, ) - access_list_mock = AsyncMock(return_value=[]) + allowed_mock = AsyncMock(return_value=[]) with ( patch.object( mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() @@ -4480,7 +4715,9 @@ class TestMCPUserEnvVarsAccessControl: mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=server) ), patch.object( - mgmt_endpoints, "get_all_mcp_servers_for_user", access_list_mock + mgmt_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + allowed_mock, ), patch.object( mgmt_endpoints, "get_user_env_vars", AsyncMock(return_value={}) @@ -4494,22 +4731,34 @@ class TestMCPUserEnvVarsAccessControl: ), ) assert result.server_id == "srv-1" - access_list_mock.assert_not_awaited() + allowed_mock.assert_not_awaited() @pytest.mark.asyncio async def test_non_admin_gets_403_not_404_for_inaccessible_server(self): - """Authorization must run before the existence check so a non-admin - cannot distinguish "server does not exist" (404) from "server exists but - you lack access" (403) and enumerate server IDs.""" - get_mcp_server_mock = AsyncMock(return_value=None) + """A non-admin cannot distinguish "server does not exist" (404) from + "server exists but you lack access" (403): both collapse to 403 so server + ids stay non-enumerable, even when neither the DB nor the registry has the + server.""" with ( patch.object( mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() ), - patch.object(mgmt_endpoints, "get_mcp_server", get_mcp_server_mock), + patch.object( + mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None) + ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + MagicMock(return_value=None), + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", AsyncMock(return_value=[]), ), ): @@ -4522,7 +4771,6 @@ class TestMCPUserEnvVarsAccessControl: ), ) assert exc.value.status_code == 403 - get_mcp_server_mock.assert_not_awaited() def test_oauth2_flow_accepted_on_create_request(): @@ -4572,3 +4820,217 @@ def test_oauth2_flow_defaults_to_none_when_omitted(): assert ( LiteLLM_MCPServerTable(server_id="srv-1", transport="http").oauth2_flow is None ) + + +class TestPerUserCredentialConfigServerResolution: + """Per-user credential and env-var endpoints must resolve config-defined MCP + servers, which live only in the in-memory registry and never get a DB row, so + a user can store their BYOK key / OAuth token / env vars against them. The + same allowed-server authorization the MCP gateway enforces also gates these + writes for non-admins. + """ + + # 32-char sha256 stable id, the shape a config.yaml server gets. + CONFIG_SERVER_ID = "3a6a3f8633340371b49562c8c4682da9" + + def _registry_only_manager(self, *, is_byok: bool = False): + """A manager mock where the server exists only in the registry (DB miss).""" + config_server = generate_mock_mcp_server_config_record( + server_id=self.CONFIG_SERVER_ID, name="Config Server" + ) + record = generate_mock_mcp_server_db_record( + server_id=self.CONFIG_SERVER_ID + ).model_copy(update={"is_byok": is_byok}) + manager = MagicMock() + manager.get_mcp_server_by_id = MagicMock( + side_effect=lambda sid: ( + config_server if sid == self.CONFIG_SERVER_ID else None + ) + ) + manager._build_mcp_server_table = MagicMock(return_value=record) + manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) + return manager + + @pytest.mark.asyncio + async def test_store_oauth_credential_resolves_config_server_for_admin(self): + """OBO token persists for a config-defined server (DB miss, registry hit). + Before the registry fallback this raised 404 "MCP Server not found".""" + manager = self._registry_only_manager() + store_mock = AsyncMock(return_value=None) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None) + ), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + patch.object(mgmt_endpoints, "store_user_oauth_credential", store_mock), + patch.object( + mgmt_endpoints, + "get_user_oauth_credential", + AsyncMock(return_value={"expires_at": None}), + ), + ): + result = await mgmt_endpoints.store_mcp_oauth_user_credential( + server_id=self.CONFIG_SERVER_ID, + payload=mgmt_endpoints.MCPOAuthUserCredentialRequest( + access_token="tok", expires_in=3600 + ), + user_api_key_dict=generate_mock_user_api_key_auth(user_id="admin"), + ) + assert result.has_credential is True + store_mock.assert_awaited_once() + manager.get_mcp_server_by_id.assert_called_once_with(self.CONFIG_SERVER_ID) + + @pytest.mark.asyncio + async def test_store_byok_credential_resolves_config_server_for_admin(self): + """BYOK key persists for a config-defined BYOK server (DB miss, registry + hit). Before the registry fallback this raised 404.""" + manager = self._registry_only_manager(is_byok=True) + store_mock = AsyncMock(return_value=None) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None) + ), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + patch.object(mgmt_endpoints, "store_user_credential", store_mock), + ): + result = await mgmt_endpoints.store_mcp_user_credential( + server_id=self.CONFIG_SERVER_ID, + payload=mgmt_endpoints.MCPUserCredentialRequest(credential="my-key"), + user_api_key_dict=generate_mock_user_api_key_auth(user_id="admin"), + ) + assert result.has_credential is True + store_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_store_oauth_credential_forbidden_for_non_admin_without_access(self): + """A non-admin storing a credential for a server not in their allowed set + gets 403 and no row is written (the store endpoints had no authz before).""" + manager = self._registry_only_manager() + manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) + store_mock = AsyncMock(return_value=None) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None) + ), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + patch.object( + mgmt_endpoints, + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object(mgmt_endpoints, "store_user_oauth_credential", store_mock), + ): + with pytest.raises(HTTPException) as exc: + await mgmt_endpoints.store_mcp_oauth_user_credential( + server_id=self.CONFIG_SERVER_ID, + payload=mgmt_endpoints.MCPOAuthUserCredentialRequest( + access_token="tok", expires_in=3600 + ), + user_api_key_dict=generate_mock_user_api_key_auth( + user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER + ), + ) + assert exc.value.status_code == 403 + store_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_store_oauth_credential_allowed_for_non_admin_with_access(self): + """A non-admin with the config server in their allowed set persists the + token; proves the non-admin authz uses the registry-aware allowed set.""" + manager = self._registry_only_manager() + manager.get_allowed_mcp_servers = AsyncMock( + return_value=[self.CONFIG_SERVER_ID] + ) + store_mock = AsyncMock(return_value=None) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None) + ), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + patch.object( + mgmt_endpoints, + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object(mgmt_endpoints, "store_user_oauth_credential", store_mock), + patch.object( + mgmt_endpoints, + "get_user_oauth_credential", + AsyncMock(return_value={"expires_at": None}), + ), + ): + result = await mgmt_endpoints.store_mcp_oauth_user_credential( + server_id=self.CONFIG_SERVER_ID, + payload=mgmt_endpoints.MCPOAuthUserCredentialRequest( + access_token="tok", expires_in=3600 + ), + user_api_key_dict=generate_mock_user_api_key_auth( + user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER + ), + ) + assert result.has_credential is True + store_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_store_env_vars_resolves_config_server_for_non_admin_with_access( + self, + ): + """Per-user env vars persist for a config server a non-admin may access. + The non-admin path previously used a DB-only access list that never + included config servers, so this 403'd before the fix.""" + env_var_server = _make_env_var_server( + server_id=self.CONFIG_SERVER_ID, + env_vars=_ENV_VARS_MIXED, + static_headers=_STATIC_HEADERS_MIXED, + ) + manager = MagicMock() + manager.get_mcp_server_by_id = MagicMock( + return_value=generate_mock_mcp_server_config_record( + server_id=self.CONFIG_SERVER_ID + ) + ) + manager._build_mcp_server_table = MagicMock(return_value=env_var_server) + manager.get_allowed_mcp_servers = AsyncMock( + return_value=[self.CONFIG_SERVER_ID] + ) + merge_mock = AsyncMock(return_value={"CORP_USERNAME": "alice"}) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None) + ), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + patch.object( + mgmt_endpoints, + "build_effective_auth_contexts", + AsyncMock(return_value=[object()]), + ), + patch.object(mgmt_endpoints, "merge_user_env_vars", merge_mock), + ): + result = await mgmt_endpoints.store_mcp_user_env_vars( + server_id=self.CONFIG_SERVER_ID, + payload=mgmt_endpoints.MCPUserEnvVarsRequest( + values={"CORP_USERNAME": "alice"} + ), + user_api_key_dict=generate_mock_user_api_key_auth( + user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER + ), + ) + merge_mock.assert_awaited_once() + _, _, _, updates, _ = merge_mock.await_args.args + assert updates == {"CORP_USERNAME": "alice"} + assert result.server_id == self.CONFIG_SERVER_ID diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index c77ac11ffc1..2b38d732e9d 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -351,6 +351,110 @@ async def test_validate_no_team_non_global_server_raises( assert "not in a team" in str(exc_info.value.detail) +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("private-server"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_no_team_proxy_admin_can_assign_private_server( + mock_access_groups, mock_allow_all +): + """Proxy admin assigning a non-global server to a teamless key — should pass (LIT-3815).""" + result = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["private-server"]}, + team_obj=None, + is_proxy_admin=True, + ) + assert result["mcp_servers"] == ["private-server"] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("private-server"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_no_team_non_admin_private_server_still_raises( + mock_access_groups, mock_allow_all +): + """The teamless override is gated on proxy admin — a non-admin still gets 403.""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["private-server"]}, + team_obj=None, + is_proxy_admin=False, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_no_team_proxy_admin_can_assign_access_group( + mock_access_groups, mock_allow_all +): + """Proxy admin assigning an access group to a teamless key — should pass (LIT-3815).""" + result = await validate_key_mcp_servers_against_team( + object_permission={"mcp_access_groups": ["group-1"]}, + team_obj=None, + is_proxy_admin=True, + ) + assert result["mcp_access_groups"] == ["group-1"] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("server-1", "server-outside"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_proxy_admin_still_bounded_by_team_scope( + mock_access_groups, mock_allow_all +): + """The override is scoped to teamless keys — an admin assigning beyond a team's scope still raises.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1", "server-outside"]}, + team_obj=team_obj, + is_proxy_admin=True, + ) + assert exc_info.value.status_code == 403 + assert "server-outside" in str(exc_info.value.detail) + + @pytest.mark.asyncio @patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index f42639cee8a..46ecb31e1c8 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -398,6 +398,83 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: app.dependency_overrides.pop(ps.user_api_key_auth, None) +def test_create_file_batch_streams_from_upload_spool(monkeypatch, llm_router: Router): + """ + Batch uploads must be passed downstream as the upload's streamable file handle + (Starlette's already-spooled file), not read into an in-memory bytes object, so + the proxy never buffers the whole payload. Non-batch uploads keep the bytes path. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + from litellm.types.llms.openai import OpenAIFileObject + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + captured: dict = {} + + async def fake_route_create_file(*, _create_file_request, **kwargs): + file_elem = _create_file_request["file"][1] + captured["file_elem"] = file_elem + if hasattr(file_elem, "read") and hasattr(file_elem, "seek"): + file_elem.seek(0) + captured["streamed_content"] = file_elem.read() + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + content = ( + b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' + b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}]}}\n' + ) + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + file_elem = captured["file_elem"] + assert not isinstance( + file_elem, (bytes, bytearray) + ), "batch upload must be a streamable handle, not in-memory bytes" + assert hasattr(file_elem, "read") and hasattr( + file_elem, "seek" + ), "batch upload must be a seekable file handle" + assert ( + captured["streamed_content"] == content + ), "the handle must stream the uploaded bytes" + + captured.clear() + resp = client.post( + "/v1/files", + files={"file": ("data.jsonl", content, "application/jsonl")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + assert isinstance( + captured["file_elem"], (bytes, bytearray) + ), "non-batch upload must stay in-memory bytes" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.flaky(retries=3, delay=2) def test_target_storage_invokes_storage_backend( mocker: MockerFixture, monkeypatch, llm_router: Router diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index e89ada5bdef..df14dc5b5dc 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -15,8 +15,6 @@ from __future__ import annotations from unittest.mock import AsyncMock, MagicMock -import pytest - from .conftest import VOLATILE_KEYS, normalize @@ -248,6 +246,193 @@ def test_config_field_info_field_not_in_db(client, auth_as, mock_prisma, monkeyp assert "not in DB" in response.json().get("detail", {}).get("error", "") +def test_config_field_info_redacts_nested_secret_for_view_only_admin( + client, auth_as, mock_prisma, monkeypatch +): + """A view-only admin reading a structured field must not receive nested + credentials. database_args carries aws_web_identity_token (a DynamoDB + role-assumption credential); it must come back redacted while non-secret + siblings like region_name stay visible.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = { + "database_args": { + "region_name": "us-east-1", + "user_table_name": "LiteLLM_UserTable", + "aws_web_identity_token": "sk-super-secret-token", + } + } + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get( + "/config/field/info", params={"field_name": "database_args"} + ) + assert response.status_code == 200 + value = response.json()["field_value"] + assert value["aws_web_identity_token"] == "REDACTED" + assert value["region_name"] == "us-east-1" + assert value["user_table_name"] == "LiteLLM_UserTable" + + +def test_config_field_info_full_admin_sees_nested_secret( + client, auth_as, mock_prisma, monkeypatch +): + """The redaction must not over-redact for a full PROXY_ADMIN, who needs + the real nested value to populate the edit form.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = { + "database_args": { + "region_name": "us-east-1", + "aws_web_identity_token": "sk-super-secret-token", + } + } + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get( + "/config/field/info", params={"field_name": "database_args"} + ) + assert response.status_code == 200 + value = response.json()["field_value"] + assert value["aws_web_identity_token"] == "sk-super-secret-token" + assert value["region_name"] == "us-east-1" + + +def test_config_field_info_redacts_top_level_scalar_for_view_only( + client, auth_as, mock_prisma, monkeypatch +): + """The top-level scalar branch must also redact for a view-only admin. + database_url carries DB credentials and is not caught by the name masker, + so it is in the explicit secret set.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {"database_url": "postgresql://admin:p4ss@db:5432/litellm"} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get( + "/config/field/info", params={"field_name": "database_url"} + ) + assert response.status_code == 200 + assert response.json()["field_value"] == "REDACTED" + + +def test_redact_general_setting_value_recurses_list_of_dicts(): + """The list branch of the recursor redacts secret leaves inside each dict + while non-secret keys survive, and a full admin gets the value untouched.""" + from litellm.proxy import proxy_server as ps + + value = [ + {"path": "/foo", "headers": {"Authorization": "Bearer sk-x"}}, + {"path": "/bar", "client_secret": "sk-y"}, + ] + redacted = ps._redact_general_setting_value( + "some_list_field", value, is_full_admin=False + ) + assert redacted[0]["headers"]["Authorization"] == "REDACTED" + assert redacted[0]["path"] == "/foo" + assert redacted[1]["client_secret"] == "REDACTED" + assert redacted[1]["path"] == "/bar" + assert ( + ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) + == value + ) + + +def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): + """Past _REDACT_SECRET_MAX_DEPTH the whole subtree is replaced with + "REDACTED" rather than returned verbatim, so a secret buried below the cap + can never leak via depth-overrun. A future refactor that flips the cap + branch to fail-open would surface here.""" + from litellm.proxy import proxy_server as ps + + # leaf and wrap keys are both NON-secret so neither the key-name + # short-circuit nor the explicit-secret set catches the leak. The cap is + # the only thing standing between the secret and the response — flip the + # cap to fail-open and the secret comes back verbatim. + nested: object = {"notes": "sk-leak-bottom"} + for _ in range(ps._REDACT_SECRET_MAX_DEPTH + 2): + nested = {"wrap": nested} + + out = ps._redact_general_setting_value( + "some_struct_field", nested, is_full_admin=False + ) + # the secret must not survive anywhere in the returned tree + assert "sk-leak-bottom" not in repr(out) + + # full admin is unaffected by the cap — the value comes back untouched + admin_out = ps._redact_general_setting_value( + "some_struct_field", nested, is_full_admin=True + ) + assert admin_out is nested + + +def test_config_list_redacts_pass_through_secret_for_view_only( + client, auth_as, mock_prisma, monkeypatch +): + """/config/list must not leak pass_through_endpoints upstream credentials + to a view-only admin. pass_through_endpoints is a known secret-bearing + field, so a non-admin gets it redacted; a full admin still sees it.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {"max_parallel_requests": 3} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr( + ps, + "general_settings", + { + "pass_through_endpoints": [ + { + "path": "/foo", + "target": "https://upstream.example.com", + "headers": {"Authorization": "Bearer sk-UPSTREAM-SECRET"}, + } + ] + }, + ) + + def _pass_through_value(body): + return next( + entry["field_value"] + for entry in body + if entry["field_name"] == "pass_through_endpoints" + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + view_resp = client.get( + "/config/list", params={"config_type": "general_settings"} + ) + assert view_resp.status_code == 200 + assert "sk-UPSTREAM-SECRET" not in view_resp.text + assert _pass_through_value(view_resp.json()) == "REDACTED" + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + admin_resp = client.get( + "/config/list", params={"config_type": "general_settings"} + ) + assert admin_resp.status_code == 200 + admin_value = _pass_through_value(admin_resp.json()) + assert admin_value[0]["headers"]["Authorization"] == "Bearer sk-UPSTREAM-SECRET" + + # --------------------------------------------------------------------------- # GET /config/list # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 33de1ede917..699606b5277 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -16,8 +16,6 @@ Pins covered: from __future__ import annotations import json -from typing import Any, AsyncIterator -from unittest.mock import AsyncMock, MagicMock import pytest @@ -26,8 +24,11 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, _fast_serialize_simple_model_response_stream, + _format_fallback_metadata_sse_event, _format_streaming_sse_chunk, _get_client_requested_model_for_streaming, + _get_streaming_fallback_metadata, + _is_positive_int_like, _restamp_streaming_chunk_model, _serialize_streaming_chunk, async_assistants_data_generator, @@ -71,6 +72,15 @@ async def _async_iter_raises(exc: Exception): raise exc +class _FakeStream: + def __init__(self, chunks, hidden_params=None): + self._chunks = chunks + self._hidden_params = hidden_params or {} + + def __aiter__(self): + return _async_iter(self._chunks) + + # --------------------------------------------------------------------------- # data_generator # --------------------------------------------------------------------------- @@ -274,6 +284,34 @@ def test_restamp_streaming_chunk_model_overrides_model_on_dict(): assert logged is True +def test_restamp_streaming_chunk_model_uses_fallback_model_from_metadata(): + chunk = _simple_chunk(model="openai/internal-fallback") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="primary-model", + request_data={"litellm_call_id": "id-1"}, + model_mismatch_logged=False, + fallback_was_attempted=True, + fallback_model_from_metadata="fallback-model", + ) + assert new_chunk.model == "fallback-model" + assert logged is True + + +def test_restamp_streaming_chunk_model_preserves_fallback_model_without_group(): + chunk = _simple_chunk(model="openai/internal-fallback") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="primary-model", + request_data={}, + model_mismatch_logged=False, + fallback_was_attempted=True, + fallback_model_from_metadata=None, + ) + assert new_chunk.model == "openai/internal-fallback" + assert logged is False + + def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged(): """For a non-BaseModel, non-dict chunk the helper returns it as-is along with the original ``model_mismatch_logged`` flag.""" @@ -288,6 +326,147 @@ def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged(): assert logged is False +def test_is_positive_int_like_invalid_and_edge_values(): + assert _is_positive_int_like(None) is False + assert _is_positive_int_like("not-a-number") is False + assert _is_positive_int_like(0) is False + assert _is_positive_int_like(-1) is False + assert _is_positive_int_like("1") is True + assert _is_positive_int_like(2) is True + + +def test_get_streaming_fallback_metadata_reads_headers(): + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + stream = _FakeStream( + [], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": "1", + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + }, + ) + assert _get_streaming_fallback_metadata(stream) == ( + True, + "fallback-model", + fallback_errors, + ) + + +def test_get_streaming_fallback_metadata_no_additional_headers(): + stream = _FakeStream([], hidden_params={}) + assert _get_streaming_fallback_metadata(stream) == (False, None, []) + + +def test_get_streaming_fallback_metadata_zero_fallback_count(): + stream = _FakeStream( + [], + hidden_params={ + "additional_headers": {"x-litellm-attempted-fallbacks": 0} + }, + ) + assert _get_streaming_fallback_metadata(stream) == (False, None, []) + + +def test_get_streaming_fallback_metadata_no_model_group_returns_none_model(): + stream = _FakeStream( + [], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + } + }, + ) + was_attempted, fallback_model, errors = _get_streaming_fallback_metadata(stream) + assert was_attempted is True + assert fallback_model is None + assert errors == [] + + +def test_restamp_streaming_chunk_model_azure_router_preserves_model(): + chunk = _simple_chunk(model="azure_ai/internal-deployment") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="azure_ai/model-router", + request_data={}, + model_mismatch_logged=False, + ) + assert new_chunk.model == "azure_ai/internal-deployment" + assert logged is False + + +def test_restamp_streaming_chunk_model_fastest_response_preserves_model(): + chunk = _simple_chunk(model="winning-model") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="gpt-4,claude-3", + request_data={"fastest_response": True}, + model_mismatch_logged=False, + ) + assert new_chunk.model == "winning-model" + assert logged is False + + +def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns(): + from pydantic import ConfigDict + + class FrozenChunk(_simple_chunk().__class__): + model_config = ConfigDict(frozen=True) + + chunk = FrozenChunk( + id="chatcmpl-test", + choices=[], + created=0, + model="openai/internal-x", + object="chat.completion.chunk", + ) + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="gpt-4", + request_data={"litellm_call_id": "test-id"}, + model_mismatch_logged=False, + ) + assert new_chunk.model == "openai/internal-x" + assert logged is True + + +def test_format_fallback_metadata_sse_event(): + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + + event = _format_fallback_metadata_sse_event( + fallback_model="fallback-model", + fallback_errors=fallback_errors, + ) + + assert isinstance(event, str) + assert event.startswith("data: ") + payload = json.loads(event.removeprefix("data: ").removesuffix("\n\n")) + assert payload["choices"] == [] + assert payload["litellm_fallback"] == { + "fallback_model": "fallback-model", + "errors": fallback_errors, + } + assert payload["id"] == "litellm-fallback-metadata" + assert payload["object"] == "chat.completion.chunk" + assert payload["model"] == "fallback-model" + assert isinstance(payload["created"], int) + + # --------------------------------------------------------------------------- # _fast_serialize_simple_model_response_stream # --------------------------------------------------------------------------- @@ -473,7 +652,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): # First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk. first = out[0] assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").rstrip(b"\n\n")) + payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) assert normalize(payload) == { "id": "", "object": "chat.completion.chunk", @@ -488,6 +667,172 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): } +@pytest.mark.asyncio +async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch): + _patch_logging_flags(monkeypatch) + + response = _FakeStream( + [_simple_chunk(model="openai/internal-fallback", content="hello")], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + } + }, + ) + out = [] + async for line in async_data_generator( + response=response, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model", "include_fallback_errors": True}, + ): + out.append(line) + + first = out[0] + assert isinstance(first, bytes) + payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + assert payload["model"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch): + _patch_logging_flags(monkeypatch) + + chunk = _simple_chunk(model="openai/internal-fallback", content="hello") + chunk._hidden_params = { + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + } + } + out = [] + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model"}, + ): + out.append(line) + + first = out[0] + assert isinstance(first, bytes) + payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + assert payload["model"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkeypatch): + """Pre-fallback chunks keep the client-requested model; once a chunk carries + fallback metadata the model latches to the fallback group for the rest of the + stream. This pins the client-visible mid-stream model change.""" + _patch_logging_flags(monkeypatch) + + primary_chunk = _simple_chunk(model="openai/internal-primary", content="hi") + fallback_chunk = _simple_chunk(model="openai/internal-fallback", content="there") + fallback_chunk._hidden_params = { + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + } + } + out = [] + async for line in async_data_generator( + response=_async_iter([primary_chunk, fallback_chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model"}, + ): + out.append(line) + + first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n")) + assert first_payload["model"] == "primary-model" + assert second_payload["model"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypatch): + _patch_logging_flags(monkeypatch) + monkeypatch.setitem(ps.general_settings, "expose_fallback_errors_to_caller", True) + + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + response = _FakeStream( + [_simple_chunk(model="openai/internal-fallback", content="hello")], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + }, + ) + out = [] + async for line in async_data_generator( + response=response, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model", "include_fallback_errors": True}, + ): + out.append(line) + + assert isinstance(out[0], bytes) + chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + assert chunk_payload["model"] == "fallback-model" + assert isinstance(out[1], str) + assert out[1].startswith("data: ") + metadata_payload = json.loads(out[1].removeprefix("data: ").removesuffix("\n\n")) + assert metadata_payload["choices"] == [] + assert metadata_payload["litellm_fallback"] == { + "fallback_model": "fallback-model", + "errors": fallback_errors, + } + assert metadata_payload["id"] == "litellm-fallback-metadata" + assert metadata_payload["object"] == "chat.completion.chunk" + assert metadata_payload["model"] == "fallback-model" + assert isinstance(metadata_payload["created"], int) + + +@pytest.mark.asyncio +async def test_async_data_generator_skips_fallback_error_event_without_opt_in( + monkeypatch, +): + _patch_logging_flags(monkeypatch) + + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + response = _FakeStream( + [_simple_chunk(model="openai/internal-fallback", content="hello")], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + }, + ) + out = [] + async for line in async_data_generator( + response=response, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model"}, + ): + out.append(line) + + assert isinstance(out[0], bytes) + payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + assert payload["model"] == "fallback-model" + + @pytest.mark.asyncio async def test_async_data_generator_mid_stream_exception_yields_error_payload( monkeypatch, diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 34405f20727..25e84fb59a7 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -14,7 +14,11 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm.proxy.proxy_server as ps -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, +) from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator from litellm.proxy.proxy_server import ( _get_proxy_model_info, @@ -1360,3 +1364,83 @@ async def test_retrieve_model_by_inaccessible_public_name_404s(monkeypatch): assert exc_info.value.status_code == 404 router.get_deployment_by_model_group_name.assert_not_called() + + +def test_get_direct_access_models_expands_all_proxy_models_sentinel(): + """A user provisioned with 'all-proxy-models' has direct access to every non-team + deployment. The sentinel must resolve via get_model_ids, not be looked up as a + literal model_name (which matches nothing). Regression for GH#22791.""" + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1", "global-id-2"] + router.get_model_list.return_value = [] + + user = LiteLLM_UserTable( + user_id="u", + models=[ps.SpecialModelNames.all_proxy_models.value], + teams=[], + ) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router) + + assert result == ["global-id-1", "global-id-2"] + router.get_model_ids.assert_called_once_with(exclude_team_models=True) + router.get_model_list.assert_not_called() + + +def test_get_direct_access_models_resolves_explicit_model_names(): + """Without the sentinel, only the user's explicitly listed models resolve to ids; + the all-proxy-models shortcut must not fire.""" + router = MagicMock() + router.get_model_list.return_value = [{"model_info": {"id": "gpt4o-id"}}] + + user = LiteLLM_UserTable(user_id="u", models=["gpt-4o"], teams=[]) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router) + + assert result == ["gpt4o-id"] + router.get_model_ids.assert_not_called() + router.get_model_list.assert_called_once_with(model_name="gpt-4o") + + +@pytest.mark.asyncio +async def test_populate_team_access_grants_all_proxy_models_user_direct_access( + monkeypatch, +): + """An internal user provisioned with 'all-proxy-models' and no teams must see proxy + models on the Models+Endpoints page. Before the fix _populate_team_access_on_models + marked direct_access=False, so _filter_models_to_user_accessible dropped every model + and the page was empty. Regression for GH#22791.""" + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1"] + + user_row = LiteLLM_UserTable( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[ps.SpecialModelNames.all_proxy_models.value], + teams=[], + ) + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) + + caller = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[] + ) + + populated = await ps._populate_team_access_on_models( + user_api_key_dict=caller, + prisma_client=prisma_client, + llm_router=router, + all_models=[global_row], + ) + visible = ps._filter_models_to_user_accessible(populated) + + assert [m["model_info"]["id"] for m in visible] == ["global-id-1"] + assert visible[0]["model_info"]["direct_access"] is True diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 56627c5be88..88dbec4020f 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1708,6 +1708,57 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=False ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_startup_exits_on_non_postgres_database_url( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + ): + """A sqlite DATABASE_URL must exit immediately, before any prisma call, + instead of stalling on a migration against the postgresql-only schema.""" + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db" + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + ): + with pytest.raises(SystemExit) as exc_info: + run_server.main( + ["--local", "--skip_server_startup"], standalone_mode=False + ) + assert exc_info.value.code == 1 + mock_setup_database.assert_not_called() + # --- Module-level helpers for worker startup hook tests --- diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index c85c5ccb39f..fb951e922fc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -857,6 +857,166 @@ def test_get_config_custom_callback_api_env_vars(monkeypatch): } +def test_get_config_returns_email_settings(monkeypatch): + """ + Regression for https://github.com/BerriAI/litellm/issues/19221 + + proxy_config.get_config() already returns environment_variables decrypted + (the DB-overlay path decrypts them, and YAML values are plaintext). The + /get/config/callbacks email block must therefore surface those values as-is + instead of decrypting a second time. The old code ran decrypt_value_helper() + on the already-plaintext value, which failed and returned None, so every + SMTP_* field came back blank on UI refresh. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + smtp_password = "super-secret-app-password" + config_data = { + "litellm_settings": {}, + "general_settings": {"alerting": ["email"]}, + "environment_variables": { + "SMTP_HOST": "smtp.resend.com", + "SMTP_PORT": "587", + "SMTP_USERNAME": "resend", + "SMTP_PASSWORD": smtp_password, + "SMTP_SENDER_EMAIL": "alerts@example.com", + "TEST_EMAIL_ADDRESS": "admin@example.com", + }, + } + + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + email_alert = next( + (a for a in response.json()["alerts"] if a["name"] == "email"), None + ) + assert email_alert is not None + variables = email_alert["variables"] + + # Non-sensitive fields round-trip verbatim (None before the fix). + assert variables["SMTP_HOST"] == "smtp.resend.com" + assert variables["SMTP_PORT"] == "587" + assert variables["SMTP_USERNAME"] == "resend" + assert variables["SMTP_SENDER_EMAIL"] == "alerts@example.com" + assert variables["TEST_EMAIL_ADDRESS"] == "admin@example.com" + + # Password is present but masked: never None, never the raw secret. + assert variables["SMTP_PASSWORD"] is not None + assert variables["SMTP_PASSWORD"] != smtp_password + assert "*" in variables["SMTP_PASSWORD"] + + +def test_get_config_returns_slack_webhook(monkeypatch): + """ + Same double-decryption regression as the email block (issue #19221): the + slack alerting block must surface the already-decrypted SLACK_WEBHOOK_URL + rather than decrypting it again into None. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + webhook_url = "https://hooks.slack.com/services/T00000/B00000/abcdefghijklmnop" + config_data = { + "litellm_settings": {}, + "general_settings": {"alerting": ["slack"]}, + "environment_variables": {"SLACK_WEBHOOK_URL": webhook_url}, + } + + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + + mock_logging = MagicMock() + mock_logging.slack_alerting_instance.alert_types = ["budget_alerts"] + mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = [ + "budget_alerts" + ] + mock_logging.slack_alerting_instance.alert_to_webhook_url = {} + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + slack_alert = next( + (a for a in response.json()["alerts"] if a["name"] == "slack"), None + ) + assert slack_alert is not None + masked_url = slack_alert["variables"]["SLACK_WEBHOOK_URL"] + + # Masked, but derived from the real URL (None before the fix). + assert masked_url is not None + assert masked_url != webhook_url + assert masked_url.startswith("http") + assert "*" in masked_url + + +def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): + """ + A webhook the admin cleared is stored as "" in environment_variables. The + slack block must surface that empty value, not silently fall back to a + SLACK_WEBHOOK_URL still present in the OS environment (which truthiness-based + `or` would do). Only a truly absent key should trigger the os.getenv lookup. + """ + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + monkeypatch.setenv( + "SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/STALE/OS/ENVVALUE" + ) + config_data = { + "litellm_settings": {}, + "general_settings": {"alerting": ["slack"]}, + "environment_variables": {"SLACK_WEBHOOK_URL": ""}, + } + + mock_router = MagicMock() + mock_router.get_settings.return_value = {} + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + + mock_logging = MagicMock() + mock_logging.slack_alerting_instance.alert_types = ["budget_alerts"] + mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = [ + "budget_alerts" + ] + mock_logging.slack_alerting_instance.alert_to_webhook_url = {} + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + + client = TestClient(app) + try: + response = client.get("/get/config/callbacks") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + slack_alert = next( + (a for a in response.json()["alerts"] if a["name"] == "slack"), None + ) + assert slack_alert is not None + assert slack_alert["variables"]["SLACK_WEBHOOK_URL"] == "" + + # Mock Prisma class MockPrisma: def __init__(self, database_url=None, proxy_logging_obj=None, http_client=None): @@ -8364,3 +8524,89 @@ def test_preserve_redacted_plugin_keys_sets_new_and_drops_orphan_placeholder(): [{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing ) assert "plugin_key" not in new_plugin[0] + + +def _config_field_info_client(monkeypatch, user_role): + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + db_record = types.SimpleNamespace( + param_value={ + "master_key": "sk-super-secret-master", + "database_url": "postgresql://user:p4ssw0rd@db:5432/litellm", + "pass_through_endpoints": [ + { + "path": "/upstream", + "target": "https://upstream.example.com", + "headers": {"Authorization": "Bearer sk-upstream-secret"}, + } + ], + "max_parallel_requests": 100, + } + ) + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=db_record) + mock_prisma = MagicMock() + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="u", user_role=user_role + ) + return TestClient(app) + + +def test_config_field_info_redacts_secrets_for_view_only_admin(monkeypatch): + """/config/field/info gates on _user_has_admin_view, which also grants + PROXY_ADMIN_VIEW_ONLY. A view-only admin reading master_key/database_url verbatim is + effectively a full admin. Secret-bearing fields must come back REDACTED for anyone who + is not a FULL PROXY_ADMIN, while non-secret fields stay readable.""" + from litellm.proxy._types import LitellmUserRoles + + client = _config_field_info_client( + monkeypatch, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + try: + for secret_field in ("master_key", "database_url", "pass_through_endpoints"): + resp = client.get("/config/field/info", params={"field_name": secret_field}) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["field_value"] == "REDACTED" + assert "secret" not in str(body["field_value"]) + assert "p4ssw0rd" not in str(body["field_value"]) + + resp = client.get( + "/config/field/info", params={"field_name": "max_parallel_requests"} + ) + assert resp.status_code == 200, resp.text + assert resp.json()["field_value"] == 100 + finally: + app.dependency_overrides.clear() + + +def test_config_field_info_returns_raw_secrets_for_full_admin(monkeypatch): + """the redaction must not over-apply. A FULL PROXY_ADMIN still + needs the real master_key value to populate the admin edit form.""" + from litellm.proxy._types import LitellmUserRoles + + client = _config_field_info_client(monkeypatch, LitellmUserRoles.PROXY_ADMIN) + try: + resp = client.get("/config/field/info", params={"field_name": "master_key"}) + assert resp.status_code == 200, resp.text + assert resp.json()["field_value"] == "sk-super-secret-master" + + resp = client.get( + "/config/field/info", params={"field_name": "pass_through_endpoints"} + ) + assert resp.status_code == 200, resp.text + assert ( + resp.json()["field_value"][0]["headers"]["Authorization"] + == "Bearer sk-upstream-secret" + ) + finally: + app.dependency_overrides.clear() diff --git a/tests/test_litellm/proxy/test_read_model_list.py b/tests/test_litellm/proxy/test_read_model_list.py new file mode 100644 index 00000000000..7703d57e21d --- /dev/null +++ b/tests/test_litellm/proxy/test_read_model_list.py @@ -0,0 +1,32 @@ +"""Tests for litellm.proxy.read_model_list (Rust AI gateway config bridge).""" + +from litellm.proxy.read_model_list import read_model_list + + +def test_read_model_list_resolves_os_environ(monkeypatch, tmp_path): + """`os.environ/` markers in the model_list are resolved via ProxyConfig.""" + monkeypatch.setenv("OPENAI_API_KEY", "sk-resolved-123") + config = tmp_path / "config.yaml" + config.write_text( + "model_list:\n" + " - model_name: gpt-realtime\n" + " litellm_params:\n" + " model: openai/gpt-realtime\n" + " api_key: os.environ/OPENAI_API_KEY\n" + ) + + model_list = read_model_list(str(config)) + + assert len(model_list) == 1 + params = model_list[0]["litellm_params"] + assert model_list[0]["model_name"] == "gpt-realtime" + assert params["model"] == "openai/gpt-realtime" + assert params["api_key"] == "sk-resolved-123" + + +def test_read_model_list_missing_key_returns_empty(tmp_path): + """A config without a model_list yields an empty list, not an error.""" + config = tmp_path / "config.yaml" + config.write_text("general_settings: {}\n") + + assert read_model_list(str(config)) == [] diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py new file mode 100644 index 00000000000..2aee0f0a4ef --- /dev/null +++ b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py @@ -0,0 +1,142 @@ +import json + +from pydantic import BaseModel + +from litellm.router_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, + get_fallback_errors_from_headers, + get_hidden_params_dict, +) + + +class StreamingWrapper: + def __init__(self): + self._hidden_params = {"additional_headers": {"x-existing": "keep"}} + + +def test_add_fallback_headers_to_streaming_wrapper(): + response = StreamingWrapper() + + result = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=1, + ) + + assert result is response + assert response._hidden_params["additional_headers"] == { + "x-existing": "keep", + "x-litellm-attempted-fallbacks": 1, + } + + +def test_add_fallback_headers_serializes_fallback_errors(): + response = StreamingWrapper() + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + + result = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=1, + fallback_errors=fallback_errors, + ) + + assert result is response + assert response._hidden_params["additional_headers"][ + "x-litellm-attempted-fallbacks" + ] == 1 + assert ( + json.loads( + response._hidden_params["additional_headers"]["x-litellm-fallback-errors"] + ) + == fallback_errors + ) + + +def test_add_retry_headers_to_streaming_wrapper(): + response = StreamingWrapper() + + result = add_retry_headers_to_response( + response=response, + attempted_retries=2, + max_retries=3, + ) + + assert result is response + assert response._hidden_params["additional_headers"] == { + "x-existing": "keep", + "x-litellm-attempted-retries": 2, + "x-litellm-max-retries": 3, + } + + +def test_get_hidden_params_dict_with_pydantic_model_hidden_params(): + class InnerHiddenParams(BaseModel): + additional_headers: dict = {} + + class Response: + def __init__(self): + self._hidden_params = InnerHiddenParams( + additional_headers={"x-custom": "value"} + ) + + result = get_hidden_params_dict(Response()) + assert result == {"additional_headers": {"x-custom": "value"}} + + +def test_get_hidden_params_dict_with_no_hidden_params(): + class PlainResponse: + pass + + assert get_hidden_params_dict(PlainResponse()) == {} + + +def test_add_fallback_headers_when_no_existing_additional_headers(): + class NoHeadersWrapper: + def __init__(self): + self._hidden_params = {} + + response = NoHeadersWrapper() + result = add_fallback_headers_to_response(response=response, attempted_fallbacks=2) + + assert result is response + assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 2 + + +def test_add_fallback_headers_returns_none_when_response_is_none(): + result = add_fallback_headers_to_response(response=None, attempted_fallbacks=1) + assert result is None + + +def test_add_fallback_headers_returns_unchanged_when_response_has_no_hidden_params(): + class PlainObject: + pass + + obj = PlainObject() + result = add_fallback_headers_to_response(response=obj, attempted_fallbacks=1) + assert result is obj + assert not hasattr(obj, "_hidden_params") + + +def test_get_fallback_errors_from_headers_existing_list_passthrough(): + errors = [{"message": "err", "type": "T", "param": None, "code": "400"}] + result = get_fallback_errors_from_headers({"x-litellm-fallback-errors": errors}) + assert result == errors + + +def test_get_fallback_errors_from_headers_invalid_json_returns_empty(): + result = get_fallback_errors_from_headers( + {"x-litellm-fallback-errors": "not-valid-json-{"} + ) + assert result == [] + + +def test_get_fallback_errors_from_headers_missing_key_returns_empty(): + result = get_fallback_errors_from_headers({}) + assert result == [] diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py new file mode 100644 index 00000000000..ca647bdce55 --- /dev/null +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -0,0 +1,139 @@ +import json + +import pytest + +from litellm.router_utils.fallback_event_handlers import run_async_fallback + + +class StreamingWrapper: + def __init__(self): + self._hidden_params = {"additional_headers": {}} + + +class FakeRouter: + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + return StreamingWrapper() + + +class AlwaysFailRouter: + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + raise RuntimeError("fallback model also failed") + + +@pytest.mark.asyncio +async def test_run_async_fallback_adds_errors_when_opted_in(): + response = await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["x-litellm-attempted-fallbacks"] == 1 + assert json.loads(additional_headers["x-litellm-fallback-errors"]) == [ + { + "message": "upstream limited request", + "type": "RuntimeError", + "param": None, + "code": None, + } + ] + + +@pytest.mark.asyncio +async def test_run_async_fallback_omits_errors_without_opt_in(): + response = await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["x-litellm-attempted-fallbacks"] == 1 + assert "x-litellm-fallback-errors" not in additional_headers + + +@pytest.mark.asyncio +async def test_run_async_fallback_raises_when_all_fallbacks_fail(): + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=AlwaysFailRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original request failed"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + +class RecordingRouter: + def __init__(self): + self.received_kwargs = None + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.received_kwargs = kwargs + return StreamingWrapper() + + +@pytest.mark.asyncio +async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): + """A nested fallback (multi-hop) must keep collecting errors, so the opt-in + flag has to reach the nested async_function_with_fallbacks call.""" + router = RecordingRouter() + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + assert router.received_kwargs.get("include_fallback_errors") is True + + +@pytest.mark.asyncio +async def test_run_async_fallback_does_not_forward_flag_without_opt_in(): + router = RecordingRouter() + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + assert "include_fallback_errors" not in router.received_kwargs + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_original_model_group(): + response = await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=["primary-model", "fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + ) + + assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 diff --git a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py new file mode 100644 index 00000000000..0d7bcbe1e53 --- /dev/null +++ b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py @@ -0,0 +1,647 @@ +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.opensandbox.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OpenSandboxSandboxConfig, +) +from litellm.utils import ProviderConfigManager + +TEST_API_BASE = "https://sandbox.test/v1" + + +def http_status_error(status_code, url="http://test"): + return httpx.HTTPStatusError( + f"status {status_code}", + request=httpx.Request("GET", url), + response=httpx.Response(status_code), + ) + + +def sse(data): + return f"data: {json.dumps(data)}" + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + def raise_for_status(self): + if self.status_code >= 400: + raise http_status_error(self.status_code) + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + def __init__( + self, + *, + create_json=None, + sandbox_states=None, + endpoint_json=None, + endpoint_responses=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + self.sandbox_states = list( + sandbox_states + or [ + { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + ] + ) + self.endpoint_json = endpoint_json or { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + } + self.endpoint_responses = ( + list(endpoint_responses) if endpoint_responses is not None else None + ) + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json, {"stream": stream})) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/code"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def get(self, url, headers=None, params=None, **kwargs): + self.calls.append(("GET", url, headers, None, params)) + if "/endpoints/44772" in url: + if self.endpoint_responses is not None and self.endpoint_responses: + response = self.endpoint_responses.pop(0) + if isinstance(response, Exception): + raise response + if isinstance(response, FakeResponse): + return response + return FakeResponse(json_data=response) + return FakeResponse(json_data=self.endpoint_json) + if "/sandboxes/" in url: + state = self.sandbox_states.pop(0) + return FakeResponse(json_data=state) + raise AssertionError(f"unexpected GET {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None, None)) + if not (200 <= self.delete_status < 300): + raise http_status_error(self.delete_status, url) + return FakeResponse(status_code=self.delete_status) + + +def test_parse_sse_lines_maps_output_result_count_and_error(): + lines = [ + sse({"type": "stdout", "text": "hello\n"}), + sse({"type": "stderr", "text": "warn\n"}), + sse({"type": "result", "results": {"text/plain": "4"}}), + sse({"type": "execution_count", "execution_count": 7}), + sse( + { + "type": "error", + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["Traceback"], + }, + } + ), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "hello\n" + assert result.stderr == "warn\n" + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 7 + assert result.error == { + "name": "ValueError", + "value": "bad", + "traceback": ["Traceback"], + } + + +def test_parse_sse_lines_skips_non_json_and_control_lines(): + lines = [ + "event: message", + "not-json", + "", + sse({"type": "stdout", "text": "ok\n"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "ok\n" + assert result.error is None + + +def test_parse_sse_lines_maps_fallback_shapes(): + lines = [ + "data:", + sse(["not-a-dict"]), + sse({"code": "BadRequest", "message": "nope"}), + sse({"type": "result", "text/plain": "4"}), + sse({"type": "error", "name": "RuntimeError", "text": "boom"}), + sse({"type": "execution_count", "execution_count": "8"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 8 + assert result.error == { + "name": "BadRequest", + "value": "nope", + "traceback": [], + } + fallback_error = OpenSandboxSandboxConfig._parse_lines( + [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] + ) + assert fallback_error.error == { + "name": "RuntimeError", + "value": "boom", + "traceback": [], + } + empty_string_error = OpenSandboxSandboxConfig._parse_lines( + [ + sse( + { + "type": "error", + "error": { + "ename": "", + "name": "FallbackName", + "evalue": "", + "value": "fallback value", + "traceback": [], + }, + } + ) + ] + ) + assert empty_string_error.error == { + "name": "", + "value": "", + "traceback": [], + } + + +def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): + def fake_secret(key): + if key == "OPEN_SANDBOX_API_KEY": + return "env-key" + if key == "OPEN_SANDBOX_API_BASE": + return TEST_API_BASE + return None + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + fake_secret, + ) + config = OpenSandboxSandboxConfig() + handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") + + assert config.validate_environment() == "env-key" + assert config.validate_environment(api_key="") == "" + assert config._api_key(api_key=None, handle=handle) == "env-key" + + handle._hidden_params = {"api_key": "stored-key"} + assert config._api_key(api_key=None, handle=handle) == "stored-key" + assert config._http(None) is not None + + body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy={"egress": [{"domain": "example.com"}]}, + secure_access=True, + ) + assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} + assert body["secureAccess"] is True + + other_body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy=None, + secure_access=False, + ) + assert body["resourceLimits"] is not other_body["resourceLimits"] + + assert config._sandbox_state(None) is None + assert config._sandbox_state({"status": "Running"}) is None + assert config._as_str_dict(None) == {} + assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( + "http://execd.local" + ) + assert config._api_base(None) == TEST_API_BASE + assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" + assert config._as_int("9") == 9 + assert config._as_int("nope") is None + assert config._as_int(None) is None + assert isinstance( + ProviderConfigManager.get_provider_sandbox_config("opensandbox"), + OpenSandboxSandboxConfig, + ) + + +def test_api_base_requires_kwarg_or_env(monkeypatch): + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + lambda key: None, + ) + + with pytest.raises(ValueError, match="api_base is required"): + OpenSandboxSandboxConfig._api_base(None) + + +@pytest.mark.asyncio +async def test_create_posts_default_body_and_omits_empty_api_key(): + client = FakeHTTPClient() + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + method, url, headers, body, _ = client.calls[0] + assert method == "POST" + assert url == f"{TEST_API_BASE}/sandboxes" + assert "OPEN-SANDBOX-API-KEY" not in headers + assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} + assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] + assert body["timeout"] == 300 + assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert handle.id == "osb_123" + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + + +@pytest.mark.asyncio +async def test_create_can_opt_into_internet_access(): + client = FakeHTTPClient() + + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + allow_internet_access=True, + client=client, + ) + + _, _, _, body, _ = client.calls[0] + assert "networkPolicy" not in body + + +@pytest.mark.asyncio +async def test_create_custom_options_poll_and_endpoint_resolution(): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + }, + sandbox_states=[ + { + "id": "osb_pending", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + } + ], + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + template="custom/image:latest", + timeout=600, + allow_internet_access=False, + api_key="osb-key", + api_base="https://sandbox.example/v1", + metadata={"suite": "unit"}, + env_vars={"PYTHONUNBUFFERED": "1"}, + resource_limits={"cpu": "500m", "memory": "512Mi"}, + resource_requests={"cpu": "250m", "memory": "256Mi"}, + entrypoint=["/bin/sh", "-lc", "sleep 3600"], + use_server_proxy=True, + client=client, + ) + + _, create_url, create_headers, body, _ = client.calls[0] + _, poll_url, poll_headers, _, _ = client.calls[1] + _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] + + assert create_url == "https://sandbox.example/v1/sandboxes" + assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert body["image"] == {"uri": "custom/image:latest"} + assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] + assert body["metadata"] == {"suite": "unit"} + assert body["env"] == {"PYTHONUNBUFFERED": "1"} + assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} + assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" + assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") + assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_params == {"use_server_proxy": True} + assert handle.id == "osb_pending" + + +@pytest.mark.asyncio +async def test_create_waits_across_pending_state(monkeypatch): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + }, + sandbox_states=[ + {"id": "osb_pending", "status": {"state": "Pending"}}, + {"id": "osb_pending", "status": {"state": "Running"}}, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + assert handle.id == "osb_pending" + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_for_terminal_state(): + client = FakeHTTPClient( + create_json={"id": "osb_failed", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_failed", "status": {"state": "Failed"}}, + ], + ) + + with pytest.raises(ValueError, match="entered Failed"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_create_times_out_waiting_for_running(): + client = FakeHTTPClient( + create_json={"id": "osb_slow", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_slow", "status": {"state": "Pending"}}, + ], + ) + + with pytest.raises(TimeoutError, match="was not Running"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=0, + poll_interval=0, + client=client, + ) + + +@pytest.mark.asyncio +async def test_create_waits_for_endpoint_resolution(monkeypatch): + client = FakeHTTPClient( + endpoint_responses=[ + http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), + { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + }, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + assert len(endpoint_calls) == 2 + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_when_endpoint_is_missing(): + client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) + + with pytest.raises(TimeoutError, match="execd endpoint.*not ready"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client + ) + + +@pytest.mark.asyncio +async def test_create_reraises_non_404_endpoint_error(): + client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_run_code_resolves_bare_id_and_posts_sse_request(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + result = await OpenSandboxSandboxConfig().arun_code( + container="osb_bare", + code="print(6*7)", + language="python", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + endpoint_call = client.calls[0] + run_call = client.calls[1] + assert endpoint_call[0] == "GET" + assert ( + endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" + ) + assert run_call[0] == "POST" + assert run_call[1] == "http://execd.local:44772/code" + assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" + assert run_call[3] == { + "code": "print(6*7)", + "context": {"language": "python"}, + } + assert run_call[4] == {"stream": True} + assert result.stdout == "42\n" + + +@pytest.mark.asyncio +async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): + client = FakeHTTPClient() + handle = ContainerHandle( + id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" + ) + handle._hidden_params = { + "execd_endpoint": "execd.example/route/44772", + "execd_headers": {}, + } + + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print(1)", client=client + ) + + assert client.calls[0][1] == "https://execd.example/route/44772/code" + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) + handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") + handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} + + with pytest.raises(ValueError, match="exceeded"): + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print('x')", client=client + ) + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + + ok = await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_gone", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + assert ok is False + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_err", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + container = await litellm.acreate_sandbox( + provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client + ) + result = await litellm.arun_code( + provider="opensandbox", + container=container, + code="print(6*7)", + api_key="", + client=client, + ) + ok = await litellm.adelete_sandbox( + provider="opensandbox", + container=container, + api_key="", + client=client, + ) + + assert container.id == "osb_123" + assert result.stdout == "42\n" + assert ok is True + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="opensandbox", + code="1/0", + api_key="", + api_base=TEST_API_BASE, + client=client, + ) + + assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] + assert client.calls[0][1].endswith("/sandboxes") + assert client.calls[1][1].endswith("/endpoints/44772") + assert client.calls[2][1].endswith("/code") + assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py new file mode 100644 index 00000000000..9ca4515239a --- /dev/null +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -0,0 +1,89 @@ +""" +Regression tests for the Cloudflare Workers AI text-generation catalog in the +model-cost map. + +The Cloudflare list was badly stale (only 4 ancient entries). These tests pin +the newly added current Workers AI models (sourced from Cloudflare's live +``/ai/models/search?task=Text Generation`` catalog) and guard against the root +``model_prices_and_context_window.json`` and the bundled +``litellm/model_prices_and_context_window_backup.json`` drifting out of sync for +the ``cloudflare/`` namespace. +""" + +import json +import os + +import pytest + +import litellm + +ROOT_MAP = os.path.join( + os.path.dirname(os.path.dirname(litellm.__file__)), + "model_prices_and_context_window.json", +) +BACKUP_MAP = os.path.join( + os.path.dirname(litellm.__file__), + "model_prices_and_context_window_backup.json", +) + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + yield + finally: + litellm.model_cost = original_model_cost + + +def _load(path: str) -> dict: + with open(path, encoding="utf-8") as f: + return json.load(f) + + +def _cloudflare_keys(data: dict) -> set: + return {k for k in data if k.startswith("cloudflare/")} + + +def test_glm_5_2_entry_is_present_and_well_formed(): + entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"] + assert entry["litellm_provider"] == "cloudflare" + assert entry["mode"] == "chat" + assert entry["supports_function_calling"] is True + assert entry["input_cost_per_token"] > 0 + assert entry["output_cost_per_token"] > 0 + + +def test_vision_model_is_flagged_supports_vision(): + entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"] + assert entry["litellm_provider"] == "cloudflare" + assert entry.get("supports_vision") is True + + +def test_additional_current_models_are_present(): + for key in ( + "cloudflare/@cf/openai/gpt-oss-120b", + "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast", + ): + entry = litellm.model_cost[key] + assert entry["litellm_provider"] == "cloudflare" + assert entry["mode"] == "chat" + assert entry["supports_function_calling"] is True + assert entry["input_cost_per_token"] > 0 + assert entry["output_cost_per_token"] > 0 + + +def test_root_and_backup_have_identical_cloudflare_keys(): + if not os.path.exists(ROOT_MAP): + pytest.skip("root cost map only ships in source checkouts") + assert _cloudflare_keys(_load(ROOT_MAP)) == _cloudflare_keys(_load(BACKUP_MAP)) + + +def test_root_and_backup_cloudflare_entries_are_byte_for_byte_equal(): + if not os.path.exists(ROOT_MAP): + pytest.skip("root cost map only ships in source checkouts") + root = {k: v for k, v in _load(ROOT_MAP).items() if k.startswith("cloudflare/")} + backup = {k: v for k, v in _load(BACKUP_MAP).items() if k.startswith("cloudflare/")} + assert root == backup diff --git a/tests/test_litellm/test_completion_timeout_resolution.py b/tests/test_litellm/test_completion_timeout_resolution.py index a76cc6f7de8..7eb79e90e60 100644 --- a/tests/test_litellm/test_completion_timeout_resolution.py +++ b/tests/test_litellm/test_completion_timeout_resolution.py @@ -63,8 +63,9 @@ def test_global_timeout_from_litellm_settings(): ) -def test_global_timeout_package_default_coerced_to_600_for_completion(): - """Package default 6000s → 600s for completion-only path.""" +def test_explicit_global_timeout_6000_is_preserved(): + """The caller passes the explicitly-configured value (or None); an explicit + 6000 must be honored, not silently coerced to 600.""" assert ( CompletionTimeout.resolve( None, @@ -73,7 +74,7 @@ def test_global_timeout_package_default_coerced_to_600_for_completion(): global_timeout=6000.0, supports_httpx_timeout=supports_httpx_timeout, ) - == 600.0 + == 6000.0 ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c2aa4a095d4..470d38caf10 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3062,6 +3062,36 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() assert credentials["custom_llm_provider"] == "bedrock" +def test_get_deployment_credentials_with_provider_includes_bucket_name(): + """ + Regression: bucket_name must survive the CredentialLiteLLMParams filter so + managed-files batch retrieval can resolve the GCS/S3 bucket. Previously it was + dropped, causing "GCS bucket_name is required" when fetching batch output files. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.5-flash", + "vertex_project": "my-project", + "vertex_location": "global", + "gcs_bucket_name": "my-batch-bucket", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="vertex-gemini" + ) + + assert credentials is not None + assert credentials["gcs_bucket_name"] == "my-batch-bucket" + assert credentials["vertex_project"] == "my-project" + assert credentials["custom_llm_provider"] == "vertex_ai" + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves @@ -4901,3 +4931,107 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) is True ) + + +class TestRouterRequestTimeoutPropagation: + """litellm_settings.request_timeout must act as an independent per-attempt timeout. + + Regression for LIT-2369: request_timeout was shadowed by router_settings.timeout, + so Bedrock (and other provider) calls fell back to the hardcoded 600s httpx + default instead of the configured value. + """ + + def _make_router(self, timeout=None, stream_timeout=None): + return litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-test", + }, + } + ], + timeout=timeout, + stream_timeout=stream_timeout, + ) + + @pytest.fixture + def explicit_request_timeout(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + try: + yield 300 + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_request_timeout_stored_independently_when_both_set( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert router.timeout == 330 + assert router.request_timeout == 300 + + def test_request_timeout_none_when_not_explicitly_configured(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + try: + router = self._make_router(timeout=330) + assert router.timeout == 330 + assert router.request_timeout is None + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_non_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 + + def test_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + # stream=True resolves through _get_stream_timeout; request_timeout must win. + assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 + + def test_explicit_stream_timeout_still_wins_over_request_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330, stream_timeout=45) + assert router._get_stream_timeout(kwargs={}, data={}) == 45 + + def test_non_stream_falls_through_to_router_timeout_without_request_timeout(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + try: + router = self._make_router(timeout=330) + assert router._get_non_stream_timeout(kwargs={}, data={}) == 330 + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_per_deployment_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 + + def test_per_request_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert ( + router._get_non_stream_timeout( + kwargs={"timeout": 60}, data={"timeout": 120} + ) + == 60 + ) diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/test_litellm/test_router_per_deployment_num_retries.py index 154ba579e4e..af2372616a6 100644 --- a/tests/test_litellm/test_router_per_deployment_num_retries.py +++ b/tests/test_litellm/test_router_per_deployment_num_retries.py @@ -4,8 +4,9 @@ GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params """ import pytest -from unittest.mock import MagicMock, patch +from unittest.mock import patch +import litellm from litellm import Router @@ -188,3 +189,133 @@ class TestPerDeploymentNumRetries: # Verify num_retries was converted from string to int assert exc.num_retries == 6 + + +class TestNumRetriesNoneGuard: + """ + Regression tests for the num_retries=None TypeError in async_function_with_retries. + + When num_retries reaches async_function_with_retries as None - e.g. a caller passes + num_retries=None explicitly (dict.get() does not fall back on an existing None value), + an auto_router/complexity_router path does not propagate it, or + Router.update_settings(num_retries=None) is used - AND the underlying call fails with a + retryable error, the comparison `if num_retries > 0:` raised: + + TypeError: '>' not supported between instances of 'NoneType' and 'int' + + This masked the real upstream error (rate limit / connection / 5xx) behind a TypeError. + Related issues: #23316, #25889, #23699, #28126. + """ + + @staticmethod + def _mock_router(num_retries=2): + return Router( + model_list=[ + { + "model_name": "mock-model", + "litellm_params": { + "model": "gpt-4o-mini", + "mock_response": "ok", + }, + } + ], + num_retries=num_retries, + ) + + def test_update_kwargs_normalises_explicit_none_to_router_default(self): + """ + _update_kwargs_before_fallbacks must normalise an explicit num_retries=None to + the router default (not leave it as None), while preserving an explicit 0. + """ + router = self._mock_router(num_retries=4) + + # explicit None -> router default + kwargs = {"num_retries": None} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 4 + + # explicit 0 is preserved (retries stay disabled) + kwargs = {"num_retries": 0} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 0 + + # absent -> router default (unchanged behaviour) + kwargs = {} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 4 + + # explicit None with router default also None -> 0 (mirrors the downstream guard) + router.num_retries = None # simulate update_settings(num_retries=None) (#28126) + kwargs = {"num_retries": None} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 0 + + @pytest.mark.asyncio + async def test_acompletion_num_retries_none_does_not_raise_typeerror(self): + """ + Per-request num_retries=None + a retryable error must NOT raise TypeError. + The router falls back to its configured num_retries and retries the (transient) + error, so the request succeeds. + """ + router = self._mock_router(num_retries=2) + with patch("asyncio.sleep", return_value=None): + response = await router.acompletion( + model="mock-model", + messages=[{"role": "user", "content": "hi"}], + num_retries=None, # the trigger + mock_testing_rate_limit_error=True, # retryable error path + ) + assert response.choices[0].message.content == "ok" + + @pytest.mark.asyncio + async def test_async_function_with_retries_none_falls_back_to_zero(self): + """ + When both the per-request value AND the router-level setting are None + (e.g. after Router.update_settings(num_retries=None), #28126), num_retries must + fall back to 0 and the real retryable error must surface - not a TypeError. + """ + router = self._mock_router(num_retries=0) + router.num_retries = None # simulate update_settings(num_retries=None) + + async def failing_fn(*args, **kwargs): + raise litellm.RateLimitError( + message="boom", model="mock-model", llm_provider="openai" + ) + + with patch("asyncio.sleep", return_value=None): + with pytest.raises(litellm.RateLimitError): + await router.async_function_with_retries( + original_function=failing_fn, + model="mock-model", + messages=[{"role": "user", "content": "hi"}], + num_retries=None, + ) + + @pytest.mark.asyncio + async def test_async_function_with_retries_none_falls_back_to_router_default(self): + """ + A None per-request num_retries falls back to the router-level setting, so retries + still happen (original_function is invoked more than once) before the real error + is raised - proving None did not silently disable retries or crash. + """ + router = self._mock_router(num_retries=3) + calls = {"n": 0} + + async def failing_fn(*args, **kwargs): + calls["n"] += 1 + raise litellm.InternalServerError( + message="boom", model="mock-model", llm_provider="openai" + ) + + with patch("asyncio.sleep", return_value=None): + with pytest.raises(litellm.InternalServerError): + await router.async_function_with_retries( + original_function=failing_fn, + model="mock-model", + messages=[{"role": "user", "content": "hi"}], + metadata={}, # populated by acompletion in the real path; log_retry needs it + num_retries=None, + ) + + # 1 initial attempt + at least 1 retry -> proves None fell back to a positive int + assert calls["n"] >= 2 diff --git a/tests/test_litellm/test_router_streaming_fallback_metadata.py b/tests/test_litellm/test_router_streaming_fallback_metadata.py new file mode 100644 index 00000000000..6ed70dc7cfe --- /dev/null +++ b/tests/test_litellm/test_router_streaming_fallback_metadata.py @@ -0,0 +1,187 @@ +import json +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.proxy.proxy_server import _should_include_fallback_errors +from litellm.router import Router +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + + +def test_apply_fallback_hidden_params_copies_from_fallback_response(): + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + chunk = litellm.ModelResponseStream( + id="test", + model="openai/internal-fallback", + choices=[], + ) + chunk._hidden_params = { + "additional_headers": {"x-existing-chunk-header": "keep"}, + "model_id": "chunk-model-id", + } + fallback_response = MagicMock() + fallback_response._hidden_params = { + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + }, + "api_base": "https://fallback.example", + } + + Router._apply_fallback_hidden_params_to_item( + fallback_item=chunk, + prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params( + fallback_response + ), + ) + + assert chunk._hidden_params["api_base"] == "https://fallback.example" + assert chunk._hidden_params["model_id"] == "chunk-model-id" + assert chunk._hidden_params["additional_headers"] == { + "x-existing-chunk-header": "keep", + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + + +def _two_group_fallback_router() -> Router: + return litellm.Router( + model_list=[ + { + "model_name": "primary-model", + "litellm_params": {"model": "openai/gpt-fake", "api_key": "sk-fake"}, + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-fake-2", "api_key": "sk-fake"}, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + +def _additional_headers(response: object) -> dict: + return get_hidden_params_dict(response).get("additional_headers", {}) + + +@pytest.mark.asyncio +async def test_include_fallback_errors_propagates_through_router(): + router = _two_group_fallback_router() + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + mock_testing_fallbacks=True, + mock_response="fallback success", + include_fallback_errors=True, + ) + + headers = _additional_headers(response) + assert headers["x-litellm-attempted-fallbacks"] == 1 + errors = json.loads(headers["x-litellm-fallback-errors"]) + assert isinstance(errors, list) and len(errors) >= 1 + assert set(errors[0].keys()) == {"message", "type", "param", "code"} + + +@pytest.mark.asyncio +async def test_router_omits_fallback_errors_without_opt_in(): + router = _two_group_fallback_router() + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + mock_testing_fallbacks=True, + mock_response="fallback success", + ) + + headers = _additional_headers(response) + assert headers["x-litellm-attempted-fallbacks"] == 1 + assert "x-litellm-fallback-errors" not in headers + + +def test_prepare_fallback_hidden_params_no_additional_headers(): + class FakeResponse: + _hidden_params = {"api_base": "http://example.com"} + + hidden_params, headers = Router._prepare_fallback_hidden_params(FakeResponse()) + assert hidden_params == {"api_base": "http://example.com"} + assert headers == {} + + +def test_apply_fallback_hidden_params_to_item_none_item(): + Router._apply_fallback_hidden_params_to_item( + None, ({"api_base": "http://fallback.example"}, {"x-custom": "value"}) + ) + + +def test_apply_fallback_hidden_params_to_item_no_existing_additional_headers(): + class FakeChunk: + _hidden_params = {"model_id": "test-id"} + + chunk = FakeChunk() + Router._apply_fallback_hidden_params_to_item( + chunk, + ( + {"api_base": "http://fallback.example"}, + {"x-litellm-attempted-fallbacks": 1}, + ), + ) + + assert chunk._hidden_params["api_base"] == "http://fallback.example" + assert chunk._hidden_params["model_id"] == "test-id" + assert chunk._hidden_params["additional_headers"] == { + "x-litellm-attempted-fallbacks": 1 + } + + +@pytest.mark.asyncio +async def test_set_response_headers_adds_model_group_to_streaming_wrapper(): + class StreamingWrapper: + def __init__(self): + self._hidden_params = {"additional_headers": {"x-existing": "keep"}} + + router = litellm.Router(model_list=[]) + response = StreamingWrapper() + + result = await router.set_response_headers( + response=response, + model_group="fallback-model", + ) + + assert result is response + assert response._hidden_params["additional_headers"] == { + "x-existing": "keep", + "x-litellm-model-group": "fallback-model", + } + + +def test_should_include_fallback_errors_gated_by_operator_setting(): + request_data: dict = {"include_fallback_errors": True} + + import litellm.proxy.proxy_server as ps + + original = ps.general_settings.copy() if isinstance(ps.general_settings, dict) else {} + try: + ps.general_settings = {} + assert _should_include_fallback_errors(request_data) is False + + ps.general_settings = {"expose_fallback_errors_to_caller": False} + assert _should_include_fallback_errors(request_data) is False + + ps.general_settings = {"expose_fallback_errors_to_caller": True} + assert _should_include_fallback_errors(request_data) is True + + ps.general_settings = {"expose_fallback_errors_to_caller": True} + assert _should_include_fallback_errors({}) is False + finally: + ps.general_settings = original diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/test_litellm/test_type_check_gate.py index 18374c5db4b..e99ad0a4f41 100644 --- a/tests/test_litellm/test_type_check_gate.py +++ b/tests/test_litellm/test_type_check_gate.py @@ -56,29 +56,58 @@ def test_paths_outside_repo_are_skipped(): def test_at_or_under_ceiling_passes(): budget = {"no-any-return": {"baseline": 5, "slack": 0}} - assert gate.evaluate({"no-any-return": 5}, budget) == [] + assert gate.evaluate({"no-any-return": 5}, {}, budget) == [] def test_one_more_error_than_ceiling_fails(): budget = {"no-any-return": {"baseline": 5, "slack": 0}} - assert gate.evaluate({"no-any-return": 6}, budget) == [ - gate.Breach("no-any-return", 6, 5) + assert gate.evaluate({"no-any-return": 6}, {}, budget) == [ + gate.Breach("no-any-return", 6, 5, 6) ] def test_slack_absorbs_small_increase_then_fails_past_it(): budget = {"arg-type": {"baseline": 5, "slack": 5}} - assert gate.evaluate({"arg-type": 10}, budget) == [] - assert gate.evaluate({"arg-type": 11}, budget) == [gate.Breach("arg-type", 11, 10)] + assert gate.evaluate({"arg-type": 10}, {}, budget) == [] + assert gate.evaluate({"arg-type": 11}, {}, budget) == [ + gate.Breach("arg-type", 11, 10, 11) + ] def test_unbudgeted_new_code_uses_default_slack(): - assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK}, {}) == [] - assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK + 1}, {}) == [ - gate.Breach("brand-new", gate.DEFAULT_SLACK + 1, gate.DEFAULT_SLACK) + assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK}, {}, {}) == [] + assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK + 1}, {}, {}) == [ + gate.Breach( + "brand-new", + gate.DEFAULT_SLACK + 1, + gate.DEFAULT_SLACK, + gate.DEFAULT_SLACK + 1, + ) ] +def test_drift_already_over_cap_in_base_is_not_blamed_on_a_flat_change(): + # The bystander case: a rule sits over its ceiling because two earlier PRs + # summed past it. A PR that branches off that base and adds nothing must pass + # -- total > cap but total == base, so the `> base` guard spares it. + budget = {"arg-type": {"baseline": 5, "slack": 5}} + assert gate.evaluate({"arg-type": 12}, {"arg-type": 12}, budget) == [] + + +def test_change_that_grows_an_over_cap_rule_is_blamed_for_only_what_it_added(): + # Over cap AND above base: blamed, and `added` is the delta vs base, not the + # whole overage, so the message points at this change's contribution. + budget = {"arg-type": {"baseline": 5, "slack": 5}} + assert gate.evaluate({"arg-type": 14}, {"arg-type": 12}, budget) == [ + gate.Breach("arg-type", 14, 10, 2) + ] + + +def test_reducing_an_over_cap_rule_below_base_passes(): + budget = {"arg-type": {"baseline": 5, "slack": 5}} + assert gate.evaluate({"arg-type": 11}, {"arg-type": 12}, budget) == [] + + def test_no_output_against_a_nonempty_budget_is_a_vacuous_run(): # A crashed type checker emits nothing; the gate must not certify it as clean. budget = {"no-untyped-def": {"baseline": 4888, "slack": 10}} diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d94a86d8e55..cf7afc1af68 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -865,6 +865,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, + "supports_speed": {"type": "boolean"}, "bedrock_output_config_effort_ceiling": { "type": "string", "enum": ["low", "medium", "high", "max", "xhigh"], diff --git a/tests/test_litellm/types/test_completion.py b/tests/test_litellm/types/test_completion.py index f24b00df3fc..cd51913c5dd 100644 --- a/tests/test_litellm/types/test_completion.py +++ b/tests/test_litellm/types/test_completion.py @@ -8,9 +8,16 @@ Usage: pytest tests/test_litellm/types/test_completion.py -v """ +import dataclasses from typing import List -from litellm.types.completion import CompletionRequest, ChatCompletionMessageParam +import pytest + +from litellm.types.completion import ( + ChatCompletionMessageParam, + CompletionRequest, + _CompletionDispatchContext, +) def test_completion_request_messages_type_validation(): @@ -146,3 +153,55 @@ def test_completion_request_with_all_params(): assert request.presence_penalty == 0.0 assert request.stream is False assert request.n == 1 + + +def _build_dispatch_context() -> _CompletionDispatchContext: + return _CompletionDispatchContext( + _azure_detection_model="gpt-4o", + acompletion=False, + api_base=None, + api_key=None, + api_version=None, + client=None, + custom_llm_provider="openai", + custom_prompt_dict={}, + extra_headers=None, + headers={}, + hf_model_name=None, + kwargs={}, + litellm_params={}, + logger_fn=None, + logging=None, # type: ignore[arg-type] + max_retries=None, + max_tokens=None, + messages=[], + metadata=None, + model="gpt-4o", + model_response=None, # type: ignore[arg-type] + optional_params={}, + organization=None, + provider_config=None, + shared_session=None, + stream=None, + temperature=None, + text_completion=False, + timeout=None, + top_p=None, + ) + + +def test_dispatch_context_is_frozen(): + """A helper must not be able to re-route the call by rebinding a dispatch + input mid-flight; this pins the frozen invariant the dispatch shape relies on.""" + ctx = _build_dispatch_context() + with pytest.raises(dataclasses.FrozenInstanceError): + ctx.model = "claude-haiku-4-5" # type: ignore[misc] + with pytest.raises(dataclasses.FrozenInstanceError): + ctx.custom_llm_provider = "anthropic" # type: ignore[misc] + + +def test_dispatch_context_uses_slots(): + """slots=True keeps the per-call context lightweight (no per-instance __dict__).""" + ctx = _build_dispatch_context() + assert not hasattr(ctx, "__dict__") + assert hasattr(type(ctx), "__slots__") diff --git a/ui/litellm-dashboard/eslint-metrics.json b/ui/litellm-dashboard/eslint-metrics.json new file mode 100644 index 00000000000..28f9fc3a6af --- /dev/null +++ b/ui/litellm-dashboard/eslint-metrics.json @@ -0,0 +1,5 @@ +{ + "@typescript-eslint/no-explicit-any": 2027, + "complexity": 128, + "max-depth": 61 +} diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 3beae6526e2..2afc145d15b 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -6882,16 +6882,16 @@ } }, "node_modules/form-data": { - "version": "4.0.5", - "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.5.tgz", - "integrity": "sha512-8RipRLol37bNs2bhoV67fiTEvdTrbMUYcFTiy3+wuuOnUog2QBHCZWXDRijWQfAkhBj2Uf5UnVaiWwA5vdd82w==", + "version": "4.0.6", + "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.6.tgz", + "integrity": "sha512-vKatAh4SlVfgbv+YtmhiRjhEMJsYpsG1Y2rMQtR+SVSbytsSD1YGzDIcrAJmdFec88u/+VoGmxnl+80gL1tRCQ==", "license": "MIT", "dependencies": { "asynckit": "^0.4.0", "combined-stream": "^1.0.8", "es-set-tostringtag": "^2.1.0", - "hasown": "^2.0.2", - "mime-types": "^2.1.12" + "hasown": "^2.0.4", + "mime-types": "^2.1.35" }, "engines": { "node": ">= 6" @@ -7248,9 +7248,9 @@ } }, "node_modules/hasown": { - "version": "2.0.3", - "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.3.tgz", - "integrity": "sha512-ej4AhfhfL2Q2zpMmLo7U1Uv9+PyhIZpgQLGT1F9miIGmiCJIoCgSmczFdrc97mWT4kVY72KA+WnnhJ5pghSvSg==", + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.4.tgz", + "integrity": "sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A==", "license": "MIT", "dependencies": { "function-bind": "^1.1.2" @@ -8149,10 +8149,20 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.1.1", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.1.1.tgz", - "integrity": "sha512-qQKT4zQxXl8lLwBtHMWwaTcGfFOZviOJet3Oy/xmGk2gZH677CJM9EvtfdSkgWcATZhj/55JZ0rmy3myCT5lsA==", + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.2.0.tgz", + "integrity": "sha512-ePWsvanv0DWuDRsW8dnt+R4jQ31SCRCQ7hhNcPXZPsoBZiemuZNYGf7adZdqX2D86j6rvKp3RpCxVTSb8WQlOw==", "dev": true, + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/puzrin" + }, + { + "type": "github", + "url": "https://github.com/sponsors/nodeca" + } + ], "license": "MIT", "dependencies": { "argparse": "^2.0.1" @@ -13726,9 +13736,9 @@ } }, "node_modules/ws": { - "version": "8.20.1", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.20.1.tgz", - "integrity": "sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==", + "version": "8.21.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", + "integrity": "sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==", "devOptional": true, "license": "MIT", "engines": { diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index c0899be8639..a8948f4be34 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -8,6 +8,7 @@ "build": "next build", "start": "next start", "lint": "eslint .", + "lint:metrics": "node scripts/update-lint-metrics.mjs", "test": "vitest", "test:dot": "vitest --reporter=dot", "test:watch": "vitest -w", @@ -83,11 +84,11 @@ }, "overrides": { "prismjs": "1.30.0", - "js-yaml": "4.1.1", + "js-yaml": "4.2.0", "glob": "13.0.0", "minimatch": "10.2.4", "lodash": "4.18.1", - "ws": "8.20.1", + "ws": "8.21.0", "braces": "3.0.3", "axios": "1.13.6", "postcss": "8.5.13", diff --git a/ui/litellm-dashboard/scripts/check-lint-budgets.mjs b/ui/litellm-dashboard/scripts/check-lint-budgets.mjs index f6208f012bb..a7aba18ae76 100644 --- a/ui/litellm-dashboard/scripts/check-lint-budgets.mjs +++ b/ui/litellm-dashboard/scripts/check-lint-budgets.mjs @@ -1,22 +1,25 @@ import { readFileSync } from "fs"; +import { countBudgetViolations, findDrift } from "./lint-budget-lib.mjs"; -const [, , reportPath, budgetsPath] = process.argv; - -const report = JSON.parse(readFileSync(reportPath, "utf8")); -const budgets = JSON.parse(readFileSync(budgetsPath, "utf8")); - -const counts = {}; -for (const file of report) { - for (const message of file.messages) { - if (message.ruleId in budgets) { - counts[message.ruleId] = (counts[message.ruleId] || 0) + 1; - } +const argv = process.argv.slice(2); +const positional = []; +const flags = {}; +for (let i = 0; i < argv.length; i += 1) { + if (argv[i] === "--check") { + flags.check = argv[(i += 1)]; + } else { + positional.push(argv[i]); } } +const [reportPath, budgetsPath] = positional; +const report = JSON.parse(readFileSync(reportPath, "utf8")); +const budgets = JSON.parse(readFileSync(budgetsPath, "utf8")); +const counts = countBudgetViolations(report, budgets); + let failed = false; for (const [rule, { max, target }] of Object.entries(budgets)) { - const count = counts[rule] || 0; + const count = counts[rule]; const note = count > max ? "OVER BUDGET" : count <= target ? "at target" : `${max - count} of headroom`; console.log(`${rule}: ${count} | max: ${max} | target: ${target} | ${note}`); if (count > max) { @@ -27,4 +30,20 @@ for (const [rule, { max, target }] of Object.entries(budgets)) { } } +if (flags.check) { + const committed = JSON.parse(readFileSync(flags.check, "utf8")); + const drift = findDrift(committed, counts); + for (const { rule, committed: was, actual } of drift) { + console.error( + `::error::${flags.check} is stale for ${rule}: committed ${was ?? "missing"}, actual ${actual ?? "not a tracked rule"}.`, + ); + } + if (drift.length > 0) { + console.error(`::error::Run \`npm run lint:metrics\` and commit ${flags.check}.`); + failed = true; + } else { + console.log(`${flags.check} is up to date.`); + } +} + process.exit(failed ? 1 : 0); diff --git a/ui/litellm-dashboard/scripts/lint-budget-lib.mjs b/ui/litellm-dashboard/scripts/lint-budget-lib.mjs new file mode 100644 index 00000000000..a43305fc4ed --- /dev/null +++ b/ui/litellm-dashboard/scripts/lint-budget-lib.mjs @@ -0,0 +1,22 @@ +export function countBudgetViolations(report, budgets) { + const counts = {}; + for (const file of report) { + for (const message of file.messages) { + if (message.ruleId in budgets) { + counts[message.ruleId] = (counts[message.ruleId] || 0) + 1; + } + } + } + return Object.fromEntries( + Object.keys(budgets) + .sort() + .map((rule) => [rule, counts[rule] || 0]), + ); +} + +export function findDrift(committed, actual) { + const rules = [...new Set([...Object.keys(actual), ...Object.keys(committed)])].sort(); + return rules + .filter((rule) => committed[rule] !== actual[rule]) + .map((rule) => ({ rule, committed: committed[rule] ?? null, actual: actual[rule] ?? null })); +} diff --git a/ui/litellm-dashboard/scripts/update-lint-metrics.mjs b/ui/litellm-dashboard/scripts/update-lint-metrics.mjs new file mode 100644 index 00000000000..16704d1f7a2 --- /dev/null +++ b/ui/litellm-dashboard/scripts/update-lint-metrics.mjs @@ -0,0 +1,25 @@ +import { execSync } from "child_process"; +import { mkdtempSync, readFileSync, writeFileSync, rmSync } from "fs"; +import { tmpdir } from "os"; +import { join } from "path"; +import { countBudgetViolations } from "./lint-budget-lib.mjs"; + +const ESLINT_EXIT_LINT_ERRORS = 1; + +const budgets = JSON.parse(readFileSync("eslint-budgets.json", "utf8")); +const dir = mkdtempSync(join(tmpdir(), "litellm-lint-")); +const reportPath = join(dir, "report.json"); + +try { + execSync(`npx eslint . -f json -o "${reportPath}"`, { stdio: "inherit" }); +} catch (err) { + if (err.status !== ESLINT_EXIT_LINT_ERRORS) throw err; +} + +const report = JSON.parse(readFileSync(reportPath, "utf8")); +rmSync(dir, { recursive: true, force: true }); + +const metrics = countBudgetViolations(report, budgets); +writeFileSync("eslint-metrics.json", JSON.stringify(metrics, null, 2) + "\n"); +console.log("Updated eslint-metrics.json"); +console.table(metrics); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx index 56b34d6a68b..f2c06387301 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx @@ -3,6 +3,7 @@ import { TextInput, Button } from "@tremor/react"; import { Select as AntdSelect, Form, Tooltip, Radio } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Providers, provider_map, providerLogoMap } from "@/components/provider_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import { MarginConfig } from "./types"; import { handleImageError } from "./provider_display_helpers"; @@ -73,7 +74,7 @@ const AddMarginForm: React.FC = ({
{`${providerEnum} handleImageError(e, providerDisplayName)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx index 61ba3194607..c4961263533 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx @@ -3,6 +3,7 @@ import { TextInput, Button } from "@tremor/react"; import { Select as AntdSelect, Form, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Providers, provider_map, providerLogoMap } from "@/components/provider_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import { DiscountConfig } from "./types"; import { handleImageError } from "./provider_display_helpers"; @@ -60,7 +61,7 @@ const AddProviderForm: React.FC = ({
{`${providerEnum} handleImageError(e, providerDisplayName)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts index cd088da09da..5489eb12487 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts @@ -1,4 +1,5 @@ import { Providers, provider_map, providerLogoMap } from "@/components/provider_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; export interface ProviderDisplayInfo { displayName: string; @@ -16,7 +17,7 @@ export const getProviderDisplayInfo = (providerValue: string): ProviderDisplayIn if (enumKey) { const displayName = Providers[enumKey as keyof typeof Providers]; - const logo = providerLogoMap[displayName]; + const logo = resolveLogoSrc(providerLogoMap[displayName]) ?? ""; return { displayName, logo, enumKey }; } diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index d777ba1b0dc..0b6e5786aaf 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -1097,3 +1097,52 @@ describe("OldTeams - delete team warning copy", () => { ); }); }); + +describe("OldTeams - LIT-2530 organization stays optional for proxy admin with a single org", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamListCall).mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 1 }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "No Org Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }], + }); + }); + + it("creates a team with no organization when exactly one organization exists", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "No Org Team" } }); + fireEvent.change(screen.getByTestId("create-team-models-select"), { target: { value: "gpt-4" } }); + + const submitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(submitButtons[submitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ team_alias: "No Org Team", organization_id: null }), + ); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index adfec4bdf6a..be9015e3730 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -262,14 +262,15 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser useEffect(() => { if (isTeamModalVisible) { const adminOrgs = getAdminOrganizations(userRole, userID, organizations); + const isOrgAdmin = userRole !== "Admin"; - // If there's exactly one organization the user is admin for, preselect it - if (adminOrgs.length === 1) { + // Org admins must scope a team to an org, so with exactly one we preselect it. + // Proxy admins can create org-less teams, so the field stays optional regardless of org count. + if (isOrgAdmin && adminOrgs.length === 1) { const org = adminOrgs[0]; form.setFieldValue("organization_id", org.organization_id); setCurrentOrgForCreateTeam(org); } else { - // Reset the organization selection for multiple orgs form.setFieldValue("organization_id", currentOrg?.organization_id || null); setCurrentOrgForCreateTeam(currentOrg); } @@ -1132,7 +1133,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser : [] } help={ - isSingleOrg + isOrgAdmin && isSingleOrg ? "You can only create teams within this organization" : isOrgAdmin ? "required" @@ -1142,7 +1143,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser