diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 00000000000..e6afd7ff530 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,19 @@ +[http] +# CI has seen transient crates.io failures from libcurl's HTTP/2 multiplexing +# during `maturin` metadata resolution. Disable multiplexing and retry more +# aggressively so editable `uv sync` builds are not failed by one flaky frame. +multiplexing = false + +[net] +retry = 5 + +# 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/.circleci/config.yml b/.circleci/config.yml index abcdbf45187..f13e9bf66f1 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -190,6 +190,8 @@ jobs: working_directory: ~/project environment: UV_PYTHON: "3.11" + CARGO_HTTP_MULTIPLEXING: "false" + CARGO_NET_RETRY: "5" steps: - checkout - run: @@ -205,6 +207,24 @@ jobs: environment: UV_HTTP_TIMEOUT: "300" command: | + $rustupInit = Join-Path $env:TEMP "rustup-init.exe" + $rustupVersion = "1.28.2" + $rustupUrl = "https://static.rust-lang.org/rustup/archive/$rustupVersion/x86_64-pc-windows-msvc/rustup-init.exe" + Invoke-WebRequest -Uri $rustupUrl -OutFile $rustupInit + $rustupExpected = "88d8258dcf6ae4f7a80c7d1088e1f36fa7025a1cfd1343731b4ee6f385121fc0" + $rustupActual = (Get-FileHash -Path $rustupInit -Algorithm SHA256).Hash.ToLower() + if ($rustupActual -ne $rustupExpected) { + throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual" + } + & $rustupInit -y --profile minimal --default-toolchain stable + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } + Remove-Item $rustupInit + $cargoBin = Join-Path $HOME ".cargo\bin" + $env:Path = "$cargoBin;$env:Path" + rustc --version + cargo --version $installer = Join-Path $env:TEMP "uv-install.ps1" Invoke-WebRequest -Uri https://astral.sh/uv/0.10.9/install.ps1 -OutFile $installer $expected = "d43ffff8d28e7d1e7d1831a212465f12b24a43c7f87f386e95e2a5915aee5d7d" @@ -222,7 +242,20 @@ jobs: if (-not (Select-String -Path $PROFILE -SimpleMatch $uvBin -Quiet)) { Add-Content -Path $PROFILE -Value "`$env:Path = `"$uvBin;`$env:Path`"" } - uv sync --frozen --group dev --python 3.11 + if (-not (Select-String -Path $PROFILE -SimpleMatch $cargoBin -Quiet)) { + Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`"" + } + for ($attempt = 1; $attempt -le 5; $attempt++) { + Write-Host "uv sync attempt $attempt/5" + uv sync --frozen --group dev --python 3.11 + if ($LASTEXITCODE -eq 0) { + break + } + if ($attempt -eq 5) { + exit $LASTEXITCODE + } + Start-Sleep -Seconds 15 + } - run: name: Run Windows-specific test command: | @@ -232,6 +265,9 @@ jobs: environment: UV_HTTP_TIMEOUT: "300" command: | + $env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path" + cargo --version + Get-ChildItem -Path "litellm\rust_bridge" -Filter "_native*" -File -ErrorAction SilentlyContinue | Remove-Item -Force uv build --wheel --out-dir dist uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py @@ -1020,7 +1056,9 @@ jobs: name: Run tests command: | mkdir -p test-results - TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py") + TEST_FILES=$(printf "%s\n%s\n" \ + "$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \ + "tests/test_litellm/ocr/test_rust_bridge.py") echo "$TEST_FILES" | circleci tests run \ --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ diff --git a/.dockerignore b/.dockerignore index a487d2a859a..6b80caeaf9f 100644 --- a/.dockerignore +++ b/.dockerignore @@ -49,6 +49,8 @@ build/ *.egg-info/ .DS_Store **/node_modules +litellm-rust/target/ +litellm/rust_bridge/_native*.so *.log .env .env.local diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 9658baeb89a..12ad124fa20 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,17 +1,17 @@ ## Relevant issues - + ## Linear ticket - + ## Pre-Submission checklist **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have added meaningful tests -- [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code) +- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem - [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review @@ -19,29 +19,13 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slack (#pr-review)](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA). -## CI (LiteLLM team) - -> **CI status guideline:** -> -> - 50-55 passing tests: main is stable with minor issues. -> - 45-49 passing tests: acceptable but needs attention -> - <= 40 passing tests: unstable; be careful with your merges and assess the risk. - -- [ ] **Branch creation CI run** - Link: - -- [ ] **CI run for the last commit** - Link: - -- [ ] **Merge / cherry-pick CI run** - Links: - ## Screenshots / Proof of Fix - + ## Type diff --git a/.github/scripts/uv_sync_with_retries.sh b/.github/scripts/uv_sync_with_retries.sh new file mode 100755 index 00000000000..85ed75af566 --- /dev/null +++ b/.github/scripts/uv_sync_with_retries.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +set -euo pipefail + +max_attempts="${UV_SYNC_MAX_ATTEMPTS:-5}" +delay_seconds="${UV_SYNC_RETRY_DELAY_SECONDS:-15}" + +export CARGO_HTTP_MULTIPLEXING="${CARGO_HTTP_MULTIPLEXING:-false}" +export CARGO_NET_RETRY="${CARGO_NET_RETRY:-5}" + +if [[ "$#" -eq 0 ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +for attempt in $(seq 1 "${max_attempts}"); do + echo "uv sync attempt ${attempt}/${max_attempts}" + status=0 + if uv sync "$@"; then + exit 0 + else + status=$? + fi + + if [[ "${attempt}" -eq "${max_attempts}" ]]; then + echo "uv sync failed after ${max_attempts} attempts" >&2 + exit "${status}" + fi + + echo "uv sync failed; retrying in ${delay_seconds}s..." + sleep "${delay_seconds}" +done diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index a42b2f8f9df..25c6d4a7019 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -73,7 +73,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index d8053c15683..439126aa1ee 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -46,7 +46,7 @@ jobs: ${{ runner.os }}-uv- - name: Install backend dependencies - run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: 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/image-scan.yml b/.github/workflows/image-scan.yml new file mode 100644 index 00000000000..90ede5a653f --- /dev/null +++ b/.github/workflows/image-scan.yml @@ -0,0 +1,65 @@ +name: Image Scan + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_branch + - "litellm_**" + paths: + - docker/Dockerfile.non_root + - uv.lock + - ui/litellm-dashboard/package-lock.json + - .github/workflows/image-scan.yml + schedule: + - cron: "41 6 * * *" + workflow_dispatch: + +permissions: {} + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + image-scan: + name: image-scan + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Download Grype v0.114.0 + run: | + curl -fsSL --retry 3 -o "$RUNNER_TEMP/grype.tar.gz" \ + https://github.com/anchore/grype/releases/download/v0.114.0/grype_0.114.0_linux_amd64.tar.gz + echo "edda0968d8827daab01d32b3cd7de192ae0915005e7bbfcfef9e68e79bc43343 $RUNNER_TEMP/grype.tar.gz" | sha256sum -c - + tar xzf "$RUNNER_TEMP/grype.tar.gz" -C "$RUNNER_TEMP" grype + chmod +x "$RUNNER_TEMP/grype" + + # Dockerfile.non_root is the rootless variant we ship. The other + # Dockerfiles share the same wolfi base and apk set, so OS-layer coverage + # is the same; matrix-scan if those variants ever diverge. + - name: Build runtime image + run: docker build -f docker/Dockerfile.non_root -t litellm-image-scan:${{ github.sha }} . + + # Scans the whole shipped artifact: OS/apk plus every language package + # baked into the image, including ones no lockfile declares (e.g. prisma's + # vendored node engine) that osv-scan cannot see. osv-scan stays the fast + # source-level gate; this is the customer's-eye-view backstop. Credential- + # free OSS, run as a pinned, checksum-verified binary; no GitHub Action + # dependency and no vendor SaaS callout. + - name: Scan image for fixable HIGH/CRITICAL CVEs + run: | + "$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \ + --only-fixed \ + --fail-on high \ + --output table diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 8094ca57467..183f12f969c 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -55,7 +55,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/.github/workflows/osv-scan.yml b/.github/workflows/osv-scan.yml index 0cd94fdd9e2..31104002dab 100644 --- a/.github/workflows/osv-scan.yml +++ b/.github/workflows/osv-scan.yml @@ -5,7 +5,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" schedule: - cron: "23 6 * * *" 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 950d6ca31a6..ff6c40ac9ae 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: @@ -50,11 +50,16 @@ jobs: run: | uv sync --frozen - - name: Check Black formatting + - name: Check ruff format + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} run: | - cd litellm - uv run --no-sync black --check --exclude '/enterprise/' . - cd .. + git diff --name-only "$BASE_SHA"...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true + if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then + echo "No changed litellm Python files to check with ruff format." + exit 0 + fi + xargs uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt" - name: Debug - Check file state 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..5b5290880c1 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: @@ -39,7 +39,7 @@ jobs: - name: Install dependencies run: | uv lock --check - uv sync --frozen --group proxy-dev --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router - name: Run MCP tests run: | 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 index 3d0a159cdc7..13e1dc4ad5e 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -9,7 +9,7 @@ on: branches: - main - litellm_internal_staging - - litellm_oss_branch + - litellm_oss_staging - "litellm_**" paths: - "litellm-rust/**" 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..4cef791a9b3 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: @@ -54,7 +54,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: 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 2226d519331..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: 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..8db218cd1fc 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: @@ -71,7 +71,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: 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..59fa5803abe 100644 --- a/.gitignore +++ b/.gitignore @@ -9,12 +9,17 @@ litellm/proxy/myenv/* litellm_uuid.txt __pycache__/ *.pyc + +# Rust bridge build artifacts (compiled, platform-specific; regenerated by maturin/cargo) +litellm/rust_bridge/_native*.so +litellm/rust_bridge/_native*.pyd +litellm-rust/target/ + bun.lockb **/.DS_Store .aider* litellm_results.jsonl secrets.toml -.gitignore litellm/proxy/litellm_secrets.toml litellm/proxy/api_log.json .idea/ @@ -36,7 +41,6 @@ litellm/tests/dynamo*.log .vscode/settings.json litellm/proxy/log.txt proxy_server_config_@.yaml -.gitignore proxy_server_config_2.yaml litellm/proxy/secret_managers/credentials.json hosted_config.yaml @@ -123,3 +127,6 @@ crash.*.log # and should be committed. .vscode .pin_list.txt + +# pytest coverage data +.coverage diff --git a/CLAUDE.md b/CLAUDE.md index b721064aaa7..eb32c2cd6da 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). 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 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: if you're adding new line(s) before the next sentence, 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 @@ -54,7 +54,7 @@ When working on a PR, keep the PR description in sync with new commits being mad Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in -Do not put names of customers or customer company names in code, PRs, and issues. The codebase is public +Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers. CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI diff --git a/Dockerfile b/Dockerfile index af49dc8d8cf..681681f28cc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,6 +21,7 @@ RUN apk add --no-cache \ gcc \ python3 \ python3-dev \ + rust \ openssl \ openssl-dev \ nodejs \ diff --git a/Makefile b/Makefile index 076eac0f4a7..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 @@ -131,7 +133,7 @@ lint-basedpyright: install-dev 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/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 f5b0a9aaf81..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, diff --git a/codecov.yaml b/codecov.yaml index 3baea13e2d3..f5acdd39136 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -3,6 +3,9 @@ codecov: notify: wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI +ignore: + - "litellm-rust/**" + # Uploads are flagged per workflow/shard (GHA) or "circleci". carryforward makes # a re-upload of a flag replace its prior session instead of accumulating a # conflicting one, and lets a commit reuse a flag from its parent when that flag diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index ab02b43d0f9..6bb925aa723 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -19,6 +19,7 @@ RUN for i in 1 2 3; do \ python3 \ python3-dev \ gcc \ + rust \ bash \ coreutils \ curl \ 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/enterprise/pyproject.toml b/enterprise/pyproject.toml index b032942427c..66f6aeb7abc 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.43" +version = "0.1.44" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.43" +version = "0.1.44" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", 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 index 1d2987e0a1a..7c723e570ef 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -2,27 +2,32 @@ 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 -The `core` and `providers` crates describe work; hosts execute work. +`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`. -- `providers/src///transformation.rs` owns the +- `core/src/providers///transformation.rs` owns the provider-specific transform. For Mistral OCR, this means - `providers/src/mistral/ocr/transformation.rs`. -- Future network execution belongs in a host/transport layer such as - `llm_http_handler`, not inside `core` or `providers`. + `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` and `providers`: +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` or `providers`: +Not allowed in `core`: - Network calls - Environment variable or secret reads - Filesystem access @@ -72,6 +77,20 @@ such as `ai-gateway`, router hosts, or standalone servers: - 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 @@ -80,7 +99,9 @@ for changes under `litellm-rust/`. ```bash cd litellm-rust cargo fmt --check -cargo clippy --workspace --all-targets -- -D warnings +# the ai-gateway binary + server code is behind the `server` feature +cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings +cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index a269a224d97..9bffe9f9ec6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -206,12 +206,24 @@ dependencies = [ "syn", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "find-msvc-tools" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -221,6 +233,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -237,12 +264,34 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.32" @@ -261,8 +310,10 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -307,6 +358,31 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -368,6 +444,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", @@ -521,6 +598,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + [[package]] name = "indoc" version = "2.0.7" @@ -544,9 +631,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "js-sys" -version = "0.3.102" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", @@ -564,14 +651,18 @@ name = "litellm-ai-gateway" version = "0.1.0" dependencies = [ "axum", + "base64", + "futures-channel", "futures-util", "litellm-core", - "litellm-providers", "pyo3", + "reqwest", "serde", "serde_json", + "sha2", "subtle", "tokio", + "tokio-tungstenite", ] [[package]] @@ -582,29 +673,19 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.18", -] - -[[package]] -name = "litellm-providers" -version = "0.1.0" -dependencies = [ - "futures-channel", - "futures-util", - "litellm-core", - "reqwest", - "serde_json", "tokio", - "tokio-tungstenite", ] [[package]] name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "litellm-ai-gateway", "litellm-core", - "litellm-providers", "pyo3", + "pyo3-async-runtimes", "serde_json", + "tokio", ] [[package]] @@ -738,6 +819,19 @@ dependencies = [ "unindent", ] +[[package]] +name = "pyo3-async-runtimes" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e" +dependencies = [ + "futures", + "once_cell", + "pin-project-lite", + "pyo3", + "tokio", +] + [[package]] name = "pyo3-build-config" version = "0.23.5" @@ -923,6 +1017,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", + "h2", "http", "http-body", "http-body-util", @@ -942,12 +1037,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots", ] @@ -1140,6 +1237,17 @@ dependencies = [ "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" @@ -1334,6 +1442,19 @@ dependencies = [ "tungstenite", ] +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "tower" version = "0.5.3" @@ -1506,9 +1627,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -1519,9 +1640,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.75" +version = "0.4.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" dependencies = [ "js-sys", "wasm-bindgen", @@ -1529,9 +1650,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1539,9 +1660,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", @@ -1552,18 +1673,31 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] [[package]] -name = "web-sys" -version = "0.3.102" +name = "wasm-streams" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" dependencies = [ "js-sys", "wasm-bindgen", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 06289e5a46f..5842ed5ba9b 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -1,9 +1,8 @@ [workspace] members = [ "crates/core", - "crates/providers", - "crates/python-bridge", "crates/ai-gateway", + "crates/python-bridge", ] resolver = "2" @@ -14,15 +13,18 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } -litellm-providers = { path = "crates/providers" } +litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" pyo3 = "0.23.5" +pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] } rand = "0.8" -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } 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 = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } 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"] } +base64 = "0.22" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 15ad1855420..1646c90ad76 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -7,6 +7,16 @@ 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 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 index 79bdc4bdb26..4055be36785 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -5,22 +5,39 @@ 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 -litellm-providers.workspace = true -axum = { workspace = true, features = ["ws"] } +# 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 -tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time"] } -serde.workspace = true serde_json.workspace = true -subtle.workspace = true +base64.workspace = true +axum = { workspace = true, features = ["ws"], optional = true } +serde.workspace = 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: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/README.md b/litellm-rust/crates/ai-gateway/README.md index 3662ce2584a..f913beff6d5 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -4,9 +4,22 @@ 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. @@ -53,6 +66,7 @@ overridden at deploy time (e.g. a Render secret file mounted at the same path). | `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. @@ -71,6 +85,18 @@ 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 diff --git a/litellm-rust/crates/ai-gateway/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs index e2dd51f656d..438a0513057 100644 --- a/litellm-rust/crates/ai-gateway/src/auth/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/auth/mod.rs @@ -12,10 +12,29 @@ 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 @@ -52,3 +71,23 @@ impl FromRequestParts for RequireMasterKey { } } } + +#[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..109b648f5db --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -0,0 +1,30 @@ +//! 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. +#[cfg(feature = "server")] +pub(crate) const DEFAULT_PROVIDER: &str = "openai"; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/README.md b/litellm-rust/crates/ai-gateway/src/integrations/README.md new file mode 100644 index 00000000000..16a162dac57 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/README.md @@ -0,0 +1,127 @@ +# LiteLLM Rust integrations + +This directory contains Rust-native equivalents of LiteLLM integration hooks. +The first supported surfaces are terminal custom loggers and pre/during-call +custom guardrails. + +## File layout + +Every integration is a folder: + +- `mod.rs` contains the implementation, trait, runner, or adapter +- `types.rs` contains the integration-local request, response, error, and future + types + +Do not add new flat integration files such as `custom_logger.rs`. Shared wire +contracts that are used by multiple integrations can stay in +`integrations/types.rs`. + +Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`. +Call-type modules, such as OCR, adapt their request and response shapes into +that generic lifecycle runner. + +## CustomLogger + +Implement `CustomLogger` when Rust code needs to observe terminal success or +failure events. Method names intentionally match Python `CustomLogger` names. + +```rust +use litellm_ai_gateway::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, +}; + +struct RecordingLogger; + +impl CustomLogger for RecordingLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let model = &model_call_details.model; + let provider = &model_call_details.custom_llm_provider; + let call_type = model_call_details.call_type.to_string(); + let request_id = model_call_details.request_id.as_deref(); + let response_object = &response_obj.object; + let duration = timing.end_time - timing.start_time; + let standard_payload = model_call_details.standard_logging_payload.as_ref(); + + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let error = model_call_details.failure_error.as_ref(); + let response_object = response_obj.map(|value| value.object.as_str()); + let duration = timing.end_time - timing.start_time; + + Ok(()) + }) + } +} +``` + +Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The +runner is a no-op when no loggers are configured, which is the expected fast +path for requests without callbacks. + +## CustomGuardrail + +Implement `CustomGuardrail` when Rust code needs to run pre-call or native +during-call checks. Method names intentionally match Python `CustomGuardrail` +entrypoints inherited from Python `CustomLogger`. + +```rust +use litellm_ai_gateway::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; + +struct BlocklistedPromptGuardrail; + +impl CustomGuardrail for BlocklistedPromptGuardrail { + fn guardrail_name(&self) -> &str { + "blocklisted-prompt" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &[GuardrailEventHook::PreCall] + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + if request.data.to_string().contains("blocked phrase") { + return Ok(GuardrailDecision::Block( + litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked( + "blocked phrase detected", + ), + )); + } + Ok(GuardrailDecision::Allow(request)) + }) + } +} +``` + +Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and +`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A +`GuardrailDecision::Mask` continues with modified request data. +`GuardrailDecision::Block` short-circuits the provider call. + +## Current boundary + +These are Rust-only primitives. Python callback and guardrail adapters are a +separate layer that should implement these Rust traits instead of changing the +runner interfaces. diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs new file mode 100644 index 00000000000..e5d4ce3a708 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs @@ -0,0 +1,468 @@ +//! Rust mirror of Python `CustomGuardrail` entrypoints used by the proxy. +//! +//! This module is intentionally Rust-only: Python/PyO3 adapters are a later +//! layer that should implement this trait rather than changing the runner. + +use std::future::Future; +use std::sync::Arc; + +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; + +pub mod types; + +pub use types::{ + GuardrailContext, GuardrailDecision, GuardrailDispatchReport, GuardrailError, + GuardrailEventHook, GuardrailFuture, GuardrailRequest, +}; + +pub trait CustomGuardrail: Send + Sync { + fn guardrail_name(&self) -> &str; + + fn supported_event_hooks(&self) -> &[GuardrailEventHook]; + + /// Python 1:1 name: `async_pre_call_hook(user_api_key_dict, cache, data, call_type)`. + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { Ok(GuardrailDecision::Allow(request)) }) + } + + /// Python 1:1 name: `async_moderation_hook(data, user_api_key_dict, call_type)`. + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { Ok(GuardrailDecision::Allow(request)) }) + } +} + +pub struct CustomGuardrailRunner { + guardrails: Vec>, +} + +impl CustomGuardrailRunner { + pub fn new(guardrails: Vec>) -> Self { + Self { guardrails } + } + + pub fn is_empty(&self) -> bool { + self.guardrails.is_empty() + } + + pub async fn run_pre_call( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + self.run_hook(GuardrailEventHook::PreCall, context, request) + .await + } + + pub async fn run_during_call( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + self.run_hook(GuardrailEventHook::DuringCall, context, request) + .await + } + + pub async fn run_before_provider( + &self, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + request: GuardrailRequest, + provider: F, + ) -> Result + where + F: FnOnce(GuardrailRequest) -> Fut, + Fut: Future>, + { + let (request, _) = self.run_hook(event_hook, context, request).await?; + provider(request).await + } + + pub async fn run_pre_call_with_failure_logging( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + logger_runner: &CustomLoggerRunner, + model_call_details: &ModelCallDetails, + timing: CallbackTiming, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + match self.run_pre_call(context, request).await { + Ok(result) => Ok(result), + Err(error) => { + let failure_details = model_call_details.clone().with_failure_error(LoggingError { + message: error.message.clone(), + kind: error.kind.clone(), + }); + let response_obj = CallbackValue::new( + "guardrail_error", + serde_json::json!({ + "message": error.message, + "kind": error.kind, + }), + ); + logger_runner + .async_log_failure_event(&failure_details, Some(&response_obj), timing) + .await; + Err(error) + } + } + } + + async fn run_hook( + &self, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + mut request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + if self.guardrails.is_empty() { + return Ok((request, GuardrailDispatchReport::default())); + } + + let mut report = GuardrailDispatchReport::default(); + for guardrail in &self.guardrails { + if !self.should_run(guardrail.as_ref(), event_hook, context) { + continue; + } + + report.invoked += 1; + let decision = match event_hook { + GuardrailEventHook::PreCall => { + guardrail + .async_pre_call_hook(context, request.clone()) + .await? + } + GuardrailEventHook::DuringCall => { + guardrail + .async_moderation_hook(context, request.clone()) + .await? + } + }; + match decision.into_request() { + Ok(next_request) => request = next_request, + Err(error) => return Err(error), + } + } + + Ok((request, report)) + } + + fn should_run( + &self, + guardrail: &dyn CustomGuardrail, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + ) -> bool { + let supports_hook = guardrail.supported_event_hooks().contains(&event_hook); + let selected = context.selected_guardrails.is_empty() + || context + .selected_guardrails + .iter() + .any(|name| name == guardrail.guardrail_name()); + supports_hook && selected + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::custom_logger::{CallType, CallbackValue, CustomLogger, LogFuture}; + use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + use serde_json::json; + use std::sync::Mutex; + + #[derive(Clone)] + enum TestDecision { + Allow, + Mask, + Block, + } + + struct RecordingCustomGuardrail { + name: String, + hooks: Vec, + decision: TestDecision, + calls: Mutex>, + } + + impl RecordingCustomGuardrail { + fn new(name: &str, hooks: Vec, decision: TestDecision) -> Self { + Self { + name: name.to_string(), + hooks, + decision, + calls: Mutex::new(Vec::new()), + } + } + + fn calls(&self) -> Vec<&'static str> { + self.calls.lock().unwrap().clone() + } + + fn decision(&self, mut request: GuardrailRequest) -> GuardrailDecision { + match self.decision { + TestDecision::Allow => GuardrailDecision::Allow(request), + TestDecision::Mask => { + request.data["masked"] = json!(true); + GuardrailDecision::Mask(request) + } + TestDecision::Block => { + GuardrailDecision::Block(GuardrailError::blocked("blocked by guardrail")) + } + } + } + } + + impl CustomGuardrail for RecordingCustomGuardrail { + fn guardrail_name(&self) -> &str { + &self.name + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push("async_pre_call_hook"); + Ok(self.decision(request)) + }) + } + + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push("async_moderation_hook"); + Ok(self.decision(request)) + }) + } + } + + #[tokio::test] + async fn pre_call_dispatches_to_async_pre_call_hook() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "pre", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]); + let context = + GuardrailContext::new(CallType::Ocr).with_selected_guardrails(vec!["pre".to_string()]); + let request = GuardrailRequest::new(json!({"messages": ["hello"]})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("guardrail allows request"); + + assert_eq!(report.invoked, 1); + assert_eq!(result.data["messages"], json!(["hello"])); + assert_eq!(guardrail.calls(), vec!["async_pre_call_hook"]); + } + + #[tokio::test] + async fn during_call_dispatches_to_async_moderation_hook() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "during", + vec![GuardrailEventHook::DuringCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]); + let context = GuardrailContext::new(CallType::Completion) + .with_selected_guardrails(vec!["during".to_string()]); + let request = GuardrailRequest::new(json!({"prompt": "hello"})); + + let (_result, report) = runner + .run_during_call(&context, request) + .await + .expect("guardrail allows request"); + + assert_eq!(report.invoked, 1); + assert_eq!(guardrail.calls(), vec!["async_moderation_hook"]); + } + + #[tokio::test] + async fn mask_decision_continues_with_updated_request() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "masker", + vec![GuardrailEventHook::PreCall], + TestDecision::Mask, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail]); + let context = GuardrailContext::new(CallType::Ocr); + let request = GuardrailRequest::new(json!({"document": "secret"})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("mask continues"); + + assert_eq!(report.invoked, 1); + assert_eq!(result.data["masked"], json!(true)); + } + + #[tokio::test] + async fn block_decision_short_circuits_and_logs_failure() { + struct RecordingFailureLogger { + errors: Mutex>, + } + + impl CustomLogger for RecordingFailureLogger { + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.errors.lock().unwrap().push( + model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()) + .unwrap_or_default(), + ); + Ok(()) + }) + } + } + + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "blocker", + vec![GuardrailEventHook::PreCall], + TestDecision::Block, + )); + let guardrail_runner = CustomGuardrailRunner::new(vec![guardrail]); + let logger = Arc::new(RecordingFailureLogger { + errors: Mutex::new(Vec::new()), + }); + let logger_runner = CustomLoggerRunner::new(vec![logger.clone()]); + let context = GuardrailContext::new(CallType::Ocr); + let details = ModelCallDetails::from_standard_logging_payload(StandardLoggingPayload { + id: "req_ocr".to_string(), + litellm_call_id: "req_ocr".to_string(), + call_type: "ocr".to_string(), + model: "mistral-ocr-latest".to_string(), + custom_llm_provider: "mistral".to_string(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: 1.0, + end_time: 1.0, + stream: false, + metadata: StandardLoggingMetadata::default(), + messages: None, + }); + + let err = guardrail_runner + .run_pre_call_with_failure_logging( + &context, + GuardrailRequest::new(json!({"document": "bad"})), + &logger_runner, + &details, + CallbackTiming::new(1.0, 2.0), + ) + .await + .expect_err("guardrail blocks request"); + + assert_eq!(err.kind, "GuardrailBlocked"); + assert_eq!( + logger.errors.lock().unwrap().as_slice(), + ["GuardrailBlocked"] + ); + } + + #[tokio::test] + async fn block_decision_short_circuits_later_guardrails_and_provider_work() { + let blocking_guardrail = Arc::new(RecordingCustomGuardrail::new( + "blocker", + vec![GuardrailEventHook::PreCall], + TestDecision::Block, + )); + let later_guardrail = Arc::new(RecordingCustomGuardrail::new( + "later", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = + CustomGuardrailRunner::new(vec![blocking_guardrail.clone(), later_guardrail.clone()]); + let provider_called = Arc::new(Mutex::new(false)); + let provider_called_for_closure = provider_called.clone(); + + let result = runner + .run_before_provider( + GuardrailEventHook::PreCall, + &GuardrailContext::new(CallType::Completion), + GuardrailRequest::new(json!({"prompt": "blocked"})), + move |_request| async move { + *provider_called_for_closure.lock().unwrap() = true; + Ok("provider response") + }, + ) + .await; + + assert!(result.is_err()); + assert_eq!(blocking_guardrail.calls(), vec!["async_pre_call_hook"]); + assert_eq!(later_guardrail.calls(), Vec::<&'static str>::new()); + assert!(!*provider_called.lock().unwrap()); + } + + #[tokio::test] + async fn run_before_provider_returns_provider_guardrail_error_directly() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "allow", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail]); + + let result = runner + .run_before_provider( + GuardrailEventHook::PreCall, + &GuardrailContext::new(CallType::Completion), + GuardrailRequest::new(json!({"prompt": "allowed"})), + |_request| async move { + Err::<&'static str, GuardrailError>(GuardrailError::blocked( + "provider-side guardrail error", + )) + }, + ) + .await; + + let err = result.expect_err("provider error is returned directly"); + assert_eq!(err.kind, "GuardrailBlocked"); + assert_eq!(err.message, "provider-side guardrail error"); + } + + #[tokio::test] + async fn no_guardrails_fast_path_dispatches_nothing() { + let runner = CustomGuardrailRunner::new(Vec::new()); + let context = GuardrailContext::new(CallType::Ocr); + let request = GuardrailRequest::new(json!({"document": "ok"})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("no guardrails allow request"); + + assert!(runner.is_empty()); + assert_eq!(report, GuardrailDispatchReport::default()); + assert_eq!(result.data["document"], json!("ok")); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs new file mode 100644 index 00000000000..825e56cc0d7 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs @@ -0,0 +1,110 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use serde_json::Value; + +use crate::integrations::custom_logger::CallType; + +pub type GuardrailFuture<'a> = + Pin> + Send + 'a>>; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GuardrailEventHook { + PreCall, + DuringCall, +} + +impl GuardrailEventHook { + pub fn as_str(&self) -> &'static str { + match self { + Self::PreCall => "pre_call", + Self::DuringCall => "during_call", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GuardrailError { + pub message: String, + pub kind: String, +} + +impl GuardrailError { + pub fn blocked(message: impl Into) -> Self { + Self { + message: message.into(), + kind: "GuardrailBlocked".to_string(), + } + } +} + +impl std::fmt::Display for GuardrailError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.kind, self.message) + } +} + +impl std::error::Error for GuardrailError {} + +#[derive(Clone, Debug)] +pub struct GuardrailContext { + pub call_type: CallType, + pub selected_guardrails: Vec, + pub metadata: HashMap, + pub user_api_key_hash: Option, + pub user_api_key_user_id: Option, + pub user_api_key_team_id: Option, + pub trace_parent: Option, +} + +impl GuardrailContext { + pub fn new(call_type: CallType) -> Self { + Self { + call_type, + selected_guardrails: Vec::new(), + metadata: HashMap::new(), + user_api_key_hash: None, + user_api_key_user_id: None, + user_api_key_team_id: None, + trace_parent: None, + } + } + + pub fn with_selected_guardrails(mut self, selected_guardrails: Vec) -> Self { + self.selected_guardrails = selected_guardrails; + self + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct GuardrailRequest { + pub data: Value, +} + +impl GuardrailRequest { + pub fn new(data: Value) -> Self { + Self { data } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum GuardrailDecision { + Allow(GuardrailRequest), + Mask(GuardrailRequest), + Block(GuardrailError), +} + +impl GuardrailDecision { + pub(super) fn into_request(self) -> Result { + match self { + Self::Allow(request) | Self::Mask(request) => Ok(request), + Self::Block(error) => Err(error), + } + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct GuardrailDispatchReport { + pub invoked: usize, +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs new file mode 100644 index 00000000000..792717dacfc --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs @@ -0,0 +1,317 @@ +//! The `CustomLogger` trait — the Rust mirror of Python +//! `litellm/integrations/custom_logger.py::CustomLogger`. +//! +//! The Python-named async terminal methods are the public Rust callback shape. + +use std::sync::Arc; + +pub mod types; + +pub use types::{ + CallType, CallbackDispatchReport, CallbackTiming, CallbackValue, LogError, LogFuture, + LoggingError, ModelCallDetails, +}; + +pub trait CustomLogger: Send + Sync { + /// Python 1:1 name: `async_log_success_event(model_call_details, response_obj, start_time, end_time)`. + fn async_log_success_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Ok(()) }) + } + + /// Python 1:1 name: `async_log_failure_event(model_call_details, response_obj, start_time, end_time)`. + fn async_log_failure_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Ok(()) }) + } +} + +pub struct CustomLoggerRunner { + loggers: Vec>, +} + +impl CustomLoggerRunner { + pub fn new(loggers: Vec>) -> Self { + Self { loggers } + } + + pub fn is_empty(&self) -> bool { + self.loggers.is_empty() + } + + pub async fn async_log_success_event( + &self, + model_call_details: &ModelCallDetails, + response_obj: &CallbackValue, + timing: CallbackTiming, + ) -> CallbackDispatchReport { + if self.loggers.is_empty() { + return CallbackDispatchReport::default(); + } + + let mut report = CallbackDispatchReport::default(); + for logger in &self.loggers { + report.invoked += 1; + if let Err(err) = logger + .async_log_success_event(model_call_details, response_obj, timing) + .await + { + report.dropped += 1; + eprintln!("litellm-ai-gateway: async_log_success_event dropped: {err}"); + } + } + report + } + + pub async fn async_log_failure_event( + &self, + model_call_details: &ModelCallDetails, + response_obj: Option<&CallbackValue>, + timing: CallbackTiming, + ) -> CallbackDispatchReport { + if self.loggers.is_empty() { + return CallbackDispatchReport::default(); + } + + let mut report = CallbackDispatchReport::default(); + for logger in &self.loggers { + report.invoked += 1; + if let Err(err) = logger + .async_log_failure_event(model_call_details, response_obj, timing) + .await + { + report.dropped += 1; + eprintln!("litellm-ai-gateway: async_log_failure_event dropped: {err}"); + } + } + report + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + use serde_json::json; + use std::sync::Mutex; + + #[derive(Clone, Debug, PartialEq)] + struct RecordedEvent { + hook: &'static str, + model: String, + provider: String, + call_type: String, + request_id: Option, + litellm_call_id: Option, + user_id: Option, + response_object: Option, + error_kind: Option, + start_time: f64, + end_time: f64, + standard_logging_model: Option, + } + + #[derive(Default)] + struct RecordingCustomLogger { + events: Mutex>, + } + + impl RecordingCustomLogger { + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } + } + + impl CustomLogger for RecordingCustomLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedEvent { + hook: "async_log_success_event", + model: model_call_details.model.clone(), + provider: model_call_details.custom_llm_provider.clone(), + call_type: model_call_details.call_type.to_string(), + request_id: model_call_details.request_id.clone(), + litellm_call_id: model_call_details.litellm_call_id.clone(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: Some(response_obj.object.clone()), + error_kind: None, + start_time: timing.start_time, + end_time: timing.end_time, + standard_logging_model: model_call_details + .standard_logging_payload + .as_ref() + .map(|payload| payload.model.clone()), + }); + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedEvent { + hook: "async_log_failure_event", + model: model_call_details.model.clone(), + provider: model_call_details.custom_llm_provider.clone(), + call_type: model_call_details.call_type.to_string(), + request_id: model_call_details.request_id.clone(), + litellm_call_id: model_call_details.litellm_call_id.clone(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: response_obj.map(|value| value.object.clone()), + error_kind: model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()), + start_time: timing.start_time, + end_time: timing.end_time, + standard_logging_model: model_call_details + .standard_logging_payload + .as_ref() + .map(|payload| payload.model.clone()), + }); + Ok(()) + }) + } + } + + fn payload(call_type: &str, model: &str, provider: &str) -> StandardLoggingPayload { + StandardLoggingPayload { + id: format!("req_{call_type}"), + litellm_call_id: format!("call_{call_type}"), + call_type: call_type.to_string(), + model: model.to_string(), + custom_llm_provider: provider.to_string(), + response_cost: 0.25, + prompt_tokens: 3, + completion_tokens: 4, + total_tokens: 7, + start_time: 10.0, + end_time: 11.5, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: Some("hash".to_string()), + user_api_key_user_id: Some("user".to_string()), + user_api_key_team_id: Some("team".to_string()), + ..Default::default() + }, + messages: Some(json!([{"role": "user", "content": "read this"}])), + } + } + + #[tokio::test] + async fn rust_custom_logger_reads_success_payload_for_ocr() { + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let details = ModelCallDetails::from_standard_logging_payload(payload( + "ocr", + "mistral-ocr-latest", + "mistral", + )); + let response = CallbackValue::new("ocr", json!({"pages": [{"markdown": "ok"}]})); + let report = runner + .async_log_success_event(&details, &response, CallbackTiming::new(10.0, 11.5)) + .await; + + assert_eq!(report.invoked, 1); + assert_eq!(report.dropped, 0); + assert_eq!( + logger.events(), + vec![RecordedEvent { + hook: "async_log_success_event", + model: "mistral-ocr-latest".to_string(), + provider: "mistral".to_string(), + call_type: "ocr".to_string(), + request_id: Some("req_ocr".to_string()), + litellm_call_id: Some("call_ocr".to_string()), + user_id: Some("user".to_string()), + response_object: Some("ocr".to_string()), + error_kind: None, + start_time: 10.0, + end_time: 11.5, + standard_logging_model: Some("mistral-ocr-latest".to_string()), + }] + ); + } + + #[tokio::test] + async fn rust_custom_logger_reads_failure_payload_for_non_ocr_call_type() { + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let details = ModelCallDetails::from_standard_logging_payload(payload( + "acompletion", + "gpt-4.1-mini", + "openai", + )) + .with_failure_error(LoggingError { + message: "provider failed".to_string(), + kind: "ProviderError".to_string(), + }); + let response = CallbackValue::new("error", json!({"message": "provider failed"})); + let report = runner + .async_log_failure_event(&details, Some(&response), CallbackTiming::new(2.0, 3.0)) + .await; + + assert_eq!(report.invoked, 1); + assert_eq!(report.dropped, 0); + assert_eq!( + logger.events(), + vec![RecordedEvent { + hook: "async_log_failure_event", + model: "gpt-4.1-mini".to_string(), + provider: "openai".to_string(), + call_type: "acompletion".to_string(), + request_id: Some("req_acompletion".to_string()), + litellm_call_id: Some("call_acompletion".to_string()), + user_id: Some("user".to_string()), + response_object: Some("error".to_string()), + error_kind: Some("ProviderError".to_string()), + start_time: 2.0, + end_time: 3.0, + standard_logging_model: Some("gpt-4.1-mini".to_string()), + }] + ); + } + + #[tokio::test] + async fn no_callback_fast_path_dispatches_nothing() { + let runner = CustomLoggerRunner::new(Vec::new()); + let details = ModelCallDetails::new("mistral-ocr-latest", "mistral", CallType::Ocr); + let response = CallbackValue::new("ocr", json!({})); + + let report = runner + .async_log_success_event(&details, &response, CallbackTiming::new(1.0, 1.5)) + .await; + + assert!(runner.is_empty()); + assert_eq!(report, CallbackDispatchReport::default()); + } + + #[test] + fn with_standard_logging_payload_keeps_top_level_fields_in_sync() { + let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion) + .with_standard_logging_payload(payload("ocr", "mistral-ocr-latest", "mistral")); + + assert_eq!(details.model, "mistral-ocr-latest"); + assert_eq!(details.custom_llm_provider, "mistral"); + assert_eq!(details.call_type, CallType::Ocr); + assert_eq!(details.request_id, Some("req_ocr".to_string())); + assert_eq!(details.litellm_call_id, Some("call_ocr".to_string())); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs new file mode 100644 index 00000000000..ba7d67bd46e --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs @@ -0,0 +1,194 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use serde_json::Value; + +use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + +pub type LogFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CallbackDispatchReport { + pub invoked: usize, + pub dropped: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CallType { + Ocr, + Realtime, + Completion, + Acompletion, + ChatCompletion, + Other(String), +} + +impl CallType { + pub fn as_str(&self) -> &str { + match self { + Self::Ocr => "ocr", + Self::Realtime => "realtime", + Self::Completion => "completion", + Self::Acompletion => "acompletion", + Self::ChatCompletion => "chat_completion", + Self::Other(value) => value.as_str(), + } + } +} + +impl From<&str> for CallType { + fn from(value: &str) -> Self { + match value { + "ocr" => Self::Ocr, + "realtime" => Self::Realtime, + "completion" => Self::Completion, + "acompletion" => Self::Acompletion, + "chat_completion" => Self::ChatCompletion, + other => Self::Other(other.to_string()), + } + } +} + +impl std::fmt::Display for CallType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallbackTiming { + pub start_time: f64, + pub end_time: f64, +} + +impl CallbackTiming { + pub fn new(start_time: f64, end_time: f64) -> Self { + Self { + start_time, + end_time, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CallbackValue { + pub object: String, + pub value: Value, +} + +impl CallbackValue { + pub fn new(object: impl Into, value: Value) -> Self { + Self { + object: object.into(), + value, + } + } +} + +#[derive(Clone, Debug)] +pub struct ModelCallDetails { + pub model: String, + pub custom_llm_provider: String, + pub call_type: CallType, + pub metadata: StandardLoggingMetadata, + pub extra_metadata: HashMap, + pub request_id: Option, + pub litellm_call_id: Option, + pub response_cost: Option, + pub standard_logging_payload: Option, + pub failure_error: Option, +} + +impl ModelCallDetails { + pub fn new( + model: impl Into, + custom_llm_provider: impl Into, + call_type: CallType, + ) -> Self { + Self { + model: model.into(), + custom_llm_provider: custom_llm_provider.into(), + call_type, + metadata: StandardLoggingMetadata::default(), + extra_metadata: HashMap::new(), + request_id: None, + litellm_call_id: None, + response_cost: None, + standard_logging_payload: None, + failure_error: None, + } + } + + pub fn from_standard_logging_payload(payload: StandardLoggingPayload) -> Self { + let request_id = Some(payload.id.clone()); + let litellm_call_id = Some(payload.litellm_call_id.clone()); + let response_cost = Some(payload.response_cost); + let metadata = payload.metadata.clone(); + Self { + model: payload.model.clone(), + custom_llm_provider: payload.custom_llm_provider.clone(), + call_type: CallType::from(payload.call_type.as_str()), + metadata, + extra_metadata: HashMap::new(), + request_id, + litellm_call_id, + response_cost, + standard_logging_payload: Some(payload), + failure_error: None, + } + } + + pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self { + self.model = payload.model.clone(); + self.custom_llm_provider = payload.custom_llm_provider.clone(); + self.call_type = CallType::from(payload.call_type.as_str()); + self.request_id = Some(payload.id.clone()); + self.litellm_call_id = Some(payload.litellm_call_id.clone()); + self.response_cost = Some(payload.response_cost); + self.metadata = payload.metadata.clone(); + self.standard_logging_payload = Some(payload); + self + } + + pub fn with_failure_error(mut self, error: LoggingError) -> Self { + self.failure_error = Some(error); + self + } +} + +#[derive(Clone, Debug)] +pub struct LoggingError { + pub message: String, + pub kind: String, +} + +#[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 {} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs new file mode 100644 index 00000000000..3dad18cb7a3 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs @@ -0,0 +1,197 @@ +//! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's +//! `/v1/rust_control_plane/logs` endpoint. +//! +//! The callback path is non-blocking: `async_log_success_event` / +//! `async_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_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH}; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, LoggingError, + ModelCallDetails, +}; +use types::{CallbackLogsRequest, EgressTunables, LogRecord}; + +pub mod types; + +/// 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 async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + if let Some(payload) = &model_call_details.standard_logging_payload { + self.enqueue(LogRecord { + status: "success".to_string(), + payload: payload.clone(), + error: None, + })?; + } + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + if let Some(payload) = &model_call_details.standard_logging_payload { + let fallback_error; + let error = match &model_call_details.failure_error { + Some(error) => error, + None => { + fallback_error = LoggingError { + message: "callback failure event".to_string(), + kind: "CallbackFailure".to_string(), + }; + &fallback_error + } + }; + self.enqueue(LogRecord { + status: "failure".to_string(), + payload: payload.clone(), + error: Some(format!("{}: {}", error.kind, error.message)), + })?; + } + Ok(()) + }) + } +} + +/// 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/litellm_python_proxy_api/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs new file mode 100644 index 00000000000..481a437747f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs @@ -0,0 +1,72 @@ +use std::time::Duration; + +use serde::Serialize; + +use crate::constants::{ + DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE, +}; +use crate::integrations::types::StandardLoggingPayload; + +#[derive(Serialize)] +pub struct CallbackLogsRequest { + pub records: Vec, +} + +#[derive(Serialize)] +pub struct CallbackLogRecord { + pub status: String, + pub standard_logging_payload: StandardLoggingPayload, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[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, + } + } +} + +pub(super) struct EgressTunables { + pub channel_capacity: usize, + pub max_batch_size: usize, + pub flush_interval: Duration, +} + +impl EgressTunables { + pub 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, + )), + } + } +} + +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) +} 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..c62f1821ef8 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs @@ -0,0 +1,12 @@ +//! Pure-Rust logging integrations. Names map 1:1 to Python +//! `litellm/integrations/`: +//! - [`custom_guardrail::CustomGuardrail`] — the guardrail callback trait +//! - [`custom_logger::CustomLogger`] — the callback trait +//! - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events +//! to the Python proxy's `/v1/rust_control_plane/logs` endpoint +//! - [`types`] — the typed `StandardLoggingPayload` wire contract + +pub mod custom_guardrail; +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..34dce93d8e0 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/types.rs @@ -0,0 +1,83 @@ +//! 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, +} + +/// 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>, +} diff --git a/litellm-rust/crates/providers/src/lib.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs similarity index 62% rename from litellm-rust/crates/providers/src/lib.rs rename to litellm-rust/crates/ai-gateway/src/io/mod.rs index 40e18961f43..3b566027646 100644 --- a/litellm-rust/crates/providers/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,5 +1,3 @@ -pub mod mistral; pub mod ocr; -pub mod openai; 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..55e02839c4e --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -0,0 +1 @@ +pub use crate::ocr::{ocr, OcrRequest}; diff --git a/litellm-rust/crates/providers/src/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs similarity index 90% rename from litellm-rust/crates/providers/src/realtime.rs rename to litellm-rust/crates/ai-gateway/src/io/realtime.rs index 398158f6dba..40a38c1579a 100644 --- a/litellm-rust/crates/providers/src/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -1,13 +1,13 @@ //! End-to-end OpenAI realtime invocation. //! -//! The host-facing entry point, mirroring `providers::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. +//! The host-facing entry point opens the WebSocket to OpenAI, then splices 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::realtime_pool`]) can pre-establish an upstream, +//! 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. @@ -26,7 +26,7 @@ use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; -use crate::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; +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"; @@ -126,6 +126,9 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult( model: &str, @@ -133,6 +136,7 @@ pub(crate) async fn splice( mut upstream_rx: UpstreamRx, prelude: Option, idle_timeout: Option, + mut observe: impl FnMut(&RealtimeEvent) + Send, mut client_in: In, mut client_out: Out, ) -> CoreResult<()> @@ -165,6 +169,10 @@ where // 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()))?; @@ -181,6 +189,7 @@ where 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) @@ -207,11 +216,13 @@ where /// 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<()> @@ -229,19 +240,22 @@ where upstream_rx, None, idle_timeout, + observe, client_in, client_out, ) .await } -/// Splice a pre-warmed upstream (taken from [`crate::realtime_pool`]) to the +/// 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::realtime_pool::WarmHandoff, + handoff: crate::io::realtime_pool::WarmHandoff, idle_timeout: Option, + observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, ) -> CoreResult<()> @@ -256,6 +270,7 @@ where handoff.rx, Some(handoff.session_created), idle_timeout, + observe, client_in, client_out, ) @@ -281,7 +296,7 @@ mod tests { /// 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-providers realtime_invokes_openai -- --ignored --nocapture` + /// `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() { @@ -303,6 +318,7 @@ mod tests { Some(&key_owned), None, None, + |_| {}, client_in, client_out, ) diff --git a/litellm-rust/crates/providers/src/realtime_pool.rs b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs similarity index 99% rename from litellm-rust/crates/providers/src/realtime_pool.rs rename to litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs index 1b1fc8112c5..bf8041f31d7 100644 --- a/litellm-rust/crates/providers/src/realtime_pool.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -6,7 +6,7 @@ //! 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 stays in `providers` (axum-free) next to the dial/splice it +//! 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 @@ -31,7 +31,7 @@ use futures_util::StreamExt; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::CoreResult; -use crate::realtime::{ +use crate::io::realtime::{ dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs, }; 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..d8ef7bb5ba1 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -0,0 +1,37 @@ +//! 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: +//! +//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks, +//! and provider I/O. Always available — no feature required. +//! - [`io`]: compatibility exports and realtime WebSocket splice helpers. +//! - 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; +pub mod ocr; + +/// 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. +mod constants; +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 index 71e4a6836ad..f9ce97801d3 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -1,22 +1,25 @@ //! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router. //! //! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment -//! (simple-shuffle) → `providers::realtime::realtime()` invokes OpenAI. The +//! (simple-shuffle) → `io::realtime::realtime()` invokes OpenAI. The //! server owns transport + config; routing lives in the `router` crate. - -mod auth; -mod gil; -#[cfg(feature = "python-config")] -mod python; -mod routes; -mod state; +//! +//! 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_providers::realtime_pool::{upstream_key, PoolConfig, RealtimePool}; -use crate::state::AppState; +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`). @@ -38,6 +41,12 @@ async fn main() { ); } + // 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 @@ -61,6 +70,7 @@ async fn main() { let state = AppState { router, master_key, + loggers: Arc::new(loggers), realtime_pool, }; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/client.rs b/litellm-rust/crates/ai-gateway/src/ocr/client.rs new file mode 100644 index 00000000000..79cc7816227 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/client.rs @@ -0,0 +1,14 @@ +use std::sync::OnceLock; +use std::time::Duration; + +const OCR_TIMEOUT_SECS: u64 = 600; + +pub(super) fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .build() + .expect("failed to build reqwest client") + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs new file mode 100644 index 00000000000..d4b4d9338e7 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -0,0 +1,447 @@ +use std::net::IpAddr; +use std::time::{Duration, Instant}; + +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use base64::Engine; +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::CoreResult; +use reqwest::Url; +use serde_json::{Map, Value}; + +use litellm_core::providers::azure_ai::ocr::transformation::{ + AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG, +}; +use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; +use litellm_core::providers::vertex_ai::ocr::transformation::{ + VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, +}; + +use super::client::http_client; + +const ERROR_BODY_MAX_CHARS: usize = 256; +const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; +const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0; +const MAX_SAFE_FETCH_REDIRECTS: usize = 10; + +pub(super) 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)") +} + +pub(super) fn ocr_provider_config( + provider: &str, + model: &str, +) -> Option<&'static dyn OcrProviderConfig> { + match provider { + "mistral" => Some(&MISTRAL_OCR_CONFIG), + "azure_ai" if is_azure_document_intelligence_model(model) => { + Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) + } + "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), + "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), + "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), + _ => None, + } +} + +fn is_azure_document_intelligence_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.contains("doc-intelligence") || model.contains("documentintelligence") +} + +pub(super) fn string_headers( + extra_headers: Option>, +) -> CoreResult> { + extra_headers + .unwrap_or_default() + .into_iter() + .map(|(key, value)| { + value + .as_str() + .map(|value| (key.clone(), value.to_string())) + .ok_or_else(|| { + CoreError::InvalidRequest(format!( + "OCR extra_headers.{key} must be a string, got {}", + litellm_core::error::json_type_name(&value) + )) + }) + }) + .collect() +} + +pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { + headers + .iter() + .any(|(key, _)| key.eq_ignore_ascii_case(name)) +} + +fn document_url_field(document: &Value) -> CoreResult> { + let Some(object) = document.as_object() else { + return Ok(None); + }; + let Some(doc_type) = object.get("type").and_then(Value::as_str) else { + return Ok(None); + }; + let field = match doc_type { + "document_url" => "document_url", + "image_url" => "image_url", + _ => return Ok(None), + }; + let Some(url) = object.get(field).and_then(Value::as_str) else { + return Ok(None); + }; + Ok(Some((field, url))) +} + +fn is_url_requiring_fetch(url: &str) -> bool { + !url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://")) +} + +fn max_document_download_bytes() -> u64 { + let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB); + (max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64 +} + +fn is_blocked_ip(ip: IpAddr) -> bool { + match ip { + IpAddr::V4(ip) => { + ip.is_private() + || ip.is_loopback() + || ip.is_link_local() + || ip.is_broadcast() + || ip.is_multicast() + || ip.is_unspecified() + } + IpAddr::V6(ip) => { + let first_segment = ip.segments()[0]; + let is_unique_local = (first_segment & 0xfe00) == 0xfc00; + let is_link_local = (first_segment & 0xffc0) == 0xfe80; + ip.is_loopback() + || ip.is_unspecified() + || ip.is_multicast() + || is_unique_local + || is_link_local + || ip + .to_ipv4_mapped() + .or_else(|| ip.to_ipv4()) + .map(|v4| is_blocked_ip(IpAddr::V4(v4))) + .unwrap_or(false) + } + } +} + +fn blocked_url_error(url: &Url) -> CoreError { + CoreError::InvalidRequest(format!( + "OCR document URL rejected by SSRF protection: {url}" + )) +} + +async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { + if !matches!(url.scheme(), "http" | "https") { + return Err(blocked_url_error(url)); + } + + let host = url.host_str().ok_or_else(|| blocked_url_error(url))?; + if let Ok(ip) = host.parse::() { + if is_blocked_ip(ip) { + return Err(blocked_url_error(url)); + } + return Ok(()); + } + + let port = url + .port_or_known_default() + .ok_or_else(|| blocked_url_error(url))?; + let addresses = tokio::net::lookup_host((host, port)) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + let mut saw_address = false; + for address in addresses { + saw_address = true; + if is_blocked_ip(address.ip()) { + return Err(blocked_url_error(url)); + } + } + if !saw_address { + return Err(blocked_url_error(url)); + } + Ok(()) +} + +fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult { + let location = response + .headers() + .get(reqwest::header::LOCATION) + .and_then(|value| value.to_str().ok()) + .ok_or_else(|| { + CoreError::InvalidResponse("OCR document redirect missing Location header".to_string()) + })?; + url.join(location) + .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}"))) +} + +async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> { + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|err| CoreError::Network(err.to_string()))?; + let mut current_url = Url::parse(url) + .map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?; + + for _ in 0..MAX_SAFE_FETCH_REDIRECTS { + validate_safe_fetch_url(¤t_url).await?; + let response = client + .get(current_url.clone()) + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + if !response.status().is_redirection() { + return Ok((current_url, response)); + } + current_url = redirect_location(&response, ¤t_url)?; + } + + Err(CoreError::InvalidRequest( + "Too many redirects while fetching OCR document URL".to_string(), + )) +} + +fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> { + if max_bytes == 0 { + return Err(CoreError::InvalidRequest(format!( + "OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}" + ))); + } + if content_length > max_bytes { + let size_mb = content_length as f64 / (1024.0 * 1024.0); + let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0); + return Err(CoreError::InvalidRequest(format!( + "OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}" + ))); + } + Ok(()) +} + +async fn read_response_with_limit( + mut response: reqwest::Response, + url: &Url, +) -> CoreResult> { + let max_bytes = max_document_download_bytes(); + if let Some(content_length) = response.content_length() { + enforce_download_size(content_length, max_bytes, url)?; + } else { + enforce_download_size(0, max_bytes, url)?; + } + + let mut bytes = Vec::new(); + let mut bytes_downloaded: u64 = 0; + while let Some(chunk) = response + .chunk() + .await + .map_err(|err| CoreError::Network(err.to_string()))? + { + bytes_downloaded += chunk.len() as u64; + enforce_download_size(bytes_downloaded, max_bytes, url)?; + bytes.extend_from_slice(&chunk); + } + Ok(bytes) +} + +pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult { + let Some((field, url)) = document_url_field(&document)? else { + return Ok(document); + }; + if !is_url_requiring_fetch(url) { + return Ok(document); + } + + let (final_url, response) = safe_get_document_url(url).await?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&body), + }); + } + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("application/octet-stream") + .to_string(); + let bytes = read_response_with_limit(response, &final_url).await?; + let data_uri = format!( + "data:{content_type};base64,{}", + BASE64_STANDARD.encode(bytes) + ); + + let mut transformed = document + .as_object() + .cloned() + .ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?; + transformed.insert(field.to_string(), Value::String(data_uri)); + Ok(Value::Object(transformed)) +} + +fn same_origin(left: &str, right: &str) -> bool { + let Ok(left) = reqwest::Url::parse(left) else { + return false; + }; + let Ok(right) = reqwest::Url::parse(right) else { + return false; + }; + left.scheme() == right.scheme() + && left.host_str() == right.host_str() + && left.port_or_known_default() == right.port_or_known_default() +} + +fn retry_after_secs(response: &reqwest::Response) -> u64 { + response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .unwrap_or(2) +} + +fn operation_status(response_json: &Value) -> CoreResult<&str> { + let status = response_json + .get("status") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("status"))?; + match status { + "succeeded" => Ok("succeeded"), + "running" | "notStarted" => Ok("running"), + "failed" => { + let message = response_json + .get("error") + .and_then(|error| error.get("message")) + .and_then(Value::as_str) + .unwrap_or("Unknown error"); + Err(CoreError::InvalidResponse(format!( + "Azure Document Intelligence analysis failed: {message}" + ))) + } + other => Err(CoreError::InvalidResponse(format!( + "Unknown operation status: {other}" + ))), + } +} + +pub(super) async fn poll_document_intelligence( + operation_url: &str, + original_url: &str, + headers: &[(String, String)], + timeout: Option, +) -> CoreResult { + if !same_origin(operation_url, original_url) { + return Err(CoreError::InvalidResponse( + "Azure Document Intelligence: rejected cross-origin polling URL".to_string(), + )); + } + + let start = Instant::now(); + let timeout = timeout.unwrap_or(Duration::from_secs( + AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, + )); + loop { + if start.elapsed() > timeout { + return Err(CoreError::Network(format!( + "Azure Document Intelligence operation polling timed out after {} seconds", + timeout.as_secs() + ))); + } + + let mut request_builder = http_client().get(operation_url); + for (key, value) in headers { + if key.eq_ignore_ascii_case("ocp-apim-subscription-key") { + request_builder = request_builder.header(key, value); + } + } + let response = request_builder + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + let retry_after = retry_after_secs(&response); + let status = response.status(); + let text = response + .text() + .await + .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 Azure DI poll response JSON: {err}")) + })?; + if operation_status(&response_json)? == "succeeded" { + return Ok(response_json); + } + tokio::time::sleep(Duration::from_secs(retry_after)).await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn blocks_private_and_metadata_ips() { + assert!(is_blocked_ip("127.0.0.1".parse().unwrap())); + assert!(is_blocked_ip("10.0.0.1".parse().unwrap())); + assert!(is_blocked_ip("169.254.169.254".parse().unwrap())); + assert!(is_blocked_ip("::1".parse().unwrap())); + assert!(is_blocked_ip("fd00::1".parse().unwrap())); + assert!(is_blocked_ip("fe80::1".parse().unwrap())); + assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap())); + assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap())); + assert!(!is_blocked_ip("8.8.8.8".parse().unwrap())); + assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap())); + } + + #[tokio::test] + async fn convert_document_url_rejects_loopback_fetch() { + let error = convert_document_url_to_data_uri(json!({ + "type": "image_url", + "image_url": "http://127.0.0.1/image.png" + })) + .await + .unwrap_err(); + + assert!(matches!( + error, + CoreError::InvalidRequest(message) + if message.contains("SSRF protection") + )); + } + + #[tokio::test] + async fn convert_document_url_leaves_data_uri_untouched() { + let document = json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abcd" + }); + + let transformed = convert_document_url_to_data_uri(document.clone()) + .await + .unwrap(); + + assert_eq!(transformed, document); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs new file mode 100644 index 00000000000..4d93c2a25db --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -0,0 +1,71 @@ +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrResponseHandling; +use litellm_core::CoreResult; +use serde_json::Value; + +use super::client::http_client; +use super::common_utils::{poll_document_intelligence, truncate_error_body}; +use super::types::ProviderOcrRequest; + +pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { + let mut request_builder = http_client().post(&request.url).json(&request.body); + for (key, value) in &request.upstream_headers { + request_builder = request_builder.header(key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); + } + + let response = request_builder + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + let status = response.status(); + if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll + && status.as_u16() == 202 + { + let operation_url = response + .headers() + .get("operation-location") + .and_then(|value| value.to_str().ok()) + .map(str::to_string) + .ok_or_else(|| { + CoreError::InvalidResponse( + "Azure Document Intelligence returned 202 but no Operation-Location header found" + .to_string(), + ) + })?; + let response_json = poll_document_intelligence( + &operation_url, + &request.url, + &request.upstream_headers, + request.timeout, + ) + .await?; + return Ok(request + .config + .transform_ocr_response(&request.model, response_json)? + .into_json()); + } + + let text = response + .text() + .await + .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(request + .config + .transform_ocr_response(&request.model, response_json)? + .into_json()) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs new file mode 100644 index 00000000000..6be74ed2714 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -0,0 +1,329 @@ +use std::future::Future; +use std::pin::Pin; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrAuthStrategy; +use litellm_core::CoreResult; +use serde_json::{json, Map, Value}; + +use super::common_utils::{ + convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, +}; +use super::types::{PreparedOcrRequest, ProviderOcrRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; +use crate::integrations::types::{ + RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, +}; + +pub(crate) struct OcrLifecycleHooks { + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, +} + +type OcrFuture<'a, T> = Pin> + Send + 'a>>; +type OcrLogFuture<'a> = Pin + Send + 'a>>; + +impl OcrLifecycleHooks { + pub(crate) fn new( + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, + ) -> Self { + Self { + logger_runner, + guardrail_runner, + request_metadata, + } + } + + async fn run_pre_call_guardrails( + &self, + request: PreparedOcrRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + + let context = guardrail_context(&self.request_metadata); + let guardrail_request = GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": request.custom_llm_provider, + "document": request.document, + "optional_params": request.optional_params, + })); + let (guardrail_request, _) = self + .guardrail_runner + .run_pre_call(&context, guardrail_request) + .await + .map_err(guardrail_error_to_core_error)?; + let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?; + Ok(PreparedOcrRequest { + document, + optional_params, + ..request + }) + } + + async fn prepare_provider_request( + &self, + request: PreparedOcrRequest, + ) -> CoreResult { + let config = ocr_provider_config(&request.custom_llm_provider, &request.model) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); + let headers = string_headers(request.extra_headers)?; + let auth_strategy = config.auth_strategy(); + let api_key = (!has_header(&headers, auth_strategy.header_name())) + .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) + .transpose()?; + let url = config.complete_url( + request.api_base.as_deref(), + &request.model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_ocr_params(&request.optional_params); + let model = request.model.clone(); + let custom_llm_provider = request.custom_llm_provider.clone(); + let document = if config.requires_data_uri_document() { + convert_document_url_to_data_uri(request.document).await? + } else { + request.document + }; + let body = config + .transform_ocr_request(&request.model, document, filtered_params)? + .data; + let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); + let body = self + .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) + .await?; + Ok(ProviderOcrRequest { + model, + config, + url, + body, + upstream_headers, + timeout: request.timeout, + }) + } + + async fn run_during_call_guardrails( + &self, + model: &str, + custom_llm_provider: &str, + url: &str, + body: Value, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(body); + } + + let context = guardrail_context(&self.request_metadata); + let guardrail_request = GuardrailRequest::new(json!({ + "model": model, + "custom_llm_provider": custom_llm_provider, + "url": url, + "body": body, + })); + let (guardrail_request, _) = self + .guardrail_runner + .run_during_call(&context, guardrail_request) + .await + .map_err(guardrail_error_to_core_error)?; + parse_ocr_during_call_guardrail_request(guardrail_request) + } + + fn standard_logging_payload( + &self, + context: &CallLifecycleContext, + timing: &CallLifecycleTiming, + ) -> StandardLoggingPayload { + StandardLoggingPayload { + id: context.litellm_call_id.clone(), + litellm_call_id: context.litellm_call_id.clone(), + call_type: context.call_type.clone(), + model: context.model.clone(), + custom_llm_provider: context.custom_llm_provider.clone(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: timing.start_time, + end_time: timing.end_time, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: self.request_metadata.user_api_key_hash.clone(), + user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(), + user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), + ..Default::default() + }, + messages: None, + } + } +} + +impl CallLifecycleHooks for OcrLifecycleHooks { + type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; + type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>; + type SuccessFuture<'a> = OcrLogFuture<'a>; + type FailureFuture<'a> = OcrLogFuture<'a>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { self.run_pre_call_guardrails(request).await }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { self.prepare_provider_request(request).await }) + } + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Value, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let response_obj = CallbackValue::new("ocr", response.clone()); + self.logger_runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ), + &response_obj, + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let logging_error = LoggingError { + message: error.to_string(), + kind: core_error_kind(error).to_string(), + }; + let response_obj = CallbackValue::new( + "error", + json!({ + "message": logging_error.message, + "kind": logging_error.kind, + }), + ); + self.logger_runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ) + .with_failure_error(logging_error), + Some(&response_obj), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } +} + +fn upstream_headers( + headers: &[(String, String)], + auth_strategy: OcrAuthStrategy, + api_key: Option<&str>, +) -> Vec<(String, String)> { + api_key + .map(|api_key| match auth_strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), + }) + .into_iter() + .chain(headers.iter().cloned()) + .collect() +} + +fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { + GuardrailContext { + call_type: CallType::Ocr, + selected_guardrails: Vec::new(), + metadata: std::collections::HashMap::new(), + user_api_key_hash: metadata.user_api_key_hash.clone(), + user_api_key_user_id: metadata.user_api_key_user_id.clone(), + user_api_key_team_id: metadata.user_api_key_team_id.clone(), + trace_parent: None, + } +} + +fn parse_ocr_pre_call_guardrail_request( + request: GuardrailRequest, +) -> CoreResult<(Value, Map)> { + let Value::Object(mut data) = request.data else { + return Err(CoreError::InvalidRequest( + "OCR pre_call guardrail must return an object".to_string(), + )); + }; + let document = data.remove("document").ok_or_else(|| { + CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string()) + })?; + let optional_params = match data.remove("optional_params") { + Some(Value::Object(params)) => params, + Some(_) => { + return Err(CoreError::InvalidRequest( + "OCR pre_call guardrail optional_params must be an object".to_string(), + )) + } + None => Map::new(), + }; + Ok((document, optional_params)) +} + +fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult { + let Value::Object(mut data) = request.data else { + return Err(CoreError::InvalidRequest( + "OCR during_call guardrail must return an object".to_string(), + )); + }; + data.remove("body").ok_or_else(|| { + CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string()) + }) +} + +fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { + CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +} + +fn core_error_kind(error: &CoreError) -> &'static str { + match error { + CoreError::Auth(_) => "AuthError", + CoreError::InvalidProvider(_) => "InvalidProvider", + CoreError::InvalidRequest(_) => "InvalidRequest", + CoreError::InvalidType { .. } => "InvalidType", + CoreError::MissingField(_) => "MissingField", + CoreError::Http { .. } => "HttpError", + CoreError::InvalidResponse(_) => "InvalidResponse", + CoreError::Network(_) => "NetworkError", + CoreError::Routing(_) => "RoutingError", + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs new file mode 100644 index 00000000000..b54ee39b21d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -0,0 +1,25 @@ +use litellm_core::call_lifecycle::CallLifecycle; +use litellm_core::CoreResult; +use serde_json::Value; + +mod client; +mod common_utils; +mod handler; +mod hooks; +mod prepare; +mod types; + +pub use types::OcrRequest; + +use handler::execute_ocr_provider_call; +use prepare::{prepare_ocr_call, PreparedOcrCall}; + +pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { + let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); + CallLifecycle::default() + .run_request(request, &hooks, execute_ocr_provider_call) + .await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs new file mode 100644 index 00000000000..5a4b350a4c4 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -0,0 +1,57 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; + +use super::hooks::OcrLifecycleHooks; +use super::types::{OcrRequest, PreparedOcrRequest}; +use crate::integrations::custom_guardrail::CustomGuardrailRunner; +use crate::integrations::custom_logger::CustomLoggerRunner; + +pub(crate) struct PreparedOcrCall { + pub(crate) request: PreparedOcrRequest, + pub(crate) hooks: OcrLifecycleHooks, +} + +pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { + let call_id = request + .litellm_call_id + .map(str::to_string) + .unwrap_or_else(new_ocr_call_id); + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .unwrap_or(CustomLlmProvider { + model: request.model, + custom_llm_provider: "mistral", + }); + let model = provider_info.model.to_string(); + let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + + PreparedOcrCall { + request: PreparedOcrRequest { + model, + custom_llm_provider, + litellm_call_id: call_id, + document: request.document, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + }, + hooks: OcrLifecycleHooks::new( + CustomLoggerRunner::new(request.callbacks), + CustomGuardrailRunner::new(request.guardrails), + request.request_metadata, + ), + } +} + +fn new_ocr_call_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + format!("ocr-{timestamp}-{sequence}") +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs new file mode 100644 index 00000000000..35747dc6985 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -0,0 +1,610 @@ +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrResponseHandling; +use serde_json::{json, Map, Value}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; +use super::{ocr, OcrRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, +}; +use crate::integrations::types::RequestMetadata; + +async fn read_http_headers(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8(request).expect("request is utf8") +} + +async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") +} + +#[derive(Clone, Debug, PartialEq)] +struct RecordedLogEvent { + hook: &'static str, + model: String, + call_type: String, + user_id: Option, + response_object: Option, + error_kind: Option, +} + +#[derive(Default)] +struct RecordingOcrLogger { + events: Mutex>, +} + +impl RecordingOcrLogger { + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } +} + +impl CustomLogger for RecordingOcrLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedLogEvent { + hook: "async_log_success_event", + model: model_call_details.model.clone(), + call_type: model_call_details.call_type.to_string(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: Some(response_obj.object.clone()), + error_kind: None, + }); + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedLogEvent { + hook: "async_log_failure_event", + model: model_call_details.model.clone(), + call_type: model_call_details.call_type.to_string(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: response_obj.map(|value| value.object.clone()), + error_kind: model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()), + }); + Ok(()) + }) + } +} + +struct RecordingOcrGuardrail { + hooks: Vec, + events: Mutex>, + block_pre_call: bool, +} + +impl RecordingOcrGuardrail { + fn new(hooks: Vec) -> Self { + Self { + hooks, + events: Mutex::new(Vec::new()), + block_pre_call: false, + } + } + + fn blocking_pre_call() -> Self { + Self { + hooks: vec![GuardrailEventHook::PreCall], + events: Mutex::new(Vec::new()), + block_pre_call: true, + } + } + + fn events(&self) -> Vec<&'static str> { + self.events.lock().unwrap().clone() + } +} + +impl CustomGuardrail for RecordingOcrGuardrail { + fn guardrail_name(&self) -> &str { + "recording-ocr-guardrail" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + mut request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("async_pre_call_hook"); + if self.block_pre_call { + return Ok(GuardrailDecision::Block(GuardrailError::blocked( + "blocked before provider", + ))); + } + request.data["document"]["guarded_pre"] = json!(true); + Ok(GuardrailDecision::Mask(request)) + }) + } + + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + mut request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("async_moderation_hook"); + request.data["body"]["guarded_during"] = json!(true); + Ok(GuardrailDecision::Mask(request)) + }) + } +} + +#[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(306); + 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, 256); +} + +#[test] +fn truncate_error_body_does_not_split_multibyte_chars() { + let body = "é".repeat(266); + let truncated = truncate_error_body(&body); + assert!(truncated.is_char_boundary(truncated.len())); +} + +#[test] +fn ocr_dispatch_supports_migrated_providers() { + assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); + assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document()); + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") + .expect("document intelligence config resolves") + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature")); + assert!(ocr_provider_config("openai", "gpt-4o").is_none()); +} + +#[test] +fn string_headers_accepts_string_values() { + let headers = json!({ + "x-trace-id": "trace-1" + }) + .as_object() + .unwrap() + .clone(); + + assert_eq!( + string_headers(Some(headers)).expect("string headers accepted"), + vec![("x-trace-id".to_string(), "trace-1".to_string())] + ); +} + +#[test] +fn auth_header_detection_is_case_insensitive() { + let headers = vec![ + ("x-trace-id".to_string(), "trace-1".to_string()), + ("authorization".to_string(), "Bearer sk-test".to_string()), + ]; + + assert!(has_header(&headers, "authorization")); + + let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; + assert!(has_header(&headers, "authorization")); + + let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; + assert!(!has_header(&headers, "authorization")); +} + +#[tokio::test] +async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let logger = Arc::new(RecordingOcrLogger::default()); + let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![ + GuardrailEventHook::PreCall, + GuardrailEventHook::DuringCall, + ])); + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: vec![logger.clone()], + guardrails: vec![guardrail.clone()], + request_metadata: RequestMetadata { + user_api_key_user_id: Some("user-1".to_string()), + ..Default::default() + }, + litellm_call_id: Some("ocr-call-1"), + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + assert_eq!( + guardrail.events(), + vec!["async_pre_call_hook", "async_moderation_hook"] + ); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_success_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: Some("user-1".to_string()), + response_object: Some("ocr".to_string()), + error_kind: None, + }] + ); + + let request = server.await.expect("server task completes"); + assert!(request.contains(r#""guarded_pre":true"#), "{request}"); + assert!(request.contains(r#""guarded_during":true"#), "{request}"); +} + +#[tokio::test] +async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let _request = read_http_request(&mut socket).await; + let response_body = "provider failed"; + let response = format!( + "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + + let logger = Arc::new(RecordingOcrLogger::default()); + let err = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: vec![logger.clone()], + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: Some("ocr-call-2"), + }) + .await + .expect_err("provider error propagates"); + + assert!(matches!(err, CoreError::Http { status: 500, .. })); + server.await.expect("server task completes"); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_failure_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: None, + response_object: Some("error".to_string()), + error_kind: Some("HttpError".to_string()), + }] + ); +} + +#[tokio::test] +async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let logger = Arc::new(RecordingOcrLogger::default()); + let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call()); + + let err = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_millis(100)), + callbacks: vec![logger.clone()], + guardrails: vec![guardrail.clone()], + request_metadata: RequestMetadata::default(), + litellm_call_id: Some("ocr-call-3"), + }) + .await + .expect_err("guardrail blocks request"); + + assert!(matches!(err, CoreError::InvalidRequest(_))); + assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_failure_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: None, + response_object: Some("error".to_string()), + error_kind: Some("InvalidRequest".to_string()), + }] + ); + let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; + assert!(accepted.is_err(), "provider socket should not be touched"); +} + +#[tokio::test] +async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_headers(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer sk-from-python".to_string()), + ); + headers.insert( + "x-trace-id".to_string(), + Value::String("trace-1".to_string()), + ); + + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-for-rust-fallback"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: Some(headers), + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let request = server.await.expect("server task completes"); + let authorization_count = request + .lines() + .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) + .count(); + assert_eq!(authorization_count, 1, "{request}"); + assert!( + request.contains("authorization: Bearer sk-from-python") + || request.contains("Authorization: Bearer sk-from-python"), + "{request}" + ); +} + +#[tokio::test] +async fn document_intelligence_poll_uses_resolved_subscription_key() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let operation_url = format!("http://{addr}/operations/1"); + + let server = tokio::spawn(async move { + let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); + let post_request = read_http_headers(&mut post_socket).await; + let post_response = format!( + "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + ); + post_socket + .write_all(post_response.as_bytes()) + .await + .expect("writes post response"); + + let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); + let poll_request = read_http_headers(&mut poll_socket).await; + let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; + let poll_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + poll_socket + .write_all(poll_response.as_bytes()) + .await + .expect("writes poll response"); + (post_request, poll_request) + }); + + let response = ocr(OcrRequest { + model: "doc-intelligence/prebuilt-read", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("di-key"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }) + .await + .expect("document intelligence request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let (post_request, poll_request) = server.await.expect("server task completes"); + assert!( + post_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{post_request}" + ); + assert!( + poll_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{poll_request}" + ); +} + +#[test] +fn string_headers_rejects_non_string_values() { + let headers = json!({ + "x-retry-count": 3 + }) + .as_object() + .unwrap() + .clone(); + + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + CoreError::InvalidRequest( + "OCR extra_headers.x-retry-count must be a string, got number".to_string() + ) + ); +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs new file mode 100644 index 00000000000..bde734a4dd1 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -0,0 +1,57 @@ +use std::sync::Arc; +use std::time::Duration; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use litellm_core::ocr::transformation::OcrProviderConfig; +use serde_json::{Map, Value}; + +use crate::integrations::custom_guardrail::CustomGuardrail; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::RequestMetadata; + +pub struct OcrRequest<'a> { + pub model: &'a str, + pub document: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, + pub callbacks: Vec>, + pub guardrails: Vec>, + pub request_metadata: RequestMetadata, + pub litellm_call_id: Option<&'a str>, +} + +pub(crate) struct PreparedOcrRequest { + pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) litellm_call_id: String, + pub(crate) document: Value, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) extra_headers: Option>, + pub(crate) optional_params: Map, + pub(crate) timeout: Option, +} + +impl CallLifecycleRequest for PreparedOcrRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "ocr", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} + +pub(crate) struct ProviderOcrRequest { + pub(crate) model: String, + pub(crate) config: &'static dyn OcrProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) timeout: Option, +} 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..c32e727de54 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -0,0 +1,388 @@ +//! `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::{ + CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; +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 async fn log_messages(&mut self, status: SessionStatus) { + self.end_time = epoch_seconds(); + let payload = self.build_payload(); + let timing = CallbackTiming::new(payload.start_time, payload.end_time); + let runner = CustomLoggerRunner::new(self.callbacks.clone()); + + match status { + SessionStatus::Success => { + let response = CallbackValue::new("realtime", serde_json::Value::Null); + let report = runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload(payload), + &response, + timing, + ) + .await; + self.dropped += report.dropped as u64; + } + SessionStatus::Failure => { + let error = LoggingError { + message: "realtime session ended in failure".to_string(), + kind: "RealtimeSessionError".to_string(), + }; + let response = CallbackValue::new( + "error", + serde_json::json!({ + "message": error.message, + "kind": error.kind, + }), + ); + let report = runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload(payload) + .with_failure_error(error), + Some(&response), + timing, + ) + .await; + self.dropped += report.dropped as u64; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::custom_logger::LogError; + use crate::integrations::custom_logger::LogFuture; + 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 async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let payload = model_call_details + .standard_logging_payload + .as_ref() + .expect("standard logging payload"); + 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(()) + }) + } + } + + #[tokio::test] + async 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).await; + 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. + #[tokio::test] + async fn failing_logger_bumps_dropped_counter() { + struct FailingLogger; + impl CustomLogger for FailingLogger { + fn async_log_success_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Err(LogError::channel_full()) }) + } + + fn async_log_failure_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { 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).await; + assert_eq!(streaming.dropped(), 1); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index 695e0c6bb39..c3f929f5f0b 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -6,8 +6,11 @@ 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; @@ -17,12 +20,29 @@ use axum::Router; use futures_util::{SinkExt, StreamExt}; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router as ModelRouter; -use litellm_providers::realtime_pool::RealtimePool; 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)) @@ -57,26 +77,63 @@ async fn handle( 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, 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(), @@ -84,5 +141,26 @@ async fn bridge( }); futures_util::pin_mut!(client_in, client_out); - let _ = service::run(&router, &pool, &model, None, client_in, client_out).await; + + // 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).await; } diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index 0cbd00d664f..d6c31edd454 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -1,6 +1,6 @@ //! Business logic: select a deployment with the (pure) core router, then call the //! provider splice. The seam between `core::router` (selection only) and -//! `providers` (the actual WebSocket I/O). +//! `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 @@ -9,12 +9,12 @@ 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; -use litellm_providers::realtime_pool::{upstream_key, RealtimePool}; /// Select a deployment for `model` and splice the client stream to the provider. /// @@ -26,6 +26,7 @@ pub async fn run( pool: &RealtimePool, model: &str, idle_timeout: Option, + observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, ) -> CoreResult<()> @@ -52,10 +53,11 @@ where params.api_base.as_deref(), ) { if let Some(handoff) = pool.take(&key) { - return litellm_providers::realtime::realtime_warm( + return crate::io::realtime::realtime_warm( provider_model, handoff, idle_timeout, + observe, client_in, client_out, ) @@ -64,11 +66,12 @@ where } // Cold path: fresh dial (the original behavior). - litellm_providers::realtime::realtime( + crate::io::realtime::realtime( provider_model, params.api_key.as_deref(), params.api_base.as_deref(), idle_timeout, + observe, client_in, client_out, ) diff --git a/litellm-rust/crates/ai-gateway/src/state.rs b/litellm-rust/crates/ai-gateway/src/state.rs index ef96037d477..3b61d8309ea 100644 --- a/litellm-rust/crates/ai-gateway/src/state.rs +++ b/litellm-rust/crates/ai-gateway/src/state.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use crate::io::realtime_pool::RealtimePool; use litellm_core::router::Router; -use litellm_providers::realtime_pool::RealtimePool; + +use crate::integrations::custom_logger::CustomLogger; /// Shared application state handed to every route handler. #[derive(Clone)] @@ -10,6 +12,8 @@ pub struct AppState { /// 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. 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/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1881bcfa602..9bd4634cc2a 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -10,3 +10,6 @@ rand.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/call_lifecycle/README.md b/litellm-rust/crates/core/src/call_lifecycle/README.md new file mode 100644 index 00000000000..692e249ef27 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/README.md @@ -0,0 +1,167 @@ +# Call lifecycle + +`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call +types migrated to Rust. It owns lifecycle ordering, phase timing, and trace +observer calls. It must not know about OCR, chat, messages, responses, +completions, provider auth, request transforms, or response normalization. + +Call-type modules own their domain behavior. For example, OCR owns document +payloads, OCR provider transforms, safe document fetch, guardrail payload shape, +callback payload shape, and provider HTTP execution. + +## Runtime order + +Every wrapped call runs in this order: + +1. `async_pre_call_hook` +2. `async_during_call_hook` +3. provider call +4. `async_log_success_event` or `async_log_failure_event` + +`async_pre_call_hook` receives the initial LiteLLM request shape. It is where +pre-call custom guardrails run. + +`async_during_call_hook` converts the initial request into the provider-ready +request. It is where provider config selection, parameter mapping, auth/header +resolution, request transforms, and during-call guardrails belong. + +The provider call receives only the provider-ready request. It should execute +I/O and call the provider response transform. + +Success and failure callbacks receive `CallLifecycleTiming`. Callback failures +must not replace the original provider or guardrail result. + +## Trace contract + +The lifecycle runner records: + +- full call start and end time +- `pre_call` phase timing +- `during_call` phase timing +- `provider_call` phase timing +- `success_callback` phase timing +- `failure_callback` phase timing + +`CallLifecycleObserver` receives phase start and end events. The default +observer is a no-op. Future OTEL support should implement this observer instead +of editing OCR, chat, messages, responses, completions, or provider modules. + +## Required shape + +Each migrated call type should use this folder shape: + +```text +litellm-rust/crates/ai-gateway/src// + mod.rs # thin public entrypoint + types.rs # public request, prepared request, provider request, response types + prepare.rs # model/provider/callback/guardrail setup + hooks.rs # CallLifecycleHooks implementation + handler.rs # provider I/O and response normalization + tests.rs # call-type lifecycle and handler tests +``` + +Provider transforms can live in `litellm-rust/crates/core/src/providers/...`. +Shared call-type helpers can live beside the call type, but generic lifecycle +code stays in this folder. + +## Core API + +The prepared request implements `CallLifecycleRequest`: + +```rust +impl CallLifecycleRequest for PreparedMessagesRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "messages", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} +``` + +The call-type hooks implement `CallLifecycleHooks`: + +```rust +impl CallLifecycleHooks< + PreparedMessagesRequest, + ProviderMessagesRequest, + MessagesResponse, +> for MessagesLifecycleHooks { + fn async_pre_call_hook(...) { + // run pre-call custom guardrails against the LiteLLM request shape + } + + fn async_during_call_hook(...) { + // map params, validate env, transform request, run during-call guardrails + } + + fn async_log_success_event(...) { + // call async_log_success_event on configured custom loggers + } + + fn async_log_failure_event(...) { + // call async_log_failure_event without swallowing the original error + } +} +``` + +The public entrypoint stays thin: + +```rust +pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { + let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?; + + CallLifecycle::default() + .run_request(request, &hooks, execute_messages_provider_call) + .await +} +``` + +Use `run_request` for new call types. Keep `run` available only for specialized +tests or existing code that already has a `CallLifecycleContext`. + +## Adding a new call type + +1. Add `/types.rs` + +Define the public request accepted by the bridge, the prepared request used by +the lifecycle runner, and the provider request consumed by the handler. + +2. Implement `CallLifecycleRequest` + +Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`. +Do not put provider-specific logic here. + +3. Add `/prepare.rs` + +Resolve model/provider once, generate or preserve `litellm_call_id`, construct +callback and guardrail runners, and return `PreparedCall`. + +4. Add `/hooks.rs` + +Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction, +provider config selection, param mapping, request transform, during-call +guardrail payload construction, and callback payload construction here. + +5. Add `/handler.rs` + +Execute the provider request and normalize the provider response. Do not repeat +provider-specific transforms here; call the provider config. + +6. Add tests + +Cover hook order, success callback payload, failure callback payload, pre-call +guardrail blocking before provider I/O, during-call body mutation, and provider +error mapping. + +## Review checklist + +- Core lifecycle has no call-type or provider-specific branches +- Public call-type entrypoint only prepares and calls `run_request` +- Provider behavior lives behind provider config/transformation code +- Hook method names map to the Python custom logger and guardrail concepts +- Phase timing is recorded once in lifecycle, not separately per call type +- Callback failures never hide the original provider or guardrail error +- Tests prove the provider socket is not touched when pre-call guardrails block diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs new file mode 100644 index 00000000000..d9b68a1b726 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -0,0 +1,414 @@ +use std::future::Future; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +use crate::{CoreError, CoreResult}; + +pub mod types; + +pub use types::{ + CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest, + CallLifecycleTiming, +}; + +pub trait CallLifecycleHooks: Send + Sync { + type PreCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a, + Resp: 'a; + + type DuringCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a, + Resp: 'a; + + type SuccessFuture<'a>: Future + Send + 'a + where + Self: 'a, + Resp: 'a; + + type FailureFuture<'a>: Future + Send + 'a + where + Self: 'a; + + fn async_pre_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::PreCallFuture<'a>; + + fn async_during_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::DuringCallFuture<'a>; + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Resp, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a>; + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a>; +} + +pub trait CallLifecycleObserver: Send + Sync { + fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {} + + fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {} +} + +#[derive(Default)] +pub struct NoopCallLifecycleObserver; + +impl CallLifecycleObserver for NoopCallLifecycleObserver {} + +pub struct CallLifecycle<'a> { + observer: &'a dyn CallLifecycleObserver, +} + +impl<'a> CallLifecycle<'a> { + pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self { + Self { observer } + } + + pub async fn run_request( + &self, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> CoreResult + where + InitialReq: CallLifecycleRequest, + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + let context = request.lifecycle_context(); + self.run(context, request, hooks, provider_call).await + } + + pub async fn run( + &self, + context: CallLifecycleContext, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> CoreResult + where + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + let call_start = epoch_seconds(); + let mut phases = Vec::new(); + + let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall); + let request = match hooks.async_pre_call_hook(&context, request).await { + Ok(request) => { + phases.push(self.finish_phase(&context, pre_call)); + request + } + Err(error) => { + phases.push(self.finish_phase(&context, pre_call)); + self.log_failure(&context, hooks, &error, call_start, &mut phases) + .await; + return Err(error); + } + }; + + let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall); + let provider_request = match hooks.async_during_call_hook(&context, request).await { + Ok(request) => { + phases.push(self.finish_phase(&context, during_call)); + request + } + Err(error) => { + phases.push(self.finish_phase(&context, during_call)); + self.log_failure(&context, hooks, &error, call_start, &mut phases) + .await; + return Err(error); + } + }; + + let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall); + let result = provider_call(provider_request).await; + phases.push(self.finish_phase(&context, provider_phase)); + + match &result { + Ok(response) => { + let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback); + let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); + hooks + .async_log_success_event(&context, response, &timing) + .await; + phases.push(self.finish_phase(&context, success_phase)); + } + Err(error) => { + self.log_failure(&context, hooks, error, call_start, &mut phases) + .await; + } + } + + result + } + + async fn log_failure( + &self, + context: &CallLifecycleContext, + hooks: &Hooks, + error: &CoreError, + call_start: f64, + phases: &mut Vec, + ) where + Hooks: CallLifecycleHooks, + { + let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback); + let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); + hooks.async_log_failure_event(context, error, &timing).await; + phases.push(self.finish_phase(context, failure_phase)); + } + + fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart { + self.observer.on_phase_start(context, phase); + PhaseStart { + phase, + start_time: epoch_seconds(), + started_at: Instant::now(), + } + } + + fn finish_phase( + &self, + context: &CallLifecycleContext, + phase_start: PhaseStart, + ) -> CallLifecyclePhaseTiming { + let timing = CallLifecyclePhaseTiming { + phase: phase_start.phase, + start_time: phase_start.start_time, + end_time: epoch_seconds(), + duration: phase_start.started_at.elapsed(), + }; + self.observer.on_phase_end(context, &timing); + timing + } +} + +impl Default for CallLifecycle<'static> { + fn default() -> Self { + static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver; + Self::new(&OBSERVER) + } +} + +struct PhaseStart { + phase: CallLifecyclePhase, + start_time: f64, + started_at: Instant, +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::pin::Pin; + use std::sync::Mutex; + + type BoxFuture<'a, T> = Pin + Send + 'a>>; + + #[derive(Default)] + struct RecordingHooks { + events: Mutex>, + } + + struct RecordingRequest(String); + + impl CallLifecycleRequest for RecordingRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1") + } + } + + impl RecordingHooks { + fn events(&self) -> Vec<&'static str> { + self.events.lock().unwrap().clone() + } + } + + impl CallLifecycleHooks for RecordingHooks { + type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; + type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type SuccessFuture<'a> = BoxFuture<'a, ()>; + type FailureFuture<'a> = BoxFuture<'a, ()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: String, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("pre_call"); + Ok(format!("{request}:pre")) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: String, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("during_call"); + Ok(format!("{request}:during")) + }) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a String, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + assert!(timing.end_time >= timing.start_time); + assert_eq!(timing.phases.len(), 3); + self.events.lock().unwrap().push("success"); + }) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + impl CallLifecycleHooks for RecordingHooks { + type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; + type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type SuccessFuture<'a> = BoxFuture<'a, ()>; + type FailureFuture<'a> = BoxFuture<'a, ()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: RecordingRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("pre_call"); + Ok(RecordingRequest(format!("{}:pre", request.0))) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: RecordingRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("during_call"); + Ok(format!("{}:during", request.0)) + }) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a String, + _timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("success"); + }) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + #[tokio::test] + async fn lifecycle_runs_hooks_around_provider_call() { + let hooks = RecordingHooks::default(); + let response = CallLifecycle::default() + .run( + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), + "request".to_string(), + &hooks, + |request| async move { + assert_eq!(request, "request:pre:during"); + Ok("response".to_string()) + }, + ) + .await + .expect("call succeeds"); + + assert_eq!(response, "response"); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); + } + + #[tokio::test] + async fn lifecycle_logs_failure_when_provider_fails() { + let hooks = RecordingHooks::default(); + let error = CallLifecycle::default() + .run( + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), + "request".to_string(), + &hooks, + |_request| async move { + Err::(CoreError::Network("provider down".to_string())) + }, + ) + .await + .expect_err("call fails"); + + assert_eq!(error, CoreError::Network("provider down".to_string())); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]); + } + + #[tokio::test] + async fn lifecycle_can_run_any_request_with_embedded_context() { + let hooks = RecordingHooks::default(); + let response = CallLifecycle::default() + .run_request( + RecordingRequest("request".to_string()), + &hooks, + |request| async move { + assert_eq!(request, "request:pre:during"); + Ok("response".to_string()) + }, + ) + .await + .expect("call succeeds"); + + assert_eq!(response, "response"); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/types.rs b/litellm-rust/crates/core/src/call_lifecycle/types.rs new file mode 100644 index 00000000000..8819c8830d2 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/types.rs @@ -0,0 +1,75 @@ +use std::time::Duration; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CallLifecycleContext { + pub call_type: String, + pub model: String, + pub custom_llm_provider: String, + pub litellm_call_id: String, +} + +impl CallLifecycleContext { + pub fn new( + call_type: impl Into, + model: impl Into, + custom_llm_provider: impl Into, + litellm_call_id: impl Into, + ) -> Self { + Self { + call_type: call_type.into(), + model: model.into(), + custom_llm_provider: custom_llm_provider.into(), + litellm_call_id: litellm_call_id.into(), + } + } +} + +pub trait CallLifecycleRequest { + fn lifecycle_context(&self) -> CallLifecycleContext; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CallLifecyclePhase { + PreCall, + DuringCall, + ProviderCall, + SuccessCallback, + FailureCallback, +} + +impl CallLifecyclePhase { + pub fn as_str(self) -> &'static str { + match self { + Self::PreCall => "pre_call", + Self::DuringCall => "during_call", + Self::ProviderCall => "provider_call", + Self::SuccessCallback => "success_callback", + Self::FailureCallback => "failure_callback", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallLifecyclePhaseTiming { + pub phase: CallLifecyclePhase, + pub start_time: f64, + pub end_time: f64, + pub duration: Duration, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CallLifecycleTiming { + pub start_time: f64, + pub end_time: f64, + pub phases: Vec, +} + +impl CallLifecycleTiming { + pub fn new(start_time: f64, end_time: f64, phases: Vec) -> Self { + Self { + start_time, + end_time, + phases, + } + } +} diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 9b29260cca4..b3e0519b772 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -13,6 +13,10 @@ pub enum CoreError { MissingField(&'static str), #[error("invalid response: {0}")] InvalidResponse(String), + #[error("invalid provider: {0}")] + InvalidProvider(String), + #[error("invalid request: {0}")] + InvalidRequest(String), #[error("{0}")] Auth(String), #[error("OCR request failed with status {status}: {body}")] diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 9d686626edc..555a04ce853 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,6 +1,9 @@ +pub mod call_lifecycle; pub mod error; pub mod ocr; +pub mod providers; pub mod realtime; pub mod router; +pub mod routing_utils; pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 7353d9d22c4..cb3e735e533 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -4,7 +4,28 @@ use crate::CoreResult; use super::types::{OcrRequestData, OcrResponseData}; -pub trait OcrProviderConfig { +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrAuthStrategy { + Bearer, + Header(&'static str), +} + +impl OcrAuthStrategy { + pub fn header_name(self) -> &'static str { + match self { + Self::Bearer => "authorization", + Self::Header(header_name) => header_name, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrResponseHandling { + Json, + AzureDocumentIntelligencePoll, +} + +pub trait OcrProviderConfig: Sync { fn supported_ocr_params(&self) -> &'static [&'static str]; fn map_ocr_params(&self, non_default_params: &Map) -> Map { @@ -29,4 +50,30 @@ pub trait OcrProviderConfig { model: &str, response_json: Value, ) -> CoreResult; + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn auth_strategy(&self) -> OcrAuthStrategy { + OcrAuthStrategy::Bearer + } + + fn requires_data_uri_document(&self) -> bool { + false + } + + fn response_handling(&self) -> OcrResponseHandling { + OcrResponseHandling::Json + } } diff --git a/litellm-rust/crates/providers/src/mistral/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs similarity index 100% rename from litellm-rust/crates/providers/src/mistral/mod.rs rename to litellm-rust/crates/core/src/providers/azure_ai/mod.rs diff --git a/litellm-rust/crates/providers/src/mistral/ocr/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs similarity index 100% rename from litellm-rust/crates/providers/src/mistral/ocr/mod.rs rename to litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs new file mode 100644 index 00000000000..060073acd47 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -0,0 +1,520 @@ +use std::collections::BTreeSet; + +use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; +use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{json, Map, Value}; + +use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; +const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; +const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; +const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; +const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96; + +const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages"]; + +pub struct AzureAiOcrConfig; +pub struct AzureDocumentIntelligenceOcrConfig; + +pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig; +pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig = + AzureDocumentIntelligenceOcrConfig; + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +fn resolve_value( + explicit: Option<&str>, + env_name: &str, + env_lookup: &dyn Fn(&str) -> Option, + missing_message: &str, +) -> CoreResult { + non_empty(explicit) + .map(str::to_string) + .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| CoreError::Auth(missing_message.to_string())) +} + +pub fn resolve_azure_ai_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_key, + AZURE_AI_API_KEY_ENV, + env_lookup, + "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params", + ) +} + +pub fn resolve_azure_ai_api_base( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_base, + AZURE_AI_API_BASE_ENV, + env_lookup, + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", + ) +} + +pub fn complete_azure_ai_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let base = resolve_azure_ai_api_base(api_base, env_lookup)?; + Ok(format!( + "{}/providers/mistral/azure/ocr", + base.trim_end_matches('/') + )) +} + +pub fn resolve_document_intelligence_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_key, + AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, + env_lookup, + "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter", + ) +} + +pub fn resolve_document_intelligence_endpoint( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_base, + AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, + env_lookup, + "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter", + ) +} + +fn encode_model_id(model: &str) -> String { + let model_id = model.rsplit('/').next().unwrap_or(model); + model_id + .bytes() + .flat_map(|byte| match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { + vec![byte as char] + } + _ => format!("%{byte:02X}").chars().collect(), + }) + .collect() +} + +fn pages_token_is_valid(token: &str) -> bool { + let mut parts = token.split('-'); + let Some(start) = parts.next() else { + return false; + }; + if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) { + return false; + } + match parts.next() { + None => true, + Some(end) => { + !end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none() + } + } +} + +fn normalize_pages_param(pages: &Value) -> CoreResult> { + match pages { + Value::String(value) => { + let normalized = value + .split(',') + .map(str::trim) + .collect::>() + .join(","); + if normalized.split(',').all(pages_token_is_valid) { + Ok(Some(normalized)) + } else { + Err(CoreError::InvalidRequest(format!( + "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." + ))) + } + } + Value::Array(values) => { + if values.is_empty() { + return Ok(None); + } + if values.iter().all(Value::is_i64) { + let mut pages = BTreeSet::new(); + for value in values { + let page = value.as_i64().expect("checked is_i64"); + if page < 0 { + return Err(CoreError::InvalidRequest( + "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), + )); + } + pages.insert(page + 1); + } + return Ok(Some( + pages + .into_iter() + .map(|page| page.to_string()) + .collect::>() + .join(","), + )); + } + if values.iter().all(Value::is_string) { + let normalized = values + .iter() + .filter_map(Value::as_str) + .map(str::trim) + .collect::>() + .join(","); + if normalized.split(',').all(pages_token_is_valid) { + return Ok(Some(normalized)); + } + return Err(CoreError::InvalidRequest(format!( + "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." + ))); + } + Err(CoreError::InvalidRequest( + "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." + .to_string(), + )) + } + _ => Err(CoreError::InvalidRequest( + "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." + .to_string(), + )), + } +} + +pub fn complete_document_intelligence_url( + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; + let mut url = format!( + "{}/documentintelligence/documentModels/{}:analyze?api-version={}", + endpoint.trim_end_matches('/'), + encode_model_id(model), + AZURE_DOCUMENT_INTELLIGENCE_API_VERSION + ); + + if let Some(pages) = optional_params.get("pages") { + if let Some(normalized) = normalize_pages_param(pages)? { + url.push_str("&pages="); + url.push_str(&normalized); + } + } + + Ok(url) +} + +fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { + let object = document.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(document), + })?; + let doc_type = object + .get("type") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("document.type"))?; + let field_name = match doc_type { + "document_url" => "document_url", + "image_url" => "image_url", + other => { + return Err(CoreError::InvalidRequest(format!( + "Invalid document type: {other}. Must be 'document_url' or 'image_url'" + ))) + } + }; + object + .get(field_name) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or(CoreError::MissingField(field_name)) +} + +fn extract_base64_from_data_uri(data_uri: &str) -> &str { + data_uri + .split_once(',') + .map(|(_, data)| data) + .unwrap_or(data_uri) +} + +fn page_markdown(page: &Map) -> String { + page.get("lines") + .and_then(Value::as_array) + .map(|lines| { + lines + .iter() + .filter_map(|line| line.get("content").and_then(Value::as_str)) + .collect::>() + .join("\n") + }) + .unwrap_or_default() +} + +fn page_dimensions(page: &Map) -> Value { + let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5); + let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0); + let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch"); + let (width, height) = if unit == "inch" { + ( + (width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, + (height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, + ) + } else { + (width as i64, height as i64) + }; + json!({ + "width": width, + "height": height, + "dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, + }) +} + +impl OcrProviderConfig for AzureAiOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_azure_ai_url(api_base, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_azure_ai_api_key(api_key, env_lookup) + } + + fn requires_data_uri_document(&self) -> bool { + true + } +} + +impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + _model: &str, + document: Value, + _optional_params: Map, + ) -> CoreResult { + let document_url = document_url_from_mistral_document(&document)?; + let mut data = Map::new(); + if document_url.starts_with("data:") { + data.insert( + "base64Source".to_string(), + Value::String(extract_base64_from_data_uri(document_url).to_string()), + ); + } else { + data.insert( + "urlSource".to_string(), + Value::String(document_url.to_string()), + ); + } + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + let status = response + .get("status") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("status"))?; + if status != "succeeded" { + return Err(CoreError::InvalidResponse(format!( + "Azure Document Intelligence analysis failed with status: {status}" + ))); + } + + let azure_pages = response + .get("analyzeResult") + .and_then(|result| result.get("pages")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + + let pages = azure_pages + .iter() + .filter_map(Value::as_object) + .map(|page| { + let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1); + json!({ + "index": page_number - 1, + "markdown": page_markdown(page), + "dimensions": page_dimensions(page), + }) + }) + .collect::>(); + + Ok(OcrResponseData { + usage_info: Some(json!({ + "pages_processed": pages.len(), + "doc_size_bytes": null, + })), + pages, + model: model.to_string(), + document_annotation: None, + object: "ocr".to_string(), + }) + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_document_intelligence_url(api_base, model, optional_params, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_document_intelligence_api_key(api_key, env_lookup) + } + + fn auth_strategy(&self) -> OcrAuthStrategy { + OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") + } + + fn response_handling(&self) -> OcrResponseHandling { + OcrResponseHandling::AzureDocumentIntelligencePoll + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn azure_ai_reuses_mistral_body_transform() { + let body = AZURE_AI_OCR_CONFIG + .transform_ocr_request( + "pixtral-12b-2409", + json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}), + serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "pixtral-12b-2409"); + assert_eq!(body["include_image_base64"], true); + assert_eq!( + body["document"]["document_url"], + "data:application/pdf;base64,abc" + ); + } + + #[test] + fn document_intelligence_url_normalizes_zero_based_pages() { + let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]); + let url = complete_document_intelligence_url( + Some("https://example.cognitiveservices.azure.com/"), + "azure_ai/doc-intelligence/prebuilt-layout", + ¶ms, + &|_| None, + ) + .expect("url builds"); + + assert_eq!( + url, + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3" + ); + } + + #[test] + fn document_intelligence_request_uses_base64_source_for_data_uri() { + let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_request( + "prebuilt-read", + json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}), + Map::new(), + ) + .expect("request transforms") + .data; + + assert_eq!(body, json!({"base64Source": "abc123"})); + } + + #[test] + fn document_intelligence_response_normalizes_pages() { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response( + "prebuilt-layout", + json!({ + "status": "succeeded", + "analyzeResult": { + "pages": [{ + "pageNumber": 2, + "width": 8.5, + "height": 11, + "unit": "inch", + "lines": [{"content": "hello"}, {"content": "world"}] + }] + } + }), + ) + .expect("response transforms"); + + assert_eq!(response.pages[0]["index"], 1); + assert_eq!(response.pages[0]["markdown"], "hello\nworld"); + assert_eq!(response.pages[0]["dimensions"]["width"], 816); + assert_eq!( + response.usage_info, + Some(json!({"pages_processed": 1, "doc_size_bytes": null})) + ); + } +} 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/providers/src/openai/realtime/mod.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs similarity index 100% rename from litellm-rust/crates/providers/src/openai/realtime/mod.rs rename to litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs diff --git a/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs similarity index 93% rename from litellm-rust/crates/providers/src/mistral/ocr/transformation.rs rename to litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index fd691177783..1a33bc1e951 100644 --- a/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -1,6 +1,6 @@ -use litellm_core::error::{json_type_name, CoreError, CoreResult}; -use litellm_core::ocr::transformation::OcrProviderConfig; -use litellm_core::ocr::types::{OcrRequestData, OcrResponseData}; +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] = &[ @@ -15,6 +15,7 @@ const SUPPORTED_OCR_PARAMS: &[&str] = &[ "extract_footer", "table_format", "confidence_scores_granularity", + "include_blocks", "id", ]; @@ -132,6 +133,24 @@ impl OcrProviderConfig for MistralOcrConfig { object: "ocr".to_string(), }) } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + Ok(complete_url(api_base)) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_api_key(api_key, env_lookup) + } } pub fn supported_ocr_params() -> &'static [&'static str] { @@ -175,6 +194,7 @@ mod tests { "extract_footer", "table_format", "confidence_scores_granularity", + "include_blocks", "id", ] ); 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..d75e750a0ba --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -0,0 +1,4 @@ +pub mod azure_ai; +pub mod mistral; +pub mod openai; +pub mod vertex_ai; diff --git a/litellm-rust/crates/providers/src/openai/mod.rs b/litellm-rust/crates/core/src/providers/openai/mod.rs similarity index 100% rename from litellm-rust/crates/providers/src/openai/mod.rs rename to litellm-rust/crates/core/src/providers/openai/mod.rs 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/providers/src/openai/realtime/transformation.rs b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs similarity index 97% rename from litellm-rust/crates/providers/src/openai/realtime/transformation.rs rename to litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs index 2e127c699e0..626e4014ff9 100644 --- a/litellm-rust/crates/providers/src/openai/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -1,6 +1,6 @@ -use litellm_core::realtime::transformation::RealtimeProviderConfig; -use litellm_core::realtime::types::{RealtimeEvent, RealtimeTransformResult}; -use litellm_core::CoreResult; +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"; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs new file mode 100644 index 00000000000..8639926c435 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -0,0 +1,435 @@ +use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{json, Map, Value}; + +use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; +const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; +const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; +const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; +const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; +const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; +const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; + +#[rustfmt::skip] +const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[ + "stream", + "temperature", + "max_tokens", + "top_p", + "n", + "stop", +]; + +pub struct VertexAiOcrConfig; +pub struct VertexAiDeepSeekOcrConfig; + +pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig; +pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; + +fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { + keys.iter() + .find_map(|key| params.get(*key).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + +pub fn is_deepseek_model(model: &str) -> bool { + model.to_ascii_lowercase().contains("deepseek") +} + +pub fn resolve_vertex_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(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| { + CoreError::Auth( + "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" + .to_string(), + ) + }) +} + +fn vertex_project( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + string_param(params, &["vertex_project", "vertex_ai_project"]) + .map(str::to_string) + .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + CoreError::InvalidRequest( + "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" + .to_string(), + ) + }) +} + +fn vertex_location( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + string_param(params, &["vertex_location", "vertex_ai_location"]) + .map(str::to_string) + .or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty())) + .or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty())) + .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) +} + +fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String { + api_base + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")) + .trim_end_matches('/') + .to_string() +} + +pub fn complete_vertex_mistral_url( + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let project = vertex_project(optional_params, env_lookup)?; + let location = vertex_location(optional_params, env_lookup); + let base = vertex_mistral_api_base(api_base, &location); + Ok(format!( + "{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict" + )) +} + +pub fn complete_vertex_deepseek_url( + api_base: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let project = vertex_project(optional_params, env_lookup)?; + let location = vertex_location(optional_params, env_lookup); + let base = api_base + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE) + .trim_end_matches('/'); + Ok(format!( + "{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions" + )) +} + +fn document_content_item(document: &Value) -> CoreResult { + let object = document.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(document), + })?; + let doc_type = object + .get("type") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("document.type"))?; + let url_field = match doc_type { + "image_url" => "image_url", + "document_url" => "document_url", + other => { + return Err(CoreError::InvalidRequest(format!( + "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" + ))) + } + }; + let url = object + .get(url_field) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or(CoreError::MissingField(url_field))?; + + Ok(json!({ + "type": "image_url", + "image_url": url, + })) +} + +fn deepseek_model_name(model: &str) -> String { + if model.starts_with("deepseek-ai/") { + model.to_string() + } else { + format!("deepseek-ai/{model}") + } +} + +fn first_choice_content(response: &Value) -> CoreResult { + response + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("message")) + .and_then(|message| message.get("content")) + .cloned() + .filter(|content| match content { + Value::String(value) => !value.is_empty(), + Value::Object(_) => true, + _ => false, + }) + .ok_or_else(|| { + CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string()) + }) +} + +fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> Value { + match content { + Value::String(content) => { + if content.trim_start().starts_with('{') { + serde_json::from_str(&content).unwrap_or_else(|_| { + json!({ + "pages": [{"index": 0, "markdown": content}], + "model": model, + "usage_info": usage.unwrap_or_else(|| json!({})), + }) + }) + } else { + json!({ + "pages": [{"index": 0, "markdown": content}], + "model": model, + "usage_info": usage.unwrap_or_else(|| json!({})), + }) + } + } + Value::Object(_) => content, + other => json!({ + "pages": [{"index": 0, "markdown": other.to_string()}], + "model": model, + "usage_info": usage.unwrap_or_else(|| json!({})), + }), + } +} + +impl OcrProviderConfig for VertexAiOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_vertex_api_key(api_key, env_lookup) + } + + fn requires_data_uri_document(&self) -> bool { + true + } +} + +impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + DEEPSEEK_SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + let mut data = Map::new(); + data.insert( + "model".to_string(), + Value::String(deepseek_model_name(model)), + ); + data.insert( + "messages".to_string(), + json!([{"role": "user", "content": [document_content_item(&document)?]}]), + ); + for (key, value) in optional_params { + if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) { + data.insert(key, value); + } + } + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + let usage = response.get("usage").cloned(); + let content = first_choice_content(&response_json)?; + let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model); + + if !ocr_data.get("pages").is_some_and(Value::is_array) { + ocr_data = json!({ + "pages": [{ + "index": 0, + "markdown": match content { + Value::String(value) => value, + other => other.to_string(), + } + }], + "model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model), + "usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})), + }); + } + + let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&ocr_data), + })?; + let pages = object + .get("pages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let usage_info = object + .get("usage_info") + .cloned() + .or_else(|| response.get("usage").cloned()); + Ok(OcrResponseData { + pages, + model: object + .get("model") + .and_then(Value::as_str) + .unwrap_or(model) + .to_string(), + document_annotation: object.get("document_annotation").cloned(), + usage_info, + object: "ocr".to_string(), + }) + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_vertex_deepseek_url(api_base, optional_params, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_vertex_api_key(api_key, env_lookup) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn vertex_mistral_url_uses_project_location_and_model() { + let params = Map::from_iter([ + ("vertex_project".to_string(), json!("proj-1")), + ("vertex_location".to_string(), json!("europe-west4")), + ]); + + let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) + .expect("url builds"); + + assert_eq!( + url, + "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + } + + #[test] + fn vertex_mistral_reuses_mistral_body_transform() { + let body = VERTEX_AI_OCR_CONFIG + .transform_ocr_request( + "mistral-ocr-maas", + json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}), + Map::new(), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "mistral-ocr-maas"); + assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc"); + } + + #[test] + fn vertex_deepseek_request_uses_ocr_endpoint_shape() { + let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG + .transform_ocr_request( + "deepseek-ocr-maas", + json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}), + Map::from_iter([("temperature".to_string(), json!(0.1))]), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"}) + ); + } + + #[test] + fn vertex_deepseek_response_wraps_markdown_content() { + let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG + .transform_ocr_response( + "deepseek-ocr-maas", + json!({ + "choices": [{"message": {"content": "# OCR text"}}], + "usage": {"prompt_tokens": 1} + }), + ) + .expect("response transforms"); + + assert_eq!( + response.pages, + vec![json!({"index": 0, "markdown": "# OCR text"})] + ); + assert_eq!(response.model, "deepseek-ocr-maas"); + assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1}))); + } +} diff --git a/litellm-rust/crates/core/src/routing_utils/README.md b/litellm-rust/crates/core/src/routing_utils/README.md new file mode 100644 index 00000000000..8585c18e421 --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/README.md @@ -0,0 +1,7 @@ +# Routing Utils + +Shared helpers for deciding how a LiteLLM model routes to an LLM provider. +Keep provider-name parsing, explicit `custom_llm_provider` handling, and model-prefix normalization here. +Do not put deployment selection or load-balancing logic here; that belongs in `router`. +Do not put provider HTTP transformation logic here; that belongs in `providers`. +Helpers in this folder should be deterministic and easy to unit test without network calls. diff --git a/litellm-rust/crates/core/src/routing_utils/mod.rs b/litellm-rust/crates/core/src/routing_utils/mod.rs new file mode 100644 index 00000000000..8336397f870 --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/mod.rs @@ -0,0 +1 @@ +pub mod provider; diff --git a/litellm-rust/crates/core/src/routing_utils/provider.rs b/litellm-rust/crates/core/src/routing_utils/provider.rs new file mode 100644 index 00000000000..6333eedebfc --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/provider.rs @@ -0,0 +1,77 @@ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CustomLlmProvider<'a> { + pub model: &'a str, + pub custom_llm_provider: &'a str, +} + +pub fn get_custom_llm_provider<'a>( + model: &'a str, + custom_llm_provider: Option<&'a str>, +) -> Option> { + if let Some(custom_llm_provider) = custom_llm_provider.filter(|provider| !provider.is_empty()) { + return Some(CustomLlmProvider { + model: strip_custom_llm_provider_prefix(model, custom_llm_provider), + custom_llm_provider, + }); + } + + let (custom_llm_provider, model) = model.split_once('/')?; + if custom_llm_provider.is_empty() || model.is_empty() { + return None; + } + Some(CustomLlmProvider { + model, + custom_llm_provider, + }) +} + +fn strip_custom_llm_provider_prefix<'a>(model: &'a str, custom_llm_provider: &str) -> &'a str { + model + .strip_prefix(custom_llm_provider) + .and_then(|model| model.strip_prefix('/')) + .unwrap_or(model) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gets_custom_llm_provider_from_model_prefix() { + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", None), + Some(CustomLlmProvider { + model: "mistral-ocr-latest", + custom_llm_provider: "mistral", + }) + ); + assert_eq!( + get_custom_llm_provider("azure_ai/doc-intelligence/prebuilt-layout", None), + Some(CustomLlmProvider { + model: "doc-intelligence/prebuilt-layout", + custom_llm_provider: "azure_ai", + }) + ); + assert_eq!(get_custom_llm_provider("mistral-ocr-latest", None), None); + assert_eq!(get_custom_llm_provider("/model", None), None); + assert_eq!(get_custom_llm_provider("provider/", None), None); + } + + #[test] + fn explicit_custom_llm_provider_strips_matching_model_prefix() { + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", Some("mistral")), + Some(CustomLlmProvider { + model: "mistral-ocr-latest", + custom_llm_provider: "mistral", + }) + ); + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", Some("vertex_ai")), + Some(CustomLlmProvider { + model: "mistral/mistral-ocr-latest", + custom_llm_provider: "vertex_ai", + }) + ); + } +} 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/providers/CLAUDE.md b/litellm-rust/crates/providers/CLAUDE.md deleted file mode 100644 index 0f7fdcda2aa..00000000000 --- a/litellm-rust/crates/providers/CLAUDE.md +++ /dev/null @@ -1,53 +0,0 @@ -# CLAUDE.md - -Rules for `litellm-rust/crates/providers`. - -## Responsibility - -`providers` owns provider-specific pure transforms. It mirrors the existing -Python provider modules closely enough that parity review is mechanical. - -Provider files should map to the Python provider tree: - -```text -providers/src///transformation.rs -``` - -For example, Mistral OCR lives at -`providers/src/mistral/ocr/transformation.rs`, matching -`litellm/llms/mistral/ocr/transformation.py`. - -Allowed: -- Provider request transforms. -- Provider response normalization. -- Supported-parameter filtering. -- Provider-specific validation that does not require I/O or secrets. - -Not allowed: -- HTTP clients or provider SDK calls. -- Environment variable reads. -- API key resolution or auth header construction. -- Logging, callbacks, spend tracking, retries, routing, cooldowns, or fallbacks. -- Panics on bad user/provider input. - -## Required Tests - -Every provider transform must include focused unit tests for: -- Supported params matching the Python provider config. -- Unknown params being dropped or transformed the same way as Python. -- Request body shape matching Python output. -- Response normalization with complete, missing, null, and extra fields. -- Bad input returning typed errors. - -For OCR specifically, assume documents can contain personal data. Tests should -prove transforms do not copy document contents into error messages. - -## Implementation Rules - -- Prefer static supported-parameter lists over allocating strings on every call. -- Keep transforms deterministic and allocation-conscious, but choose clarity over - premature micro-optimization for tiny parameter lists. -- Use typed errors from `core`; avoid stringly-typed error plumbing. -- Add comments only when they explain Python-parity decisions or provider quirks. -- Put route-level provider dispatch in a route file such as `providers/src/ocr.rs`. - Do not move provider-specific transform logic into the Python bridge. diff --git a/litellm-rust/crates/providers/Cargo.toml b/litellm-rust/crates/providers/Cargo.toml deleted file mode 100644 index c5b41424d66..00000000000 --- a/litellm-rust/crates/providers/Cargo.toml +++ /dev/null @@ -1,18 +0,0 @@ -[package] -name = "litellm-providers" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-core.workspace = true -reqwest.workspace = true -serde_json.workspace = true -tokio.workspace = true -tokio-tungstenite.workspace = true -futures-util.workspace = true - -[dev-dependencies] -serde_json.workspace = true -futures-channel = "0.3" diff --git a/litellm-rust/crates/providers/src/ocr.rs b/litellm-rust/crates/providers/src/ocr.rs deleted file mode 100644 index dcd56a5f0b4..00000000000 --- a/litellm-rust/crates/providers/src/ocr.rs +++ /dev/null @@ -1,127 +0,0 @@ -//! 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 crate::mistral::ocr::transformation as mistral; -use crate::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/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/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 80b6478daac..83e163c38f1 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -6,11 +6,13 @@ license.workspace = true repository.workspace = true [lib] -name = "litellm_python_bridge" +name = "_native" crate-type = ["cdylib"] [dependencies] litellm-core.workspace = true -litellm-providers.workspace = true +litellm-ai-gateway = { workspace = true, default-features = false } pyo3 = { workspace = true, features = ["extension-module"] } +pyo3-async-runtimes.workspace = true serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/build.rs b/litellm-rust/crates/python-bridge/build.rs new file mode 100644 index 00000000000..0f7293007b2 --- /dev/null +++ b/litellm-rust/crates/python-bridge/build.rs @@ -0,0 +1,6 @@ +fn main() { + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") { + println!("cargo:rustc-cdylib-link-arg=-undefined"); + println!("cargo:rustc-cdylib-link-arg=dynamic_lookup"); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 15e93f7b00c..946a99f990c 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,7 +1,7 @@ use std::time::Duration; +use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest}; use litellm_core::error::CoreError; -use litellm_providers::ocr::run_ocr; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; @@ -9,6 +9,13 @@ use serde_json::{Map, Value}; mod gil; +type MarshaledOcrInputs = ( + Value, + Option>, + Map, + Option, +); + fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { let json = py.import("json")?; let encoded: String = json.call_method1("dumps", (value,))?.extract()?; @@ -22,59 +29,96 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { 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()) - } + CoreError::InvalidProvider(_) + | CoreError::InvalidRequest(_) + | 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. +fn optional_object_to_map( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult> { + match value { + Some(value) => match py_to_json(py, value.bind(py))? { + Value::Object(map) => Ok(map), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + }, + None => Ok(Map::new()), + } +} + +fn optional_timeout(timeout_seconds: Option) -> Option { + timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }) +} + +fn marshal_inputs( + py: Python<'_>, + document: Py, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + let document = py_to_json(py, document.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + + Ok((document, extra_headers, optional_params, timeout)) +} + #[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] fn ocr( py: Python<'_>, model: String, document: Py, api_key: Option, api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let document = py_to_json(py, document.bind(py))?; + let (document, extra_headers, optional_params, timeout) = marshal_inputs( + py, + document, + extra_headers, + optional_params, + timeout_seconds, + )?; - 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, + pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { + model: &model, document, - api_key.as_deref(), - api_base.as_deref(), + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, optional_params, timeout, - ) + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + })) }); match result { @@ -83,8 +127,50 @@ fn ocr( } } -/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe -/// how often the bridge has dropped the GIL for blocking work. +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn aocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (document, extra_headers, optional_params, timeout) = marshal_inputs( + py, + document, + extra_headers, + optional_params, + timeout_seconds, + )?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = run_ocr(OcrRequest { + model: &model, + document, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await + .map_err(core_error_to_pyerr)?; + + Python::with_gil(|py| json_to_py(py, value)) + }) +} + #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); @@ -93,8 +179,9 @@ fn gil_stats(py: Python<'_>) -> PyResult> { } #[pymodule] -fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { +fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(aocr, module)?)?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; Ok(()) } diff --git a/litellm/__init__.py b/litellm/__init__.py index d0513f77b35..5ae5942f32f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -390,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 ) @@ -416,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' @@ -485,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 @@ -507,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 ####### @@ -551,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 @@ -1406,7 +1394,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * -from .ocr.rust_bridge import use_litellm_rust +from .rust_bridge.ocr import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * @@ -1457,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 ) 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/_embedding_router.py b/litellm/caching/_embedding_router.py new file mode 100644 index 00000000000..1ec898012e9 --- /dev/null +++ b/litellm/caching/_embedding_router.py @@ -0,0 +1,45 @@ +"""Shared selection of the embedding path for semantic caches. + +Both the Redis and qdrant semantic caches need the same decision: when the +configured embedding model is a proxy Router deployment, embeddings must run +through the Router so per-deployment auth (e.g. Bedrock aws_role_name) is +applied. Otherwise fall back to a direct litellm embedding call. + +This module is dependency-injected: callers pass the proxy ``llm_router`` and +``llm_model_list`` in, so the decision logic is unit-testable without importing +``litellm.proxy.proxy_server``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from litellm.router import Router + + +def resolve_embedding_router( + embedding_model: str, + llm_router: Router | None, + llm_model_list: list[dict[str, Any]] | None, +) -> Router | None: + """Return ``llm_router`` iff it serves ``embedding_model`` as a deployment.""" + if llm_router is None: + return None + router_model_names: list[str] = ( + [m["model_name"] for m in llm_model_list if "model_name" in m] + if llm_model_list is not None + else [] + ) + if embedding_model in router_model_names: + return llm_router + return None + + +def build_router_embedding_metadata( + request_metadata: dict[str, Any] | None, +) -> dict[str, Any]: + """Forward the caller's full metadata, flagged as a semantic-cache embedding.""" + metadata: dict[str, Any] = dict(request_metadata or {}) + metadata["semantic-cache-embedding"] = True + return metadata diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index cb122e90102..5f2269d1945 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -574,8 +574,9 @@ class Cache: if prompt_kwarg in kwargs: cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] - if isinstance(kwargs.get("metadata"), dict): - cache_lookup_kwargs["metadata"] = {} + metadata = kwargs.get("metadata") + if isinstance(metadata, dict): + cache_lookup_kwargs["metadata"] = dict(metadata) return cache_lookup_kwargs 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/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 68d3b8c20b3..504ef8a54eb 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.types.utils import EmbeddingResponse +from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router from .base_cache import BaseCache @@ -219,37 +220,50 @@ class QdrantSemanticCache(BaseCache): cached_key = payload.get(self.CACHE_KEY_FIELD_NAME) return cached_key is not None and str(cached_key) == str(key) - async def _get_async_embedding(self, prompt: str, **kwargs) -> Any: - llm_model_list = None - llm_router = None - + def _get_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> EmbeddingResponse: + """Embed via the proxy Router when it serves the model, else direct.""" try: - from litellm.proxy.proxy_server import ( - llm_model_list as proxy_llm_model_list, - llm_router as proxy_llm_router, - ) - - llm_model_list = proxy_llm_model_list - llm_router = proxy_llm_router + from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: - pass + llm_model_list = None + llm_router = None - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list ) - if llm_router is not None and self.embedding_model in router_model_names: - user_api_key = kwargs.get("metadata", {}).get("user_api_key", "") - return await llm_router.aembedding( + if router is not None: + return router.embedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, - metadata={ - "user_api_key": user_api_key, - "semantic-cache-embedding": True, - "trace_id": kwargs.get("metadata", {}).get("trace_id", None), - }, + metadata=build_router_embedding_metadata(metadata), + ) + return litellm.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + ) + + async def _get_async_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> EmbeddingResponse: + try: + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list + ) + if router is not None: + return await router.aembedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), ) return await litellm.aembedding( @@ -269,11 +283,7 @@ class QdrantSemanticCache(BaseCache): # create an embedding for prompt embedding_response = cast( EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding @@ -312,11 +322,7 @@ class QdrantSemanticCache(BaseCache): # convert to embedding embedding_response = cast( EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding @@ -388,7 +394,9 @@ class QdrantSemanticCache(BaseCache): # get the prompt messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, **kwargs) + embedding_response = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # get the embedding embedding = embedding_response["data"][0]["embedding"] @@ -424,7 +432,9 @@ class QdrantSemanticCache(BaseCache): messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, **kwargs) + embedding_response = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # get the embedding embedding = embedding_response["data"][0]["embedding"] diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index ba07511448a..5c9f74e4c18 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -435,6 +435,7 @@ class RedisCache(BaseCache): _redis_client = self.redis_client start_time = time.time() set_ttl = self.get_ttl(ttl=ttl) + key = self.check_and_fix_namespace(key=key) try: start_time = time.time() result: int = _redis_client.incr(name=key, amount=value) # type: ignore @@ -498,6 +499,7 @@ class RedisCache(BaseCache): ) return [] + pattern = self.check_and_fix_namespace(key=pattern) async for key in _redis_client.scan_iter(match=pattern + "*", count=count): # type: ignore keys.append(key) if len(keys) >= count: @@ -538,6 +540,11 @@ class RedisCache(BaseCache): Register a Lua script with Redis asynchronously. Works with both standalone Redis and Redis Cluster. + The returned callable namespaces every key it is invoked with, so Lua + scripts hit the same prefixed keys as get/set/increment. Without this, + scripts would operate on raw keys while the rest of the cache uses the + namespace, leaving rate-limit and lock keys outside the configured prefix. + Args: script (str): The Lua script to register @@ -548,7 +555,15 @@ class RedisCache(BaseCache): _redis_client = self.init_async_client() # For standalone Redis if hasattr(_redis_client, "register_script"): - return _redis_client.register_script(script) # type: ignore + registered_script = _redis_client.register_script(script) # type: ignore + + async def namespaced_script( + keys: list[str], args: list[Any], client: Any = None + ) -> Any: + keys = [self.check_and_fix_namespace(key=key) for key in keys] + return await registered_script(keys=keys, args=args, client=client) + + return namespaced_script # For Redis Cluster elif hasattr(_redis_client, "script_load"): # Load the script and get its SHA @@ -556,6 +571,7 @@ class RedisCache(BaseCache): # Return a callable that uses evalsha async def script_callable(keys: List[str], args: List[Any]) -> Any: + keys = [self.check_and_fix_namespace(key=key) for key in keys] return _redis_client.evalsha(script_sha, len(keys), *keys, *args) # type: ignore return script_callable @@ -1257,6 +1273,7 @@ class RedisCache(BaseCache): async def delete_cache_keys(self, keys): # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` _redis_client: Any = self.init_async_client() + keys = [self.check_and_fix_namespace(key=key) for key in keys] # keys is a list, unpack it so it gets passed as individual elements to delete await _redis_client.delete(*keys) @@ -1322,10 +1339,12 @@ class RedisCache(BaseCache): async def async_delete_cache(self, key: str): # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) # keys is str return await _redis_client.delete(key) def delete_cache(self, key): + key = self.check_and_fix_namespace(key=key) self.redis_client.delete(key) async def _pipeline_increment_helper( @@ -1432,6 +1451,7 @@ class RedisCache(BaseCache): try: # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl` _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) ttl = await _redis_client.ttl(key) if ttl <= -1: # -1 means the key does not exist, -2 key does not exist return None @@ -1460,6 +1480,7 @@ class RedisCache(BaseCache): int: The length of the list after the push operation """ _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) start_time = time.time() try: response = await _redis_client.rpush(key, *values) @@ -1499,7 +1520,8 @@ class RedisCache(BaseCache): ) -> List[int]: """Helper function for pipeline rpush operations""" for rpush_op in rpush_list: - pipe.rpush(rpush_op["key"], *rpush_op["values"]) + key = self.check_and_fix_namespace(key=rpush_op["key"]) + pipe.rpush(key, *rpush_op["values"]) results = await pipe.execute() # Preserve positional correspondence — raise on per-command errors for r in results: @@ -1586,6 +1608,7 @@ class RedisCache(BaseCache): **kwargs, ) -> Union[Any, List[Any]]: _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) start_time = time.time() print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") try: @@ -1658,17 +1681,19 @@ class RedisCache(BaseCache): if major_version >= 7: for lpop_op in lpop_list: - pipe.lpop(lpop_op["key"], lpop_op["count"]) + key = self.check_and_fix_namespace(key=lpop_op["key"]) + pipe.lpop(key, lpop_op["count"]) raw_results = await pipe.execute() else: # For Redis < 7, LPOP doesn't support count param. # Issue `count` individual LPOP commands per key, all in one pipeline. counts: List[int] = [] for lpop_op in lpop_list: + key = self.check_and_fix_namespace(key=lpop_op["key"]) count = lpop_op["count"] or 1 counts.append(count) for _ in range(count): - pipe.lpop(lpop_op["key"]) + pipe.lpop(key) flat_results = await pipe.execute() # Re-group the flat results back into per-key lists 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/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index cce4b75795f..e79392bb7f0 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -16,12 +16,13 @@ import os from typing import Any, Dict, List, Optional, Tuple, cast import litellm -from litellm._logging import print_verbose +from litellm._logging import print_verbose, verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.types.utils import EmbeddingResponse +from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router from .base_cache import BaseCache @@ -67,9 +68,6 @@ class RedisSemanticCache(BaseCache): Exception: If similarity_threshold is not provided or required Redis connection information is missing """ - from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped] - from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] - if index_name is None: index_name = self.DEFAULT_REDIS_INDEX_NAME @@ -107,15 +105,42 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Redis semantic-cache redis_url: {redis_url}") - # Initialize the Redis vectorizer and cache - cache_vectorizer = CustomTextVectorizer(self._get_embedding) + # Defer redisvl index construction until first use. redisvl's + # CustomTextVectorizer eagerly embeds a probe string at construction; + # building lazily ensures that probe runs after llm_router is wired so + # per-deployment auth (e.g. Bedrock aws_role_name) is applied. + self._index_name = index_name + self._redis_url = redis_url + self._llmcache = None - self.llmcache = self._init_semantic_cache( - semantic_cache_cls=SemanticCache, - index_name=index_name, - redis_url=redis_url, - cache_vectorizer=cache_vectorizer, - ) + @property + def llmcache(self) -> object: + if getattr(self, "_llmcache", None) is None: + self._llmcache = self._build_llmcache() + return self._llmcache + + @llmcache.setter + def llmcache(self, value: object) -> None: + self._llmcache = value + + def _build_llmcache(self) -> object: + # CustomTextVectorizer probes its embedding dimension at construction by + # embedding "dimension test", so the first cache request issues one extra + # billable embedding on top of the request's own. + from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped] + from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] + + try: + cache_vectorizer = CustomTextVectorizer(self._get_embedding) + return self._init_semantic_cache( + semantic_cache_cls=SemanticCache, + index_name=self._index_name, + redis_url=self._redis_url, + cache_vectorizer=cache_vectorizer, + ) + except Exception as e: + verbose_logger.error(f"Redis semantic-cache index build failed: {e}") + raise @classmethod def _cache_key_filterable_field(cls) -> Dict[str, str]: @@ -285,27 +310,43 @@ class RedisSemanticCache(BaseCache): return dict_method() return value - def _get_embedding(self, prompt: str) -> List[float]: + def _get_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> List[float]: """ - Generate an embedding vector for the given prompt using the configured embedding model. - - Args: - prompt: The text to generate an embedding for - - Returns: - List[float]: The embedding vector + Routes through the proxy Router when the embedding model is a Router + deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies, + mirroring ``_get_async_embedding``; otherwise embeds directly. """ - # Create an embedding from prompt - embedding_response = cast( - EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + try: + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list ) - embedding = embedding_response["data"][0]["embedding"] - return embedding + if router is not None: + embedding_response = cast( + EmbeddingResponse, + router.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), + ), + ) + else: + embedding_response = cast( + EmbeddingResponse, + litellm.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + ), + ) + return embedding_response["data"][0]["embedding"] def _get_cache_logic(self, cached_response: Any) -> Any: """ @@ -357,7 +398,12 @@ class RedisSemanticCache(BaseCache): value_str = str(value) - store_kwargs: Dict[str, Any] = { + prompt_embedding = self._get_embedding( + prompt, metadata=kwargs.get("metadata") + ) + + store_kwargs: dict[str, Any] = { + "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -393,8 +439,12 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Dict[str, Any] = { + prompt_embedding = self._get_embedding( + prompt, metadata=kwargs.get("metadata") + ) + check_kwargs: dict[str, Any] = { "prompt": prompt, + "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), } results = self.llmcache.check(**check_kwargs) @@ -435,49 +485,42 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _get_async_embedding(self, prompt: str, **kwargs) -> List[float]: + async def _get_async_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> List[float]: """ Asynchronously generate an embedding for the given prompt. Args: prompt: The text to generate an embedding for - **kwargs: Additional arguments that may contain metadata + metadata: Request metadata forwarded to the Router embedding call Returns: List[float]: The embedding vector """ - from litellm.proxy.proxy_server import llm_model_list, llm_router - - # Route the embedding request through the proxy if appropriate - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] - ) - try: - if llm_router is not None and self.embedding_model in router_model_names: - # Use the router for embedding generation - user_api_key = kwargs.get("metadata", {}).get("user_api_key", "") - embedding_response = await llm_router.aembedding( + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list + ) + try: + if router is not None: + embedding_response = await router.aembedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, - metadata={ - "user_api_key": user_api_key, - "semantic-cache-embedding": True, - "trace_id": kwargs.get("metadata", {}).get("trace_id", None), - }, + metadata=build_router_embedding_metadata(metadata), ) else: - # Generate embedding directly embedding_response = await litellm.aembedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, ) - - # Extract and return the embedding vector return embedding_response["data"][0]["embedding"] except Exception as e: print_verbose(f"Error generating async embedding: {str(e)}") @@ -504,9 +547,11 @@ class RedisSemanticCache(BaseCache): value_str = str(value) # Generate embedding for the value (response) to cache - prompt_embedding = await self._get_async_embedding(prompt, **kwargs) + prompt_embedding = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) - store_kwargs: Dict[str, Any] = { + store_kwargs: dict[str, Any] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -544,11 +589,13 @@ class RedisSemanticCache(BaseCache): return None # Generate embedding for the prompt - prompt_embedding = await self._get_async_embedding(prompt, **kwargs) + prompt_embedding = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Dict[str, Any] = { + check_kwargs: dict[str, Any] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), 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 212d34357f8..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) @@ -293,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 @@ -882,31 +886,29 @@ openai_compatible_providers: List = [ "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/google_genai/main.py b/litellm/google_genai/main.py index bdbb483dcf6..d35601c7e13 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -49,6 +49,7 @@ class GenerateContentSetupResult(BaseModel): custom_llm_provider: str generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] generate_content_config_dict: Dict[str, Any] + native_request_fields: dict[str, object] litellm_params: GenericLiteLLMParams litellm_logging_obj: LiteLLMLoggingObj litellm_call_id: Optional[str] @@ -152,6 +153,7 @@ class GenerateContentHelper: request_body={}, # Will be handled by adapter generate_content_provider_config=None, # type: ignore generate_content_config_dict=dict(config or {}), + native_request_fields={}, litellm_params=litellm_params, litellm_logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, @@ -171,6 +173,12 @@ class GenerateContentHelper: system_instruction = kwargs.get("systemInstruction") or kwargs.get( "system_instruction" ) + # Native top-level REST fields arrive as loose kwargs and are otherwise dropped. + native_request_fields: dict[str, object] = { + field: kwargs[field] + for field in generate_content_provider_config.get_generate_content_request_top_level_fields() + if field in kwargs + } request_body = ( generate_content_provider_config.transform_generate_content_request( model=model, @@ -201,12 +209,29 @@ class GenerateContentHelper: request_body=request_body, generate_content_provider_config=generate_content_provider_config, generate_content_config_dict=generate_content_config_dict, + native_request_fields=native_request_fields, litellm_params=litellm_params, litellm_logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, ) +def _merge_native_request_fields( + native_request_fields: dict[str, object], + extra_body: dict[str, object] | None, +) -> dict[str, object] | None: + """ + Merge native top-level request fields into ``extra_body`` so the HTTP handler + forwards them verbatim onto the outgoing request body. An explicit ``extra_body`` + value wins on conflict. Returns ``None`` only when there is genuinely nothing to + forward (no native fields and no caller-supplied ``extra_body``), preserving the + prior behavior without discarding an explicit ``extra_body={}``. + """ + if not native_request_fields and extra_body is None: + return None + return {**native_request_fields, **(extra_body or {})} + + @client async def agenerate_content( model: str, @@ -350,7 +375,9 @@ def generate_content( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), @@ -447,7 +474,9 @@ async def agenerate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=True, client=kwargs.get("client"), @@ -503,6 +532,11 @@ def generate_content_stream( **kwargs, ) + # Extract systemInstruction from kwargs to pass to handler + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) + # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: if "stream" in kwargs: @@ -531,12 +565,15 @@ def generate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), stream=True, litellm_metadata=kwargs.get("litellm_metadata", {}), + system_instruction=system_instruction, ) except Exception as e: 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/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/opentelemetry.py b/litellm/integrations/opentelemetry.py index 6b50ef49b49..c652141a2b6 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -189,6 +189,37 @@ def _normalize_team_metadata_keys(value: Any) -> List[str]: return [str(item).strip() for item in value if str(item).strip()] +_FREEZE_MAX_DEPTH = 16 + +HashableScope = Union[ + str, + int, + float, + bool, + bytes, + None, + tuple["HashableScope", ...], + frozenset["HashableScope"], +] + + +def _freeze_for_dedupe(value: object, _depth: int = 0) -> HashableScope: + if _depth >= _FREEZE_MAX_DEPTH: + return repr(value) + if isinstance(value, (list, tuple)): + return tuple(_freeze_for_dedupe(item, _depth + 1) for item in value) + if isinstance(value, set): + return frozenset(_freeze_for_dedupe(item, _depth + 1) for item in value) + if isinstance(value, dict): + return frozenset( + (_freeze_for_dedupe(key, _depth + 1), _freeze_for_dedupe(item, _depth + 1)) + for key, item in value.items() + ) + if isinstance(value, (str, int, float, bytes)) or value is None: + return value + return repr(value) + + @dataclass class OpenTelemetryConfig: exporter: Union[str, SpanExporter] = "console" @@ -1073,10 +1104,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): can be re-read with mutated entries between calls, so dedupe must be at entry granularity. Scope: the entry's stable identity. - ``scope`` parts can be any hashable identity. The marker is stored - in ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it - is request-local (kwargs is shared across the sync/async callbacks - and lifecycle hooks for one request). + ``scope`` parts may include unhashable containers (list, dict, set); + they are normalized into a hashable shape via ``_freeze_for_dedupe`` + before keying the marker dict. The marker is stored in + ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it is + request-local (kwargs is shared across the sync/async callbacks and + lifecycle hooks for one request). """ litellm_params = kwargs.get("litellm_params") if not isinstance(litellm_params, dict): @@ -1098,7 +1131,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): spans_logged = {} _otel_internal["spans_logged"] = spans_logged - dedupe_key = (self.__class__.__name__, id(self), *scope) + dedupe_key = ( + self.__class__.__name__, + id(self), + *(_freeze_for_dedupe(part) for part in scope), + ) if spans_logged.get(dedupe_key) is True: return False 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/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 4f4c046ec59..6f901ce0db9 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 8a02aa72b00..4a09e3f2b4d 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -621,7 +621,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": ( @@ -714,9 +718,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/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e05c8b3577a..2fa78900447 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 @@ -2953,7 +2968,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: @@ -3297,9 +3314,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) @@ -3767,7 +3782,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: @@ -4612,7 +4649,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: @@ -5786,7 +5823,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..d2c97a5d6e7 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 @@ -481,6 +485,7 @@ class PromptTokensDetailsResult(TypedDict): character_count: int image_count: int video_length_seconds: float + audio_length_seconds: float def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -531,6 +536,13 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + audio_length_seconds = ( + cast( + Optional[float], + getattr(usage.prompt_tokens_details, "audio_length_seconds", 0), + ) + or 0.0 + ) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -542,6 +554,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: character_count=character_count, image_count=image_count, video_length_seconds=float(video_length_seconds), + audio_length_seconds=float(audio_length_seconds), ) @@ -663,6 +676,14 @@ def _calculate_input_cost( prompt_tokens_details["video_length_seconds"], ) + ### AUDIO LENGTH COST + if prompt_tokens_details["audio_length_seconds"]: + prompt_cost += calculate_cost_component( + model_info, + "input_cost_per_audio_per_second", + prompt_tokens_details["audio_length_seconds"], + ) + return prompt_cost @@ -739,6 +760,7 @@ def generic_cost_per_token( character_count=0, image_count=0, video_length_seconds=0.0, + audio_length_seconds=0.0, ) if usage.prompt_tokens_details: prompt_tokens_details = _parse_prompt_tokens_details(usage) 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..4bdda6de2c8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -466,7 +466,7 @@ def get_format_from_file_id(file_id: Optional[str]) -> Optional[str]: def update_messages_with_model_file_ids( messages: List[AllMessageValues], - model_id: str, + model_id: str | None, model_file_id_mapping: Dict[str, Dict[str, str]], ) -> List[AllMessageValues]: """ @@ -519,7 +519,7 @@ def update_messages_with_model_file_ids( if file_id: provider_file_id = ( model_file_id_mapping.get(file_id, {}).get(model_id) - if model_file_id_mapping + if model_file_id_mapping and model_id is not None else None ) if ( @@ -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/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 b14e12de7cd..3135b5f831a 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -54,9 +54,9 @@ class SensitiveDataMasker: # 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..1b1b652eaa0 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"] @@ -585,6 +585,17 @@ class ChunkProcessor: # # Update usage information if needed prompt_tokens = 0 completion_tokens = 0 + # Anthropic's `message_start` SSE event carries usage.output_tokens=1 as a + # cursor/placeholder; the real value only arrives in `message_delta`. + # If a stream is cancelled before `message_delta` lands, the last-wins + # accumulator below leaves completion_tokens stuck at 1 — which then + # bypasses the `completion_tokens or token_counter(...)` fallback in + # calculate_usage() because 1 is truthy. Count the completion-bearing + # usage events so `_reset_anthropic_cursor_completion_tokens` can tell a + # legitimate single-token reply (Anthropic emits 1 in BOTH message_start + # AND message_delta, so >=2 events is positive evidence message_delta + # arrived) from a stale lone cursor. + completion_usage_updates = 0 ## anthropic prompt caching information ## cache_creation_input_tokens: Optional[int] = None cache_read_input_tokens: Optional[int] = None @@ -617,6 +628,7 @@ class ChunkProcessor: and usage_chunk_dict["completion_tokens"] > 0 ): completion_tokens = usage_chunk_dict["completion_tokens"] + completion_usage_updates += 1 if usage_chunk_dict["cache_creation_input_tokens"] is not None and ( usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None @@ -667,6 +679,12 @@ class ChunkProcessor: prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] + completion_tokens = self._reset_anthropic_cursor_completion_tokens( + chunks=chunks, + completion_tokens=completion_tokens, + completion_usage_updates=completion_usage_updates, + ) + return UsagePerChunk( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, @@ -678,6 +696,47 @@ class ChunkProcessor: prompt_tokens_details=prompt_tokens_details, ) + @staticmethod + def _reset_anthropic_cursor_completion_tokens( + chunks: list[dict[str, Any] | ModelResponse], + completion_tokens: int, + completion_usage_updates: int, + ) -> int: + """Reset a stale Anthropic ``message_start`` cursor placeholder to 0. + + See the ``completion_usage_updates`` comment in + ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale + cursor when either it is > 1 (definitely not a placeholder) or we saw + >= 2 completion-bearing usage events (positive evidence ``message_delta`` + arrived). Otherwise — the only completion update we ever saw was the + Anthropic ``message_start`` cursor (=1) — reset to 0 so + ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates + from the actually-received completion text instead of trusting the + placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the + heuristic (which encodes Anthropic's specific message_start SSE shape) + does not silently affect other providers that may legitimately report + ``completion_tokens=1`` from a single usage event. + """ + saw_non_cursor_completion = ( + completion_tokens > 1 or completion_usage_updates >= 2 + ) + if saw_non_cursor_completion: + return completion_tokens + + custom_llm_provider: Optional[str] = None + if chunks: + first_chunk = chunks[0] + if isinstance(first_chunk, dict): + hp = first_chunk.get("_hidden_params") + else: + hp = getattr(first_chunk, "_hidden_params", None) + if isinstance(hp, dict): + custom_llm_provider = hp.get("custom_llm_provider") + + if custom_llm_provider == "anthropic" and completion_tokens == 1: + return 0 + return completion_tokens + def calculate_usage( self, chunks: List[Union[Dict[str, Any], ModelResponse]], @@ -720,15 +779,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/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 39b1cc742d4..92a1510f3a1 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -21,16 +21,39 @@ else: LiteLLMLoggingObj = Any +OPENAI_STYLE_IMAGE_MODEL_PREFIXES: tuple[str, ...] = ("openai/",) + + class AimlImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api.aimlapi.com" IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" + @staticmethod + def _is_openai_style_model(model: str) -> bool: + """ + OpenAI image models routed through AI/ML API (e.g. ``openai/gpt-image-2``) + use the upstream OpenAI request schema, not the flux-style schema used by + the rest of the AI/ML catalog. + """ + return model.startswith(OPENAI_STYLE_IMAGE_MODEL_PREFIXES) + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: """ https://api.aimlapi.com/v1/images/generations """ + if self._is_openai_style_model(model): + return [ + "n", + "size", + "quality", + "response_format", + "output_format", + "background", + "moderation", + "output_compression", + ] return ["n", "response_format", "size"] def map_openai_params( @@ -41,39 +64,38 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): drop_params: bool, ) -> dict: supported_params = self.get_supported_openai_params(model) + is_openai_style = self._is_openai_style_model(model) for k in non_default_params.keys(): - if k not in optional_params.keys(): - if k in supported_params: - # Map OpenAI params to AI/ML params - if k == "n": - optional_params["num_images"] = non_default_params[k] - elif k == "response_format": - optional_params["output_format"] = non_default_params[k] - elif k == "size": - # Map OpenAI size format to AI/ML image_size - size_value = non_default_params[k] - if isinstance(size_value, str): - # Handle standard OpenAI sizes like "1024x1024" - if "x" in size_value: - width, height = map(int, size_value.split("x")) - optional_params["image_size"] = { - "width": width, - "height": height, - } - else: - # Pass through predefined sizes - optional_params["image_size"] = size_value - else: - optional_params["image_size"] = size_value - else: - optional_params[k] = non_default_params[k] - elif drop_params: - pass + if k in optional_params.keys(): + continue + if k not in supported_params: + if drop_params: + continue + raise ValueError( + f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." + ) + + if is_openai_style: + optional_params[k] = non_default_params[k] + continue + + if k == "n": + optional_params["num_images"] = non_default_params[k] + elif k == "response_format": + optional_params["output_format"] = non_default_params[k] + elif k == "size": + size_value = non_default_params[k] + if isinstance(size_value, str) and "x" in size_value: + width, height = map(int, size_value.split("x")) + optional_params["image_size"] = { + "width": width, + "height": height, + } else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) + optional_params["image_size"] = size_value + else: + optional_params[k] = non_default_params[k] return optional_params @@ -131,10 +153,13 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): headers: dict, ) -> dict: """ - Transform the image generation request to the AI/ML flux image generation request body + Transform the image generation request to the AI/ML image generation request body https://api.aimlapi.com/v1/images/generations """ + if self._is_openai_style_model(model): + return {"model": model, "prompt": prompt, **optional_params} + aiml_image_generation_request_body: AimlImageGenerationRequestParams = ( AimlImageGenerationRequestParams( prompt=prompt, diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index fd67a7fbaf1..7c4e9386d5f 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -43,7 +43,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig): api_base: Optional[str] = None, ) -> dict: """Validate and prepare environment-specific headers and parameters.""" - auth_header = self.anthropic_model_info.get_auth_header(api_key) + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params" 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..b1212f93059 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -364,6 +364,7 @@ class AnthropicChatCompletion(BaseLLM): messages=messages, optional_params={**optional_params, "is_vertex_request": is_vertex_request}, litellm_params=litellm_params, + api_base=api_base, ) config = ProviderConfigManager.get_provider_chat_config( @@ -624,7 +625,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 822b75b37f4..2b2264e6da9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -750,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 @@ -2144,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 0e41ef619ba..d27187cf2e7 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 ( @@ -498,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( @@ -543,6 +548,19 @@ class AnthropicModelInfo(BaseLLMModelInfo): return list(set(betas)) + @staticmethod + def _make_api_key_auth_header( + api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False + ) -> dict: + if use_bearer_for_custom_base and ( + api_base + and "api.anthropic.com" not in api_base + and not api_key.startswith("sk-ant-") + ): + value = api_key if api_key.startswith("Bearer ") else f"Bearer {api_key}" + return {"authorization": value} + return {"x-api-key": api_key} + def get_anthropic_headers( self, api_key: Optional[str] = None, @@ -562,6 +580,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): user_anthropic_beta_headers: Optional[List[str]] = None, code_execution_tool_used: bool = False, container_with_skills_used: bool = False, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, ) -> dict: betas = set() # Anthropic no longer requires the prompt-caching beta header @@ -610,7 +630,11 @@ class AnthropicModelInfo(BaseLLMModelInfo): elif auth_token and not api_key: headers["authorization"] = f"Bearer {auth_token}" elif api_key: - headers["x-api-key"] = api_key + headers.update( + self._make_api_key_auth_header( + api_key, api_base, use_bearer_for_custom_base + ) + ) if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) @@ -639,6 +663,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> Dict: + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + use_bearer_for_custom_base: bool = bool( + isinstance(litellm_params, dict) + and litellm_params.get("use_bearer_for_custom_base", False) + ) # Check for Anthropic OAuth token in headers headers, api_key = optionally_handle_anthropic_oauth( headers=headers, api_key=api_key @@ -694,6 +724,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): effort_used=effort_used, code_execution_tool_used=code_execution_tool_used, container_with_skills_used=container_with_skills_used, + api_base=api_base, + use_bearer_for_custom_base=use_bearer_for_custom_base, ) headers = {**headers, **anthropic_headers} @@ -729,18 +761,24 @@ class AnthropicModelInfo(BaseLLMModelInfo): return auth_token or get_secret_str("ANTHROPIC_AUTH_TOKEN") @staticmethod - def get_auth_header(api_key: Optional[str] = None) -> Optional[dict]: + def get_auth_header( + api_key: str | None = None, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, + ) -> dict | None: """Resolve Anthropic credentials and return the appropriate auth header dict. - Checks ANTHROPIC_API_KEY first (-> x-api-key), then - ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). + Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on + use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). Returns None if neither is available. """ resolved_key = AnthropicModelInfo.get_api_key(api_key) if resolved_key is not None: if is_anthropic_oauth_key(resolved_key): return {"authorization": f"Bearer {resolved_key}"} - return {"x-api-key": resolved_key} + return AnthropicModelInfo._make_api_key_auth_header( + resolved_key, api_base, use_bearer_for_custom_base + ) auth_token = AnthropicModelInfo.get_auth_token() if auth_token is not None: return {"authorization": f"Bearer {auth_token}"} @@ -754,7 +792,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): self, api_key: Optional[str] = None, api_base: Optional[str] = None ) -> List[str]: api_base = AnthropicModelInfo.get_api_base(api_base) - auth_header = AnthropicModelInfo.get_auth_header(api_key) + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if api_base is None or auth_header is None: raise ValueError( "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint." @@ -999,6 +1037,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..a92834b5d71 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="", ) @@ -207,32 +205,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if "delta" not in merged_chunk: merged_chunk["delta"] = {} - uncached_input_tokens = chunk.usage.prompt_tokens or 0 - if ( - hasattr(chunk.usage, "prompt_tokens_details") - and chunk.usage.prompt_tokens_details - ): - cached_tokens = ( - getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0 - ) - uncached_input_tokens -= cached_tokens + from .transformation import LiteLLMAnthropicMessagesAdapter - usage_dict: UsageDelta = { - "input_tokens": uncached_input_tokens, - "output_tokens": chunk.usage.completion_tokens or 0, - } - if ( - hasattr(chunk.usage, "_cache_creation_input_tokens") - and chunk.usage._cache_creation_input_tokens > 0 - ): - usage_dict["cache_creation_input_tokens"] = ( - chunk.usage._cache_creation_input_tokens - ) - if ( - hasattr(chunk.usage, "_cache_read_input_tokens") - and chunk.usage._cache_read_input_tokens > 0 - ): - usage_dict["cache_read_input_tokens"] = chunk.usage._cache_read_input_tokens + usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + chunk.usage + ) merged_chunk["usage"] = usage_dict if self.applied_edits and "context_management" not in merged_chunk: merged_chunk["context_management"] = ContextManagementResponse( diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 75a8acdfcc3..1fd593f63c8 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, @@ -1402,6 +1399,95 @@ class LiteLLMAnthropicMessagesAdapter: return "tool_use" return "end_turn" + @staticmethod + def _positive_int(value: object) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int) and value > 0: + return value + if isinstance(value, float) and value.is_integer() and value > 0: + return int(value) + return 0 + + @classmethod + def _first_positive_usage_value( + cls, usage: Usage, field_names: tuple[str, ...] + ) -> int: + for field_name in field_names: + value = cls._positive_int(getattr(usage, field_name, None)) + if value > 0: + return value + return 0 + + @classmethod + def _first_positive_prompt_tokens_detail_value( + cls, usage: Usage, field_names: tuple[str, ...] + ) -> int: + prompt_tokens_details = getattr(usage, "prompt_tokens_details", None) + if prompt_tokens_details is None: + return 0 + + for field_name in field_names: + if isinstance(prompt_tokens_details, dict): + value = cls._positive_int(prompt_tokens_details.get(field_name)) + else: + value = cls._positive_int( + getattr(prompt_tokens_details, field_name, None) + ) + if value > 0: + return value + return 0 + + @classmethod + def _get_cache_read_input_tokens(cls, usage: Usage) -> int: + explicit_value = cls._first_positive_usage_value( + usage, ("cache_read_input_tokens", "_cache_read_input_tokens") + ) + if explicit_value > 0: + return explicit_value + return cls._first_positive_prompt_tokens_detail_value(usage, ("cached_tokens",)) + + @classmethod + def _get_cache_creation_input_tokens(cls, usage: Usage) -> int: + explicit_value = cls._first_positive_usage_value( + usage, ("cache_creation_input_tokens", "_cache_creation_input_tokens") + ) + if explicit_value > 0: + return explicit_value + return cls._first_positive_prompt_tokens_detail_value( + usage, ("cache_creation_tokens", "cache_write_tokens") + ) + + @classmethod + def _translate_openai_usage_to_anthropic_usage_delta( + cls, usage: Usage + ) -> UsageDelta: + cache_read_input_tokens = cls._get_cache_read_input_tokens(usage) + cache_creation_input_tokens = cls._get_cache_creation_input_tokens(usage) + input_tokens = max( + (usage.prompt_tokens or 0) + - cache_read_input_tokens + - cache_creation_input_tokens, + 0, + ) + + usage_delta = UsageDelta( + input_tokens=input_tokens, + output_tokens=usage.completion_tokens or 0, + ) + if cache_creation_input_tokens > 0: + usage_delta["cache_creation_input_tokens"] = cache_creation_input_tokens + if cache_read_input_tokens > 0: + usage_delta["cache_read_input_tokens"] = cache_read_input_tokens + return usage_delta + + @classmethod + def _translate_openai_usage_to_anthropic_usage(cls, usage: Usage) -> AnthropicUsage: + return cast( + AnthropicUsage, + cls._translate_openai_usage_to_anthropic_usage_delta(usage), + ) + def translate_openai_response_to_anthropic( self, response: ModelResponse, @@ -1433,35 +1519,17 @@ class LiteLLMAnthropicMessagesAdapter: ) # extract usage usage: Usage = getattr(response, "usage") - uncached_input_tokens = usage.prompt_tokens or 0 - cached_tokens = 0 - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - cached_tokens = ( - getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 - ) - uncached_input_tokens -= cached_tokens - - anthropic_usage = AnthropicUsage( - input_tokens=uncached_input_tokens, - output_tokens=usage.completion_tokens or 0, - ) - if ( - hasattr(usage, "_cache_creation_input_tokens") - and usage._cache_creation_input_tokens > 0 - ): - anthropic_usage["cache_creation_input_tokens"] = ( - usage._cache_creation_input_tokens - ) - if cached_tokens > 0: - anthropic_usage["cache_read_input_tokens"] = cached_tokens + anthropic_usage = self._translate_openai_usage_to_anthropic_usage(usage) if polyfill_result is not None and polyfill_result.iterations_usage is not None: message_iteration: UsageIteration = { "type": "message", - "input_tokens": uncached_input_tokens, + "input_tokens": anthropic_usage["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 +1569,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": {}, } @@ -1647,39 +1713,15 @@ class LiteLLMAnthropicMessagesAdapter: else: litellm_usage_chunk = None if litellm_usage_chunk is not None: - uncached_input_tokens = litellm_usage_chunk.prompt_tokens or 0 - cached_tokens = 0 - if ( - hasattr(litellm_usage_chunk, "prompt_tokens_details") - and litellm_usage_chunk.prompt_tokens_details - ): - cached_tokens = ( - getattr( - litellm_usage_chunk.prompt_tokens_details, - "cached_tokens", - 0, - ) - or 0 - ) - uncached_input_tokens -= cached_tokens - - usage_delta = UsageDelta( - input_tokens=uncached_input_tokens, - output_tokens=litellm_usage_chunk.completion_tokens or 0, + usage_delta = self._translate_openai_usage_to_anthropic_usage_delta( + litellm_usage_chunk ) - if ( - hasattr(litellm_usage_chunk, "_cache_creation_input_tokens") - and litellm_usage_chunk._cache_creation_input_tokens > 0 - ): - usage_delta["cache_creation_input_tokens"] = ( - litellm_usage_chunk._cache_creation_input_tokens - ) - if cached_tokens > 0: - usage_delta["cache_read_input_tokens"] = cached_tokens 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 7b10a447bc8..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) 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/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/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 56296df94a1..170cb086bb0 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -84,7 +84,7 @@ class AnthropicFilesHandler: # Get Anthropic API credentials api_base = self.anthropic_model_info.get_api_base(api_base) - auth_header = self.anthropic_model_info.get_auth_header(api_key) + auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError("Missing Anthropic API Key") diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index ea9bf00f505..7ffb6beb4c7 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -95,7 +95,9 @@ class AnthropicFilesConfig(BaseFilesConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - auth_header = AnthropicModelInfo.get_auth_header(api_key) + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter." diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 4ea768b02af..2bfdf7ef4ba 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -38,10 +38,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): # Get API key from litellm_params if available api_key = None + api_base = None if litellm_params is not None: api_key = litellm_params.api_key + api_base = litellm_params.api_base - auth_header = AnthropicModelInfo.get_auth_header(api_key) + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API" 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/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index d4144a75718..cc65ad706ab 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -11,7 +11,7 @@ The operation location must be polled until the analysis completes. import asyncio import re import time -from typing import Any, Dict, Optional +from typing import Any, Dict from urllib.parse import quote import httpx @@ -35,6 +35,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.secret_managers.main import get_secret_str +AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" + class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ @@ -54,6 +56,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def __init__(self) -> None: super().__init__() + def get_api_key_env_var(self) -> str | None: + return AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR + def get_supported_ocr_params(self, model: str) -> list: """ Get supported OCR parameters for Azure Document Intelligence. @@ -144,9 +149,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -156,7 +161,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") + api_key = get_secret_str(AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR) if api_key is None: raise ValueError( @@ -182,10 +187,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py index f661ddb9ebc..ee35fc28994 100644 --- a/litellm/llms/azure_ai/ocr/transformation.py +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -2,7 +2,7 @@ Azure AI OCR transformation implementation. """ -from typing import Dict, Optional +from typing import Dict from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.image_handling import ( @@ -13,6 +13,8 @@ from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestDat from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.secret_managers.main import get_secret_str +AZURE_AI_OCR_API_KEY_ENV_VAR = "AZURE_AI_API_KEY" + class AzureAIOCRConfig(MistralOCRConfig): """ @@ -30,13 +32,16 @@ class AzureAIOCRConfig(MistralOCRConfig): def __init__(self) -> None: super().__init__() + def get_api_key_env_var(self) -> str | None: + return AZURE_AI_OCR_API_KEY_ENV_VAR + def validate_environment( self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -46,7 +51,7 @@ class AzureAIOCRConfig(MistralOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = get_secret_str("AZURE_AI_API_KEY") + api_key = get_secret_str(AZURE_AI_OCR_API_KEY_ENV_VAR) if api_key is None: raise ValueError( @@ -72,10 +77,10 @@ class AzureAIOCRConfig(MistralOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ 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/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index e8b3bf1a576..8fb7eb9fde0 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -62,6 +62,19 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): "get_supported_generate_content_optional_params is not implemented" ) + def get_generate_content_request_top_level_fields(self) -> tuple[str, ...]: + """ + Native Google ``GenerateContentRequest`` fields that sit at the top level + (siblings of ``generationConfig``) rather than inside it. The proxy forwards + these verbatim from a native request so ``generateContent`` is a drop-in for + Google's REST API. + + Excludes ``contents``, ``model`` and ``tools`` (dedicated params), + ``systemInstruction`` (dedicated extraction) and ``generationConfig`` (mapped + to ``config``). + """ + return ("safetySettings", "toolConfig", "cachedContent", "labels") + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 263e0c094ce..a2946c62506 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -2,7 +2,7 @@ Base OCR transformation configuration. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Union import httpx from pydantic import PrivateAttr @@ -25,16 +25,16 @@ DocumentType = Dict[str, str] class OCRPageDimensions(LiteLLMPydanticObjectBase): """Page dimensions from OCR response.""" - dpi: Optional[int] = None - height: Optional[int] = None - width: Optional[int] = None + dpi: int | None = None + height: int | None = None + width: int | None = None class OCRPageImage(LiteLLMPydanticObjectBase): """Image extracted from OCR page.""" - image_base64: Optional[str] = None - bbox: Optional[Dict[str, Any]] = None + image_base64: str | None = None + bbox: Dict[str, Any] | None = None model_config = {"extra": "allow"} @@ -44,8 +44,8 @@ class OCRPage(LiteLLMPydanticObjectBase): index: int markdown: str - images: Optional[List[OCRPageImage]] = None - dimensions: Optional[OCRPageDimensions] = None + images: List[OCRPageImage] | None = None + dimensions: OCRPageDimensions | None = None model_config = {"extra": "allow"} @@ -53,9 +53,9 @@ class OCRPage(LiteLLMPydanticObjectBase): class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" - pages_processed: Optional[int] = None - credits: Optional[float] = None - doc_size_bytes: Optional[int] = None + pages_processed: int | None = None + credits: float | None = None + doc_size_bytes: int | None = None model_config = {"extra": "allow"} @@ -68,8 +68,8 @@ class OCRResponse(LiteLLMPydanticObjectBase): pages: List[OCRPage] model: str - document_annotation: Optional[Any] = None - usage_info: Optional[OCRUsageInfo] = None + document_annotation: Any | None = None + usage_info: OCRUsageInfo | None = None object: str = "ocr" model_config = {"extra": "allow"} @@ -81,8 +81,8 @@ class OCRResponse(LiteLLMPydanticObjectBase): class OCRRequestData(LiteLLMPydanticObjectBase): """OCR request data structure.""" - data: Optional[Union[Dict, bytes]] = None - files: Optional[Dict[str, Any]] = None + data: Union[Dict, bytes] | None = None + files: Dict[str, Any] | None = None class BaseOCRConfig: @@ -101,6 +101,12 @@ class BaseOCRConfig: """ return [] + def get_api_key_env_var(self) -> str | None: + """ + Return the provider-specific API key environment variable name, if any. + """ + return None + def map_ocr_params( self, non_default_params: dict, @@ -114,9 +120,9 @@ class BaseOCRConfig: self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -127,10 +133,10 @@ class BaseOCRConfig: def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ 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 9fca7bc61af..29e97068100 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -225,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 @@ -321,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 @@ -1300,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 diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5e97394f459..69c0a8529b4 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -904,9 +904,12 @@ class BedrockModelInfo(BaseLLMModelInfo): "mantle/": "mantle", } - # Check explicit routes first + # Check explicit routes first. Match each prefix only as a leading path + # segment so the `bedrock_mantle/` provider prefix is never mistaken for + # the `mantle/` invoke route (which would mangle + # `bedrock_mantle/openai.gpt-5.5` into `bedrock_openai.gpt-5.5`). for prefix, route_type in route_mappings.items(): - if prefix in model: + if BedrockModelInfo._model_has_route_prefix(model, prefix): return route_type # Check for nova spec prefixes (nova/ and nova-2/) @@ -930,14 +933,14 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Check if the model is an explicit converse route. """ - return "converse/" in model + return BedrockModelInfo._model_has_route_prefix(model, "converse/") @staticmethod def _explicit_claude_platform_route(model: str) -> bool: """ Check if the model is an explicit Claude Platform on AWS route. """ - return "claude_platform/" in model + return BedrockModelInfo._model_has_route_prefix(model, "claude_platform/") @staticmethod def get_claude_platform_model(model: str) -> str: @@ -967,42 +970,58 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Check if the model is an explicit invoke route. """ - return "invoke/" in model + return BedrockModelInfo._model_has_route_prefix(model, "invoke/") @staticmethod def _explicit_agent_route(model: str) -> bool: """ Check if the model is an explicit agent route. """ - return "agent/" in model + return BedrockModelInfo._model_has_route_prefix(model, "agent/") @staticmethod def _explicit_agentcore_route(model: str) -> bool: """ Check if the model is an explicit agentcore route. """ - return "agentcore/" in model + return BedrockModelInfo._model_has_route_prefix(model, "agentcore/") + + @staticmethod + def _model_has_route_prefix(model: str, prefix: str) -> bool: + """Whether a route prefix (e.g. ``mantle/``) appears as a leading path segment. + + A route token is only valid at the start of the model id or immediately + after a ``/``. A plain substring check matches the ``bedrock_mantle/`` + provider prefix against the ``mantle/`` route, so the body model gets + mangled to ``bedrock_openai.gpt-5.5``; anchoring to a segment boundary + keeps the bare model id intact. + + ``f"/{prefix}" in model`` matches the token as a segment at any path + depth, not just the second segment; that is intentional and acceptable + for these short, unambiguous route tokens. + """ + return model.startswith(prefix) or f"/{prefix}" in model @staticmethod def _explicit_mantle_route(model: str) -> bool: """ Check if the model is an explicit mantle route (bedrock-mantle endpoint). """ - return "mantle/" in model + return BedrockModelInfo._model_has_route_prefix(model, "mantle/") @staticmethod def _explicit_converse_like_route(model: str) -> bool: """ Check if the model is an explicit converse like route. """ - return "converse_like/" in model + return BedrockModelInfo._model_has_route_prefix(model, "converse_like/") @staticmethod def _explicit_async_invoke_route(model: str) -> bool: """ Check if the model is an explicit async invoke route. """ - return "async_invoke/" in model + return BedrockModelInfo._model_has_route_prefix(model, "async_invoke/") @staticmethod def _explicit_openai_route(model: str) -> bool: @@ -1010,7 +1029,7 @@ class BedrockModelInfo(BaseLLMModelInfo): Check if the model is an explicit openai route. Used for Bedrock imported models that use OpenAI Chat Completions format. """ - return "openai/" in model + return BedrockModelInfo._model_has_route_prefix(model, "openai/") @staticmethod def get_bedrock_provider_config_for_messages_api( 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/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/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 1000ab12803..33601b47609 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -72,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. @@ -719,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() @@ -780,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() @@ -841,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() @@ -888,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() @@ -1200,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() @@ -1235,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() @@ -1274,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 @@ -1312,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() 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 948c90f9f99..d33ec295e94 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,13 +1,14 @@ +import asyncio import json import ssl from functools import lru_cache -from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, + Iterator, List, Literal, Optional, @@ -16,6 +17,7 @@ from typing import ( cast, get_type_hints, ) +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx # type: ignore from openai.types.file_deleted import FileDeleted @@ -27,8 +29,8 @@ 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.realtime_streaming import RealTimeStreaming 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 ( BaseAnthropicMessagesConfig, @@ -107,6 +109,7 @@ from litellm.types.llms.openai import ( 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 @@ -132,7 +135,6 @@ 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, @@ -3438,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 ): @@ -3519,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, }, @@ -3596,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 ): @@ -3639,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", 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 e153d00e6ab..a8cb40ac6db 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -1755,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/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/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/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index 21e0e27a314..07a67815f6e 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -2,7 +2,7 @@ Mistral OCR transformation implementation. """ -from typing import Any, Dict, Optional +from typing import Any, Dict import httpx @@ -15,6 +15,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.secret_managers.main import get_secret_str +MISTRAL_OCR_API_KEY_ENV_VAR = "MISTRAL_API_KEY" + class MistralOCRConfig(BaseOCRConfig): """ @@ -42,6 +44,7 @@ class MistralOCRConfig(BaseOCRConfig): - extract_footer: Whether to extract document footer - table_format: Table output format ("markdown" or "html") - confidence_scores_granularity: Confidence score level ("word" or "page") + - include_blocks: Whether to return paragraph-level bounding boxes and typed content blocks (OCR 4) - id: Request identifier """ return [ @@ -56,9 +59,13 @@ class MistralOCRConfig(BaseOCRConfig): "extract_footer", "table_format", "confidence_scores_granularity", + "include_blocks", "id", ] + def get_api_key_env_var(self) -> str | None: + return MISTRAL_OCR_API_KEY_ENV_VAR + def map_ocr_params( self, non_default_params: dict, @@ -85,9 +92,9 @@ class MistralOCRConfig(BaseOCRConfig): self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -95,7 +102,7 @@ class MistralOCRConfig(BaseOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = get_secret_str("MISTRAL_API_KEY") + api_key = get_secret_str(MISTRAL_OCR_API_KEY_ENV_VAR) if api_key is None: raise ValueError( @@ -113,10 +120,10 @@ class MistralOCRConfig(BaseOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ 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/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/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/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/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index b92f7ca1aff..4a95f519646 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -149,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..3bdcbd25949 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 @@ -1039,9 +1046,9 @@ def _gemini_convert_messages_with_history( if invocation.get("tool_type"): tr_dict["toolType"] = invocation["tool_type"] tr_part: Dict[str, Any] = {"toolResponse": tr_dict} - if "thought_signature" in invocation: + if "response_thought_signature" in invocation: tr_part["thoughtSignature"] = invocation[ - "thought_signature" + "response_thought_signature" ] assistant_content.append(tr_part) # 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..423a5dd5d17 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 @@ -1633,13 +1634,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): resp = tool_responses_by_id.pop(call_id, None) if resp is not None: merged["response"] = resp.get("response") - # Keep response signature if call didn't have one - if "thought_signature" not in merged and "thought_signature" in resp: - merged["thought_signature"] = resp["thought_signature"] + if "thought_signature" in resp: + merged["response_thought_signature"] = resp["thought_signature"] invocations.append(merged) # Any orphan responses (shouldn't happen, but be safe) for resp_id, resp_entry in tool_responses_by_id.items(): + if "thought_signature" in resp_entry: + resp_entry["response_thought_signature"] = resp_entry[ + "thought_signature" + ] invocations.append(resp_entry) return invocations if invocations else None @@ -2195,7 +2199,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 +2476,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 +2544,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 +3323,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/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 165dac24903..53f5fb464df 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -274,6 +274,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore @@ -377,6 +378,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index ba6e6f0c056..27ca1bd92a4 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,10 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from typing import Dict, List, Optional, Tuple +from collections.abc import Mapping +from typing import Dict, List, Optional, Sequence, Tuple + +from pydantic import TypeAdapter, ValidationError from litellm.types.llms.vertex_ai import ( BlobType, @@ -13,10 +16,17 @@ from litellm.types.llms.vertex_ai import ( FileDataType, GeminiEmbeddingInput, PartType, + PromptTokensDetails, + UsageMetadata, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + PromptTokensDetailsWrapper, + Usage, +) from litellm.utils import get_formatted_prompt, token_counter SUPPORTED_EMBEDDING_MIME_TYPES = { @@ -294,11 +304,133 @@ def transform_openai_input_gemini_embed_content( return request_body +_IMAGE_MIME_TYPES = frozenset({"image/png", "image/jpeg"}) +_VIDEO_TOKENS_PER_SECOND = 258.0 +_AUDIO_TOKENS_PER_SECOND = 32.0 +_usage_metadata_adapter = TypeAdapter(UsageMetadata) + + +def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata]: + if not isinstance(raw_usage_metadata, dict): + return None + try: + return _usage_metadata_adapter.validate_python(raw_usage_metadata) + except ValidationError: + return None + + +def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: + if isinstance(input, str): + return (input,) + return tuple( + sub + for element in input + for sub in (element if isinstance(element, list) else [element]) + ) + + +def _is_image_element( + element: str, + resolved_files: Mapping[str, Mapping[str, str]], +) -> bool: + if element.startswith("data:") and ";base64," in element: + try: + mime_type, _ = _parse_data_url(element) + except ValueError: + return False + return mime_type in _IMAGE_MIME_TYPES + if _is_gcs_url(element): + try: + return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES + except ValueError: + return False + if _is_file_reference(element): + file_info = resolved_files.get(element) + return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES + return False + + +def _count_input_images( + input: GeminiEmbeddingInput, + resolved_files: Mapping[str, Mapping[str, str]], +) -> int: + return sum( + 1 + for element in _flatten_input(input) + if _is_image_element(element, resolved_files) + ) + + +def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: + return sum( + detail["tokenCount"] for detail in details if detail["modality"] == modality + ) + + +def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage: + if _is_multimodal_input(input): + return Usage(prompt_tokens=0, total_tokens=0) + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + return Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens) + + +def _usage_from_embed_content_response( + input: GeminiEmbeddingInput, + model: str, + raw_usage_metadata: object, + resolved_files: Mapping[str, Mapping[str, str]], +) -> Usage: + usage_metadata = _parse_usage_metadata(raw_usage_metadata) + if usage_metadata is None: + return _fallback_usage(input, model) + + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens + + details: Sequence[PromptTokensDetails] = ( + usage_metadata.get("promptTokensDetails") or () + ) + text_tokens = _tokens_for_modality(details, "TEXT") + audio_tokens = _tokens_for_modality(details, "AUDIO") + video_tokens = _tokens_for_modality(details, "VIDEO") + image_count = _count_input_images(input, resolved_files) + + video_length_seconds = ( + video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 + ) + audio_length_seconds = ( + audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 + ) + + # generic_cost_per_token rewrites text_tokens to the full prompt minus + # other modalities when both text_tokens and image_count are zero. For + # video, that misallocates video tokens to text; a 1-token floor sidesteps + # the rewrite and keeps billing on input_cost_per_video_per_second. + needs_video_text_floor = ( + video_length_seconds > 0 and text_tokens == 0 and image_count == 0 + ) + resolved_text_tokens = 1 if needs_video_text_floor else text_tokens + + return Usage( + prompt_tokens=prompt_tokens, + total_tokens=total_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=resolved_text_tokens, + audio_tokens=audio_tokens, + image_count=image_count, + video_length_seconds=video_length_seconds, + audio_length_seconds=audio_length_seconds, + ), + ) + + def process_embed_content_response( input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, response_json: dict, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). @@ -308,6 +440,8 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint + resolved_files: Mapping of file references (files/abc) to {mime_type, uri}, + used to bill resolved image references at the per-image rate Returns: EmbeddingResponse with single embedding @@ -327,14 +461,11 @@ def process_embed_content_response( model_response.data = [openai_embedding] model_response.model = model - - if _is_multimodal_input(input): - prompt_tokens = 0 - else: - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) - model_response.usage = Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + model_response.usage = _usage_from_embed_content_response( + input=input, + model=model, + raw_usage_metadata=response_json.get("usageMetadata"), + resolved_files=resolved_files or {}, ) return model_response diff --git a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py index 516ee03ba55..a98311d04eb 100644 --- a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py +++ b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py @@ -3,7 +3,7 @@ Vertex AI DeepSeek OCR transformation implementation. """ import json -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict import httpx @@ -18,6 +18,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY" + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj else: @@ -28,21 +30,24 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): """ Vertex AI DeepSeek OCR transformation configuration. - Vertex AI DeepSeek OCR uses the chat completion API format through the openapi endpoint. - This transformation converts OCR requests to chat completion format and vice versa. + This transformation converts standard LiteLLM OCR requests to the + Vertex AI DeepSeek OCR OpenAPI endpoint shape and normalizes the response. """ def __init__(self) -> None: super().__init__() self.vertex_base = VertexBase() + def get_api_key_env_var(self) -> str | None: + return VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR + def validate_environment( self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -50,6 +55,13 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): Vertex AI uses Bearer token authentication with access token from credentials. """ + if api_key is not None: + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + **headers, + } + # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} @@ -77,18 +89,15 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ Get complete URL for Vertex AI DeepSeek OCR endpoint. - Vertex AI endpoint format: - https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions - Args: api_base: Vertex AI API base URL (optional) model: Model name (e.g., "deepseek-ai/deepseek-ocr-maas") @@ -123,8 +132,6 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # Ensure no trailing slash api_base = api_base.rstrip("/") - # Vertex AI DeepSeek OCR endpoint format - # Format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/endpoints/openapi/chat/completions return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" def transform_ocr_request( @@ -136,9 +143,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRRequestData: """ - Transform OCR request to chat completion format for Vertex AI DeepSeek OCR. + Transform OCR request for Vertex AI DeepSeek OCR. - Converts OCR document format to chat completion messages format: + Converts OCR document format to the Vertex AI DeepSeek OCR payload: - Input: {"type": "image_url", "image_url": "gs://..."} - Output: {"model": "deepseek-ai/deepseek-ocr-maas", "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "gs://..."}]}]} @@ -150,7 +157,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs: Additional arguments Returns: - OCRRequestData with JSON data in chat completion format + OCRRequestData with JSON data for the DeepSeek OCR endpoint """ verbose_logger.debug( "Vertex AI DeepSeek OCR transform_ocr_request (sync) called" @@ -173,7 +180,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'" ) - # Build chat completion message content + # Build DeepSeek OCR message content content_item = {} if image_url: content_item = {"type": "image_url", "image_url": image_url} @@ -181,25 +188,21 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # For document URLs, we use image_url type as well (Vertex AI supports both) content_item = {"type": "image_url", "image_url": document_url} - # Build chat completion request + # Build DeepSeek OCR request data = { "model": "deepseek-ai/" + model, "messages": [{"role": "user", "content": [content_item]}], } # Add optional parameters (stream, temperature, etc.) - # Filter out OCR-specific params that don't apply to chat completion - chat_completion_params = {} + deepseek_ocr_params = {} for key, value in optional_params.items(): - # Include common chat completion params if key in ["stream", "temperature", "max_tokens", "top_p", "n", "stop"]: - chat_completion_params[key] = value + deepseek_ocr_params[key] = value - data.update(chat_completion_params) + data.update(deepseek_ocr_params) - verbose_logger.debug( - "Vertex AI DeepSeek OCR: Transformed request to chat completion format" - ) + verbose_logger.debug("Vertex AI DeepSeek OCR: Transformed request") return OCRRequestData(data=data, files=None) @@ -212,7 +215,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRRequestData: """ - Transform OCR request to chat completion format for Vertex AI DeepSeek OCR (async). + Transform OCR request for Vertex AI DeepSeek OCR (async). Same as sync version - no async-specific logic needed. @@ -224,7 +227,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs: Additional arguments Returns: - OCRRequestData with JSON data in chat completion format + OCRRequestData with JSON data for the DeepSeek OCR endpoint """ return self.transform_ocr_request( model=model, @@ -242,12 +245,11 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRResponse: """ - Transform chat completion response to OCR format. + Transform Vertex AI DeepSeek OCR response to OCR format. - Vertex AI DeepSeek OCR returns chat completion format: + Vertex AI DeepSeek OCR returns an OpenAPI response: { "id": "...", - "object": "chat.completion", "choices": [{ "message": { "role": "assistant", @@ -274,16 +276,16 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): try: response_json = raw_response.json() - # Extract content from chat completion response + # Extract OCR content from provider response choices = response_json.get("choices", []) if not choices: - raise ValueError("No choices in chat completion response") + raise ValueError("No choices in DeepSeek OCR response") message = choices[0].get("message", {}) content = message.get("content", "") if not content: - raise ValueError("No content in chat completion response") + raise ValueError("No content in DeepSeek OCR response") # Try to parse content as JSON (OCR result might be JSON string) ocr_data = None @@ -376,7 +378,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRResponse: """ - Async transform chat completion response to OCR format. + Async transform Vertex AI DeepSeek OCR response to OCR format. Same as sync version - no async-specific logic needed. diff --git a/litellm/llms/vertex_ai/ocr/transformation.py b/litellm/llms/vertex_ai/ocr/transformation.py index cbf15803132..a725762b3c5 100644 --- a/litellm/llms/vertex_ai/ocr/transformation.py +++ b/litellm/llms/vertex_ai/ocr/transformation.py @@ -2,7 +2,7 @@ Vertex AI Mistral OCR transformation implementation. """ -from typing import Dict, Optional +from typing import Dict from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.image_handling import ( @@ -14,6 +14,8 @@ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +VERTEX_AI_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY" + class VertexAIOCRConfig(MistralOCRConfig): """ @@ -32,13 +34,16 @@ class VertexAIOCRConfig(MistralOCRConfig): super().__init__() self.vertex_base = VertexBase() + def get_api_key_env_var(self) -> str | None: + return VERTEX_AI_OCR_API_KEY_ENV_VAR + def validate_environment( self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -46,6 +51,13 @@ class VertexAIOCRConfig(MistralOCRConfig): Vertex AI uses Bearer token authentication with access token from credentials. """ + if api_key is not None: + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + **headers, + } + # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} @@ -73,10 +85,10 @@ class VertexAIOCRConfig(MistralOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 8a92e7ec4a5..a633ef3298a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -35,27 +35,21 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert Validate the environment for the request """ + # Work on a local copy — router shallow-copies litellm_params so the caller's + # headers dict may be the shared deployment extra_headers object. + headers = dict(headers) vertex_ai_project = VertexBase.safe_get_vertex_ai_project(litellm_params) vertex_ai_location = VertexBase.safe_get_vertex_ai_location(litellm_params) - project_id: Optional[str] = None - if "Authorization" not in headers: - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( - litellm_params - ) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params) + access_token, project_id = self._ensure_access_token( + credentials=vertex_credentials, + project_id=vertex_ai_project, + custom_llm_provider="vertex_ai", + ) + headers["Authorization"] = f"Bearer {access_token}" - access_token, project_id = self._ensure_access_token( - credentials=vertex_credentials, - project_id=vertex_ai_project, - custom_llm_provider="vertex_ai", - ) - - headers["Authorization"] = f"Bearer {access_token}" - else: - # Authorization already in headers, but we still need project_id - project_id = vertex_ai_project - - # Always calculate api_base if not provided, regardless of Authorization header + # Calculate api_base if not provided if api_base is None: api_base = self.get_complete_vertex_url( custom_api_base=api_base, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 960d3483848..78669f1e789 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -194,9 +194,11 @@ class VertexAIPartnerModels(VertexBase): encoding=encoding, ) elif "claude" in model: - if headers is None: - headers = {} - headers.update({"Authorization": "Bearer {}".format(access_token)}) + # Build a new dict so we never mutate the shared deployment extra_headers object. + headers = { + **(headers or {}), + "Authorization": "Bearer {}".format(access_token), + } optional_params.update( { 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/main.py b/litellm/main.py index cabb070dbb1..983807027ed 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -633,6 +633,7 @@ async def acompletion( try: # Use a partial function to pass your keyword arguments + kwargs.pop("acompletion", None) func = partial(completion, **completion_kwargs, **kwargs) # Add the context to the function @@ -5063,6 +5064,12 @@ def completion( # type: ignore ######### unpacking kwargs ##################### args = locals() + # Set by the responses->completion fallback so completion() does not bridge + # back to the Responses API: that round-trip mutually recurses forever for a + # model whose model_cost mode is "responses" but whose provider has no + # Responses API config (get_provider_responses_api_config -> None). + skip_responses_api_bridge = kwargs.pop("_skip_responses_api_bridge", False) + skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False) if not skip_mcp_handler and tools: from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp @@ -5358,7 +5365,7 @@ def completion( # type: ignore messages = update_messages_with_model_file_ids( messages=messages, - model_id=kwargs.get("model_info", {}).get("id", None), + model_id=(kwargs.get("model_info") or {}).get("id", None), model_file_id_mapping=cast( Dict[str, Dict[str, str]], kwargs.get("model_file_id_mapping") or {}, @@ -5448,8 +5455,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" @@ -5562,7 +5570,10 @@ def completion( # type: ignore # detection when the deployment name differs from the model name. _azure_detection_model = base_model or model - if responses_api_model_info.get("mode") == "responses": + if ( + responses_api_model_info.get("mode") == "responses" + and not skip_responses_api_bridge + ): from litellm.completion_extras import responses_api_bridge optional_params, rs_val = ( @@ -7587,8 +7598,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 @@ -7659,7 +7670,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 @@ -8065,7 +8078,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, @@ -8406,7 +8424,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 @@ -8418,9 +8438,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: @@ -8877,9 +8895,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 6ebac7efc8d..d44fc654a56 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -273,6 +273,19 @@ "/v1/images/generations" ] }, + "aiml/openai/gpt-image-2": { + "litellm_provider": "aiml", + "metadata": { + "notes": "OpenAI gpt-image-2 via AI/ML API - flagship multimodal image generation and editing model with reasoning and 2K output. output_cost_per_image is AI/ML's published medium-quality rate; like the other aiml image entries it is billed as a flat per-image price" + }, + "mode": "image_generation", + "output_cost_per_image": 0.054, + "source": "https://docs.aimlapi.com/api-references/image-models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "amazon.nova-canvas-v1:0": { "litellm_provider": "bedrock", "max_input_tokens": 2600, @@ -25550,8 +25563,18 @@ }, "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", - "ocr_cost_per_page": 0.001, - "annotation_cost_per_page": 0.003, + "ocr_cost_per_page": 0.004, + "annotation_cost_per_page": 0.005, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, + "mistral/mistral-ocr-4-0": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.004, + "annotation_cost_per_page": 0.005, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -25770,7 +25793,7 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "mistral/mistral-medium-latest": { + "mistral/mistral-medium-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, @@ -25778,12 +25801,45 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, + "mistral/mistral-medium-2604": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-latest": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium-3-1-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", @@ -25810,6 +25866,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 3a9ef8db804..93b3a892659 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -4,13 +4,12 @@ Main OCR function for LiteLLM. import asyncio import base64 -import contextvars import mimetypes import os import re -from functools import partial +from dataclasses import dataclass from io import IOBase -from typing import Any, Callable, Coroutine, Dict, Optional, Union, cast +from typing import Any, Callable, Coroutine, Union, cast import httpx @@ -20,7 +19,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.rust_bridge import ocr as rust_ocr_bridge from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -29,91 +28,287 @@ 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) +@dataclass +class _PreparedOCRRequest: + model: str + document: dict[str, Any] + api_key: str | None + api_base: str | None + custom_llm_provider: str + extra_headers: dict[str, object] | None + provider_config: BaseOCRConfig + optional_params: dict[str, object] + litellm_params: dict[str, object] + effective_timeout: Union[float, httpx.Timeout] + litellm_logging_obj: LiteLLMLoggingObj -def _run_rust_ocr( - rust_ocr: RustOcr, - logging_obj: LiteLLMLoggingObj, - provider_config: BaseOCRConfig, - resolve_api_key: Callable[[str], Optional[str]], +@dataclass +class _PreparedRustOCRCall: + api_key: str | None + api_base: str | None + headers: dict[str, object] + optional_params: dict[str, object] + + +_RUST_OCR_PROVIDERS = { + "mistral", + "azure_ai", + "vertex_ai", +} + + +def _prepare_ocr_request( 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. + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + timeout: Union[float, httpx.Timeout] | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + kwargs: dict[str, Any], +) -> _PreparedOCRRequest: + litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) + litellm_call_id = cast(str | None, kwargs.get("litellm_call_id", None)) - 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={}, + if not isinstance(document, dict): + raise ValueError( + f"document must be a dict with 'type' and URL/file field, got {type(document)}" + ) + + doc_type = document.get("type") + + if doc_type == "file": + document = convert_file_document_to_url_document(document) + doc_type = document.get("type") + + if doc_type not in ["document_url", "image_url"]: + raise ValueError( + f"Invalid document type: {doc_type}. " + "Must be 'document_url', 'image_url', or 'file'" + ) + + ( + model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( model=model, - api_key=resolved_api_key, + custom_llm_provider=custom_llm_provider, api_base=api_base, - litellm_params=litellm_params, + api_key=api_key, ) - resolved_complete_url = provider_config.get_complete_url( - api_base=api_base, + + if dynamic_api_key: + api_key = dynamic_api_key + if dynamic_api_base: + api_base = dynamic_api_base + + ocr_provider_config = ProviderConfigManager.get_provider_ocr_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) + + if ocr_provider_config is None: + raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}") + + verbose_logger.debug(f"OCR call - model: {model}, provider: {custom_llm_provider}") + + litellm_params = GenericLiteLLMParams(**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) + + optional_params = ocr_provider_config.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + ) + + verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") + + effective_timeout = timeout or request_timeout + + litellm_logging_obj.update_from_kwargs( + kwargs=kwargs, model=model, optional_params=optional_params, - litellm_params=litellm_params, + litellm_params={ + "litellm_call_id": litellm_call_id, + "api_base": api_base, + }, + custom_llm_provider=custom_llm_provider, ) - logging_obj.pre_call( + + return _PreparedOCRRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + provider_config=ocr_provider_config, + optional_params=cast(dict[str, object], optional_params), + litellm_params=dict(litellm_params), + effective_timeout=effective_timeout, + litellm_logging_obj=litellm_logging_obj, + ) + + +def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: + return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS + + +def _rust_bridge_optional_params( + prepared_request: _PreparedOCRRequest, + resolve_secret: Callable[[str], str | None], +) -> dict[str, object]: + optional_params = dict(prepared_request.optional_params) + if prepared_request.custom_llm_provider == "vertex_ai": + vertex_project = ( + prepared_request.litellm_params.get("vertex_project") + or prepared_request.litellm_params.get("vertex_ai_project") + or litellm.vertex_project + or resolve_secret("VERTEXAI_PROJECT") + ) + vertex_location = ( + prepared_request.litellm_params.get("vertex_location") + or prepared_request.litellm_params.get("vertex_ai_location") + or litellm.vertex_location + or resolve_secret("VERTEXAI_LOCATION") + or resolve_secret("VERTEX_LOCATION") + ) + if vertex_project is not None: + optional_params["vertex_project"] = vertex_project + if vertex_location is not None: + optional_params["vertex_location"] = vertex_location + return optional_params + + +def _rust_bridge_api_base( + prepared_request: _PreparedOCRRequest, + resolve_secret: Callable[[str], str | None], +) -> str | None: + if prepared_request.api_base is not None: + return prepared_request.api_base + if prepared_request.custom_llm_provider == "azure_ai": + model = prepared_request.model.lower() + if "doc-intelligence" in model or "documentintelligence" in model: + return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") + return resolve_secret("AZURE_AI_API_BASE") + return None + + +def _prepare_rust_ocr_call( + prepared_request: _PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> _PreparedRustOCRCall: + provider_config = prepared_request.provider_config + api_key_env_var = provider_config.get_api_key_env_var() + resolved_api_key = prepared_request.api_key or ( + resolve_api_key(api_key_env_var) if api_key_env_var is not None else None + ) + resolved_headers = provider_config.validate_environment( + headers=prepared_request.extra_headers or {}, + model=prepared_request.model, + api_key=resolved_api_key, + api_base=prepared_request.api_base, + litellm_params=prepared_request.litellm_params, + ) + resolved_complete_url = provider_config.get_complete_url( + api_base=prepared_request.api_base, + model=prepared_request.model, + optional_params=prepared_request.optional_params, + litellm_params=prepared_request.litellm_params, + ) + rust_api_base = _rust_bridge_api_base(prepared_request, resolve_api_key) + rust_optional_params = _rust_bridge_optional_params( + prepared_request, resolve_api_key + ) + prepared_request.litellm_logging_obj.pre_call( input="OCR document processing", api_key=resolved_api_key, additional_args={ "complete_input_dict": { - "model": model, - "document": document, - **optional_params, + "model": prepared_request.model, + "document": prepared_request.document, + **rust_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, - ) + return _PreparedRustOCRCall( + api_key=resolved_api_key, + api_base=rust_api_base, + headers=cast(dict[str, object], resolved_headers), + optional_params=rust_optional_params, ) +def _run_rust_ocr( + prepared_request: _PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> OCRResponse | None: + if rust_ocr_bridge.load_rust_ocr() is None: + return None + prepared = _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ) + rust_response = rust_ocr_bridge.ocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, + ) + if rust_response is None: + return None + return OCRResponse.model_validate(rust_response) + + +async def _run_rust_aocr( + prepared_request: _PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> OCRResponse | None: + if rust_ocr_bridge.load_rust_aocr() is None: + return None + prepared = _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ) + rust_response = await rust_ocr_bridge.aocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, + ) + if rust_response is None: + return None + return OCRResponse.model_validate(rust_response) + + @client async def aocr( model: str, - document: Dict[str, Any], - api_key: Optional[str] = None, - api_base: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + document: dict[str, Any], + api_key: str | None = None, + api_base: str | None = None, + timeout: Union[float, httpx.Timeout] | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, Any] | None = None, **kwargs, ) -> OCRResponse: """ @@ -174,19 +369,18 @@ async def aocr( ) ``` """ - local_vars = locals() + completion_kwargs: dict[str, object] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } try: - loop = asyncio.get_event_loop() - kwargs["aocr"] = True - - # Get custom llm provider - if custom_llm_provider is None: - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, api_base=api_base - ) - - func = partial( - ocr, + prepared = _prepare_ocr_request( model=model, document=document, api_key=api_key, @@ -194,17 +388,45 @@ async def aocr( timeout=timeout, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - **kwargs, + kwargs=kwargs, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update( + {"model": model, "custom_llm_provider": custom_llm_provider} ) - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + from litellm.secret_managers.main import get_secret_str - if asyncio.iscoroutine(init_response): - response = await init_response - else: - response = init_response + rust_response = await _run_rust_aocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is None: + verbose_logger.debug( + "Async Rust OCR bridge unavailable; falling back to Python path" + ) + else: + return rust_response + + response = base_llm_http_handler.ocr( + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=True, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + if asyncio.iscoroutine(response): + response = await response if response is None: raise ValueError( @@ -217,20 +439,145 @@ async def aocr( model=model, custom_llm_provider=custom_llm_provider, original_exception=e, - completion_kwargs=local_vars, + completion_kwargs=completion_kwargs, extra_kwargs=kwargs, ) +################################################# +# Public utilities — used by the SDK and the proxy +################################################# + +_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$") + +_MIME_TYPE_MAP = { + ".pdf": "application/pdf", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".bmp": "image/bmp", +} + + +def get_mime_type(file_path: str) -> str: + """ + Determine MIME type from file path extension. + + Falls back to mimetypes.guess_type, then to 'application/octet-stream'. + """ + ext = os.path.splitext(file_path)[1].lower() + mime = _MIME_TYPE_MAP.get(ext) + if mime: + return mime + guessed, _ = mimetypes.guess_type(file_path) + return guessed or "application/octet-stream" + + +def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]: + """ + Convert a file-type document dict to a document_url-type document dict + with an inline base64 data URI. + + Accepts document dicts like: + {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path + {"type": "file", "file": } # file-like object (BinaryIO) + {"type": "file", "file": b"raw bytes"} # raw bytes + + Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or + ``open(path, "rb")`` instead. See the str check below for the rationale. + + Returns: + {"type": "document_url", "document_url": "data:;base64,"} + or {"type": "image_url", "image_url": "data:;base64,"} + """ + file_input = document.get("file") + if file_input is None: + raise ValueError( + "document with type='file' must include a 'file' field containing " + "a pathlib.Path, file-like object, or bytes" + ) + + file_bytes: bytes + mime_type: str = "application/octet-stream" + file_name: str | None = None + + if isinstance(file_input, str): + # Bare strings are rejected here. The OCR ``document`` accepts a + # ``{"type": "file", "file": }`` shape, and when this helper + # runs in a proxy request handler ```` is attacker-controlled. + # Opening it as a path is an arbitrary local file read on the proxy + # host, which is then base64-encoded and forwarded to the OCR + # provider — an exfiltration primitive. + raise ValueError( + "OCR file input does not accept bare str values. Pass bytes, " + "a pathlib.Path, or a file-like object. To OCR a local file " + "from a path, call open(path, 'rb') yourself." + ) + if isinstance(file_input, os.PathLike): + # os.PathLike (pathlib.Path and custom __fspath__ classes) is a + # Python-level type that HTTP form values can't fabricate. + file_path = str(file_input) + if not os.path.isfile(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + mime_type = get_mime_type(file_path) + file_name = os.path.basename(file_path) + with open(file_path, "rb") as f: + file_bytes = f.read() + elif isinstance(file_input, bytes): + file_bytes = file_input + elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): + if hasattr(file_input, "name"): + file_name = getattr(file_input, "name", None) + if file_name: + mime_type = get_mime_type(file_name) + file_bytes = file_input.read() + if isinstance(file_bytes, str): + file_bytes = file_bytes.encode("utf-8") + else: + raise ValueError( + f"Unsupported file input type: {type(file_input)}. " + "Expected pathlib.Path, bytes, or a file-like object." + ) + + if not file_bytes: + raise ValueError("File is empty or could not be read") + + if "mime_type" in document: + mime_type = document["mime_type"] + + if not _MIME_PATTERN.match(mime_type): + raise ValueError(f"Invalid MIME type: {mime_type}") + + base64_data = base64.b64encode(file_bytes).decode("utf-8") + data_uri = f"data:{mime_type};base64,{base64_data}" + + if mime_type.startswith("image/"): + verbose_logger.debug( + f"OCR file input: Converted file to image_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "image_url", "image_url": data_uri} + + verbose_logger.debug( + f"OCR file input: Converted file to document_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "document_url", "document_url": data_uri} + + @client def ocr( model: str, - document: Dict[str, Any], - api_key: Optional[str] = None, - api_base: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + document: dict[str, Any], + api_key: str | None = None, + api_base: str | None = None, + timeout: Union[float, httpx.Timeout] | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, Any] | None = None, **kwargs, ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: """ @@ -295,131 +642,62 @@ def ocr( print(f"Page {page.index}: {page.markdown}") ``` """ - local_vars = locals() + completion_kwargs: dict[str, object] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } try: - 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 - - # Validate document parameter format - if not isinstance(document, dict): - raise ValueError( - f"document must be a dict with 'type' and URL/file field, got {type(document)}" - ) - - doc_type = document.get("type") - - # Handle file type: convert to document_url/image_url with base64 data URI - if doc_type == "file": - document = convert_file_document_to_url_document(document) - doc_type = document.get("type") - - if doc_type not in ["document_url", "image_url"]: - raise ValueError( - f"Invalid document type: {doc_type}. " - "Must be 'document_url', 'image_url', or 'file'" - ) - - ( - model, - custom_llm_provider, - dynamic_api_key, - dynamic_api_base, - ) = litellm.get_llm_provider( + completion_kwargs["aocr"] = _is_async + prepared = _prepare_ocr_request( model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, + document=document, api_key=api_key, - ) - - # Update with dynamic values if available - if dynamic_api_key: - api_key = dynamic_api_key - if dynamic_api_base: - api_base = dynamic_api_base - - ocr_provider_config: Optional[BaseOCRConfig] = ( - ProviderConfigManager.get_provider_ocr_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) - ) - - if ocr_provider_config is None: - raise ValueError( - f"OCR is not supported for provider: {custom_llm_provider}" - ) - - verbose_logger.debug( - f"OCR call - model: {model}, provider: {custom_llm_provider}" - ) - - litellm_params = GenericLiteLLMParams(**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) - - optional_params = ocr_provider_config.map_ocr_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - ) - - verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") - - effective_timeout = timeout or request_timeout - - litellm_logging_obj.update_from_kwargs( + api_base=api_base, kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params={ - "litellm_call_id": litellm_call_id, - "api_base": api_base, - }, custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout=timeout, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update( + {"model": model, "custom_llm_provider": custom_llm_provider} ) - # 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: + if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + from litellm.secret_managers.main import get_secret_str + + rust_response = _run_rust_ocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response 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), - ) + return rust_response response = base_llm_http_handler.ocr( - model=model, - document=document, - optional_params=optional_params, - timeout=effective_timeout, - logging_obj=litellm_logging_obj, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, aocr=_is_async, - headers=extra_headers, - provider_config=ocr_provider_config, - litellm_params=dict(litellm_params), + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, ) return response @@ -428,131 +706,6 @@ def ocr( model=model, custom_llm_provider=custom_llm_provider, original_exception=e, - completion_kwargs=local_vars, + completion_kwargs=completion_kwargs, extra_kwargs=kwargs, ) - - -################################################# -# Public utilities — used by the SDK and the proxy -################################################# - -_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$") - -_MIME_TYPE_MAP = { - ".pdf": "application/pdf", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".webp": "image/webp", - ".tiff": "image/tiff", - ".tif": "image/tiff", - ".bmp": "image/bmp", -} - - -def get_mime_type(file_path: str) -> str: - """ - Determine MIME type from file path extension. - - Falls back to mimetypes.guess_type, then to 'application/octet-stream'. - """ - ext = os.path.splitext(file_path)[1].lower() - mime = _MIME_TYPE_MAP.get(ext) - if mime: - return mime - guessed, _ = mimetypes.guess_type(file_path) - return guessed or "application/octet-stream" - - -def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str, str]: - """ - Convert a file-type document dict to a document_url-type document dict - with an inline base64 data URI. - - Accepts document dicts like: - {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path - {"type": "file", "file": } # file-like object (BinaryIO) - {"type": "file", "file": b"raw bytes"} # raw bytes - - Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or - ``open(path, "rb")`` instead. See the str check below for the rationale. - - Returns: - {"type": "document_url", "document_url": "data:;base64,"} - or {"type": "image_url", "image_url": "data:;base64,"} - """ - file_input = document.get("file") - if file_input is None: - raise ValueError( - "document with type='file' must include a 'file' field containing " - "a pathlib.Path, file-like object, or bytes" - ) - - file_bytes: bytes - mime_type: str = "application/octet-stream" - file_name: Optional[str] = None - - if isinstance(file_input, str): - # Bare strings are rejected here. The OCR ``document`` accepts a - # ``{"type": "file", "file": }`` shape, and when this helper - # runs in a proxy request handler ```` is attacker-controlled. - # Opening it as a path is an arbitrary local file read on the proxy - # host, which is then base64-encoded and forwarded to the OCR - # provider — an exfiltration primitive. - raise ValueError( - "OCR file input does not accept bare str values. Pass bytes, " - "a pathlib.Path, or a file-like object. To OCR a local file " - "from a path, call open(path, 'rb') yourself." - ) - if isinstance(file_input, os.PathLike): - # os.PathLike (pathlib.Path and custom __fspath__ classes) is a - # Python-level type that HTTP form values can't fabricate. - file_path = str(file_input) - if not os.path.isfile(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - mime_type = get_mime_type(file_path) - file_name = os.path.basename(file_path) - with open(file_path, "rb") as f: - file_bytes = f.read() - elif isinstance(file_input, bytes): - file_bytes = file_input - elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): - if hasattr(file_input, "name"): - file_name = getattr(file_input, "name", None) - if file_name: - mime_type = get_mime_type(file_name) - file_bytes = file_input.read() - if isinstance(file_bytes, str): - file_bytes = file_bytes.encode("utf-8") - else: - raise ValueError( - f"Unsupported file input type: {type(file_input)}. " - "Expected pathlib.Path, bytes, or a file-like object." - ) - - if not file_bytes: - raise ValueError("File is empty or could not be read") - - if "mime_type" in document: - mime_type = document["mime_type"] - - if not _MIME_PATTERN.match(mime_type): - raise ValueError(f"Invalid MIME type: {mime_type}") - - base64_data = base64.b64encode(file_bytes).decode("utf-8") - data_uri = f"data:{mime_type};base64,{base64_data}" - - if mime_type.startswith("image/"): - verbose_logger.debug( - f"OCR file input: Converted file to image_url data URI " - f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" - ) - return {"type": "image_url", "image_url": data_uri} - else: - verbose_logger.debug( - f"OCR file input: Converted file to document_url data URI " - f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" - ) - return {"type": "document_url", "document_url": data_uri} diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py deleted file mode 100644 index 61f9e8ca69a..00000000000 --- a/litellm/ocr/rust_bridge.py +++ /dev/null @@ -1,74 +0,0 @@ -""" -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 c3ebc81541d..aa4aef74e13 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -544,7 +544,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 hasattr(response, "iter_bytes") and is_streaming_request: 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..b7ac6a8c325 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]}..." ) @@ -624,7 +624,9 @@ class MCPRequestHandler: Permission hierarchy (all rules are intersections): 1. Get allowed servers from key permissions - 2. Get allowed servers from team permissions (key inherits from team, or intersection) + 2. Get allowed servers from team permissions (key inherits from team, or + intersection; or inherits nothing when require_key_mcp_access_defined + is enabled, making the team a ceiling rather than a default) 3. Get allowed servers from end_user permissions (intersected if set) 4. Get allowed servers from agent permissions (intersected if set) 5. Get allowed servers from org permissions — org acts as a ceiling: if the org @@ -677,7 +679,16 @@ class MCPRequestHandler: if not team_set: base = key_set # no team restriction elif not key_set: - base = team_set # key has no own perms → inherits team + # A key that grants no MCP servers of its own inherits the + # team's by default. With require_key_mcp_access_defined the + # team is a ceiling rather than a default, so the key must + # grant servers explicitly (or via an access group) to reach + # any — it inherits none. + base = ( + set() + if general_settings.get("require_key_mcp_access_defined", False) + else team_set + ) else: base = key_set & team_set # both restrict → intersect @@ -1383,9 +1394,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/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 index 7bcdb3e6529..969bbf01ec8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -7,9 +7,9 @@ no precedence cascade. It is wildcard-free with an `assert_never` tail, so addin an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly at runtime instead of returning `None`. -This skeleton ships every arm as a `not_implemented` stub. Each mode's real body, with its -injected seam, lands in its own follow-up PR; until then the arm returns a typed error rather -than silently producing no credential. Pure v2: no imports from v1. +`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 @@ -17,8 +17,13 @@ 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 ( @@ -26,11 +31,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthorizationCodeConfig, AuthSpecKind, AwsSigV4Config, + Byok, ClientCredentialsConfig, CredError, NoneConfig, PassthroughConfig, ServerSpec, + SharedKey, Subject, TokenExchangeConfig, ) @@ -40,7 +47,7 @@ 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 skeleton needs none, since every arm is a stub. + is built; the live `none` and `api_key`-shared arms read from the config and need none. """ async def resolve_credentials( @@ -48,9 +55,9 @@ class UpstreamCredentialProvider: ) -> Result[httpx.Auth, CredError]: match server.config: case NoneConfig(): - return _not_implemented(AuthSpecKind.none) - case ApiKeyConfig(): - return _not_implemented(AuthSpecKind.api_key) + return Ok(NoOpAuth()) + case ApiKeyConfig() as config: + return self._api_key(config) case PassthroughConfig(): return _not_implemented(AuthSpecKind.passthrough) case ClientCredentialsConfig(): @@ -63,6 +70,22 @@ class UpstreamCredentialProvider: 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( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2149f079a3d..7b4f1e13a52 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -79,6 +79,7 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.server import ( ListMCPToolsRestAPIResponseObject, + MCPInfo, MCPServer, _tool_name_matches, execute_mcp_tool, @@ -238,14 +239,24 @@ if MCP_AVAILABLE: ) return {} - def _create_tool_response_objects(tools, server_mcp_info): - """Helper function to create tool response objects.""" + def _create_tool_response_objects(tools, server: MCPServer): + """Helper function to create tool response objects. + + Enriches the server's ``mcp_info`` with ``server_id`` and ``alias`` so + REST clients can map the internal ``server_name`` to the user-facing + alias without needing access to the ``mcp_routes``-gated server listing. + """ + enriched_mcp_info: MCPInfo = { + **(server.mcp_info or {}), + "server_id": server.server_id, + "alias": server.alias, + } return [ ListMCPToolsRestAPIResponseObject( name=tool.name, description=tool.description, inputSchema=tool.inputSchema, - mcp_info=server_mcp_info, + mcp_info=enriched_mcp_info, ) for tool in tools ] @@ -405,7 +416,7 @@ if MCP_AVAILABLE: ) if not apply_tool_filters: - return _create_tool_response_objects(tools, server.mcp_info) + return _create_tool_response_objects(tools, server) # Always apply allowed_tools/disallowed_tools so the blacklist is # enforced even when no allowlist is set (matches the SSE/HTTP path). @@ -436,7 +447,7 @@ if MCP_AVAILABLE: if _tool_name_matches(tool.name, allowed_tools_for_server) ] - return _create_tool_response_objects(tools, server.mcp_info) + return _create_tool_response_objects(tools, server) async def _resolve_allowed_mcp_servers_for_tool_call( user_api_key_dict: UserAPIKeyAuth, @@ -587,6 +598,8 @@ if MCP_AVAILABLE: "mcp_info": { "server_name": "zapier", "logo_url": "https://www.zapier.com/logo.png", + "server_id": "a1b2c3d4-...", + "alias": "zapier_prod", } } ], 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 5bde80ecc90..4035387c74e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1053,9 +1053,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 @@ -2362,6 +2362,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.", @@ -3177,6 +3182,7 @@ class SpendLogsMetadata(TypedDict): dict ] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] + litellm_call_id: Optional[str] applied_guardrails: Optional[List[str]] mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] @@ -3995,9 +4001,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_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/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/client/exceptions.py b/litellm/proxy/client/exceptions.py index fffd1b78b7e..c4089381e30 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -2,18 +2,30 @@ from typing import Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + + +def _redact_orig_exception( + orig_exception: Union[requests.exceptions.HTTPError, str], +) -> Union[requests.exceptions.HTTPError, str]: + if isinstance(orig_exception, requests.exceptions.HTTPError): + return requests.exceptions.HTTPError( + redact_string(str(orig_exception)), response=orig_exception.response + ) + return redact_string(str(orig_exception)) + class UnauthorizedError(Exception): """Exception raised when the API returns a 401 Unauthorized response.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) class NotFoundError(Exception): """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index d8687cbad16..845b49d1581 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + from .exceptions import UnauthorizedError @@ -314,6 +316,9 @@ class KeysManagementClient: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: + redacted_message = redact_string(str(e)) if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + raise UnauthorizedError(e) from None + raise requests.exceptions.HTTPError( + redacted_message, response=e.response + ) from None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0135d625bc9..0f7d40c1181 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -516,12 +516,13 @@ def _override_openai_response_model( LiteLLM internally prefixes some provider/deployment model identifiers (e.g. `hosted_vllm/...`). That internal identifier should not be returned to clients in the OpenAI `model` field. - Note: This is intentionally verbose. A model mismatch is a useful signal that an internal - model identifier is being stamped/preserved somewhere in the request/response pipeline. - We log mismatches as warnings (and then restamp to the client-requested value) so these - paths stay observable for maintainers/operators without breaking client compatibility. + Note: This is intentionally verbose at debug level. A model mismatch is a useful signal that an + internal model identifier is being stamped/preserved somewhere in the request/response pipeline. + We log mismatches as debug (and then restamp to the client-requested value) so these paths stay + observable for maintainers without breaking client compatibility or alarming operators. - Errors are reserved for cases where the proxy cannot read/override the response model field. + Responses that omit an OpenAI-style `model` field are left unchanged (silent return), + including dict responses with no `model` key. Exceptions: 1. If a fallback occurred (indicated by x-litellm-attempted-fallbacks header), @@ -577,6 +578,8 @@ def _override_openai_response_model( return if isinstance(response_obj, dict): + if "model" not in response_obj: + return downstream_model = response_obj.get("model") if downstream_model != requested_model: verbose_proxy_logger.debug( @@ -589,11 +592,6 @@ def _override_openai_response_model( return if not hasattr(response_obj, "model"): - verbose_proxy_logger.error( - "%s: cannot override response model; missing `model` attribute. response_type=%s", - log_context, - type(response_obj), - ) return downstream_model = getattr(response_obj, "model", None) @@ -608,7 +606,7 @@ def _override_openai_response_model( try: setattr(response_obj, "model", requested_model) except Exception as e: - verbose_proxy_logger.error( + verbose_proxy_logger.debug( "%s: failed to override response.model=%r on response_type=%s. error=%s", log_context, requested_model, @@ -1094,9 +1092,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 @@ -1539,7 +1537,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": streaming_headers = custom_headers 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/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/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8a432eb2f42..928d00109fa 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1849,7 +1849,10 @@ async def test_model_connection( "responses", "ocr", ] - ] = fastapi.Body("chat", description="The mode to test the model with"), + ] = fastapi.Body( + None, + description="The mode to test the model with. If not provided, auto-detected from model capabilities.", + ), litellm_params: Dict = fastapi.Body( None, description="Parameters for litellm.completion, litellm.embedding for the health check", 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/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/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index f28bcc2bcd4..bc2da33672f 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,8 +1,27 @@ +import math from typing import TYPE_CHECKING, Any, Dict, Optional, Union from fastapi import HTTPException, status from pydantic import BaseModel + +# Defined above the `litellm.proxy.*` imports so the name is bound even when +# this module is imported first through the proxy import cycle (CodeQL: +# module-level cyclic import). Depends only on `math` + `HTTPException`. +def validate_finite_spend(spend: float | None) -> None: + """Reject NaN/±inf spend before it reaches the DB / spend counter. + + A non-finite spend would otherwise slip past `spend >= max_budget` + enforcement, since any comparison with NaN (and `-inf >= max_budget`) + is False, letting the entity keep spending past its configured budget. + """ + if spend is not None and not math.isfinite(spend): + raise HTTPException( + status_code=400, + detail={"error": f"spend must be a finite number. Received: {spend}"}, + ) + + from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( 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..99e276cfbb4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _user_has_admin_view, require_caller_user_id_for_non_admin, + validate_finite_spend, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -746,7 +747,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 +1054,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 @@ -1256,6 +1257,25 @@ def _check_user_update_authz( ) +async def _invalidate_user_spend_counter_if_changed( + non_default_values: dict[str, Any], +) -> None: + """Invalidate the cross-pod spend counter after a direct ``spend`` change. + + A direct ``spend`` change must also invalidate the cross-pod spend counter + enforcement reads; the DB write alone leaves a warm counter at the stale + value. ``non_default_values["user_id"]`` is populated in every branch of the + caller (incl. the email-new-user insert path, whose response is a bare model + and not safely subscriptable). + """ + if non_default_values.get("spend") is not None: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter( + counter_key=f"spend:user:{non_default_values['user_id']}" + ) + + async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, @@ -1336,6 +1356,9 @@ async def _update_single_user_helper( existing_metadata=existing_metadata or {}, ) + # Reject NaN/±inf spend before it can reach the DB / spend counter. + validate_finite_spend(non_default_values.get("spend")) + # Perform the update response: Optional[Dict[str, Any]] = None @@ -1384,6 +1407,8 @@ async def _update_single_user_helper( litellm_proxy_admin_name=litellm_proxy_admin_name, ) + await _invalidate_user_spend_counter_if_changed(non_default_values) + if response is None: raise HTTPException( status_code=400, @@ -1705,7 +1730,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..9f316256ac9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -69,6 +69,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _set_object_metadata_field, _team_member_has_permission, + validate_finite_spend, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, @@ -229,10 +230,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 +383,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 +877,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 +966,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 +2053,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 +2169,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 +2192,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, @@ -2204,6 +2211,9 @@ async def _validate_update_key_data( user_api_key_cache: Any, ) -> None: """Validate permissions and constraints for key update.""" + # Reject NaN/±inf spend before it can reach the DB / spend counter. + validate_finite_spend(data.spend) + _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value _check_allowed_routes_caller_permission( @@ -2263,12 +2273,14 @@ async def _validate_update_key_data( # existing admin-only budget semantics). budget_limits uses # model_fields_set because an explicit null/[] clears the field # and must gate the same as setting or changing it. + # - spend gates on presence alone (not a value diff): the DB spend + # lags the live cross-pod counter, so letting an "unchanged" spend + # through the non-admin path would let a key owner / team member + # overwrite the live counter below real usage and silently weaken + # enforcement. _is_budget_change = ( (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) - or ( - data.spend is not None - and data.spend != getattr(existing_key_row, "spend", None) - ) + or data.spend is not None or "budget_limits" in data.model_fields_set ) @@ -2422,6 +2434,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( @@ -2602,15 +2615,24 @@ async def update_key_fn( ) if data.spend is not None: - try: - from litellm.proxy.proxy_server import _invalidate_spend_counter + from litellm.proxy.proxy_server import spend_counter_cache - token_to_invalidate = _hash_token_if_needed(key) - await _invalidate_spend_counter( - counter_key=f"spend:key:{token_to_invalidate}" - ) - except Exception: - pass + counter_key = f"spend:key:{_hash_token_if_needed(key)}" + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=data.spend, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=data.spend, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to update spend counter %s in Redis after key spend update: %s. " + "Budget checks may use stale value until counter expires.", + counter_key, + redis_err, + ) asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -2963,12 +2985,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 +3999,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 +4689,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 +6013,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 +6128,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 f896047a219..4ab308990b5 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2540,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 01449d8ab26..689b574d5bc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2098,7 +2098,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 d0281885482..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' diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4c36b42615e..3c91f1bc1d7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -346,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, ) @@ -670,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 @@ -690,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") @@ -760,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() @@ -1077,6 +1114,33 @@ _OPENAPI_HTTP_METHODS = { # `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO # and cache endpoint files. _ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} +_DB_LITELLM_PARAM_ENV_REF_KEYS = frozenset( + { + "api_key", + "client_secret", + "vertex_credentials", + "vertex_ai_credentials", + "aws_access_key_id", + "aws_secret_access_key", + } +) + + +def _db_model_is_team_scoped(model: object) -> bool: + model_info = getattr(model, "model_info", None) + if isinstance(model_info, BaseModel): + return getattr(model_info, "team_id", None) is not None + if isinstance(model_info, str): + try: + model_info = json.loads(model_info) + except (TypeError, ValueError): + model_info = None + if isinstance(model_info, dict) and model_info.get("team_id") is not None: + return True + if getattr(model_info, "team_id", None) is not None: + return True + model_name = getattr(model, "model_name", None) + return isinstance(model_name, str) and model_name.startswith("model_name_") def _strip_operation_id_method_suffix(operation_id: str) -> str: @@ -1914,9 +1978,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 @@ -3223,13 +3287,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 @@ -4058,7 +4118,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( @@ -4133,7 +4195,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) @@ -4410,7 +4500,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 @@ -5045,8 +5138,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, @@ -5190,6 +5282,24 @@ class ProxyConfig: deleted_deployments += 1 return deleted_deployments + def _resolve_db_litellm_param( + self, key: str, value: object, resolve_env_refs: bool = True + ) -> object: + if not isinstance(value, str): + return value + + decrypted_value = decrypt_value_helper( + value=value, key=key, return_original_value=True + ) + if ( + resolve_env_refs + and key in _DB_LITELLM_PARAM_ENV_REF_KEYS + and isinstance(decrypted_value, str) + and decrypted_value.startswith("os.environ/") + ): + return get_secret(decrypted_value) + return decrypted_value + def _add_deployment(self, db_models: list) -> int: """ Iterate through db models @@ -5207,15 +5317,13 @@ class ProxyConfig: ## ADD MODEL LOGIC for m in db_models: _litellm_params = m.litellm_params + resolve_env_refs = not _db_model_is_team_scoped(m) if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - if isinstance(v, str): - # decrypt value - returns original value if decryption fails or no key is set - _value = decrypt_value_helper( - value=v, key=k, return_original_value=True - ) - _litellm_params[k] = _value + _litellm_params[k] = self._resolve_db_litellm_param( + key=k, value=v, resolve_env_refs=resolve_env_refs + ) _litellm_params = LiteLLM_Params(**_litellm_params) else: @@ -5243,15 +5351,15 @@ class ProxyConfig: _model_list: list = [] for m in new_models: _litellm_params = m.litellm_params + resolve_env_refs = not _db_model_is_team_scoped(m) if isinstance(_litellm_params, BaseModel): _litellm_params = _litellm_params.model_dump() if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - decrypted_value = decrypt_value_helper( - value=v, key=k, return_original_value=True + _litellm_params[k] = self._resolve_db_litellm_param( + key=k, value=v, resolve_env_refs=resolve_env_refs ) - _litellm_params[k] = decrypted_value _litellm_params = LiteLLM_Params(**_litellm_params) else: verbose_proxy_logger.error( @@ -6581,10 +6689,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) @@ -6896,7 +7004,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() @@ -7874,7 +7998,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 @@ -8707,7 +8832,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 @@ -8884,7 +9015,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 @@ -9035,9 +9172,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 @@ -9219,9 +9356,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 @@ -9435,9 +9572,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 @@ -9470,9 +9610,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 @@ -12430,7 +12570,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: @@ -13186,7 +13331,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): @@ -13528,7 +13679,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: @@ -14920,7 +15076,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( @@ -15173,7 +15336,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 }, ) @@ -15382,6 +15548,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"}, @@ -15428,9 +15595,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 @@ -15547,7 +15714,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 }, ) @@ -15615,9 +15785,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) @@ -15661,10 +15831,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() @@ -15730,18 +15904,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 ) @@ -15772,15 +15942,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( @@ -16638,6 +16802,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/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/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8d89ff4a1ff..f6e0303e349 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -79,6 +79,7 @@ def _get_spend_logs_metadata( cold_storage_object_key: Optional[str] = None, litellm_overhead_time_ms: Optional[float] = None, cost_breakdown: Optional[CostBreakdown] = None, + litellm_call_id: Optional[str] = None, ) -> SpendLogsMetadata: if metadata is None: return SpendLogsMetadata( @@ -109,6 +110,7 @@ def _get_spend_logs_metadata( attempted_retries=None, max_retries=None, cost_breakdown=None, + litellm_call_id=litellm_call_id, ) verbose_proxy_logger.debug( "getting payload for SpendLogs, available keys in metadata: " @@ -133,6 +135,7 @@ def _get_spend_logs_metadata( clean_metadata["cold_storage_object_key"] = cold_storage_object_key clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms clean_metadata["cost_breakdown"] = cost_breakdown + clean_metadata["litellm_call_id"] = litellm_call_id return clean_metadata @@ -383,6 +386,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if standard_logging_payload is not None else None ), + litellm_call_id=cast( + Optional[str], + kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), + ), ) special_usage_fields = ["completion_tokens", "prompt_tokens", "total_tokens"] 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..1187640bf0c 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: @@ -61,6 +59,7 @@ class LiteLLMCompletionTransformationHandler: completion_args = {} completion_args.update(kwargs) completion_args.update(litellm_completion_request) + completion_args["_skip_responses_api_bridge"] = True litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper @@ -69,12 +68,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 @@ -111,6 +108,7 @@ class LiteLLMCompletionTransformationHandler: acompletion_args = {} acompletion_args.update(kwargs) acompletion_args.update(litellm_completion_request) + acompletion_args["_skip_responses_api_bridge"] = True litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper @@ -119,12 +117,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 2c46baaada5..3edcbd430f1 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -307,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 a9ff709716a..8afef50759f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -431,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"] = ( @@ -486,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"] = {} @@ -526,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 @@ -548,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 @@ -615,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 {} @@ -2967,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 @@ -3504,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): @@ -3613,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( @@ -4414,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 @@ -6725,9 +6723,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: @@ -6761,9 +6757,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: @@ -6824,9 +6818,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 @@ -7611,7 +7607,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) @@ -7671,9 +7671,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): @@ -9596,7 +9594,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 ( @@ -9616,7 +9615,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 ( 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/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/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/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py new file mode 100644 index 00000000000..3da5b98449b --- /dev/null +++ b/litellm/rust_bridge/__init__.py @@ -0,0 +1,9 @@ +"""LiteLLM Rust bridge package.""" + +from litellm.rust_bridge.loader import ( + get_native_bridge, + native_bridge_available, +) +from litellm.rust_bridge.ocr import use_litellm_rust + +__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"] diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py new file mode 100644 index 00000000000..3ae0bf9307c --- /dev/null +++ b/litellm/rust_bridge/loader.py @@ -0,0 +1,28 @@ +"""Loader for the packaged LiteLLM Rust extension.""" + +from __future__ import annotations + +from types import ModuleType + +_BRIDGE_SENTINEL = object() +_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL + + +def get_native_bridge() -> ModuleType | None: + """Return the packaged Rust extension, or ``None`` when unavailable.""" + global _cached_bridge + if _cached_bridge is not _BRIDGE_SENTINEL: + return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None + + try: + from litellm.rust_bridge import _native + except ImportError: + _cached_bridge = None + return None + _cached_bridge = _native + return _native + + +def native_bridge_available() -> bool: + """Whether the packaged Rust extension is importable.""" + return get_native_bridge() is not None diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py new file mode 100644 index 00000000000..36a088b6b1a --- /dev/null +++ b/litellm/rust_bridge/ocr.py @@ -0,0 +1,159 @@ +"""Thin Python wrapper for the native Rust OCR bridge.""" + +from __future__ import annotations + +import os +from typing import Any, Awaitable, Final, Protocol, Union, cast + +import httpx + + +class RustOcr(Protocol): + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise NotImplementedError + + +class RustAocr(Protocol): + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: + raise NotImplementedError + + +class _Unset: + pass + + +_UNSET: Final[_Unset] = _Unset() + + +def _env_enables_rust_ocr() -> bool: + return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +_rust_ocr_enabled = _env_enables_rust_ocr() +_rust_ocr_impl: RustOcr | None = None +_rust_aocr_impl: RustAocr | None = None + + +def use_litellm_rust( + enabled: bool = True, + *, + ocr: RustOcr | None | _Unset = _UNSET, + aocr: RustAocr | None | _Unset = _UNSET, +) -> None: + global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl + _rust_ocr_enabled = enabled + if not isinstance(ocr, _Unset): + _rust_ocr_impl = ocr + if not isinstance(aocr, _Unset): + _rust_aocr_impl = aocr + + +def rust_ocr_enabled() -> bool: + return _rust_ocr_enabled + + +def load_rust_ocr() -> RustOcr | None: + if _rust_ocr_impl is not None: + return _rust_ocr_impl + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + if native_bridge is None: + return None + return cast(RustOcr, native_bridge.ocr) + + +def load_rust_aocr() -> RustAocr | None: + if _rust_aocr_impl is not None: + return _rust_aocr_impl + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + if native_bridge is None: + return None + return cast(RustAocr, getattr(native_bridge, "aocr", None)) + + +def _timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) + + +def ocr( + *, + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_ocr = load_rust_ocr() + if rust_ocr is None: + return None + return rust_ocr( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + optional_params=optional_params, + timeout_seconds=_timeout_to_seconds(timeout), + ) + + +async def aocr( + *, + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_aocr = load_rust_aocr() + if rust_aocr is None: + return None + return await rust_aocr( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + optional_params=optional_params, + timeout_seconds=_timeout_to_seconds(timeout), + ) 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/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/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/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 24d6e84fba7..c3f99b1d18b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1539,6 +1539,9 @@ class PromptTokensDetailsWrapper( video_length_seconds: Optional[float] = None """Length of videos sent to the model. Used for Vertex AI multimodal embeddings.""" + audio_length_seconds: Optional[float] = None + """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + cache_creation_tokens: Optional[int] = None """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching.""" @@ -1553,6 +1556,8 @@ class PromptTokensDetailsWrapper( del self.image_count if self.video_length_seconds is None: del self.video_length_seconds + if self.audio_length_seconds is None: + del self.audio_length_seconds if self.web_search_requests is None: del self.web_search_requests if self.cache_creation_tokens is None: diff --git a/litellm/utils.py b/litellm/utils.py index 946b260f98a..6964e2158f9 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( @@ -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 ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d8e001e4753..c2d4e23fcfb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -273,6 +273,19 @@ "/v1/images/generations" ] }, + "aiml/openai/gpt-image-2": { + "litellm_provider": "aiml", + "metadata": { + "notes": "OpenAI gpt-image-2 via AI/ML API - flagship multimodal image generation and editing model with reasoning and 2K output. output_cost_per_image is AI/ML's published medium-quality rate; like the other aiml image entries it is billed as a flat per-image price" + }, + "mode": "image_generation", + "output_cost_per_image": 0.054, + "source": "https://docs.aimlapi.com/api-references/image-models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "amazon.nova-canvas-v1:0": { "litellm_provider": "bedrock", "max_input_tokens": 2600, @@ -16153,6 +16166,47 @@ "supports_service_tier": true, "supports_image_size": false }, + "gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -16194,6 +16248,44 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, @@ -17765,6 +17857,49 @@ "supports_service_tier": true, "supports_image_size": false }, + "gemini/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "rpm": 1000, + "tpm": 4000000, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17808,6 +17943,48 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini/gemini-3.1-flash-image": { + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.045, + "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-flash-image-preview": { "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -25567,8 +25744,18 @@ }, "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", - "ocr_cost_per_page": 0.001, - "annotation_cost_per_page": 0.003, + "ocr_cost_per_page": 0.004, + "annotation_cost_per_page": 0.005, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, + "mistral/mistral-ocr-4-0": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.004, + "annotation_cost_per_page": 0.005, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -25787,7 +25974,7 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "mistral/mistral-medium-latest": { + "mistral/mistral-medium-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, @@ -25795,12 +25982,45 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, + "mistral/mistral-medium-2604": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-latest": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium-3-1-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", @@ -25827,6 +26047,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -35582,6 +35803,21 @@ "tpm": 8000000, "supports_image_size": false }, + "vertex_ai/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" + }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -35597,6 +35833,19 @@ "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, + "vertex_ai/gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, diff --git a/pyproject.toml b/pyproject.toml index 1cb39153e83..6e99d81f8f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ proxy = [ "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.26.0,<2.0", "litellm-proxy-extras==0.4.74", - "litellm-enterprise==0.1.43", + "litellm-enterprise==0.1.44", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "polars>=1.38.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", @@ -235,8 +234,24 @@ healthcheck = [ ] [build-system] -requires = ["uv_build==0.11.8"] -build-backend = "uv_build" +requires = ["maturin==1.9.4"] +build-backend = "maturin" + +[tool.maturin] +manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml" +module-name = "litellm.rust_bridge._native" +python-source = "." +bindings = "pyo3" +exclude = [ + "litellm/proxy/enterprise", + "litellm/proxy/enterprise/**", + "**/__pycache__", + "**/__pycache__/**", + "**/.pytest_cache", + "**/.pytest_cache/**", + "**/.ruff_cache", + "**/.ruff_cache/**", +] [tool.uv] constraint-dependencies = [ @@ -254,18 +269,6 @@ litellm-enterprise = { workspace = true } [tool.uv.workspace] members = ["enterprise", "litellm-proxy-extras"] -[tool.uv.build-backend] -module-root = "" -source-exclude = [ - "litellm/proxy/enterprise", - "**/__pycache__", - "**/__pycache__/**", - "**/.pytest_cache", - "**/.pytest_cache/**", - "**/.ruff_cache", - "**/.ruff_cache/**", -] - [tool.isort] profile = "black" @@ -329,4 +332,3 @@ pytest_add_cli_args = [ [tool.coverage.run] source = ["litellm"] relative_files = true - diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ae46f020de1..10c820324ea 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,7 +1,7 @@ { "ANN001": { "baseline": 2865, - "slack": 50 + "slack": 287 }, "ANN002": { "baseline": 64, @@ -9,19 +9,19 @@ }, "ANN003": { "baseline": 759, - "slack": 30 + "slack": 76 }, "ANN201": { "baseline": 1944, - "slack": 50 + "slack": 194 }, "ANN202": { "baseline": 858, - "slack": 30 + "slack": 86 }, "ANN204": { "baseline": 658, - "slack": 20 + "slack": 66 }, "ANN205": { "baseline": 117, @@ -33,7 +33,7 @@ }, "ANN401": { "baseline": 1886, - "slack": 50 + "slack": 189 }, "ASYNC230": { "baseline": 11, @@ -231,10 +231,6 @@ "baseline": 6, "slack": 3 }, - "PLR0913": { - "baseline": 1813, - "slack": 50 - }, "PLR1704": { "baseline": 3, "slack": 3 diff --git a/ruff-strict.toml b/ruff-strict.toml index 1caa3567872..d58885fe848 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -2,7 +2,7 @@ extend = "ruff.toml" [lint] preview = true -select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR0913", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] +select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] extend-select = [] [lint.mccabe] 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/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 861d65489e8..df9815d6557 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -1,16 +1,19 @@ #!/usr/bin/env python3 -"""Non-gating ratchet guard: budget ceilings may only fall, never rise. +"""Non-gating ratchet guard: budget baselines and ceilings may only fall, never rise. Every `*-budget.json` file (ruff-strict, type-discipline, basedpyright-code) is a -one-way ratchet: each rule's ceiling is `baseline + slack`, and the whole point is -to drive that number DOWN over time. This check compares every budget file against -its own content at the merge-base with the target branch and fails (exits 1, red) if: +one-way ratchet: each rule's ceiling is `baseline + slack`, and both the recorded +`baseline` (the live violation count) and that ceiling are meant to be driven DOWN +over time. This check compares every budget file against its own content at the +merge-base with the target branch and fails (exits 1, red) if: - * a rule's ceiling went up, + * a rule's ceiling (`baseline + slack`) went up, + * a rule's `baseline` went up, even if `slack` was lowered to keep the ceiling + flat (a higher baseline bakes in more accepted debt and must be acknowledged), * a rule was dropped from a budget (its ceiling effectively became infinite), or * an entire budget file was deleted. -New rules and lowered/equal ceilings are fine. +New rules and lowered/equal baselines and ceilings are fine. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -66,7 +69,12 @@ def _load_head(rel: str) -> dict | None: def _ref_is_commit(ref: str) -> bool: - return _run(["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"]).returncode == 0 + return ( + _run( + ["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"] + ).returncode + == 0 + ) def _load_base(rel: str, ref: str) -> dict | None: @@ -81,13 +89,55 @@ def _load_base(rel: str, ref: str) -> dict | None: return json.loads(proc.stdout) +def _baselines(budget: dict) -> dict[str, int]: + """Map each rule to its recorded baseline; skip malformed specs.""" + return { + rule: int(spec.get("baseline", 0)) + for rule, spec in budget.items() + if isinstance(spec, dict) + } + + def _caps(budget: dict) -> dict[str, int]: """Map each rule to its ceiling (baseline + slack); skip malformed specs.""" - caps: dict[str, int] = {} - for rule, spec in budget.items(): - if isinstance(spec, dict): - caps[rule] = int(spec.get("baseline", 0)) + int(spec.get("slack", 0)) - return caps + return { + rule: int(spec.get("baseline", 0)) + int(spec.get("slack", 0)) + for rule, spec in budget.items() + if isinstance(spec, dict) + } + + +def _regression_detail( + rule: str, + base_caps: dict[str, int], + head_caps: dict[str, int], + base_baselines: dict[str, int], + head_baselines: dict[str, int], +) -> str | None: + """Why `rule` regressed vs base, or None when it held flat or fell. + + A dropped rule is terminal; otherwise a raised ceiling and a raised baseline are + independent loosenings (the latter catches a baseline bump masked by a slack cut), + so both reasons are reported when both apply. + """ + base_cap = base_caps[rule] + if rule not in head_caps: + return f"rule dropped (ceiling {base_cap} -> removed)" + reasons = tuple( + message + for raised, message in ( + ( + head_caps[rule] > base_cap, + f"ceiling raised {base_cap} -> {head_caps[rule]}", + ), + ( + head_baselines[rule] > base_baselines[rule], + f"baseline raised {base_baselines[rule]} -> {head_baselines[rule]}", + ), + ) + if raised + ) + return "; ".join(reasons) or None def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: @@ -96,18 +146,17 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr if head is None: return [Regression(rel, "*", "budget file was deleted (every ceiling removed)")] - base_caps = _caps(base) - head_caps = _caps(head) + base_caps, head_caps = _caps(base), _caps(head) + base_baselines, head_baselines = _baselines(base), _baselines(head) return [ - Regression( - rel, - rule, - f"rule dropped (ceiling {base_cap} -> removed)" - if rule not in head_caps - else f"ceiling raised {base_cap} -> {head_caps[rule]}", + Regression(rel, rule, detail) + for rule in sorted(base_caps) + if ( + detail := _regression_detail( + rule, base_caps, head_caps, base_baselines, head_baselines + ) ) - for rule, base_cap in sorted(base_caps.items()) - if rule not in head_caps or head_caps[rule] > base_cap + is not None ] @@ -141,7 +190,9 @@ def main() -> int: regressions.extend(regressions_for(rel, base, head)) if regressions: - print(f"FAIL: budget ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):") + print( + f"FAIL: budget baseline(s)/ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):" + ) for reg in regressions: print(f" {reg.budget} {reg.rule}: {reg.detail}") print( 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 1d11d676207..6d0e12314f0 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -51,6 +51,7 @@ IGNORE_FUNCTIONS = [ "_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. + "_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap. ] diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 681cd536259..60fbd505d67 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -24,6 +24,12 @@ EXCLUDED_GUARD_ONLY_VARS = { "MAVVRIK_FOCUS_FREQUENCY", } +# Temporary/internal rollout flags are intentionally not added to the public +# environment settings docs until the feature is ready for broad use. +EXCLUDED_ROLLOUT_FLAGS = { + "LITELLM_USE_RUST_OCR", +} + EXCLUDED_TERMINAL_VARS = { "TERM", "TERM_PROGRAM", @@ -71,6 +77,7 @@ for root, dirs, files in os.walk(repo_base): for match in getenv_matches if match not in EXCLUDED_TERMINAL_VARS and match not in EXCLUDED_GUARD_ONLY_VARS + and match not in EXCLUDED_ROLLOUT_FLAGS ) # Extract only the key part, excluding terminal vars # Find all keys using litellm.get_secret() 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..e059ac62429 --- /dev/null +++ b/tests/e2e/gateway/litellm-config.yml @@ -0,0 +1,198 @@ +# 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 + + - model_name: rust-ocr-mistral + litellm_params: + model: mistral/mistral-ocr-latest + api_key: os.environ/MISTRAL_API_KEY + + - model_name: rust-ocr-azure-ai + litellm_params: + model: azure_ai/mistral-document-ai-2505 + api_base: os.environ/AZURE_API_BASE + api_key: os.environ/AZURE_API_KEY + + - model_name: rust-ocr-azure-document-intelligence + litellm_params: + model: azure_ai/doc-intelligence/prebuilt-layout + api_base: os.environ/AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT + api_key: os.environ/AZURE_DOCUMENT_INTELLIGENCE_API_KEY + + - model_name: rust-ocr-vertex-mistral + litellm_params: + model: vertex_ai/mistral-ocr-2505 + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + + - model_name: rust-ocr-vertex-deepseek + litellm_params: + model: vertex_ai/deepseek-ocr-maas + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + + +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/gateway/test_ocr_rust_e2e.py b/tests/e2e/gateway/test_ocr_rust_e2e.py new file mode 100644 index 00000000000..6ce59b2b5ac --- /dev/null +++ b/tests/e2e/gateway/test_ocr_rust_e2e.py @@ -0,0 +1,148 @@ +""" +Gateway E2E smoke for Rust-backed OCR. + +Start the proxy with: + +LITELLM_USE_RUST_OCR=1 litellm --config tests/e2e/gateway/litellm-config.yml --port 4000 +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import httpx +import pytest +import yaml + +TEST_PDF_URL = ( + "https://cdn.jsdelivr.net/gh/BerriAI/litellm" + "@d769e81c90d453240c61fc572cdb27fae06a89d0" + "/tests/llm_translation/fixtures/dummy.pdf" +) +TEST_IMAGE_URL = ( + "https://cdn.jsdelivr.net/gh/BerriAI/litellm" + "@d769e81c90d453240c61fc572cdb27fae06a89d0" + "/tests/image_gen_tests/test_image.png" +) + +RUST_OCR_GATEWAY_CASES = [ + pytest.param( + "rust-ocr-mistral", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="mistral", + ), + pytest.param( + "rust-ocr-azure-ai", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="azure_ai", + ), + pytest.param( + "rust-ocr-azure-document-intelligence", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="azure_document_intelligence", + ), + pytest.param( + "rust-ocr-vertex-mistral", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="vertex_mistral", + ), + pytest.param( + "rust-ocr-vertex-deepseek", + { + "type": "image_url", + "image_url": os.getenv("RUST_OCR_IMAGE_URL", TEST_IMAGE_URL), + }, + id="vertex_deepseek", + ), +] + +CONFIG_PATH = Path(__file__).with_name("litellm-config.yml") + + +@dataclass(frozen=True) +class OcrGateway: + base_url: str + master_key: str + + def model_names(self) -> set[str]: + with httpx.Client( + timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120")) + ) as client: + response = client.get( + f"{self.base_url.rstrip('/')}/model/info", + headers={"Authorization": f"Bearer {self.master_key}"}, + ) + assert response.status_code == 200, response.text + return { + model["model_name"] + for model in response.json().get("data", []) + if "model_name" in model + } + + def ocr(self, model: str, document: dict[str, str]) -> httpx.Response: + with httpx.Client( + timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120")) + ) as client: + return client.post( + f"{self.base_url.rstrip('/')}/v1/ocr", + headers={"Authorization": f"Bearer {self.master_key}"}, + json={"model": model, "document": document}, + ) + + +@dataclass(frozen=True) +class OcrResources: + gateway: OcrGateway + + +@pytest.fixture +def resources() -> OcrResources: + proxy_url = os.getenv("LITELLM_PROXY_URL") + if not proxy_url: + pytest.skip( + "Start a Rust OCR proxy and set LITELLM_PROXY_URL, e.g. http://localhost:4000" + ) + return OcrResources( + gateway=OcrGateway( + base_url=proxy_url, + master_key=os.getenv("LITELLM_MASTER_KEY", "sk-1234"), + ) + ) + + +def _assert_ocr_response_shape(response_json: dict[str, Any]) -> None: + assert response_json["object"] == "ocr" + assert response_json["model"] + assert isinstance(response_json["pages"], list) + assert len(response_json["pages"]) > 0 + assert "index" in response_json["pages"][0] + assert "markdown" in response_json["pages"][0] + + +class TestRustOcrGateway: + def test_rust_ocr_models_are_on_gateway_config(self) -> None: + config = yaml.safe_load(CONFIG_PATH.read_text()) + configured_models = { + model_config["model_name"] for model_config in config["model_list"] + } + + expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES} + assert expected_models.issubset(configured_models) + + def test_running_gateway_loaded_rust_ocr_models( + self, resources: OcrResources + ) -> None: + expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES} + assert expected_models.issubset(resources.gateway.model_names()) + + @pytest.mark.parametrize(("model", "document"), RUST_OCR_GATEWAY_CASES) + def test_rust_ocr_model_gateway_response( + self, resources: OcrResources, model: str, document: dict[str, str] + ) -> None: + response = resources.gateway.ocr(model, document) + + assert response.status_code == 200, response.text + _assert_ocr_response_shape(response.json()) 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/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 873777189c9..f0a73325afa 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -386,6 +386,62 @@ async def test_aiml_image_generation_with_dynamic_api_key(): assert captured_json_data["model"] == "flux-pro/v1.1" +@pytest.mark.asyncio +async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape(): + """End-to-end check that ``aiml/openai/gpt-image-2`` keeps the upstream + OpenAI request shape (``size``/``n``/``response_format``) instead of + being remapped to the AI/ML flux schema (``image_size``/``num_images``/ + ``output_format``), and hits the correct upstream model name. + """ + from unittest.mock import MagicMock, patch + import json as _json + + mock_aiml_response = { + "created": 1703658209, + "data": [{"url": "https://example.com/gpt-image-2.png"}], + } + + captured = {} + + def capture_post_call(*args, **kwargs): + captured["url"] = kwargs.get("url") or (args[0] if args else None) + captured["headers"] = kwargs.get("headers", {}) + captured["json"] = kwargs.get("json", {}) + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_aiml_response + mock_response.text = _json.dumps(mock_aiml_response) + return mock_response + + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: + mock_post.side_effect = capture_post_call + + await litellm.aimage_generation( + prompt="A T-Rex relaxing on a beach", + model="aiml/openai/gpt-image-2", + api_key="test-key-mocked-no-credits-needed", + size="1024x1536", + quality="high", + response_format="b64_json", + n=1, + ) + + assert captured["url"] is not None + assert "api.aimlapi.com" in captured["url"] + assert "/v1/images/generations" in captured["url"] + + body = captured["json"] + assert body["model"] == "openai/gpt-image-2" + assert body["prompt"] == "A T-Rex relaxing on a beach" + assert body["size"] == "1024x1536" + assert body["quality"] == "high" + assert body["response_format"] == "b64_json" + assert body["n"] == 1 + assert "image_size" not in body + assert "num_images" not in body + assert "output_format" not in body + + @pytest.mark.asyncio async def test_azure_image_generation_request_body(): """Azure deployment URL selects the model; JSON body omits ``model`` (#26316).""" 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/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py index f4cc9735177..c37a2e3f65d 100644 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ b/tests/logging_callback_tests/test_gcs_pub_sub.py @@ -31,6 +31,7 @@ verbose_logger.setLevel(logging.DEBUG) ignored_keys = [ "request_id", + "metadata.litellm_call_id", "session_id", "startTime", "endTime", 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..bb13a7ce8cc 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 @@ -1881,10 +1883,11 @@ def test_get_server_auth_header_no_auth_headers(): def test_create_tool_response_objects(): - """Test _create_tool_response_objects function.""" + """Test _create_tool_response_objects enriches mcp_info with server_id and alias.""" from litellm.proxy._experimental.mcp_server.rest_endpoints import ( _create_tool_response_objects, ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer from mcp.types import Tool as MCPTool # Create mock tools @@ -1901,20 +1904,32 @@ def test_create_tool_response_objects(): ), ] - server_mcp_info = { - "server_name": "zapier", + server = MCPServer( + server_id="a1b2c3d4", + name="zapier_internal", + alias="zapier", + transport="http", + mcp_info={ + "server_name": "zapier_internal", + "logo_url": "https://zapier.com/logo.png", + }, + ) + + result = _create_tool_response_objects(mock_tools, server) + + expected_mcp_info = { + "server_name": "zapier_internal", "logo_url": "https://zapier.com/logo.png", + "server_id": "a1b2c3d4", + "alias": "zapier", } - - result = _create_tool_response_objects(mock_tools, server_mcp_info) - assert len(result) == 2 assert result[0].name == "send_email" assert result[0].description == "Send an email" - assert result[0].mcp_info == server_mcp_info + assert result[0].mcp_info == expected_mcp_info assert result[1].name == "create_event" assert result[1].description == "Create a calendar event" - assert result[1].mcp_info == server_mcp_info + assert result[1].mcp_info == expected_mcp_info @pytest.mark.asyncio @@ -1928,6 +1943,8 @@ async def test_get_tools_for_single_server(): # Create a mock server (pin allowlist fields; MagicMock auto-attrs are truthy) mock_server = MagicMock() mock_server.mcp_info = {"server_name": "zapier"} + mock_server.server_id = "zapier_id" + mock_server.alias = "zapier_alias" mock_server.allowed_tools = None mock_server.disallowed_tools = None @@ -1961,7 +1978,11 @@ async def test_get_tools_for_single_server(): # Verify the result assert len(result) == 1 assert result[0].name == "send_email" - assert result[0].mcp_info == {"server_name": "zapier"} + assert result[0].mcp_info == { + "server_name": "zapier", + "server_id": "zapier_id", + "alias": "zapier_alias", + } @pytest.mark.asyncio 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/test_litellm/caching/test_embedding_router.py b/tests/test_litellm/caching/test_embedding_router.py new file mode 100644 index 00000000000..550095a112a --- /dev/null +++ b/tests/test_litellm/caching/test_embedding_router.py @@ -0,0 +1,67 @@ +import os +import sys +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching._embedding_router import ( + build_router_embedding_metadata, + resolve_embedding_router, +) + + +def test_resolve_returns_router_when_model_is_a_deployment(): + router = MagicMock() + assert ( + resolve_embedding_router("sem-embed", router, [{"model_name": "sem-embed"}]) + is router + ) + + +def test_resolve_returns_none_when_model_not_in_router(): + router = MagicMock() + assert ( + resolve_embedding_router("sem-embed", router, [{"model_name": "other"}]) is None + ) + + +def test_resolve_returns_none_when_router_is_none(): + assert ( + resolve_embedding_router("sem-embed", None, [{"model_name": "sem-embed"}]) + is None + ) + + +def test_resolve_returns_none_when_model_list_is_none(): + router = MagicMock() + assert resolve_embedding_router("sem-embed", router, None) is None + + +def test_resolve_skips_entries_missing_model_name(): + router = MagicMock() + model_list = [ + {"litellm_params": {"model": "bedrock/x"}}, + {"model_name": "sem-embed"}, + ] + assert resolve_embedding_router("sem-embed", router, model_list) is router + assert resolve_embedding_router("other", router, [{"litellm_params": {}}]) is None + + +def test_build_metadata_preserves_request_fields_and_adds_flag(): + md = build_router_embedding_metadata( + {"user_api_key": "sk-x", "user_api_key_team_id": "team-1", "trace_id": "t-1"} + ) + assert md == { + "user_api_key": "sk-x", + "user_api_key_team_id": "team-1", + "trace_id": "t-1", + "semantic-cache-embedding": True, + } + + +def test_build_metadata_handles_none_and_does_not_mutate_input(): + original = {"user_api_key": "sk-x"} + md = build_router_embedding_metadata(original) + assert md == {"user_api_key": "sk-x", "semantic-cache-embedding": True} + assert original == {"user_api_key": "sk-x"} + assert build_router_embedding_metadata(None) == {"semantic-cache-embedding": True} diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index 949e6ccc292..67d4e2d9892 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -1,5 +1,6 @@ import os import sys +import types from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -806,3 +807,104 @@ def test_qdrant_semantic_cache_large_vector_size(): ) create_payload = put_call.kwargs["json"] assert create_payload["vectors"]["size"] == 4096 + + +def _router_proxy_module(router, model_name): + mod = types.ModuleType("litellm.proxy.proxy_server") + mod.llm_router = router + mod.llm_model_list = [{"model_name": model_name}] + return mod + + +def test_qdrant_sync_get_cache_routes_through_router(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.qdrant_api_base = "http://test.qdrant.local" + cache.collection_name = "test_collection" + cache.headers = {"Content-Type": "application/json", "api-key": "test_key"} + cache.similarity_threshold = 0.8 + cache.sync_client = MagicMock() + search_response = MagicMock() + search_response.status_code = 200 + search_response.json.return_value = {"result": []} + cache.sync_client.post.return_value = search_response + + router = MagicMock() + router.embedding = MagicMock( + return_value={"data": [{"embedding": [0.3, 0.3, 0.3]}]} + ) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + with patch("litellm.embedding") as direct_embed: + result = cache.get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) + + assert result is None + router.embedding.assert_called_once() + assert router.embedding.call_args.kwargs["model"] == "sem-embed" + direct_embed.assert_not_called() + + +def test_qdrant_sync_set_cache_falls_back_to_direct(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "text-embedding-ada-002" + cache.qdrant_api_base = "http://test.qdrant.local" + cache.collection_name = "test_collection" + cache.headers = {"Content-Type": "application/json", "api-key": "test_key"} + cache.sync_client = MagicMock() + put_response = MagicMock() + put_response.status_code = 200 + cache.sync_client.put.return_value = put_response + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = None + fake_proxy.llm_model_list = None + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch( + "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.1, 0.1]}]} + ) as direct_embed: + cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ) + + direct_embed.assert_called_once() + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_forwards_full_metadata(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + await cache._get_async_embedding( + "hello", + metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"}, + ) + + md = router.aembedding.call_args.kwargs["metadata"] + assert md["user_api_key"] == "sk-x" + assert md["user_api_key_team_id"] == "team-1" + assert md["semantic-cache-embedding"] is True diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 78192400fb0..37c938bb38d 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -3,7 +3,6 @@ import sys from unittest.mock import MagicMock, patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") @@ -94,6 +93,42 @@ async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( mock_redis_instance.expire.assert_not_awaited() +@pytest.mark.parametrize("namespace", [None, "litellm"]) +@pytest.mark.asyncio +async def test_async_delete_cache_applies_namespace( + namespace, monkeypatch, redis_no_ping +): + """async_delete_cache must prefix keys with the namespace, matching every + other cache operation. Without this, Redis NOPERM errors occur when an + ACL restricts DEL to the litellm:* pattern.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_delete_cache(key="3997c4abcdef") + + expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" + mock_redis_instance.delete.assert_awaited_once_with(expected_key) + + +@pytest.mark.parametrize("namespace", [None, "litellm"]) +def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): + """delete_cache must prefix keys with the namespace, matching every other + cache operation.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_client = MagicMock() + redis_cache.redis_client = mock_redis_client + + redis_cache.delete_cache(key="3997c4abcdef") + + expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" + mock_redis_client.delete.assert_called_once_with(expected_key) + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") @@ -517,3 +552,177 @@ async def test_async_lpop_with_float_redis_version( # Verify the method completed without error assert result is not None + + +# LIT-3374: the namespace must be applied uniformly across every key-taking +# Redis operation, not just get/set/increment. Before the fix these paths wrote +# or read raw keys, so with a namespace configured the prefixed keys other +# operations created were silently missed. + + +@pytest.mark.parametrize( + "namespace, raw_keys, expected_keys", + [ + (None, ["{k:v}:tokens", "{k:v}:requests"], ["{k:v}:tokens", "{k:v}:requests"]), + ( + "litellm_sandbox", + ["{k:v}:tokens", "{k:v}:requests"], + ["litellm_sandbox:{k:v}:tokens", "litellm_sandbox:{k:v}:requests"], + ), + ], +) +@pytest.mark.asyncio +async def test_async_register_script_namespaces_keys( + namespace, raw_keys, expected_keys, monkeypatch, redis_no_ping +): + """The callable returned by async_register_script (used by the rate limiter + Lua scripts, pod-lock release, and budget limiters) must namespace every key + it is invoked with. The hash tag is preserved so cluster slotting is intact.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + + registered_script = AsyncMock(return_value="ok") + mock_redis_instance = MagicMock() + mock_redis_instance.register_script = MagicMock(return_value=registered_script) + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + script = redis_cache.async_register_script("return 1") + result = await script(keys=raw_keys, args=[60]) + + assert result == "ok" + registered_script.assert_awaited_once_with( + keys=expected_keys, args=[60], client=None + ) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_delete_cache_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_delete_cache("k") + mock_redis_instance.delete.assert_awaited_once_with(expected) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_delete_cache_keys_namespaces_keys( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.delete_cache_keys(["k"]) + mock_redis_instance.delete.assert_awaited_once_with(expected) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_get_ttl_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + mock_redis_instance.ttl = AsyncMock(return_value=42) + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + ttl = await redis_cache.async_get_ttl("k") + assert ttl == 42 + mock_redis_instance.ttl.assert_awaited_once_with(expected) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_lpop_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + mock_redis_instance.lpop = AsyncMock(return_value=b"value") + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_lpop(key="k") + mock_redis_instance.lpop.assert_awaited_once_with(expected, None) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_rpush_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + mock_redis_instance.rpush = AsyncMock(return_value=1) + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_rpush("k", ["v"]) + mock_redis_instance.rpush.assert_awaited_once_with(expected, "v") + + +@pytest.mark.parametrize("namespace, expected_match", [(None, "k*"), ("ns", "ns:k*")]) +@pytest.mark.asyncio +async def test_async_scan_iter_namespaces_pattern( + namespace, expected_match, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + + captured = {} + + def scan_iter(match, count): + captured["match"] = match + + async def gen(): + for _ in (): + yield _ + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_scan_iter(pattern="k") + assert captured["match"] == expected_match + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +def test_increment_cache_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_client = MagicMock() + mock_client.incr.return_value = 5 + mock_client.ttl.return_value = 100 + redis_cache.redis_client = mock_client + redis_cache.increment_cache(key="k", value=1) + mock_client.incr.assert_called_once_with(name=expected, amount=1) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_client = MagicMock() + redis_cache.redis_client = mock_client + redis_cache.delete_cache(key="k") + mock_client.delete.assert_called_once_with(expected) diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 13f9d00136d..1d3129d6467 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -104,6 +104,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch): # Verify llmcache.check was called redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -138,10 +139,16 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): ] ) - with patch.object( - redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", + with ( + patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, + ), + patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ), ): metadata = {} result = redis_semantic_cache.get_cache( @@ -176,16 +183,21 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8) redis_semantic_cache.llmcache.store = MagicMock() - redis_semantic_cache.set_cache( - key="test_key", - value={"content": "Paris"}, - messages=[{"content": "What is the capital of France?"}], - ttl=60, - ) + with patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, + ): + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ttl=60, + ) redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ttl=60, ) @@ -299,10 +311,11 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat monkeypatch.setenv("REDIS_PASSWORD", "test_password") with pytest.raises(ValueError, match="connection failed"): - RedisSemanticCache( + cache = RedisSemanticCache( similarity_threshold=0.8, index_name="existing_index", ) + _ = cache.llmcache def test_redis_semantic_cache_reraises_unexpected_index_error(): @@ -534,6 +547,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input(): return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} ) redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) redis_semantic_cache.set_cache( key="test_key", @@ -544,6 +558,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input(): redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ) @@ -564,6 +579,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input(): } ] ) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) with patch.object( redis_semantic_cache, @@ -581,6 +597,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input(): assert metadata["semantic-similarity"] == pytest.approx(0.9) redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -594,6 +611,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} ) redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) redis_semantic_cache.set_cache( key="test_key", @@ -616,6 +634,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?\nAnswer briefly.", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ) @@ -740,6 +759,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) redis_semantic_cache.llmcache = MagicMock() redis_semantic_cache.llmcache.check = MagicMock(return_value=[]) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) with patch.object( redis_semantic_cache, @@ -757,6 +777,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): assert metadata["semantic-similarity"] == 0.0 redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -870,6 +891,63 @@ async def test_redis_semantic_cache_async_paths_set_similarity_on_misses(): ) +def test_redis_get_embedding_routes_through_router(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]}) + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": "sem-embed"}] + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch("litellm.embedding") as direct_embed: + vec = cache._get_embedding("hello", metadata={"user_api_key": "sk-x"}) + + assert vec == [0.5, 0.6] + router.embedding.assert_called_once() + assert router.embedding.call_args.kwargs["model"] == "sem-embed" + assert router.embedding.call_args.kwargs["input"] == "hello" + assert router.embedding.call_args.kwargs["cache"] == { + "no-store": True, + "no-cache": True, + } + assert router.embedding.call_args.kwargs["metadata"] == { + "user_api_key": "sk-x", + "semantic-cache-embedding": True, + } + direct_embed.assert_not_called() + + +def test_redis_get_embedding_falls_back_to_direct(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "text-embedding-ada-002" + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = None + fake_proxy.llm_model_list = None + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch( + "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]} + ) as direct_embed: + vec = cache._get_embedding("hello") + + assert vec == [0.1, 0.2] + direct_embed.assert_called_once() + + def test_cache_get_cache_passes_responses_input_to_backend_cache(): from litellm.caching.caching import Cache @@ -893,7 +971,7 @@ def test_cache_get_cache_passes_responses_input_to_backend_cache(): ) -def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): +def test_cache_get_cache_filters_non_lookup_kwargs_from_backend_cache(): from litellm.caching.caching import Cache cache = Cache.__new__(Cache) @@ -927,7 +1005,11 @@ def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): forwarded_kwargs = cache.cache.get_cache.call_args.kwargs assert forwarded_kwargs == { "input": "What is the capital of France?", - "metadata": {"semantic-similarity": 0.7}, + "metadata": { + "user_api_key": "sk-secret", + "trace_id": "trace-id", + "semantic-similarity": 0.7, + }, } assert forwarded_kwargs["metadata"] is not metadata cache._get_cache_logic.assert_called_once_with( @@ -988,3 +1070,166 @@ def test_cache_get_cache_passes_responses_input_to_dynamic_cache(): cached_result={"content": "Paris"}, max_age=float("inf"), ) + + +def test_redis_sync_set_cache_passes_precomputed_vector(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.llmcache = MagicMock() + cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + cache._get_ttl = MagicMock(return_value=None) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ) + + cache._get_embedding.assert_called_once() + cache.llmcache.store.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_sync_get_cache_passes_precomputed_vector(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.similarity_threshold = 0.8 + cache.llmcache = MagicMock() + cache.llmcache.check = MagicMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + with patch.object( + cache, "_get_cache_key_filter_expression", return_value="cache-key-filter" + ): + result = cache.get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) + + assert result == {"content": "Paris"} + cache._get_embedding.assert_called_once() + cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_async_embedding_forwards_full_metadata(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": "sem-embed"}] + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + await cache._get_async_embedding( + "hello", + metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"}, + ) + + md = router.aembedding.call_args.kwargs["metadata"] + assert md["user_api_key"] == "sk-x" + assert md["user_api_key_team_id"] == "team-1" # FAILS today: team_id is dropped + assert md["semantic-cache-embedding"] is True + + +def test_redis_init_defers_redisvl_construction(monkeypatch): + semantic_cache_mock = MagicMock() + custom_vectorizer_mock = MagicMock() + + with patch.dict( + "sys.modules", + { + "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), + "redisvl.utils.vectorize": MagicMock( + CustomTextVectorizer=custom_vectorizer_mock + ), + }, + ): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "test_password") + + cache = RedisSemanticCache(similarity_threshold=0.8) + + semantic_cache_mock.assert_not_called() + custom_vectorizer_mock.assert_not_called() + + first = cache.llmcache + semantic_cache_mock.assert_called_once() + custom_vectorizer_mock.assert_called_once() + + second = cache.llmcache + assert first is second + semantic_cache_mock.assert_called_once() + + +def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch): + built_cache = MagicMock() + semantic_cache_mock = MagicMock( + side_effect=[ConnectionError("redis down"), built_cache] + ) + custom_vectorizer_mock = MagicMock() + + with patch.dict( + "sys.modules", + { + "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), + "redisvl.utils.vectorize": MagicMock( + CustomTextVectorizer=custom_vectorizer_mock + ), + }, + ): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "test_password") + + cache = RedisSemanticCache(similarity_threshold=0.8) + + with pytest.raises(ConnectionError, match="redis down"): + _ = cache.llmcache + + assert cache.llmcache is built_cache + assert semantic_cache_mock.call_count == 2 + + +def test_redis_llmcache_setter_supported(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + sentinel = MagicMock() + cache.llmcache = sentinel + assert cache.llmcache is sentinel 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/google_genai/test_google_genai_main.py b/tests/test_litellm/google_genai/test_google_genai_main.py index 5854e4b55af..8f56b4e4bc0 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -2,6 +2,7 @@ """ Test to verify the Google GenAI generate_content adapter functionality """ + import json import os import sys @@ -42,4 +43,228 @@ async def test_agenerate_content_stream(): stream=True, ) mock_post.assert_called_once() - mock_post.call_args.kwargs["stream"] == True + assert mock_post.call_args.kwargs["stream"] is True + + +def _mock_gemini_post_response(): + """A minimal stand-in for a successful Gemini generateContent HTTP response.""" + from unittest.mock import MagicMock + + resp = MagicMock() + resp.status_code = 200 + resp.headers = {} + resp.json.return_value = { + "candidates": [ + { + "content": {"parts": [{"text": "hi"}], "role": "model"}, + "finishReason": "STOP", + } + ] + } + return resp + + +NATIVE_TOP_LEVEL_FIELD_CASES = [ + ( + "safetySettings", + [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}], + ), + ("toolConfig", {"functionCallingConfig": {"mode": "AUTO"}}), + ("cachedContent", "cachedContents/abc123"), + ("labels", {"team": "search"}), +] + + +@pytest.mark.parametrize("field_name, field_value", NATIVE_TOP_LEVEL_FIELD_CASES) +def test_native_top_level_field_forwarded_to_request_body(field_name, field_value): + """ + Regression for https://github.com/BerriAI/litellm/issues/12671 + + Google's native generateContent body carries fields like safetySettings at the + top level (siblings of generationConfig). The proxy spreads them as loose kwargs + into generate_content. They must reach Google's request body at the top level and + must NOT be silently dropped nor nested under generationConfig. + """ + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + **{field_name: field_value}, + ) + + assert mock_post.called, "expected the request to reach the HTTP client" + body = mock_post.call_args.kwargs["json"] + assert ( + body[field_name] == field_value + ), f"{field_name} should be forwarded to Google at the top level" + assert field_name not in body.get("generationConfig", {}), ( + f"{field_name} must be a top-level sibling of generationConfig, " + "not nested inside it" + ) + + +@pytest.mark.asyncio +async def test_native_safety_settings_forwarded_async(): + """The async path (used by the proxy's :generateContent route) must also forward + native top-level fields.""" + from unittest.mock import AsyncMock, patch + + from litellm.google_genai.main import agenerate_content + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + AsyncHTTPHandler, + "post", + new_callable=AsyncMock, + return_value=_mock_gemini_post_response(), + ) as mock_post: + await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert "safetySettings" not in body.get("generationConfig", {}) + + +def test_native_fields_coexist_with_generation_config(): + """Forwarding native top-level fields must not regress the already-working + generationConfig path; both must land in their correct positions.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + generationConfig={"temperature": 0, "responseMimeType": "application/json"}, + ) + + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + generation_config = body["generationConfig"] + assert generation_config["temperature"] == 0 + assert generation_config["responseMimeType"] == "application/json" + + +def test_explicit_extra_body_overrides_native_top_level_field(): + """An explicit extra_body value takes precedence over the same top-level field.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + native = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}] + override = [ + {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"} + ] + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=native, + extra_body={"safetySettings": override}, + ) + + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == override + + +def test_native_fields_and_system_instruction_forwarded_on_sync_stream(): + """The sync streaming path (generate_content_stream) must forward native top-level + fields AND systemInstruction. The PR changed the merge here and newly added the + systemInstruction kwarg; without coverage a regression on either ships green.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content_stream + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + system_instruction = {"parts": [{"text": "Be terse"}]} + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content_stream( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + systemInstruction=system_instruction, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert body["systemInstruction"] == system_instruction + assert "safetySettings" not in body.get("generationConfig", {}) + + +@pytest.mark.asyncio +async def test_native_fields_forwarded_on_async_stream(): + """The async streaming path (agenerate_content_stream) backs the proxy's + :streamGenerateContent route and must forward native top-level fields too.""" + from unittest.mock import AsyncMock, patch + + from litellm.google_genai.main import agenerate_content_stream + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + AsyncHTTPHandler, + "post", + new_callable=AsyncMock, + return_value=_mock_gemini_post_response(), + ) as mock_post: + await agenerate_content_stream( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert "safetySettings" not in body.get("generationConfig", {}) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index e47e437a131..7ffd09b931f 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -4690,6 +4690,109 @@ class TestOpenTelemetrySpanDedupe(unittest.TestCase): self.assertTrue(otel._emit_once(kwargs, "success")) self.assertFalse(otel._emit_once(kwargs, "success")) + def test_emit_once_accepts_list_valued_scope_part(self): + """Regression for LIT-3428 / LIT-3764: a list-valued ``guardrail_mode`` + (the shape Presidio expands to with ``output_parse_pii: true``) must + not raise ``TypeError: unhashable type: 'list'`` when building the + dedupe key. Pre-fix, this call crashed inside ``dict.get``.""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + self.assertTrue( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]) + ) + self.assertFalse( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]), + "Same list scope must dedupe to False on the second call", + ) + + def test_emit_once_distinct_list_scopes_dont_collide(self): + """Two different list-valued scopes on the same handler/kwargs must + each emit exactly once. Catches a regression where every list collapses + to the same key (e.g. ``str(list)`` collisions on near-identical input).""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call"])) + self.assertTrue( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]), + "Distinct list scopes must produce distinct dedupe keys", + ) + self.assertFalse(otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call"])) + self.assertFalse( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]) + ) + + def test_emit_once_accepts_dict_and_set_scope_parts(self): + """``guardrail_mode`` can also arrive as a ``GuardrailMode`` TypedDict + (i.e. a plain dict at runtime). Sets are not produced today but flow + through the same normalization. Both must hash without raising.""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + self.assertTrue( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"tags": ["pre", "post"]}) + ) + self.assertFalse( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"tags": ["pre", "post"]}) + ) + self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"a", "b"})) + + def test_emit_once_handles_self_referential_scope_without_recursion_error(self): + """``_freeze_for_dedupe`` caps recursion at ``_FREEZE_MAX_DEPTH`` and + falls back to ``repr`` past the cap, so a self-referential container + in scope must not crash ``_emit_once``. ``guardrail_mode`` cannot + construct such input today, but the cap is the bound that justifies + recursion on the logging hot path.""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + cyclic: list = [] + cyclic.append(cyclic) + self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, cyclic)) + self.assertFalse(otel._emit_once(kwargs, "guardrail", "pii", 1.0, cyclic)) + + def test_create_guardrail_span_does_not_raise_on_list_mode(self): + """End-to-end regression for LIT-3428: ``_create_guardrail_span`` + must produce exactly one span (not raise ``TypeError``) when the + guardrail entry's ``guardrail_mode`` is a list.""" + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + + otel = OpenTelemetry(tracer_provider=tracer_provider) + otel.tracer = tracer_provider.get_tracer(__name__) + + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": {"custom_llm_provider": "openai", "metadata": {}}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "hidden_params": {}, + "guardrail_information": [ + { + "guardrail_name": "presidio-pii", + "guardrail_mode": ["pre_call", "post_call"], + "guardrail_response": "ok", + "start_time": 1.0, + "end_time": 2.0, + } + ], + }, + } + + otel._create_guardrail_span(kwargs=kwargs, context=None) + otel._create_guardrail_span(kwargs=kwargs, context=None) + + guardrail_spans = [ + s for s in span_exporter.get_finished_spans() if s.name == "guardrail" + ] + self.assertEqual( + len(guardrail_spans), + 1, + "List-valued guardrail_mode must emit exactly one guardrail span " + "across repeated lifecycle entrypoints", + ) + def test_handle_success_emits_single_litellm_request_span_on_double_call(self): """Sync + async callback paths firing for the same kwargs must result in exactly one litellm_request span.""" 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 7f3d5a959a1..a5d1934237f 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 @@ -951,6 +951,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "character_count": 0, "image_count": 0, "video_length_seconds": 0.0, + "audio_length_seconds": 0.0, } model_info: ModelInfo = {} 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_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f0db0409bd7..5a8ac0313a7 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -11,7 +11,9 @@ sys.path.insert( import time +import litellm from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import set_callbacks from litellm.types.utils import ModelResponse, TextCompletionResponse @@ -3408,6 +3410,140 @@ def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_resp assert result.usage.prompt_tokens == 4 # type: ignore[attr-defined] +class _SuccessCapturingLogger(CustomLogger): + """Records the success payload. success_payload is populated only in + async_log_success_event, so it stays None when the buggy no-op + async_log_stream_event path runs for streaming.""" + + def __init__(self): + super().__init__() + self.success_payload = None + self.success_calls = 0 + self.stream_event_calls = 0 + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_calls += 1 + self.success_payload = kwargs.get("standard_logging_object") + + async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): + self.stream_event_calls += 1 + + +def _responses_stream_sse_bytes(): + """A full Responses stream: an opened message item, two text deltas, then the + terminal response.completed carrying usage. Exercises mid-stream delta handling + in addition to end-of-stream success logging.""" + import json + + events = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "msg-1", + "type": "message", + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "hello ", + "sequence_number": 2, + }, + { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "world", + "sequence_number": 3, + }, + { + "type": "response.completed", + "sequence_number": 4, + "response": _responses_api_response_with_text("hello world").model_dump(), + }, + ] + return [f"data: {json.dumps(e)}\n\n".encode("utf-8") for e in events] + + +def _fake_streaming_responses_http_response(): + sse_chunks = _responses_stream_sse_bytes() + + async def aiter_bytes(*args, **kwargs): + for chunk in sse_chunks: + yield chunk + + resp = MagicMock() + resp.status_code = 200 + resp.headers = {} + resp.aiter_bytes = aiter_bytes + return resp + + +def _chunk_text(chunk): + if isinstance(chunk, (bytes, bytearray)): + return chunk.decode("utf-8", "ignore") + return str(chunk) + + +async def _drain_until_logged(logger, max_iter=30): + for _ in range(max_iter): + if logger.success_payload is not None: + break + await asyncio.sleep(0.1) + + +@pytest.mark.asyncio +async def test_streaming_anthropic_messages_openai_bridge_fires_success_logging( + monkeypatch, +): + """Regression for #28595 / #28943. The existing tests above call + _handle_anthropic_messages_response_logging directly; they do not cover the + streaming wiring that originally broke. Drive a real streaming + anthropic_messages call routed to the OpenAI Responses backend (upstream SSE + mocked) and assert the bridge surfaces delta chunks and fires success logging + exactly once with real cost. On the broken version the stream ran but only the + no-op async_log_stream_event was called, so success_payload stayed None and the + SpendLogs row never landed.""" + logger = _SuccessCapturingLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + + chunks = [] + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=AsyncMock(return_value=_fake_streaming_responses_http_response()), + ): + stream = await litellm.anthropic_messages( + model="openai/gpt-4o", + api_key="sk-test-28595", + messages=[{"role": "user", "content": "ping"}], + max_tokens=16, + stream=True, + ) + async for chunk in stream: # logging fires on stream end; must drain fully + chunks.append(chunk) + + await _drain_until_logged(logger) + + assert chunks, "stream yielded no chunks" + assert any("content_block_delta" in _chunk_text(c) for c in chunks), ( + "no delta chunks surfaced; the streaming text deltas were not forwarded" + ) + assert logger.success_payload is not None, ( + "async_log_success_event never fired for streaming /v1/messages -> openai " + "Responses bridge; the no-op stream path dropped the spend row" + ) + assert logger.success_calls == 1, "bridge call must log success exactly once" + assert logger.success_payload["response_cost"] > 0 + assert logger.success_payload["call_type"] == "anthropic_messages" + + def test_failure_handler_records_recovered_partial_spend(logging_obj): """A stream interrupted mid-flight still billed the provider for the chunks already delivered. When the router stashes that recovered usage as diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py new file mode 100644 index 00000000000..453c7490d98 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -0,0 +1,321 @@ +""" +Regression tests for the Anthropic message_start cursor=1 bug in +ChunkProcessor._calculate_usage_per_chunk. + +Background +---------- +Anthropic streams a `message_start` event that carries +`usage.output_tokens=1` as a placeholder ("cursor"). The real cumulative +output count only arrives in the final `message_delta` event. When a +stream is cancelled before `message_delta` lands (very common for +thinking models on long-tail prompts), the last-wins accumulator in +ChunkProcessor leaves completion_tokens stuck at 1. Because 1 is +truthy, the `completion_tokens or token_counter(text=...)` fallback in +calculate_usage() never fires, and the request is billed for 1 output +token even when several thousand tokens of text were actually streamed. + +These tests pin the post-fix behavior: completion_tokens should reset +to 0 when the only update we saw was the cursor, allowing the +text-based fallback to estimate from the real completion text. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +def _make_chunk( + *, + content: str = "", + usage: Usage = None, + finish_reason: str = None, + custom_llm_provider: str = "anthropic", +) -> ModelResponseStream: + chunk = ModelResponseStream( + id="msg_test", + created=1738900000, + model="claude-sonnet-4-6", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=finish_reason, + index=0, + delta=Delta(content=content, role="assistant"), + ) + ], + usage=usage, + ) + # The cursor reset is now gated on provider; populate the same field the + # real streaming_handler sets (see litellm/litellm_core_utils/streaming_handler.py). + chunk._hidden_params = {"custom_llm_provider": custom_llm_provider} + return chunk + + +class TestAnthropicCursorBug: + """The core regression: completion_tokens=1 cursor must not leak through.""" + + def test_only_message_start_cursor_resets_completion_to_zero(self): + """ + Stream cancelled before message_delta — only the message_start cursor + (output_tokens=1) was seen. Per-chunk accumulator must reset to 0 so + token_counter fallback can estimate from completion text. + """ + # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + # Several content_block_delta chunks (no usage attached) + text_chunks = [ + _make_chunk(content="Hello"), + _make_chunk(content=" world"), + _make_chunk(content=" this is partial."), + ] + chunks = [message_start, *text_chunks] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 1024 + # The cursor value of 1 must NOT leak through — should be reset to 0 + # so the text-based fallback estimates the real completion length. + assert result["completion_tokens"] == 0, ( + "completion_tokens=1 from message_start cursor leaked through. " + "Should reset to 0 when only cursor was seen, so token_counter " + "fallback in calculate_usage() can estimate from completion text." + ) + + def test_message_start_plus_message_delta_uses_delta_value(self): + """ + Normal complete stream: message_start cursor=1, then message_delta=3847. + Last-wins must give 3847 (the real value). + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]] + # message_delta with the real cumulative output_tokens + message_delta = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=3847, total_tokens=4871), + finish_reason="stop", + ) + chunks = [message_start, *text_chunks, message_delta] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 1024 + assert result["completion_tokens"] == 3847 + + def test_calculate_usage_falls_back_to_token_counter_for_cursor_only(self): + """ + End-to-end via calculate_usage(): cursor-only stream + real completion + text should produce a token-counter estimate, NOT 1. + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + # ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark) + text_chunks = [ + _make_chunk(content="Based on your question, I think the answer is "), + _make_chunk(content="forty-two. Here is my reasoning: "), + ] + chunks = [message_start, *text_chunks] + completion_output = ( + "Based on your question, I think the answer is forty-two. " + "Here is my reasoning: " + ) + + processor = ChunkProcessor(chunks=chunks, messages=[]) + usage = processor.calculate_usage( + chunks=chunks, + model="claude-sonnet-4-6", + completion_output=completion_output, + messages=[], + ) + + # Should be a token_counter estimate of the text, not the cursor 1 + assert usage.completion_tokens > 1, ( + f"Expected token_counter estimate of completion text, got " + f"completion_tokens={usage.completion_tokens} (likely stuck at cursor)" + ) + + def test_cache_fields_preserved_from_message_start(self): + """cache_read / cache_creation come from message_start and must survive.""" + message_start_usage = Usage( + prompt_tokens=1024, completion_tokens=1, total_tokens=1025 + ) + # Anthropic puts these in message_start + message_start_usage.cache_read_input_tokens = 512 + message_start_usage.cache_creation_input_tokens = 128 + message_start = _make_chunk(usage=message_start_usage) + + chunks = [message_start, _make_chunk(content="hi")] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["cache_read_input_tokens"] == 512 + assert result["cache_creation_input_tokens"] == 128 + + def test_openai_streaming_unaffected(self): + """ + OpenAI's only usage chunk is the penultimate one (with + stream_options.include_usage=true), and it carries the real value + directly. Our cursor fix must not break this path — output > 1 + means saw_non_cursor_completion=True so no reset happens. + """ + # Simulate OpenAI: content chunks first, then ONE usage chunk at the end + text_chunks = [_make_chunk(content=t) for t in ["The", " answer", " is 42"]] + usage_chunk = _make_chunk( + usage=Usage(prompt_tokens=42, completion_tokens=15, total_tokens=57), + finish_reason="stop", + ) + chunks = [*text_chunks, usage_chunk] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 42 + assert result["completion_tokens"] == 15 + + def test_single_token_completion_legitimate_case(self): + """ + Edge case: a stream that legitimately completes with output_tokens=1 + (e.g., model returns just "Yes."). Without saw_non_cursor_completion + we'd reset to 0 and fall through to token_counter — but token_counter + on a 1-token string also gives ~1, so billing is still approximately + correct. This test pins that the result is sane (1 or 0). + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21) + ) + text_chunk = _make_chunk(content="Yes.") + # Anthropic's message_delta also gives output_tokens=1 in this case + message_delta = _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21), + finish_reason="stop", + ) + chunks = [message_start, text_chunk, message_delta] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + usage = processor.calculate_usage( + chunks=chunks, + model="claude-sonnet-4-6", + completion_output="Yes.", + messages=[], + ) + + # Two completion-bearing usage events (message_start AND message_delta + # both with output_tokens=1) is positive evidence that message_delta + # arrived — saw_non_cursor_completion goes True via the count >= 2 + # branch and the reset is suppressed. Result: completion_tokens stays + # at the legitimate value of 1. + assert usage.completion_tokens == 1, ( + f"Legitimate single-token completion should bill exactly 1 token " + f"(message_start + message_delta both saw output_tokens=1, " + f"confirming message_delta arrived), got {usage.completion_tokens}" + ) + + def test_anthropic_cache_only_chunks_after_message_start_still_resets(self): + """ + Cache-only chunks (cache_read_input_tokens > 0 but completion_tokens=0) + following message_start should not be mistaken for completion progress. + The cursor=1 from message_start stays the only completion update; reset + must fire so token_counter estimates from completion text instead of + billing the placeholder. + """ + message_start_usage = Usage( + prompt_tokens=1024, completion_tokens=1, total_tokens=1025 + ) + message_start_usage.cache_read_input_tokens = 4096 + message_start = _make_chunk(usage=message_start_usage) + # Subsequent chunks with cache fields but no completion_tokens + cache_chunk_usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + cache_chunk_usage.cache_read_input_tokens = 4096 + cache_chunk = _make_chunk(content="partial", usage=cache_chunk_usage) + # No message_delta — stream was cancelled + chunks = [message_start, cache_chunk] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["cache_read_input_tokens"] == 4096 + assert result["completion_tokens"] == 0, ( + "cache chunks alone don't count as completion progress — only " + "completion_tokens > 0 in a usage event proves real output happened. " + "Reset to 0 forces token_counter fallback." + ) + + +class TestProviderGuard: + """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic + providers, even if they happen to report completion_tokens=1.""" + + def test_non_anthropic_provider_completion_tokens_one_not_reset(self): + """ + Some non-Anthropic provider legitimately reports completion_tokens=1 + in its single usage chunk. Without the provider guard the cursor + heuristic would silently reset it to 0 and bill via token_counter, + producing a different (often inflated) number than what the provider + actually charged. + """ + chunks = [ + _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + finish_reason="stop", + custom_llm_provider="openai", + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 1, ( + "Non-Anthropic providers must not be subject to the message_start " + "cursor reset — their completion_tokens=1 is the real value." + ) + + def test_unknown_provider_completion_tokens_one_not_reset(self): + """No custom_llm_provider on hidden_params (older path or custom + plugin) — heuristic must not fire.""" + chunk = _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ) + # Explicitly clear hidden_params to simulate the unknown-provider case + chunk._hidden_params = {} + processor = ChunkProcessor(chunks=[chunk], messages=[]) + result = processor._calculate_usage_per_chunk(chunks=[chunk]) + assert result["completion_tokens"] == 1 + + +class TestNonAnthropicStreamingIntact: + """Make sure providers without cursor pattern still work.""" + + def test_completion_tokens_above_one_never_resets(self): + """Any chunk reporting completion_tokens > 1 sets saw_non_cursor + and prevents the reset.""" + chunks = [ + _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + + def test_no_usage_chunks_leaves_zero(self): + """Stream with zero usage info → completion_tokens stays 0 + (token_counter fallback will handle it).""" + chunks = [_make_chunk(content="hi"), _make_chunk(content=" there")] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["prompt_tokens"] == 0 + assert result["completion_tokens"] == 0 diff --git a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py new file mode 100644 index 00000000000..7485f2121df --- /dev/null +++ b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py @@ -0,0 +1,147 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + +import litellm + +litellm.model_cost = litellm.get_model_cost_map(url="") + +from litellm.llms.aiml.image_generation.cost_calculator import ( + cost_calculator as aiml_cost_calculator, +) +from litellm.llms.aiml.image_generation.transformation import ( + AimlImageGenerationConfig, +) +from litellm.types.utils import ImageObject, ImageResponse + + +def test_openai_style_model_supports_full_openai_param_surface(): + params = AimlImageGenerationConfig().get_supported_openai_params( + "openai/gpt-image-2" + ) + assert { + "n", + "size", + "quality", + "response_format", + "output_format", + "background", + "moderation", + "output_compression", + } == set(params) + + +def test_flux_style_model_keeps_legacy_param_surface(): + assert AimlImageGenerationConfig().get_supported_openai_params("flux-pro/v1.1") == [ + "n", + "response_format", + "size", + ] + + +def test_openai_style_request_passes_params_through_unchanged(): + """gpt-image-2 must receive OpenAI-shaped fields (size string, n, response_format) verbatim; + the flux-style remapping to ``num_images``/``image_size``/``output_format`` would break the upstream call. + """ + config = AimlImageGenerationConfig() + mapped = config.map_openai_params( + non_default_params={ + "n": 1, + "size": "1024x1536", + "quality": "high", + "response_format": "b64_json", + "output_format": "png", + }, + optional_params={}, + model="openai/gpt-image-2", + drop_params=False, + ) + + body = config.transform_image_generation_request( + model="openai/gpt-image-2", + prompt="A cute baby sea otter", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + + assert body == { + "model": "openai/gpt-image-2", + "prompt": "A cute baby sea otter", + "n": 1, + "size": "1024x1536", + "quality": "high", + "response_format": "b64_json", + "output_format": "png", + } + + +def test_flux_style_request_still_remaps_to_legacy_fields(): + config = AimlImageGenerationConfig() + mapped = config.map_openai_params( + non_default_params={ + "n": 2, + "size": "1024x1024", + "response_format": "png", + }, + optional_params={}, + model="flux-pro/v1.1", + drop_params=False, + ) + + body = config.transform_image_generation_request( + model="flux-pro/v1.1", + prompt="hello", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + + assert body["model"] == "flux-pro/v1.1" + assert body["prompt"] == "hello" + assert body["num_images"] == 2 + assert body["image_size"] == {"width": 1024, "height": 1024} + assert body["output_format"] == "png" + assert "n" not in body + assert "size" not in body + assert "response_format" not in body + + +def test_openai_style_unsupported_param_raises_without_drop_params(): + with pytest.raises(ValueError): + AimlImageGenerationConfig().map_openai_params( + non_default_params={"image_size": {"width": 1024, "height": 1024}}, + optional_params={}, + model="openai/gpt-image-2", + drop_params=False, + ) + + +def test_openai_style_unsupported_param_dropped_with_drop_params(): + mapped = AimlImageGenerationConfig().map_openai_params( + non_default_params={"image_size": {"width": 1024, "height": 1024}}, + optional_params={}, + model="openai/gpt-image-2", + drop_params=True, + ) + assert mapped == {} + + +def test_cost_calculator_uses_aiml_pricing_for_gpt_image_2(): + """Regression: pricing must come from the ``aiml/openai/gpt-image-2`` entry, + not the upstream OpenAI token-based entry. + """ + response = ImageResponse( + data=[ + ImageObject(b64_json=None, url="https://example.com/1.png"), + ImageObject(b64_json=None, url="https://example.com/2.png"), + ] + ) + assert aiml_cost_calculator( + model="openai/gpt-image-2", image_response=response + ) == pytest.approx(0.054 * 2) 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..45820b9833f 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( @@ -2146,6 +2175,283 @@ def test_translate_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_ assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 +def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with_integral_floats(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.prompt_tokens_details = { + "cached_tokens": 30.0, + "cache_write_tokens": 20.0, + } + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 70 + assert anthropic_usage["output_tokens"] == 50 + assert anthropic_usage["cache_read_input_tokens"] == 30 + assert anthropic_usage["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.prompt_tokens_details = { + "cached_tokens": 30.5, + "cache_creation_tokens": 20.25, + } + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 120 + assert anthropic_usage["output_tokens"] == 50 + assert "cache_read_input_tokens" not in anthropic_usage + assert "cache_creation_input_tokens" not in anthropic_usage + + +def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.cache_read_input_tokens = True + usage.cache_creation_input_tokens = True + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 120 + assert anthropic_usage["output_tokens"] == 50 + assert "cache_read_input_tokens" not in anthropic_usage + assert "cache_creation_input_tokens" not in anthropic_usage + + +def test_translate_openai_response_to_anthropic_cache_creation_from_prompt_tokens_details(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="gpt-4o-2024-08-06", + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_response_to_anthropic_cache_tokens_from_usage_fields(): + usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170) + usage.cache_read_input_tokens = 30 + usage.cache_creation_input_tokens = 20 + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_response_to_anthropic_cache_tokens_from_private_usage_fields(): + usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + response.usage._cache_read_input_tokens = 30 + response.usage._cache_creation_input_tokens = 20 + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_details(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_hidden_params_usage(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + ) + response._hidden_params = {"usage": usage} + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_applied_edits(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + applied_edits=[{"type": "compact_20260112"}], + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + assert message_delta["context_management"]["applied_edits"][0]["type"] == ( + "compact_20260112" + ) + + # ===================================================================== # Web Search Tool Transformation Tests # ===================================================================== diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index f74c5b61300..d67de0dcaf8 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -24,6 +24,7 @@ from litellm.types.utils import ( Message, ModelResponse, ModelResponseStream, + PromptTokensDetailsWrapper, StreamingChoices, Usage, ) @@ -88,6 +89,59 @@ def test_fake_stream_usage_preserved(): assert message_delta["usage"]["input_tokens"] == 10 +def test_delayed_usage_chunk_preserves_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=5, + total_tokens=125, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + chunks = [ + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Two."), + finish_reason=None, + ) + ], + ), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + ), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason=None, + ) + ], + usage=usage, + ), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o") + events = list(wrapper) + + message_delta = next( + event for event in events if event.get("type") == "message_delta" + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 5 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 8bc39a6d85e..d9d9474d6f3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -30,7 +30,9 @@ from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, Function, + PromptTokensDetailsWrapper, StreamingChoices, + Usage, ) @@ -107,6 +109,34 @@ def _input_json_deltas(events: List[dict]) -> List[str]: ] +def test_held_stop_reason_usage_merge_preserves_openai_cache_token_details(): + """OpenAI-compatible usage chunks carry cache reads in prompt_tokens_details.""" + wrapper = AnthropicStreamWrapper(completion_stream=iter([]), model="claude-x") + wrapper.holding_stop_reason_chunk = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + + usage_chunk = MagicMock() + usage_chunk.usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + + merged_chunk = wrapper._merge_usage_into_held_stop_reason_chunk(usage_chunk) + + assert merged_chunk["usage"]["input_tokens"] == 70 + assert merged_chunk["usage"]["output_tokens"] == 50 + assert merged_chunk["usage"]["cache_read_input_tokens"] == 30 + assert merged_chunk["usage"]["cache_creation_input_tokens"] == 20 + + def test_first_text_delta_after_tool_use_is_not_dropped_sync(): """A tool_use -> text transition (text resuming after a tool call) carries the resumed text's first token in the trigger chunk. Without the fix it was 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/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a09e55d4ed7..b6c055f4390 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -159,6 +159,59 @@ class TestGetAnthropicHeaders: assert "authorization" not in headers assert "anthropic-dangerous-direct-browser-access" not in headers + def test_custom_api_base_uses_bearer_header(self): + """Custom api_base and non-standard API key should produce Authorization: Bearer header when opted in.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key="my-custom-ollama-token", + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + use_bearer_for_custom_base=True, + ) + + assert headers["authorization"] == "Bearer my-custom-ollama-token" + assert "x-api-key" not in headers + + def test_custom_api_base_uses_bearer_header_already_starts_with_bearer(self): + """If the key already starts with Bearer and Bearer opt-in is enabled, use it directly.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key="Bearer my-custom-ollama-token", + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + use_bearer_for_custom_base=True, + ) + + assert headers["authorization"] == "Bearer my-custom-ollama-token" + assert "x-api-key" not in headers + + def test_custom_api_base_uses_x_api_key_when_standard_key(self): + """If the key is standard sk-ant- key, use x-api-key even with custom api_base.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key=FAKE_REGULAR_KEY, + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + ) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in headers + def test_oauth_includes_standard_headers(self): """OAuth path should still include standard Anthropic headers.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -242,6 +295,46 @@ class TestValidateEnvironmentOAuth: assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY assert "authorization" not in updated_headers + + def test_custom_api_base_via_param(self): + """validate_environment uses Bearer when use_bearer_for_custom_base is set in litellm_params.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"use_bearer_for_custom_base": True}, + api_key="custom-api-key", + api_base="https://custom-gateway.com", + ) + + assert updated_headers["authorization"] == "Bearer custom-api-key" + assert "x-api-key" not in updated_headers + + def test_custom_api_base_via_litellm_params(self): + """validate_environment uses Bearer when api_base and use_bearer_for_custom_base are in litellm_params.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"api_base": "https://custom-gateway.com", "use_bearer_for_custom_base": True}, + api_key="custom-api-key", + api_base=None, + ) + + assert updated_headers["authorization"] == "Bearer custom-api-key" + assert "x-api-key" not in updated_headers assert "anthropic-dangerous-direct-browser-access" not in updated_headers @@ -1004,6 +1097,20 @@ class TestGetAuthHeader: result = AnthropicModelInfo.get_auth_header() assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + def test_custom_api_base_get_auth_header_uses_bearer(self): + """Non-standard API key and custom api_base returns Bearer when use_bearer_for_custom_base=True.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key="my-custom-key", api_base="https://custom-gateway.com", use_bearer_for_custom_base=True) + assert result == {"authorization": "Bearer my-custom-key"} + + def test_custom_api_base_get_auth_header_uses_x_api_key_when_standard(self): + """Standard sk-ant- key with custom api_base should still return x-api-key.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key=FAKE_REGULAR_KEY, api_base="https://custom-gateway.com") + assert result == {"x-api-key": FAKE_REGULAR_KEY} + class TestGetApiBaseFallbackChain: """Tests for AnthropicModelInfo.get_api_base() fallback to ANTHROPIC_BASE_URL.""" @@ -1329,6 +1436,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/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 6298eeb25e9..6b869076044 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -263,3 +263,123 @@ def test_bundled_bedrock_opus_model_info_declares_output_config_effort_ceiling( model_info = GetModelCostMap.load_local_model_cost_map()[model] assert model_info["bedrock_output_config_effort_ceiling"] == expected_ceiling + + +def test_route_prefix_matched_as_path_segment_not_substring(): + """Route tokens like ``mantle/`` must match only at a path-segment boundary. + + The ``bedrock_mantle/`` provider prefix contains the substring ``mantle/``; + a substring match misroutes ``bedrock_mantle/openai.gpt-5.5`` to the Claude + Mythos mantle config, whose request transform strips ``mantle/`` and mangles + the body model into ``bedrock_openai.gpt-5.5``. These assertions fail under + the old substring matching and pass once matching is anchored to ``startswith`` + or a ``/`` boundary. + """ + # The bedrock_mantle/ provider prefix must NOT be read as the mantle/ route. + assert ( + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" + ) + assert ( + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "invoke" + ) + assert ( + BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") + is False + ) + + # A genuine mantle route still resolves, via the startswith branch... + assert ( + BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") + == "mantle" + ) + # ...and via the mid-path "/mantle/" branch (after the bedrock/ provider prefix). + assert ( + BedrockModelInfo.get_bedrock_route( + "bedrock/mantle/anthropic.claude-mythos-preview" + ) + == "mantle" + ) + + +def test_model_has_route_prefix_exercises_both_branches(): + """``_model_has_route_prefix`` matches on ``startswith`` or a ``/`` boundary only.""" + # startswith branch + assert ( + BedrockModelInfo._model_has_route_prefix( + "mantle/anthropic.claude-mythos-preview", "mantle/" + ) + is True + ) + # f"/{prefix}" boundary branch + assert ( + BedrockModelInfo._model_has_route_prefix( + "bedrock/mantle/anthropic.claude-mythos-preview", "mantle/" + ) + is True + ) + # neither branch: the token only appears glued to another segment + assert ( + BedrockModelInfo._model_has_route_prefix( + "bedrock_mantle/openai.gpt-5.5", "mantle/" + ) + is False + ) + + +@pytest.mark.parametrize( + "route_method, token", + [ + (BedrockModelInfo._explicit_converse_route, "converse"), + (BedrockModelInfo._explicit_converse_like_route, "converse_like"), + (BedrockModelInfo._explicit_invoke_route, "invoke"), + (BedrockModelInfo._explicit_async_invoke_route, "async_invoke"), + (BedrockModelInfo._explicit_agent_route, "agent"), + (BedrockModelInfo._explicit_agentcore_route, "agentcore"), + (BedrockModelInfo._explicit_claude_platform_route, "claude_platform"), + (BedrockModelInfo._explicit_openai_route, "openai"), + ], + ids=[ + "converse", + "converse_like", + "invoke", + "async_invoke", + "agent", + "agentcore", + "claude_platform", + "openai", + ], +) +def test_explicit_route_helpers_match_token_only_as_path_segment(route_method, token): + """Each migrated ``_explicit_*_route`` matches its token only as a path segment. + + A leading segment (start of the id or right after a ``/``) matches; the token + glued onto a preceding segment does not. Reverting any method to the old + ``"/" in model`` substring check makes the non-segment case return True + and fails this test. + """ + # leading-segment forms match + assert route_method(f"{token}/some-model") is True + assert route_method(f"bedrock/{token}/some-model") is True + # the token only as a non-segment substring must not match + assert route_method(f"x{token}/y") is False + + +def test_explicit_invoke_route_does_not_match_async_invoke(): + """``invoke/`` must not substring-match ``async_invoke/`` models. + + This is the concrete improvement of the segment-boundary migration: the old + ``"invoke/" in model`` check wrongly classified async-invoke models as the + invoke route. + """ + async_invoke_model = "async_invoke/twelvelabs.marengo-embed-2-7-v1:0" + assert BedrockModelInfo._explicit_invoke_route(async_invoke_model) is False + assert ( + BedrockModelInfo._explicit_invoke_route(f"bedrock/{async_invoke_model}") + is False + ) + # ...while async_invoke/ is still detected as its own route. + assert BedrockModelInfo._explicit_async_invoke_route(async_invoke_model) is True + assert ( + BedrockModelInfo._explicit_async_invoke_route(f"bedrock/{async_invoke_model}") + is True + ) diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py new file mode 100644 index 00000000000..b98b0c22e50 --- /dev/null +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -0,0 +1,40 @@ +""" +Cost tests for Mistral OCR models against the real litellm cost map +(no monkeypatching of get_model_info). These regress the pricing entries +for mistral-ocr-4-0 and mistral-ocr-latest, which now both resolve to +OCR 4 at $4 / 1000 pages. +""" + +import pytest + +import litellm +from litellm.cost_calculator import completion_cost +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo + +OCR4_COST_PER_PAGE = 0.004 + + +def _ocr_response(model: str, pages_processed: int) -> OCRResponse: + return OCRResponse( + pages=[OCRPage(index=i, markdown=f"page {i}") for i in range(pages_processed)], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed), + ) + + +@pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) +def test_model_info_ocr4_price(model: str) -> None: + info = litellm.get_model_info(model=f"mistral/{model}", custom_llm_provider="mistral") + assert info["ocr_cost_per_page"] == OCR4_COST_PER_PAGE + + +@pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: + cost = completion_cost( + completion_response=_ocr_response(model, pages_processed), + model=f"mistral/{model}", + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py index 97461561a05..8ce8b777f71 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py @@ -5,6 +5,7 @@ Tests the supported OCR parameters and their mapping behaviour. No real API calls are made — all tests are fully mocked/local. """ +import httpx import pytest from litellm.llms.mistral.ocr.transformation import MistralOCRConfig @@ -40,9 +41,7 @@ class TestGetSupportedOcrParams: "bbox_annotation_format", "document_annotation_format", ]: - assert ( - param in supported - ), f"Previously supported param '{param}' is missing" + assert param in supported, f"Previously supported param '{param}' is missing" class TestMapOcrParams: @@ -93,12 +92,11 @@ class TestNewSupportedParams: "table_format", "confidence_scores_granularity", "document_annotation_prompt", + "include_blocks", "id", ], ) - def test_new_param_in_supported_list( - self, config: MistralOCRConfig, param_name: str - ) -> None: + def test_new_param_in_supported_list(self, config: MistralOCRConfig, param_name: str) -> None: supported = config.get_supported_ocr_params(model=MODEL) assert param_name in supported @@ -114,12 +112,11 @@ class TestNewParamsMapOcr: ("confidence_scores_granularity", "word"), ("confidence_scores_granularity", "page"), ("document_annotation_prompt", "Extract all invoice line items"), + ("include_blocks", True), ("id", "req-123"), ], ) - def test_new_param_passed_through( - self, config: MistralOCRConfig, param_name: str, param_value: str - ) -> None: + def test_new_param_passed_through(self, config: MistralOCRConfig, param_name: str, param_value: str) -> None: result = config.map_ocr_params( non_default_params={param_name: param_value}, optional_params={}, @@ -144,12 +141,11 @@ class TestTransformOcrRequest: ("document_annotation_prompt", "Extract all invoice line items"), ("id", "req-123"), ("extract_header", True), + ("include_blocks", True), ("pages", [0, 1]), ], ) - def test_param_included_in_request_body( - self, config: MistralOCRConfig, param_name: str, param_value - ) -> None: + def test_param_included_in_request_body(self, config: MistralOCRConfig, param_name: str, param_value) -> None: result = config.transform_ocr_request( model=MODEL, document=self.SAMPLE_DOCUMENT, @@ -176,3 +172,65 @@ class TestTransformOcrRequest: ) for key, value in optional_params.items(): assert result.data[key] == value + + +class TestTransformOcrResponseOcr4Fields: + """OCR 4 adds blocks, confidence_scores, tables, hyperlinks, header and footer + to each page. These must survive transform_ocr_response so callers actually + receive the new structured output rather than having it silently dropped.""" + + def _response(self, page: dict) -> httpx.Response: + return httpx.Response( + 200, + json={ + "pages": [page], + "model": "mistral-ocr-4-0", + "usage_info": {"pages_processed": 1}, + }, + ) + + def test_blocks_and_confidence_scores_preserved(self, config: MistralOCRConfig) -> None: + page = { + "index": 0, + "markdown": "# Invoice", + "blocks": [ + { + "type": "title", + "top_left_x": 10, + "top_left_y": 20, + "bottom_right_x": 300, + "bottom_right_y": 60, + "content": "Invoice", + } + ], + "confidence_scores": {"page": 0.98}, + } + + result = config.transform_ocr_response( + model="mistral-ocr-4-0", + raw_response=self._response(page), + logging_obj=None, + ) + + assert result.pages[0].blocks == page["blocks"] + assert result.pages[0].confidence_scores == page["confidence_scores"] + + def test_ocr4_fields_survive_model_dump(self, config: MistralOCRConfig) -> None: + page = { + "index": 0, + "markdown": "table page", + "tables": [{"rows": 2, "cols": 3}], + "hyperlinks": ["https://example.com"], + "header": "Acme Corp", + "footer": "Page 1", + } + + result = config.transform_ocr_response( + model="mistral-ocr-4-0", + raw_response=self._response(page), + logging_obj=None, + ) + + dumped_page = result.model_dump()["pages"][0] + for field in ("tables", "hyperlinks", "header", "footer"): + assert dumped_page[field] == page[field] 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/gemini/test_context_circulation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py index 6d913ad5d1d..1422531edbf 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py @@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini.transformation import ( ) from litellm.types.llms.vertex_ai import HttpxPartType - # --- Response extraction tests --- @@ -63,6 +62,7 @@ class TestExtractServerSideToolInvocations: assert result[0]["args"] == {"queries": ["weather Buenos Aires"]} assert result[0]["response"] == {"weather": "Sunny, 20°C"} assert result[0]["thought_signature"] == "sig_call_1" + assert result[0]["response_thought_signature"] == "sig_resp_1" def test_returns_none_when_no_server_side_tools(self): """No toolCall/toolResponse parts → returns None.""" @@ -145,6 +145,59 @@ class TestExtractServerSideToolInvocations: assert result[0]["id"] == "exec1" assert "response" not in result[0] + def test_extracts_tool_call_and_response_with_different_signatures(self): + """Case where toolCall and toolResponse have different signatures.""" + parts: List[HttpxPartType] = [ + { + "thoughtSignature": "sig_call_1", + "toolCall": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + }, + }, + { + "thoughtSignature": "sig_resp_1", + "toolResponse": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "response": {"weather": "Sunny, 20°C"}, + }, + }, + ] + + result = VertexGeminiConfig._extract_server_side_tool_invocations(parts) + + assert result is not None + assert len(result) == 1 + assert result[0]["tool_type"] == "GOOGLE_SEARCH_WEB" + assert result[0]["id"] == "abc123" + assert result[0]["args"] == {"queries": ["weather Buenos Aires"]} + assert result[0]["response"] == {"weather": "Sunny, 20°C"} + assert result[0]["thought_signature"] == "sig_call_1" + assert result[0]["response_thought_signature"] == "sig_resp_1" + + def test_orphan_response_signature_extraction(self): + """Orphan toolResponse is captured and has response_thought_signature set.""" + parts: List[HttpxPartType] = [ + { + "thoughtSignature": "sig_resp_1", + "toolResponse": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "orphan123", + "response": {"result": "some response"}, + }, + }, + ] + + result = VertexGeminiConfig._extract_server_side_tool_invocations(parts) + + assert result is not None + assert len(result) == 1 + assert result[0]["id"] == "orphan123" + assert result[0]["response"] == {"result": "some response"} + assert result[0]["response_thought_signature"] == "sig_resp_1" + # --- Input re-injection tests --- @@ -200,6 +253,76 @@ class TestReInjectServerSideToolInvocations: "weather": "Sunny, 20°C" } + def test_roundtrip_single_invocation_with_different_signatures(self): + """Server-side invocations with different signatures for call and response.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "It's sunny in Buenos Aires.", + "provider_specific_fields": { + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny, 20°C"}, + "thought_signature": "sig_call", + "response_thought_signature": "sig_resp", + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + contents = _gemini_convert_messages_with_history(messages) + + model_turn = [c for c in contents if c["role"] == "model"] + assert len(model_turn) == 1 + + parts = model_turn[0]["parts"] + tool_call_parts = [p for p in parts if "toolCall" in p] + tool_response_parts = [p for p in parts if "toolResponse" in p] + + assert len(tool_call_parts) == 1 + assert tool_call_parts[0]["thoughtSignature"] == "sig_call" + + assert len(tool_response_parts) == 1 + assert tool_response_parts[0]["thoughtSignature"] == "sig_resp" + + def test_roundtrip_orphan_response_signature(self): + """Orphan response signature is preserved and re-injected into toolResponse part.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "It's sunny in Buenos Aires.", + "provider_specific_fields": { + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "orphan123", + "response": {"result": "some response"}, + "response_thought_signature": "sig_orphan_resp", + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + contents = _gemini_convert_messages_with_history(messages) + + model_turn = [c for c in contents if c["role"] == "model"] + assert len(model_turn) == 1 + + parts = model_turn[0]["parts"] + tool_response_parts = [p for p in parts if "toolResponse" in p] + + assert len(tool_response_parts) == 1 + assert tool_response_parts[0]["thoughtSignature"] == "sig_orphan_resp" + def test_no_invocations_no_extra_parts(self): """Without server_side_tool_invocations, no extra parts are added.""" messages = [ diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index bb4e6c67e9e..86b3f0976ab 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -10,9 +10,11 @@ Covers: import pytest +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, _is_multimodal_input, + process_embed_content_response, process_response, transform_openai_input_gemini_content, transform_openai_input_gemini_embed_content, @@ -72,7 +74,9 @@ class TestBuildPartForInput: assert part["file_data"]["file_uri"] == GCS_URL def test_file_reference_resolved(self): - resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} + resolved = { + "files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"} + } part = _build_part_for_input("files/abc", resolved_files=resolved) assert part["file_data"] is not None assert part["file_data"]["mime_type"] == "image/jpeg" @@ -94,7 +98,9 @@ class TestTransformOpenaiInputGeminiContent: def test_multiple_texts(self): result = transform_openai_input_gemini_content( - input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={} + input=["hello", "world"], + model="gemini-embedding-2-preview", + optional_params={}, ) assert len(result["requests"]) == 2 assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" @@ -109,7 +115,10 @@ class TestTransformOpenaiInputGeminiContent: ) assert len(result["requests"]) == 2 # First request is text - assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" + assert ( + result["requests"][0]["content"]["parts"][0]["text"] + == "The food was delicious" + ) # Second request is image assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None @@ -288,3 +297,207 @@ class TestProcessResponse: model="gemini-embedding-2-preview", optional_params={}, ) + + +class TestProcessEmbedContentResponseUsage: + """Gemini Embedding 2 embedContent usageMetadata must drive spend. + + Regression for multimodal calls recording prompt_tokens=0 / spend=$0. + """ + + MODEL = "gemini-embedding-2" + + def test_multimodal_image_preserves_usage_metadata(self): + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 258 + assert result.usage.total_tokens == 258 + assert result.usage.prompt_tokens_details.image_count == 1 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_text_modality_detail_populated(self): + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + }, + } + result = process_embed_content_response( + input="a short caption", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 12 + assert result.usage.prompt_tokens_details.text_tokens == 12 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_video_modality_derives_seconds_and_text_floor(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 516, + "totalTokenCount": 516, + "promptTokensDetails": [{"modality": "VIDEO", "tokenCount": 516}], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 516 + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( + 2.0 + ) + assert result.usage.prompt_tokens_details.text_tokens == 1 + + def test_missing_usage_metadata_does_not_estimate_from_base64(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 0 + assert result.usage.total_tokens == 0 + + def test_missing_usage_metadata_text_falls_back_to_token_counter(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input="hello world this is plain text", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens > 0 + + def test_file_reference_image_billed_per_image_not_text(self): + """files/... image refs must bill per-image, not at the text token rate.""" + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=["files/img123"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/img123": { + "mime_type": "image/png", + "uri": "https://example.com/img123", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(0.00012) + + def test_file_reference_non_image_not_counted_as_image(self): + """A files/... ref resolving to a non-image mime must not be image-counted.""" + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 64, + "totalTokenCount": 64, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}], + }, + } + result = process_embed_content_response( + input=["files/clip1"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/clip1": { + "mime_type": "audio/mpeg", + "uri": "https://example.com/clip1", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 0 + assert result.usage.prompt_tokens_details.audio_tokens == 64 + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( + 2.0 + ) + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(2.0 * 0.00016) + + def test_video_plus_audio_does_not_double_bill_text(self): + """Video+audio responses must not get video tokens reassigned to text.""" + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 580, + "totalTokenCount": 580, + "promptTokensDetails": [ + {"modality": "VIDEO", "tokenCount": 516}, + {"modality": "AUDIO", "tokenCount": 64}, + ], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.text_tokens == 1 + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( + 2.0 + ) + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( + 2.0 + ) + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + # 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4 + assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 6f4bb4e59c2..d64b5c8d742 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -1,10 +1,11 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels from litellm.types.router import GenericLiteLLMParams @@ -233,40 +234,36 @@ def test_both_compact_and_context_management_headers_added(): ), f"anthropic-beta should contain 'context-management-2025-06-27', got: {updated_headers['anthropic-beta']}" -def test_validate_environment_with_authorization_header_calculates_api_base(): - """Test that api_base is calculated even when Authorization header is already present""" +def test_validate_environment_always_refreshes_token_ignoring_stale_bearer(): + """Regression: stale Authorization in shared deployment extra_headers must not + skip token refresh on /v1/messages — _ensure_access_token is always called.""" config = VertexAIPartnerModelsAnthropicMessagesConfig() - # Simulate scenario where Authorization is already in headers (e.g., from cached extra_headers) - headers = {"Authorization": "Bearer existing-token"} + headers = {"Authorization": "Bearer EXPIRED"} litellm_params = { "vertex_project": "test-project", "vertex_location": "us-central1", - "extra_headers": {"anthropic-beta": "context-1m-2025-08-07"}, } - optional_params = {} - with patch.object( - config, "get_complete_vertex_url", return_value="https://mock-vertex-url" - ) as mock_get_url: + with ( + patch.object( + config, "_ensure_access_token", return_value=("fresh-token", "test-project") + ) as mock_ensure, + patch.object( + config, "get_complete_vertex_url", return_value="https://mock-vertex-url" + ), + ): updated_headers, api_base = config.validate_anthropic_messages_environment( headers=headers, model="claude-sonnet-4", messages=[], - optional_params=optional_params, + optional_params={}, litellm_params=litellm_params, api_base=None, ) - # Verify that api_base was calculated even though Authorization was already present - assert ( - api_base == "https://mock-vertex-url" - ), f"api_base should be calculated even with Authorization header. Got: {api_base}" - assert mock_get_url.called, "get_complete_vertex_url should be called" - - # Verify Authorization header is still present - assert ( - "Authorization" in updated_headers - ), "Authorization header should be preserved" + mock_ensure.assert_called_once() + assert updated_headers["Authorization"] == "Bearer fresh-token" + assert api_base == "https://mock-vertex-url" def test_transform_anthropic_messages_request_removes_scope_from_cache_control(): @@ -371,3 +368,77 @@ def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instanc assert first_config is second_config finally: ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() + + +def test_validate_environment_does_not_mutate_caller_headers(): + """Regression: beta headers (e.g. web-search) must not leak into the caller's + headers dict — which may be the shared deployment extra_headers object.""" + config = VertexAIPartnerModelsAnthropicMessagesConfig() + caller_headers: dict = {} + + with ( + patch.object( + config, "_ensure_access_token", return_value=("token", "test-project") + ), + patch.object( + config, "get_complete_vertex_url", return_value="https://mock-url" + ), + ): + config.validate_anthropic_messages_environment( + headers=caller_headers, + model="claude-sonnet-4", + messages=[], + optional_params={ + "tools": [{"type": "web_search_20250305", "name": "web_search"}] + }, + litellm_params={ + "vertex_ai_project": "p", + "vertex_ai_location": "us-central1", + }, + api_base=None, + ) + + assert ( + caller_headers == {} + ), "validate_anthropic_messages_environment must not mutate the caller's headers dict" + + +def test_vertex_claude_completion_does_not_mutate_shared_extra_headers(): + """Regression: router shallow-copies litellm_params so extra_headers is a shared + reference. Verify that the chat/completions path builds a new headers dict instead + of calling .update() on the shared object.""" + handler = VertexAIPartnerModels() + shared_extra_headers = {} # simulates deployment["litellm_params"]["extra_headers"] + + mock_response = MagicMock() + + with ( + patch.object( + handler, "_ensure_access_token", return_value=("ya29.fresh", "proj") + ), + patch.object( + handler, "get_complete_vertex_url", return_value="https://mock-url" + ), + patch( + "litellm.llms.anthropic.chat.AnthropicChatCompletion.completion", + return_value=mock_response, + ), + ): + handler.completion( + model="claude-haiku-4-5@20251001", + messages=[{"role": "user", "content": "hi"}], + model_response=MagicMock(), + print_verbose=lambda *a, **k: None, + encoding=None, + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=shared_extra_headers, + timeout=30, + litellm_params={}, + ) + + assert ( + shared_extra_headers == {} + ), "extra_headers must not be mutated by completion()" diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e028064e4c..acad249a2bb 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,8 +1,9 @@ -"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" +"""Tests for the optional Rust-backed OCR path.""" import importlib -import sys +import builtins import types +from typing import Any import httpx import pytest @@ -14,12 +15,16 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse # 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") +rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") +rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" -DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} +DOCUMENT: dict[str, object] = { + "type": "document_url", + "document_url": "https://example.com/doc.pdf", +} -FAKE_OCR_RESPONSE = { +FAKE_OCR_RESPONSE: dict[str, object] = { "pages": [{"index": 0, "markdown": "hello world"}], "model": "mistral-ocr-2505-completion", "document_annotation": None, @@ -28,21 +33,35 @@ FAKE_OCR_RESPONSE = { } +class CapturedException(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" - def __init__(self): - self.calls = [] + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] def __call__( - self, model, document, api_key, api_base, optional_params, timeout_seconds - ): + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: self.calls.append( { "model": model, "document": document, "api_key": api_key, "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, "optional_params": optional_params, "timeout_seconds": timeout_seconds, } @@ -50,13 +69,81 @@ class RecordingBridge: return dict(FAKE_OCR_RESPONSE) +class RecordingAsyncBridge: + """A fake async ``RustAocr`` callable that records the args it was handed.""" + + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_OCR_RESPONSE) + + +class RaisingBridge: + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise RuntimeError("bridge failed") + + +class RaisingAsyncBridge: + async def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise RuntimeError("bridge failed") + + class RecordingLogging: """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" - def __init__(self): - self.pre_call_kwargs = None + def __init__(self) -> None: + self.pre_call_kwargs: dict[str, object] | None = None - def pre_call(self, *, input, api_key, additional_args): + def pre_call( + self, + *, + input: str, + api_key: str | None, + additional_args: dict[str, object], + ) -> None: self.pre_call_kwargs = { "input": input, "api_key": api_key, @@ -67,21 +154,71 @@ class RecordingLogging: 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 __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None: + self.api_key_env_var = api_key_env_var - def get_complete_url(self, *, api_base, model, optional_params, litellm_params): + def get_api_key_env_var(self) -> str: + return self.api_key_env_var + + def validate_environment( + self, + *, + headers: dict[str, object], + model: str, + api_key: str | None, + api_base: str | None, + litellm_params: dict[str, object], + ) -> dict[str, object]: + return {"Authorization": f"Bearer {api_key}", **headers} + + def get_complete_url( + self, + *, + api_base: str | None, + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + ) -> str: return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" +def build_prepared_request( + *, + logging_obj: RecordingLogging | None = None, + provider_config: FakeOCRConfig | None = None, + model: str = "mistral-ocr-latest", + document: dict[str, object] = DOCUMENT, + api_key: str | None = "sk-test", + api_base: str | None = None, + custom_llm_provider: str = "mistral", + extra_headers: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + litellm_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = 12.5, +) -> Any: + return ocr_main._PreparedOCRRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + provider_config=provider_config or FakeOCRConfig(), + optional_params=optional_params or {}, + litellm_params=litellm_params or {}, + effective_timeout=timeout, + litellm_logging_obj=logging_obj or RecordingLogging(), + ) + + @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @pytest.fixture @@ -92,6 +229,14 @@ def fake_bridge(): return bridge +@pytest.fixture +def fake_async_bridge(): + """Enable the async Rust path with an injected recording bridge.""" + bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, aocr=bridge) + return bridge + + def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False litellm.use_litellm_rust() @@ -100,12 +245,61 @@ def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False +def test_env_var_enables_rust_ocr(monkeypatch): + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + assert rust_bridge._env_enables_rust_ocr() is True + + 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_native_bridge_loader_returns_none_when_extension_absent(monkeypatch): + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "litellm.rust_bridge" and "_native" in fromlist: + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + assert rust_bridge_loader.get_native_bridge() is None + + +def test_native_bridge_loader_caches_absent_extension(monkeypatch): + real_import = builtins.__import__ + attempts = 0 + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + nonlocal attempts + if name == "litellm.rust_bridge" and "_native" in fromlist: + attempts += 1 + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + assert rust_bridge_loader.get_native_bridge() is None + assert rust_bridge_loader.get_native_bridge() is None + assert attempts == 1 + + +def test_native_bridge_available_reflects_loader(monkeypatch): + fake_module = types.ModuleType("litellm.rust_bridge._native") + monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module) + + assert rust_bridge_loader.native_bridge_available() is True + + +def test_load_rust_aocr_returns_injected_impl(): + bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, aocr=bridge) + assert rust_bridge.load_rust_aocr() is bridge + + def test_toggle_without_ocr_arg_preserves_injected_impl(): """Regression: routine enable/disable calls must not clobber a prior injection. @@ -114,97 +308,172 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): a caller toggled the flag without re-passing ``ocr=``. """ bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + async_bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge) litellm.use_litellm_rust(False) assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge.load_rust_aocr() is async_bridge litellm.use_litellm_rust(True) assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge.load_rust_aocr() is async_bridge -def test_explicit_ocr_none_clears_injected_impl(): +def test_explicit_ocr_none_clears_injected_impl(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + async_bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge) - litellm.use_litellm_rust(True, ocr=None) + litellm.use_litellm_rust(True, ocr=None, aocr=None) assert rust_bridge.load_rust_ocr() is None + assert rust_bridge.load_rust_aocr() is None -def test_load_rust_ocr_none_when_extension_absent(): +def test_load_rust_ocr_none_when_extension_absent(monkeypatch): """With no injected impl and no compiled wheel, the loader returns None so the caller degrades to the Python path instead of raising ImportError.""" + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI assert rust_bridge.load_rust_ocr() is None + assert rust_bridge.load_rust_aocr() is None def test_load_rust_ocr_uses_compiled_extension(monkeypatch): - """With no injected impl but a compiled ``litellm_python_bridge`` importable, + """With no injected impl but a packaged ``litellm.rust_bridge._native`` importable, the loader returns the extension's ``ocr`` callable. The native wheel isn't - built in CI, so stand in a fake module via ``sys.modules``.""" - fake_module = types.ModuleType("litellm_python_bridge") + built in CI, so stand in a fake module via the bridge loader.""" + fake_module = types.ModuleType("litellm.rust_bridge._native") fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) + fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: fake_module, + ) litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension assert rust_bridge.load_rust_ocr() is fake_module.ocr + assert rust_bridge.load_rust_aocr() is fake_module.aocr def test_timeout_to_seconds_handles_float_timeout_and_none(): - assert 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 + assert rust_bridge._timeout_to_seconds(12.5) == 12.5 + assert rust_bridge._timeout_to_seconds(None) is None + assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 -def test_run_rust_ocr_forwards_args_and_wraps_response(): +def test_bridge_wrapper_forwards_prepared_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, + litellm.use_litellm_rust(True, ocr=bridge) + response = rust_bridge.ocr( 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, + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True, "pages": [0]}, + timeout=12.5, ) - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" + assert response == FAKE_OCR_RESPONSE call = bridge.calls[0] assert call == { "model": "mistral-ocr-latest", "document": DOCUMENT, "api_key": "sk-test", "api_base": "https://proxy.internal", + "custom_llm_provider": "mistral", + "extra_headers": { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + }, + "optional_params": {"include_image_base64": True, "pages": [0]}, + "timeout_seconds": 12.5, + } + + +@pytest.mark.asyncio +async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): + bridge = RecordingAsyncBridge() + + litellm.use_litellm_rust(True, aocr=bridge) + response = await rust_bridge.aocr( + model="mistral-ocr-maas", + document=DOCUMENT, + api_key=None, + api_base=None, + custom_llm_provider="vertex_ai", + extra_headers=None, + optional_params={"vertex_project": "project-1"}, + timeout=httpx.Timeout(30.0, read=42.0), + ) + + assert response == FAKE_OCR_RESPONSE + assert bridge.calls[0] == { + "model": "mistral-ocr-maas", + "document": DOCUMENT, + "api_key": None, + "api_base": None, + "custom_llm_provider": "vertex_ai", + "extra_headers": None, + "optional_params": {"vertex_project": "project-1"}, + "timeout_seconds": 42.0, + } + + +def test_run_rust_ocr_prepares_request_and_wraps_response(): + bridge = RecordingBridge() + logging_obj = RecordingLogging() + litellm.use_litellm_rust(True, ocr=bridge) + + response = ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + logging_obj=logging_obj, + api_base="https://proxy.internal", + extra_headers={"x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True}, + timeout=12.5, + ), + resolve_api_key=lambda _name: None, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + assert bridge.calls[0] == { + "model": "mistral-ocr-latest", + "document": DOCUMENT, + "api_key": "sk-test", + "api_base": "https://proxy.internal", + "custom_llm_provider": "mistral", + "extra_headers": { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + }, "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() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, - logging_obj=RecordingLogging(), - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request(api_key=None, timeout=None), 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" @@ -212,46 +481,148 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() - resolver_calls = [] + litellm.use_litellm_rust(True, ocr=bridge) - def _resolver(name): - resolver_calls.append(name) - return "sk-from-vault" + def _resolver(name: str) -> str | None: + raise AssertionError(f"resolver should not be called for {name}") ocr_main._run_rust_ocr( - rust_ocr=bridge, - logging_obj=RecordingLogging(), - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request( + api_key="sk-explicit", + timeout=None, + ), 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_uses_provider_api_key_env_var(): + bridge = RecordingBridge() + resolver_calls = [] + litellm.use_litellm_rust(True, ocr=bridge) + + def _resolver(name): + resolver_calls.append(name) + return "sk-provider-env" + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), + model="provider-ocr-model", + api_key=None, + timeout=None, + ), + resolve_api_key=_resolver, + ) + + assert resolver_calls == ["PROVIDER_OCR_API_KEY"] + assert bridge.calls[0]["api_key"] == "sk-provider-env" + + +def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + custom_llm_provider="vertex_ai", + model="mistral-ocr-maas", + litellm_params={ + "vertex_project": "project-1", + "vertex_location": "us-central1", + "vertex_credentials": "redacted", + }, + optional_params={"include_image_base64": True}, + timeout=None, + ), + resolve_api_key=lambda _name: None, + ) + + assert bridge.calls[0]["optional_params"] == { + "include_image_base64": True, + "vertex_project": "project-1", + "vertex_location": "us-central1", + } + + +def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + def _resolver(name: str) -> str | None: + return { + "VERTEXAI_PROJECT": "project-from-secret", + "VERTEXAI_LOCATION": "us-east5", + }.get(name) + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + custom_llm_provider="vertex_ai", + model="mistral-ocr-maas", + timeout=None, + ), + resolve_api_key=_resolver, + ) + + assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" + assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" + + +def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_base=None, + timeout=None, + ), + resolve_api_key=lambda name: ( + "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None + ), + ) + + assert bridge.calls[0]["api_base"] == "https://azure.example.com" + + +def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + custom_llm_provider="azure_ai", + model="doc-intelligence/prebuilt-layout", + api_base=None, + timeout=None, + ), + resolve_api_key=lambda name: ( + "https://document-intelligence.example.com" + if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" + else None + ), + ) + + assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" 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() + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=RecordingBridge(), - logging_obj=logging_obj, - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request( + logging_obj=logging_obj, + api_base="https://api.mistral.ai/v1", + extra_headers={"x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True}, + timeout=None, + ), resolve_api_key=lambda _name: None, - 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 @@ -260,9 +631,11 @@ def test_run_rust_ocr_runs_pre_call_logging(): 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"} + assert additional_args["headers"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } def test_ocr_routes_to_rust_when_enabled(fake_bridge): @@ -270,6 +643,7 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): model=MODEL, document=DOCUMENT, api_key="sk-test", + extra_headers={"x-trace-id": "trace-1"}, include_image_base64=True, ) @@ -277,14 +651,108 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): 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["custom_llm_provider"] == "mistral" + assert call["extra_headers"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } assert call["optional_params"].get("include_image_base64") is True +def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): + response = litellm.ocr( + model="azure_ai/pixtral-12b-2409", + document=DOCUMENT, + api_key="sk-test", + api_base="https://example.services.ai.azure.com", + ) + + assert isinstance(response, OCRResponse) + assert len(fake_bridge.calls) == 1 + assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409" + assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" + + +def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): + response = litellm.ocr( + model=MODEL, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + api_key="sk-test", + ) + + assert isinstance(response, OCRResponse) + document = fake_bridge.calls[0]["document"] + assert document["type"] == "document_url" + assert document["document_url"].startswith("data:application/pdf;base64,") + + +def test_ocr_exception_type_uses_resolved_provider_context( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + def fake_exception_type(**kwargs: object) -> CapturedException: + captured.update(kwargs) + return CapturedException("wrapped") + + monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) + litellm.use_litellm_rust(True, ocr=RaisingBridge()) + + with pytest.raises(CapturedException): + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured["model"] == "mistral-ocr-latest" + assert captured["custom_llm_provider"] == "mistral" + + +@pytest.mark.asyncio +async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): + response = await litellm.aocr( + model=MODEL, + document=DOCUMENT, + api_key="sk-test", + extra_headers={"x-trace-id": "trace-1"}, + include_image_base64=True, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + assert len(fake_async_bridge.calls) == 1 + call = fake_async_bridge.calls[0] + assert call["model"] == "mistral-ocr-latest" + assert call["document"] == DOCUMENT + assert call["api_key"] == "sk-test" + assert call["custom_llm_provider"] == "mistral" + assert call["extra_headers"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } + assert call["optional_params"].get("include_image_base64") is True + + +@pytest.mark.asyncio +async def test_aocr_exception_type_uses_resolved_provider_context( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + def fake_exception_type(**kwargs: object) -> CapturedException: + captured.update(kwargs) + return CapturedException("wrapped") + + monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) + litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge()) + + with pytest.raises(CapturedException): + await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured["model"] == "mistral-ocr-latest" + assert captured["custom_llm_provider"] == "mistral" + + def test_ocr_forwards_timeout_to_rust(fake_bridge): """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s client ceiling doesn't silently override shorter deadlines.""" @@ -294,12 +762,10 @@ def test_ocr_forwards_timeout_to_rust(fake_bridge): 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") + from litellm.constants import request_timeout + assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) @@ -317,6 +783,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): """Rust enabled but no bridge available (no injected impl, no compiled wheel): ocr() must degrade to the Python HTTP handler instead of raising.""" + monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI captured = {} @@ -331,3 +798,26 @@ def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): assert captured.get("called") is True # Python path was used assert isinstance(response, OCRResponse) + + +def test_ocr_provider_configs_expose_api_key_env_vars(): + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, + ) + from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig + from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( + VertexAIDeepSeekOCRConfig, + ) + from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig + + assert BaseOCRConfig().get_api_key_env_var() is None + assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" + assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" + assert ( + AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() + == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" + ) + assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" + assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 20fd5c1d86a..b3b0e8adcf6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -170,6 +170,104 @@ class TestMCPRequestHandler: mock_key_servers.assert_called_once_with(user_api_key_auth) mock_team_servers.assert_called_once_with(user_api_key_auth) + @pytest.mark.parametrize( + "require_key_mcp_access_defined,expected", + [ + # Default (flag off): a key with no MCP scope of its own inherits + # the team's servers. + (False, ["team_server1", "team_server2"]), + # Flag on: the team is a ceiling, not a default — the key inherits + # nothing and must grant servers explicitly. + (True, []), + ], + ) + async def test_require_key_mcp_access_defined_gates_team_inheritance( + self, require_key_mcp_access_defined, expected + ): + """The require_key_mcp_access_defined general setting flips an empty key + from inheriting its team's MCP servers (default) to inheriting none.""" + auth = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"require_key_mcp_access_defined": require_key_mcp_access_defined}, + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + + assert sorted(result) == sorted(expected) + + @pytest.mark.parametrize( + "key_servers,grants,expected,scenario", + [ + # Explicit key subset is still honored under the flag (intersected + # with the team ceiling) — the flag only removes empty-key inheritance. + (["team_server1"], [], ["team_server1"], "explicit_subset_survives"), + # An access-group grant is the escape hatch: it surfaces even though + # the key inherits nothing from the team. + ([], ["granted_server"], ["granted_server"], "access_group_grant_survives"), + ], + ) + async def test_require_key_mcp_access_defined_preserves_explicit_grants( + self, key_servers, grants, expected, scenario + ): + """With require_key_mcp_access_defined on, a key still reaches servers it + grants explicitly or via an access group — only blanket team inheritance + is removed.""" + auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + access_group_ids=["grp"] if grants else [], + ) + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=key_servers, + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=grants, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"require_key_mcp_access_defined": True}, + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + + assert sorted(result) == sorted(expected) + @pytest.mark.parametrize("team_servers", [[], ["team_server1", "team_server2"]]) async def test_no_mcp_servers_sentinel_returns_empty(self, team_servers): """A key scoped to the no-mcp-servers sentinel resolves to zero servers, 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 index 7885617aa46..75be6dfc157 100644 --- 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 @@ -1,57 +1,104 @@ -"""Tests for the resolver dispatch skeleton. +"""Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. -Every mode must reach its own arm and, until that arm is built, return a typed -`not_implemented` CredError rather than silently producing no credential. Parametrizing over -one config per mode also guards reachability: if a `case` were dropped, that mode would fall to -the `assert_never` tail and raise here instead of returning the stub. +`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, - AuthSpecKind, AwsSigV4Config, + Byok, ClientCredentialsConfig, Error, NoneConfig, + NoOpAuth, + Ok, PassthroughConfig, ServerSpec, SharedKey, + StaticHeaderAuth, Subject, TokenExchangeConfig, UpstreamCredentialProvider, ) -_ONE_CONFIG_PER_MODE = [ - (AuthSpecKind.none, NoneConfig()), - (AuthSpecKind.api_key, ApiKeyConfig(key_source=SharedKey(value=SecretStr("k")))), - (AuthSpecKind.passthrough, PassthroughConfig()), - (AuthSpecKind.client_credentials, ClientCredentialsConfig()), - (AuthSpecKind.token_exchange, TokenExchangeConfig()), - (AuthSpecKind.authorization_code, AuthorizationCodeConfig()), - (AuthSpecKind.aws_sigv4, AwsSigV4Config(region="us-east-1")), +_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("kind, config", _ONE_CONFIG_PER_MODE) -async def test_every_mode_reaches_its_arm_and_returns_not_implemented(kind, config): - spec = ServerSpec( - server_id="s", resource="https://upstream.example.com", config=config +@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) ) - subject = Subject(tenant_id="", subject_id="") - - result = await UpstreamCredentialProvider().resolve_credentials(subject, spec) - assert isinstance(result, Error) assert result.error.tag == "not_implemented" - assert kind.value in result.error.summary - - -def test_all_seven_modes_are_covered(): - # Guards that the parametrization (and therefore the dispatch) spans every AuthSpecKind, so a - # newly added mode without a test row is caught here rather than slipping through. - assert {kind for kind, _ in _ONE_CONFIG_PER_MODE} == set(AuthSpecKind) 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_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 47c9396f121..27f9a311250 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1800,3 +1800,71 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(RuntimeError("weird")) assert "weird" not in message assert "proxy logs" in message.lower() + + +class TestToolResponseMcpInfoEnrichment: + """The REST tools/list response must expose the user-facing alias and the + server_id alongside the internal server_name so clients (agent builder UIs) + can map the internal config key to a friendly name without needing the + mcp_routes-gated server listing. + """ + + def test_enriches_mcp_info_with_alias_and_server_id(self): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + server = MCPServer( + server_id="a1b2c3d4", + name="mcpAtlassian", + alias="atlassian", + server_name="mcpAtlassian", + transport=MCPTransport.http, + mcp_info={"server_name": "mcpAtlassian"}, + ) + tools = [ + MCPTool( + name="get_issue", + description="Fetch a Jira issue", + inputSchema={"type": "object"}, + ) + ] + + result = rest_endpoints._create_tool_response_objects(tools, server) + + assert result[0].mcp_info == { + "server_name": "mcpAtlassian", + "server_id": "a1b2c3d4", + "alias": "atlassian", + } + + def test_alias_none_is_explicit_in_mcp_info(self): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + server = MCPServer( + server_id="server-uuid", + name="no_alias_server", + alias=None, + server_name="no_alias_server", + transport=MCPTransport.http, + mcp_info={"server_name": "no_alias_server"}, + ) + tools = [ + MCPTool( + name="ping", + description="Ping", + inputSchema={"type": "object"}, + ) + ] + + result = rest_endpoints._create_tool_response_objects(tools, server) + + assert result[0].mcp_info == { + "server_name": "no_alias_server", + "server_id": "server-uuid", + "alias": None, + } 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_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/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 136408e01ac..620daefb39e 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,5 +1,6 @@ import os import sys +import traceback import pytest import requests @@ -11,7 +12,7 @@ sys.path.insert( import responses -from litellm.proxy.client.exceptions import UnauthorizedError +from litellm.proxy.client.exceptions import NotFoundError, UnauthorizedError from litellm.proxy.client.keys import KeysManagementClient @@ -420,3 +421,96 @@ def test_info_server_error(client): ) with pytest.raises(requests.exceptions.HTTPError): client.info(key="test-key") + + +LEAKY_KEY = "sk-1234567890abcdefghijklmnop" + + +def _render_full_traceback(exc: BaseException) -> str: + return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) + + +@responses.activate +def test_info_not_found_redacts_key_everywhere(client): + """A 404 must not echo the raw key embedded in the request URL. + + Covers str(exc) and the rendered traceback, since the chain through + __cause__ / __context__ is what logging.exception() and the default + excepthook print. + """ + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=404, + json={"error": {"message": "Key not found", "code": "404"}}, + ) + with pytest.raises(requests.exceptions.HTTPError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert exc.response is not None + assert exc.response.status_code == 404 + assert exc.request is not None + # Known residual: the live request URL still carries the key, since the + # response is preserved so callers keep status_code / text. str(exc) and the + # traceback are scrubbed; the URL-borne key is the root issue tracked in + # LIT-4013 (move the lookup key out of the query string server-side). + assert LEAKY_KEY in exc.response.request.url + + +@responses.activate +def test_info_unauthorized_redacts_key_everywhere(client): + """A 401 surfaced as UnauthorizedError must not echo the raw key in the + message, the retained original, or the rendered traceback chain.""" + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=401, + json={"error": "Unauthorized"}, + ) + with pytest.raises(UnauthorizedError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in str(exc.orig_exception) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert isinstance(exc.orig_exception, requests.exceptions.HTTPError) + assert exc.orig_exception.response is not None + assert exc.orig_exception.response.status_code == 401 + + +def _http_error_with_key(prefix: str, status: int) -> requests.exceptions.HTTPError: + resp = requests.Response() + resp.status_code = status + return requests.exceptions.HTTPError( + f"{prefix} for url: http://x/key/info?key={LEAKY_KEY}", response=resp + ) + + +def test_unauthorized_error_redacts_wrapped_key(): + """UnauthorizedError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = UnauthorizedError( + _http_error_with_key("401 Client Error: Unauthorized", 401) + ) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 401 + + +def test_not_found_error_redacts_wrapped_key(): + """NotFoundError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = NotFoundError(_http_error_with_key("404 Client Error: Not Found", 404)) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 404 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_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 7a8d04507dc..dd9cbcf5232 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -570,3 +570,41 @@ class TestRequireCallerUserIdForNonAdmin: assert exc_info.value.status_code == 403 assert "Service-account keys" in str(exc_info.value.detail) + + +class TestValidateFiniteSpend: + """`validate_finite_spend` rejects NaN/±inf so a non-finite spend cannot + bypass `spend >= max_budget` enforcement (NaN/-inf compare false).""" + + def test_none_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + assert validate_finite_spend(None) is None + + def test_finite_value_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + assert validate_finite_spend(0.0) is None + assert validate_finite_spend(12.5) is None + # Negative spend is intentionally allowed. Admins may set a negative + # spend counter to grant an entity extra allowance for the current + # budget period only (e.g. a large one-time spend grant), effectively + # raising their headroom without raising the recurring budget ceiling. + # Future changes should continue to allow negative spend counters. + assert validate_finite_spend(-50.0) is None + + @pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) + def test_non_finite_is_rejected(self, bad): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_finite_spend(bad) + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index b4602e0ad8b..27e82df90c1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3102,6 +3102,86 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): assert result is not None +@pytest.mark.asyncio +async def test_admin_user_update_spend_invalidates_counter(mocker): + """A direct /user/update spend change must invalidate the cross-pod + spend counter so enforcement re-reads the new DB value.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.update_data = mocker.AsyncMock( + return_value={"user_id": "target-user", "spend": -25.0} + ) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mock_invalidate = mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + + # Use a negative spend: this also implicitly validates that negative spend + # is allowed, which is desirable. Admins may grant an entity extra + # allowance for the current budget period only (a one-time spend grant) + # without raising the recurring budget ceiling. Future changes should + # continue allowing negative spend counters. + user_request = UpdateUserRequest(user_id="target-user", spend=-25) + admin_caller = UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=admin_caller + ) + mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + + +@pytest.mark.asyncio +async def test_user_update_rejects_non_finite_spend(mocker): + """NaN/inf spend is rejected before any DB write or counter invalidation.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.update_data = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mock_invalidate = mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + + user_request = UpdateUserRequest(user_id="target-user", spend=float("nan")) + admin_caller = UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with pytest.raises(HTTPException) as exc: + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=admin_caller + ) + assert exc.value.status_code == 400 + mock_prisma_client.update_data.assert_not_called() + mock_invalidate.assert_not_awaited() + + @pytest.mark.asyncio async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): """Regression for LIT-3889. diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index b8ec8a8a388..97397fb06be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6639,9 +6639,9 @@ async def test_reset_key_spend_success(monkeypatch): @pytest.mark.asyncio -async def test_update_key_spend_invalidates_counter(monkeypatch): +async def test_update_key_spend_updates_counter(monkeypatch): """ - Test that updating a key's spend via update_key_fn immediately invalidates the spend counter. + Test that updating a key's spend via update_key_fn immediately updates the spend counter. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, @@ -6676,15 +6676,17 @@ async def test_update_key_spend_invalidates_counter(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr("litellm.store_audit_logs", False) + mock_spend_counter_cache = MagicMock() + mock_spend_counter_cache.redis_cache = MagicMock() + mock_spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_cache", + mock_spend_counter_cache, + ) - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" - ) as mock_delete_cache, - patch( - "litellm.proxy.proxy_server._invalidate_spend_counter" - ) as mock_invalidate, - ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ) as mock_delete_cache: mock_delete_cache.return_value = None user_api_key_dict = UserAPIKeyAuth( @@ -6704,7 +6706,12 @@ async def test_update_key_spend_invalidates_counter(monkeypatch): ) mock_delete_cache.assert_awaited_once() - mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") + mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with( + key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + ) + mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( + key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + ) @pytest.mark.asyncio @@ -9718,6 +9725,39 @@ class TestKeyOwnerPrivilegeEscalation: assert exc_info.value.status_code == 403 mock_check.assert_called_once() + @pytest.mark.asyncio + async def test_creator_cannot_reset_own_spend_to_stale_value(self): + """Submitting `spend` equal to the stale DB value must still require + admin. The DB spend lags the live cross-pod counter, so an + "unchanged" spend on the non-admin path would let the creator + overwrite the live counter below real usage. Any explicit `spend` + is a budget change, regardless of value match.""" + existing = self._make_existing_key(created_by="creator-123") + existing.spend = 0.0 + # spend equals the stale DB value (0.0) — the old `!=` gate skipped + # the admin check here. + data = UpdateKeyRequest(key="sk-test", spend=0.0) + auth = self._make_auth(user_id="creator-123") + + mock_check = AsyncMock( + side_effect=HTTPException(status_code=403, detail="Not authorized") + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + mock_check, + ): + with pytest.raises(HTTPException): + await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=auth, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + mock_check.assert_called_once() + @pytest.mark.asyncio async def test_assigned_user_blocked_from_model_escalation(self): data = UpdateKeyRequest(key="sk-test", models=["gpt-4", "claude-opus"]) 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_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 592232f45f5..196ff045208 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -10,8 +10,8 @@ from __future__ import annotations import os from types import SimpleNamespace -from typing import Any, Dict, List, Optional -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Any, Dict +from unittest.mock import AsyncMock, MagicMock import pytest @@ -887,6 +887,84 @@ def test_ProxyConfig__add_deployment_invalid_litellm_params_skips(monkeypatch): assert pc._add_deployment(db_models=[bad]) == 0 +def test_ProxyConfig__add_deployment_resolves_env_refs_after_db_decrypt(monkeypatch): + monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret") + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: value, + ) + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="env-model", + model_info={"id": "model-1"}, + litellm_params={ + "model": "openai/gpt-4o-mini", + "api_key": "os.environ/LITELLM_DB_MODEL_API_KEY", + "api_base": "os.environ/LITELLM_MASTER_KEY", + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.api_key == "resolved-secret" + assert deployment.litellm_params.api_base == "os.environ/LITELLM_MASTER_KEY" + + +def test_ProxyConfig__add_deployment_keeps_team_env_refs_literal(monkeypatch): + def fail_on_call(secret_name, *args, **kwargs): + raise AssertionError("team DB models should not resolve env refs") + + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: value, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call) + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="model_name_team-1_abc", + model_info={"id": "model-1", "team_id": "team-1"}, + litellm_params={ + "model": "openai/gpt-4o-mini", + "api_key": "os.environ/LITELLM_MASTER_KEY", + "api_base": "https://attacker.example", + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.api_key == "os.environ/LITELLM_MASTER_KEY" + assert deployment.litellm_params.api_base == "https://attacker.example" + + +def test_ProxyConfig__resolve_db_litellm_param_skips_non_string_values(monkeypatch): + def fail_on_call(value, key, return_original_value): + raise AssertionError("decrypt_value_helper should only receive strings") + + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + fail_on_call, + ) + pc = ProxyConfig() + + assert pc._resolve_db_litellm_param(key="tpm", value=100) == 100 + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- @@ -919,6 +997,71 @@ def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch): } +def test_ProxyConfig_decrypt_model_list_from_db_resolves_env_refs_after_db_decrypt( + monkeypatch, +): + monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret") + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: ( + "os.environ/LITELLM_DB_MODEL_API_KEY" + if key == "api_key" + else "os.environ/LITELLM_MASTER_KEY" if key == "api_base" else value + ), + ) + pc = ProxyConfig() + m = SimpleNamespace( + model_id="model-1", + model_name="env-model", + model_info={"id": "model-1"}, + litellm_params={ + "api_key": "encrypted-env-ref", + "api_base": "encrypted-api-base-env-ref", + "model": "openai/gpt-4o-mini", + }, + blocked=False, + ) + + out = pc.decrypt_model_list_from_db(new_models=[m]) + + assert out[0]["litellm_params"]["api_key"] == "resolved-secret" + assert out[0]["litellm_params"]["api_base"] == "os.environ/LITELLM_MASTER_KEY" + + +def test_ProxyConfig_decrypt_model_list_from_db_keeps_team_env_refs_literal_after_db_decrypt( + monkeypatch, +): + def fail_on_call(secret_name, *args, **kwargs): + raise AssertionError("team DB models should not resolve env refs") + + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: ( + "os.environ/LITELLM_MASTER_KEY" if key == "api_key" else value + ), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call) + pc = ProxyConfig() + m = SimpleNamespace( + model_id="model-1", + model_name="model_name_team-1_abc", + model_info={"id": "model-1", "team_id": "team-1"}, + litellm_params={ + "api_key": "encrypted-env-ref", + "api_base": "https://attacker.example", + "model": "openai/gpt-4o-mini", + }, + blocked=False, + ) + + out = pc.decrypt_model_list_from_db(new_models=[m]) + + assert out[0]["litellm_params"]["api_key"] == "os.environ/LITELLM_MASTER_KEY" + assert out[0]["litellm_params"]["api_base"] == "https://attacker.example" + + def test_ProxyConfig_decrypt_model_list_from_db_invalid_params_skips(): pc = ProxyConfig() bad = SimpleNamespace( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 0b583129591..6d0b4509b45 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -342,6 +342,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat ignored_keys = [ "request_id", + "metadata.litellm_call_id", "session_id", "startTime", "endTime", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e305054d075..874e0654a1f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2120,3 +2120,112 @@ def test_get_logging_payload_failure_without_recovered_usage_is_zero(): ) assert payload["total_tokens"] == 0 + + +def test_get_logging_payload_sets_litellm_call_id_for_correlation(): + """LIT-3868: a successful spend log must carry the x-litellm-call-id (the + trace id) in its metadata, distinct from request_id, which stays the + provider response id. Without this there is no way to correlate a DB row + with its trace for a successful call. + """ + provider_response_id = "chatcmpl-e6e6f3e9-c392-404e-9a71-5361c79d8470" + trace_call_id = "c6a77556-19ce-4406-b287-53f5fb4b2b55" + + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "id": provider_response_id, + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + metadata = json.loads(payload["metadata"]) + + assert payload["request_id"] == provider_response_id + assert metadata["litellm_call_id"] == trace_call_id + assert metadata["litellm_call_id"] != payload["request_id"] + + +def test_get_logging_payload_litellm_call_id_falls_back_to_litellm_params(): + """litellm_call_id may only be present in litellm_params; it must still land + in the spend log metadata so correlation works on that path too. + """ + trace_call_id = "fallback-7a1c-42d9-9f0e-2b6c5d4e3f21" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_params": { + "litellm_call_id": trace_call_id, + "metadata": {"user_api_key": "sk-test"}, + }, + } + response_obj = { + "id": "chatcmpl-abc123", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + + +def test_get_logging_payload_litellm_call_id_when_response_has_no_id(): + """When the provider returns no id, request_id falls back to the call id, so + request_id and the metadata call id hold the same value and correlation + still resolves. + """ + trace_call_id = "noid-5b2e-4c7a-9d10-3f8a1c2b4e6d" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + assert payload["request_id"] == trace_call_id + + +def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): + """On a cache hit request_id is suffixed to stay unique, but the metadata + litellm_call_id stays the raw trace id so the row still points at its trace. + """ + trace_call_id = "cache-9a1c-42d9-9f0e-2b6c5d4e3f21" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "cache_hit": True, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "id": "chatcmpl-cache-src", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + assert "_cache_hit" in payload["request_id"] + assert json.loads(payload["metadata"])["litellm_call_id"] != payload["request_id"] diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 364861d6e31..3d3a2cc1013 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2022,6 +2022,52 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == requested_model + def test_skips_model_override_when_response_has_no_model_attribute(self): + from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult + + response_obj = SearchResponse( + results=[SearchResult(title="t", url="http://x.com", snippet="s")], + object="search", + ) + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-search-tool", + log_context="test_context", + ) + + assert not hasattr(response_obj, "model") + + def test_skips_model_override_for_dict_without_model_key(self): + response_obj = { + "object": "search", + "results": [{"title": "t", "url": "http://x.com", "snippet": "s"}], + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-search-tool", + log_context="test_context", + ) + + assert "model" not in response_obj + + def test_override_model_swallows_setattr_failure(self): + class ReadOnlyModelResponse: + @property + def model(self) -> str: + return "downstream-model" + + response_obj = ReadOnlyModelResponse() + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-model", + log_context="test_context", + ) + + assert response_obj.model == "downstream-model" + class TestIsAzureModelRouterRequest: """Tests for _is_azure_model_router_request helper""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index a203fcc7ec0..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): diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py new file mode 100644 index 00000000000..04db7192364 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py @@ -0,0 +1,73 @@ +"""Regression tests for the responses -> completion fallback bridge guard. + +When the Responses API falls back to chat completions (no native responses +config), it must tag the forwarded ``litellm.completion`` / ``litellm.acompletion`` +call with ``_skip_responses_api_bridge=True`` so ``completion()`` does not bridge +the request straight back to the Responses API and mutually recurse forever. + +Both fallback paths are covered: the sync ``response_api_handler`` (``_is_async`` +False) and the async ``async_response_api_handler`` (``_is_async`` True). The +module-level ``litellm.completion`` / ``litellm.acompletion`` are patched to +capture the forwarded kwargs; if the flag-setting line is removed the captured +kwargs lack the flag and these tests fail. +""" + +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.responses.litellm_completion_transformation.handler import ( + LiteLLMCompletionTransformationHandler, +) + + +class _StopForwarding(Exception): + """Raised by the mocked (a)completion once the forwarded kwargs are captured.""" + + +def test_sync_fallback_tags_skip_responses_api_bridge(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + def fake_completion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.completion", fake_completion): + with pytest.raises(_StopForwarding): + handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=False, + ) + + assert captured.get("_skip_responses_api_bridge") is True + + +@pytest.mark.asyncio +async def test_async_fallback_tags_skip_responses_api_bridge(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.acompletion", fake_acompletion): + coro = handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=True, + ) + with pytest.raises(_StopForwarding): + await coro + + assert captured.get("_skip_responses_api_bridge") is True diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 9f19944fdba..77cee8a485c 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -1,7 +1,8 @@ """Tests for scripts/budget_ratchet_check.py. -The guard's whole contract is "ceilings may only fall": a raised ceiling, a dropped -rule, or a deleted file is a regression, while a lowered/equal ceiling, a brand-new +The guard's contract is "baselines and ceilings may only fall": a raised ceiling, a +raised baseline (even when slack is cut to keep the ceiling flat), a dropped rule, or +a deleted file is a regression, while a lowered/equal baseline and ceiling, a brand-new rule, or a brand-new budget file is fine. Each branch is pinned here. """ @@ -10,7 +11,9 @@ import subprocess import sys from pathlib import Path -_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" +_MODULE_PATH = ( + Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" +) _spec = importlib.util.spec_from_file_location("budget_ratchet_check", _MODULE_PATH) ratchet = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(ratchet) @@ -35,10 +38,30 @@ def test_raised_ceiling_is_a_regression(): def test_lowered_or_equal_ceiling_is_clean(): base = {"LIT006": _spec_of(1013, 10)} + # baseline drops, slack flat -> ceiling falls assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1000, 10)}) == [] + # nothing changes assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 10)}) == [] - # slack traded for baseline at the same ceiling is fine - assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)}) == [] + # slack cut while baseline holds -> ceiling falls, baseline flat + assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 0)}) == [] + + +def test_raised_baseline_is_a_regression_even_when_ceiling_held_flat(): + # baseline 1013 -> 1023 with slack cut 10 -> 0 keeps the ceiling at 1023, but a + # higher baseline bakes in more accepted debt and must still surface as a regression + base = {"LIT006": _spec_of(1013, 10)} + regs = ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)}) + assert [r.rule for r in regs] == ["LIT006"] + assert "baseline raised 1013 -> 1023" in regs[0].detail + assert "ceiling raised" not in regs[0].detail + + +def test_raised_baseline_and_ceiling_report_both_reasons(): + base = {"LIT006": _spec_of(1013, 10)} + regs = ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1100, 10)}) + assert [r.rule for r in regs] == ["LIT006"] + assert "ceiling raised 1023 -> 1110" in regs[0].detail + assert "baseline raised 1013 -> 1100" in regs[0].detail def test_dropped_rule_is_a_regression(): diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py index 496b0276b87..7cc05d6e30a 100644 --- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py +++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py @@ -3,19 +3,45 @@ from pathlib import Path import pytest +import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -@pytest.mark.parametrize("model", ["mistral/mistral-medium-3-5"]) -def test_mistral_medium_3_5_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) +MEDIUM_3_5_MODELS = ( + "mistral/mistral-medium-3-5", + "mistral/mistral-medium-2604", + "mistral/mistral-medium-latest", +) - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" +SYNCED_MODELS = MEDIUM_3_5_MODELS + ( + "mistral/mistral-medium-2508", + "mistral/mistral-medium-3-1-2508", +) + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force get_model_info to resolve against the in-repo cost map instead of the + remote one fetched at import time, which still carries the pre-merge pricing.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS) +def test_medium_3_5_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" assert info["litellm_provider"] == "mistral" assert info["mode"] == "chat" @@ -27,10 +53,11 @@ def test_mistral_medium_3_5_model_info(model): assert info["max_output_tokens"] == 262144 assert info["max_tokens"] == 262144 + assert info["supports_reasoning"] is True + assert info["supports_vision"] is True assert info["supports_function_calling"] is True assert info["supports_response_schema"] is True assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True assert info["supports_assistant_prefill"] is True routed_model, provider, _, _ = get_llm_provider(model=model) @@ -38,18 +65,32 @@ def test_mistral_medium_3_5_model_info(model): assert provider == "mistral" -def test_mistral_medium_3_5_backup_matches_main(): - """Ensure the bundled model cost map stays in sync with the canonical file.""" - repo_root = Path(__file__).parents[2] - main_path = repo_root / "model_prices_and_context_window.json" - backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json" +def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map): + """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must + return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values.""" + info = litellm.get_model_info(model="mistral/mistral-medium-latest") - with open(main_path) as f: - main_cost = json.load(f) - with open(backup_path) as f: - backup_cost = json.load(f) + assert info["input_cost_per_token"] == 1.5e-06 + assert info["output_cost_per_token"] == 7.5e-06 + assert info["max_input_tokens"] == 262144 + assert info["supports_reasoning"] is True - for model in ("mistral/mistral-medium-3-5",): - assert backup_cost.get(model) == main_cost.get( - model - ), f"{model} differs between main and backup model cost maps" + +def test_mistral_medium_2508_keeps_medium_3_1_specs(): + """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing.""" + info = _load(MAIN_PATH).get("mistral/mistral-medium-2508") + assert info is not None, "mistral/mistral-medium-2508 missing from cost map" + + assert info["input_cost_per_token"] == 4e-07 + assert info["output_cost_per_token"] == 2e-06 + assert info["max_input_tokens"] == 131072 + assert info.get("supports_reasoning") is not True + + +@pytest.mark.parametrize("model", SYNCED_MODELS) +def test_backup_matches_main(model): + """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3e2150848a7..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 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.json b/ui/litellm-dashboard/package.json index 23da0bcd636..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", 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/SearchTools/CreateSearchTools.tsx b/ui/litellm-dashboard/src/components/SearchTools/CreateSearchTools.tsx index cf355961d11..a8ed1a5e005 100644 --- a/ui/litellm-dashboard/src/components/SearchTools/CreateSearchTools.tsx +++ b/ui/litellm-dashboard/src/components/SearchTools/CreateSearchTools.tsx @@ -3,8 +3,8 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { useQuery } from "@tanstack/react-query"; import { Button, TextInput } from "@tremor/react"; import { Form, Input, Modal, Select, Tooltip, Typography } from "antd"; -import Image from "next/image"; import React, { useState } from "react"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import NotificationsManager from "../molecules/notifications_manager"; import { createSearchTool, fetchAvailableSearchProviders } from "../networking"; import SearchConnectionTest from "./SearchConnectionTest"; @@ -13,7 +13,7 @@ import { AvailableSearchProvider, SearchTool } from "./types"; const { TextArea } = Input; // Search provider logos folder path (matches existing provider logo pattern) -const searchProviderLogosFolder = "../ui/assets/logos/"; +const searchProviderLogosFolder = "/ui/assets/logos/"; // Helper function to get logo path for a search provider const getSearchProviderLogo = (providerName: string): string => { @@ -28,12 +28,13 @@ interface SearchProviderLabelProps { const SearchProviderLabel: React.FC = ({ providerName, displayName }) => (
- = ({ visible, onClose, accessTok value={info.agent_type} label={
- + {info.agent_type_display_name}
} >
- {info.agent_type_display_name} + {info.agent_type_display_name}
{info.agent_type_display_name}
{info.description &&
{info.description}
} @@ -942,7 +947,9 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok - {selectedLogo && currentStep < 1 && Agent} + {selectedLogo && currentStep < 1 && ( + Agent + )}

Add New Agent

} diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index 0334c55f66c..3c6b3829fef 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -7,7 +7,7 @@ interface CallbackConfig { description: string; } -const asset_logos_folder = "../ui/assets/logos/"; +const asset_logos_folder = "/ui/assets/logos/"; export const CALLBACK_CONFIGS: CallbackConfig[] = [ { diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 41a2dd7f67c..4ac4242b8ab 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -20,6 +20,7 @@ import { shouldRenderLLMJudgeFields, shouldRenderPIIConfigSettings, } from "./guardrail_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import GuardrailOptionalParams from "./guardrail_optional_params"; import GuardrailProviderFields from "./guardrail_provider_fields"; import LLMJudgeFields from "./llm_judge/LLMJudgeFields"; @@ -689,7 +690,7 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a
{guardrailLogoMap[value] && ( = ({ visible, onClose, a
{guardrailLogoMap[value] && ( = ({
{guardrailLogoMap[value] && ( = ({ src, name }) => { const [hasError, setHasError] = useState(false); @@ -29,7 +30,7 @@ const LogoWithFallback: React.FC<{ src: string; name: string }> = ({ src, name } return ( setHasError(true)} diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts index c49eedaac23..f81277f13c3 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts @@ -16,7 +16,7 @@ export interface GuardrailCardInfo { providerKey?: string; } -const ASSET_PREFIX = "../ui/assets/logos/"; +const ASSET_PREFIX = "/ui/assets/logos/"; export const LITELLM_CONTENT_FILTER_CARDS: GuardrailCardInfo[] = [ { diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_detail.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_detail.tsx index 00daeeb8f01..c92486bbad9 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_detail.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_detail.tsx @@ -2,6 +2,7 @@ import React, { useState } from "react"; import { Button } from "antd"; import { ArrowLeftOutlined } from "@ant-design/icons"; import AddGuardrailForm from "./add_guardrail_form"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import { GUARDRAIL_PRESETS } from "./guardrail_garden_configs"; import { GuardrailCardInfo } from "./guardrail_garden_data"; @@ -60,7 +61,7 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {/* ── Header block (Vertex-style) ── */}
{ diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx index 837d0cf83fc..3ac9fe4087a 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx @@ -1,3 +1,5 @@ +import { resolveLogoSrc } from "@/lib/assetPaths"; + // Legacy enum - keeping for backward compatibility export enum GuardrailProviders { PresidioPII = "Presidio PII", @@ -113,7 +115,7 @@ export const shouldRenderLLMJudgeFields = (provider: string | null) => { return guardrail_provider_map[provider] === "llm_as_a_judge"; }; -const asset_logos_folder = "../ui/assets/logos/"; +const asset_logos_folder = "/ui/assets/logos/"; export const guardrailLogoMap: Record = { "Zscaler AI Guard": `${asset_logos_folder}zscaler.svg`, @@ -163,9 +165,9 @@ export const getGuardrailLogoAndName = (guardrailValue: string): { logo: string; // Get the display name from current GuardrailProviders and logo from map const currentProviders = getGuardrailProviders(); const displayName = currentProviders[enumKey as keyof typeof currentProviders]; - const logo = guardrailLogoMap[displayName as keyof typeof guardrailLogoMap]; + const logo = resolveLogoSrc(guardrailLogoMap[displayName as keyof typeof guardrailLogoMap]) ?? ""; - return { logo: logo || "", displayName: displayName || guardrailValue }; + return { logo, displayName: displayName || guardrailValue }; }; /** Tri-state UI value for `litellm_params.skip_system_message_in_guardrail` (inherit = use global). */ diff --git a/ui/litellm-dashboard/src/components/logging_settings_view.tsx b/ui/litellm-dashboard/src/components/logging_settings_view.tsx index aeb83f255bc..5124d98da5d 100644 --- a/ui/litellm-dashboard/src/components/logging_settings_view.tsx +++ b/ui/litellm-dashboard/src/components/logging_settings_view.tsx @@ -2,6 +2,7 @@ import React from "react"; import { Tag } from "antd"; import { CogIcon, BanIcon } from "@heroicons/react/outline"; import { callbackInfo, callback_map, reverse_callback_map } from "./callback_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; interface LoggingConfig { callback_name: string; @@ -68,7 +69,7 @@ export function LoggingSettingsView({
{loggingConfigs.map((config, index) => { const displayName = getLoggingDisplayName(config.callback_name); - const logoUrl = callbackInfo[displayName]?.logo; + const logoUrl = resolveLogoSrc(callbackInfo[displayName]?.logo); return (
{ // Handle both display names and internal values const displayName = reverse_callback_map[callbackName] || callbackName; - const logoUrl = callbackInfo[displayName]?.logo; + const logoUrl = resolveLogoSrc(callbackInfo[displayName]?.logo); return (
= ({ value, onChange }) => {value && (
Selected logo { @@ -96,7 +97,7 @@ const MCPLogoSelector: React.FC = ({ value, onChange }) => style={{ width: 40, height: 40 }} > {logo.name} handleImgError(logo.url)} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx index 6d0ef76b11b..1a05edf7c2c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolsetsTab.tsx @@ -12,6 +12,17 @@ import { MCPToolset, MCPToolsetTool } from "./types"; const { Text: AntdText } = Typography; +// Display-only. Toolsets persist {server_id, bare tool_name}; the gateway serves +// each tool prefixed as "{server-prefix}-{tool}". Render that qualified form so +// the same tool name on different servers stays distinguishable. This mirrors the +// backend default MCP_TOOL_PREFIX_SEPARATOR; overriding that env var only changes +// this cosmetic label, never what is stored or how tools are matched. +const MCP_TOOL_PREFIX_SEPARATOR = "-"; + +function displayToolName(serverPrefix: string | undefined, toolName: string): string { + return serverPrefix ? `${serverPrefix}${MCP_TOOL_PREFIX_SEPARATOR}${toolName}` : toolName; +} + interface MCPToolsetsTabProps { accessToken: string | null; userRole: string | null; @@ -138,6 +149,10 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset const [saving, setSaving] = useState(false); const [serverSearch, setServerSearch] = useState(""); const { data: mcpServers = [] } = useMCPServers(); + const serverPrefixById = React.useMemo( + () => new Map(mcpServers.map((s) => [s.server_id, s.alias || s.server_name || s.server_id])), + [mcpServers], + ); React.useEffect(() => { if (open) { @@ -254,7 +269,7 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset >
- {tool.tool_name} + {displayToolName(serverPrefixById.get(tool.server_id), tool.tool_name)} {tool.server_id.slice(0, 8)}…
@@ -282,8 +297,9 @@ function toolsetColumns( isAdmin: boolean, onEdit: (t: MCPToolset) => void, onDelete: (id: string) => void, - proxyBaseUrl: string, + serverPrefixById: Map, ): ColumnDef[] { + const proxyBaseUrl = getProxyBaseUrl(); return [ { header: "Toolset ID", @@ -334,7 +350,7 @@ function toolsetColumns( key={i} className="inline-flex items-center px-1.5 py-0.5 rounded bg-purple-50 border border-purple-200 text-purple-700 text-xs" > - {t.tool_name} + {displayToolName(serverPrefixById.get(t.server_id), t.tool_name)} ))} {tools.length > 4 && +{tools.length - 4} more} @@ -431,6 +447,7 @@ function ToolsetUsageGuide() { export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { const queryClient = useQueryClient(); const { data: toolsets = [], isLoading } = useMCPToolsets(); + const { data: mcpServers = [] } = useMCPServers(); const [createOpen, setCreateOpen] = useState(false); const [editToolset, setEditToolset] = useState(null); const [deleteId, setDeleteId] = useState(null); @@ -466,8 +483,11 @@ export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { } }; - const proxyBaseUrl = getProxyBaseUrl(); - const columns = toolsetColumns(isAdmin, setEditToolset, setDeleteId, proxyBaseUrl); + const serverPrefixById = React.useMemo( + () => new Map(mcpServers.map((s) => [s.server_id, s.alias || s.server_name || s.server_id])), + [mcpServers], + ); + const columns = toolsetColumns(isAdmin, setEditToolset, setDeleteId, serverPrefixById); return (
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx index e3e84cb7434..a8fbe1c1fbc 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx @@ -1,6 +1,7 @@ import React from "react"; import { Button, TextInput } from "@tremor/react"; import { MCPTool, InputSchema, InputSchemaProperty } from "./types"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import { Form, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import NotificationsManager from "../molecules/notifications_manager"; @@ -301,7 +302,7 @@ export function ToolTestPanel({ {tool.mcp_info.logo_url && ( // eslint-disable-next-line @next/next/no-img-element {`${tool.mcp_info.server_name} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index ddcc9f65d38..b45c902b9d7 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -29,8 +29,9 @@ import NotificationsManager from "../molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { useTestMCPConnection } from "@/hooks/useTestMCPConnection"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { resolveLogoSrc } from "@/lib/assetPaths"; -const asset_logos_folder = "../ui/assets/logos/"; +const asset_logos_folder = "/ui/assets/logos/"; export const mcpLogoImg = `${asset_logos_folder}mcp_logo.png`; interface CreateMCPServerProps { @@ -586,7 +587,7 @@ const CreateMCPServer: React.FC = ({ )} MCP Logo = ({
MCP Logo = ({ > {server.icon_url ? ( {server.title} {tool.mcp_info.logo_url && ( {`${tool.mcp_info.server_name} diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx index 2324cbdd0e9..b86a379d3d1 100644 --- a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx @@ -4,6 +4,7 @@ import type { UploadProps } from "antd/es/upload"; import React, { useState } from "react"; import ProviderSpecificFields from "../add_model/provider_specific_fields"; import { Providers, providerLogoMap } from "../provider_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import { resetCredentialFormOnProviderChange } from "./credential_form_helpers"; const { Link } = Typography; @@ -67,7 +68,7 @@ const AddCredentialsModal: React.FC = ({ open, onCance
{`${providerEnum} { diff --git a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx index f504ba7a78a..d087edc1069 100644 --- a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx @@ -5,6 +5,7 @@ import { useEffect, useState } from "react"; import ProviderSpecificFields from "../add_model/provider_specific_fields"; import { CredentialItem } from "../networking"; import { Providers, providerLogoMap } from "../provider_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import { resetCredentialFormOnProviderChange } from "./credential_form_helpers"; const { Link } = Typography; @@ -100,7 +101,7 @@ export default function EditCredentialsModal({
{`${providerEnum} { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 3d15aea8fdd..f5bae832e64 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -32,6 +32,9 @@ import NotificationsManager from "./molecules/notifications_manager"; import type { MCPUserEnvVarsStatus } from "./mcp_tools/types"; import { createApiClient, deriveErrorMessage } from "@/lib/http/client"; import { resolveApiBase } from "@/lib/http/resolveApiBase"; +import { serverRootPath, setServerRootPath } from "@/lib/serverRootPath"; + +export { serverRootPath }; export { deriveErrorMessage }; export { ApiError } from "@/lib/http/client"; @@ -46,8 +49,6 @@ const resolveDefaultBase = (fallback: string | null): string | null => ? "http://localhost:4000" : fallback; const defaultProxyBaseUrl = resolveDefaultBase(null); -const defaultServerRootPath = "/"; -export let serverRootPath = defaultServerRootPath; const WORKER_URL_KEY = "litellm_worker_url"; // If a worker URL is in localStorage, use it as the initial proxyBaseUrl. // This survives page navigation and the sessionStorage.clear() in user_dashboard. @@ -92,7 +93,7 @@ const updateProxyBaseUrl = (serverRootPath: string, receivedProxyBaseUrl: string }; const updateServerRootPath = (receivedServerRootPath: string) => { - serverRootPath = receivedServerRootPath; + setServerRootPath(receivedServerRootPath); }; export const getProxyBaseUrl = (): string => { diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index acc77d1263b..d870c278db0 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -436,3 +436,27 @@ describe("provider_info_helpers", () => { }); }); }); + +describe("getProviderLogoAndName under a custom server_root_path", () => { + afterEach(() => { + vi.resetModules(); + vi.doUnmock("@/lib/serverRootPath"); + }); + + // Regression: under SERVER_ROOT_PATH=/litellm the logo must be requested at + // /litellm/ui/assets/logos/... A bare /ui/... path is served off the root and + // 404s behind the reverse proxy. + it("prefixes the server root path onto the resolved logo", async () => { + vi.resetModules(); + vi.doMock("@/lib/serverRootPath", () => ({ serverRootPath: "/litellm" })); + const { getProviderLogoAndName } = await import("./provider_info_helpers"); + expect(getProviderLogoAndName("openai").logo).toBe("/litellm/ui/assets/logos/openai_small.svg"); + }); + + it("leaves the logo at /ui/... when mounted at the root", async () => { + vi.resetModules(); + vi.doMock("@/lib/serverRootPath", () => ({ serverRootPath: "/" })); + const { getProviderLogoAndName } = await import("./provider_info_helpers"); + expect(getProviderLogoAndName("openai").logo).toBe("/ui/assets/logos/openai_small.svg"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index a19a6700a59..e7b7cbf29ee 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -1,3 +1,5 @@ +import { resolveLogoSrc } from "@/lib/assetPaths"; + export enum Providers { A2A_Agent = "A2A Agent", AI21 = "Ai21", @@ -317,7 +319,7 @@ export const getProviderLogoAndName = (providerValue: string): { logo: string; d // Handle special case for "gemini" provider value if (providerValue.toLowerCase() === "gemini") { const displayName = Providers.Google_AI_Studio; - const logo = providerLogoMap[displayName]; + const logo = resolveLogoSrc(providerLogoMap[displayName]) ?? ""; return { logo, displayName }; } @@ -334,7 +336,7 @@ export const getProviderLogoAndName = (providerValue: string): { logo: string; d // Get the display name from Providers enum and logo from map const displayName = Providers[enumKey as keyof typeof Providers]; - const logo = providerLogoMap[displayName as keyof typeof providerLogoMap]; + const logo = resolveLogoSrc(providerLogoMap[displayName as keyof typeof providerLogoMap]) ?? ""; return { logo, displayName }; }; diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index ed45d95222d..bff0dbb35b8 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -22,6 +22,7 @@ import React, { useEffect, useState } from "react"; import { Button as Button2, Form, Input, Modal, Select, Typography } from "antd"; import EmailSettings from "./email_settings"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import NotificationsManager from "./molecules/notifications_manager"; const { Title, Paragraph } = Typography; @@ -53,7 +54,7 @@ interface genericCallbackParams { litellm_callback_params: string[] | null; // known required params for this callback } -const assetsLogoFolder = "../ui/assets/logos/"; +const assetsLogoFolder = "/ui/assets/logos/"; interface DynamicParamsFieldsProps { params: string[]; @@ -156,10 +157,11 @@ const CallbackSelector: React.FC = ({ > {callbackConfigs.map((callbackConfig) => { const logo = callbackConfig.logo; - const logoSrc = + const logoSrc = resolveLogoSrc( logo && (logo.includes("/") || logo.startsWith("data:") || logo.startsWith("http")) ? logo - : `${assetsLogoFolder}${logo}`; + : `${assetsLogoFolder}${logo}`, + ); return ( diff --git a/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx index 56d2e23a90c..49602fd5b05 100644 --- a/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx +++ b/ui/litellm-dashboard/src/components/team/LoggingSettings.tsx @@ -6,6 +6,7 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, Card, TextInput } from "@tremor/react"; import { PlusIcon, TrashIcon, CogIcon, BanIcon } from "@heroicons/react/outline"; import { callbackInfo, callback_map, mapDisplayToInternalNames } from "../callback_info_helpers"; +import { resolveLogoSrc } from "@/lib/assetPaths"; import NumericalInput from "../shared/numerical_input"; const { Option } = Select; @@ -178,7 +179,7 @@ const LoggingSettings: React.FC = ({ optionLabelProp="label" > {allCallbacks.map((callbackName) => { - const logo = callbackInfo[callbackName]?.logo; + const logo = resolveLogoSrc(callbackInfo[callbackName]?.logo); const description = callbackInfo[callbackName]?.description; return (