diff --git a/.circleci/config.yml b/.circleci/config.yml index 6e368a3debe..aa851f829e4 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2648,6 +2648,19 @@ jobs: name: Start mock LLM server command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true + - run: + name: Start mock Presidio server + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py + background: true + - run: + name: Wait for mock Presidio server + command: | + for i in $(seq 1 30); do + if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi + sleep 1 + done + echo "Mock Presidio server never answered /health on port 8091" >&2 + exit 1 - run: name: Start LiteLLM proxy environment: @@ -2778,6 +2791,19 @@ jobs: name: Start mock LLM server command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py background: true + - run: + name: Start mock Presidio server + command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py + background: true + - run: + name: Wait for mock Presidio server + command: | + for i in $(seq 1 30); do + if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi + sleep 1 + done + echo "Mock Presidio server never answered /health on port 8091" >&2 + exit 1 - run: name: Start LiteLLM proxy under a server root path environment: diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml index d8256917805..9a362b4c85c 100644 --- a/.github/workflows/auto-close-duplicates.yml +++ b/.github/workflows/auto-close-duplicates.yml @@ -22,7 +22,7 @@ on: permissions: {} jobs: - test: + sweep-tests: if: github.event_name == 'pull_request' runs-on: ubuntu-latest timeout-minutes: 5 diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml new file mode 100644 index 00000000000..61a56f1f47e --- /dev/null +++ b/.github/workflows/cost-map-guard.yml @@ -0,0 +1,45 @@ +name: Cost map guard + +on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed + pull_request_target: + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + cost-map-guard: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + - name: Fetch the pull request head and its merge base + id: revisions + env: + GH_TOKEN: ${{ github.token }} + BASE_SHA: ${{ github.event.pull_request.base.sha }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')" + git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" + echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT" + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + - name: Run the guard + env: + MERGE_BASE: ${{ steps.revisions.outputs.merge_base }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + HEAD_REF: ${{ github.event.pull_request.head.ref }} + run: | + uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF" diff --git a/.github/workflows/publish-basedpyright-base-counts.yml b/.github/workflows/publish-basedpyright-base-counts.yml index cd443a8e9db..27d4682dbd9 100644 --- a/.github/workflows/publish-basedpyright-base-counts.yml +++ b/.github/workflows/publish-basedpyright-base-counts.yml @@ -1,6 +1,6 @@ name: Publish basedpyright base counts -# Every commit on litellm_internal_staging is some branch's future merge-base. +# Every commit on main or litellm_internal_staging can become a future merge-base. # Publishing its per-rule basedpyright counts as an artifact lets # scripts/type_check_gate.py download them in seconds instead of paying a # 60-110s second basedpyright pass on every fresh worktree or moved merge-base. @@ -10,13 +10,13 @@ name: Publish basedpyright base counts on: push: branches: + - main - litellm_internal_staging workflow_dispatch: inputs: ref: - description: "Ref to compute and publish base counts for" + description: "Ref to compute and publish base counts for (defaults to the workflow run's commit)" required: false - default: litellm_internal_staging permissions: contents: read diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml deleted file mode 100644 index 74e6be69604..00000000000 --- a/.github/workflows/report-rust-release-wheel.yml +++ /dev/null @@ -1,129 +0,0 @@ -name: Report LiteLLM Rust release wheel - -on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs - workflow_run: - workflows: - - LiteLLM Rust - types: - - completed - -permissions: {} - -concurrency: - group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }} - cancel-in-progress: false - -jobs: - report-release-wheel: - name: report release wheel - if: >- - github.event.workflow_run.event == 'pull_request' && - github.event.workflow_run.path == '.github/workflows/test-rust.yml' && - github.event.workflow_run.head_repository.full_name == github.repository && - github.event.workflow_run.pull_requests[0].number != null - runs-on: ubuntu-latest - timeout-minutes: 5 - permissions: - pull-requests: write - - steps: - - name: Link release wheel report on PR - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 - env: - COMMENT_MARKER: "" - with: - script: | - const marker = process.env.COMMENT_MARKER; - const workflowRun = context.payload.workflow_run; - const allowedConclusions = new Set([ - "action_required", - "cancelled", - "failure", - "neutral", - "skipped", - "stale", - "startup_failure", - "success", - "timed_out", - ]); - if ( - !allowedConclusions.has(workflowRun.conclusion) || - workflowRun.event !== "pull_request" || - workflowRun.path !== ".github/workflows/test-rust.yml" || - workflowRun.head_repository?.full_name !== - `${context.repo.owner}/${context.repo.repo}` || - workflowRun.pull_requests?.length !== 1 - ) { - throw new Error("unexpected source workflow"); - } - const pullRequest = workflowRun.pull_requests[0]; - const pullRequestNumber = pullRequest.number; - const headSha = workflowRun.head_sha; - const runId = workflowRun.id; - if ( - !Number.isSafeInteger(pullRequestNumber) || - pullRequestNumber <= 0 || - !Number.isSafeInteger(runId) || - runId <= 0 || - !/^[0-9a-f]{40}$/.test(headSha) || - pullRequest.head?.sha !== headSha - ) { - throw new Error("invalid source workflow metadata"); - } - const runUrl = - `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + - `/actions/runs/${runId}`; - const result = - workflowRun.conclusion === "success" - ? "successfully" - : `with \`${workflowRun.conclusion}\``; - const body = [ - marker, - "## LiteLLM Rust workflow", - "", - `Workflow completed ${result} for \`${headSha}\``, - "", - `[View workflow run](${runUrl})`, - ].join("\n"); - const comments = await github.paginate(github.rest.issues.listComments, { - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: pullRequestNumber, - per_page: 100, - }); - const existing = comments.find( - (comment) => - comment.user?.login === "github-actions[bot]" && - comment.body?.startsWith(marker), - ); - const currentPullRequest = ( - await github.rest.pulls.get({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: pullRequestNumber, - }) - ).data; - if ( - currentPullRequest.state !== "open" || - currentPullRequest.head.repo?.full_name !== - `${context.repo.owner}/${context.repo.repo}` || - currentPullRequest.head.sha !== headSha - ) { - core.info("source workflow no longer matches the current pull request head"); - return; - } - if (existing) { - await github.rest.issues.updateComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: existing.id, - body, - }); - } else { - await github.rest.issues.createComment({ - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: pullRequestNumber, - body, - }); - } diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml index 1daaadeabe2..18d5e3de1eb 100644 --- a/.github/workflows/sync-together-ai-models.yml +++ b/.github/workflows/sync-together-ai-models.yml @@ -13,10 +13,12 @@ jobs: sync_together_ai_models: if: github.repository == 'BerriAI/litellm' runs-on: ubuntu-latest + env: + BASE_BRANCH: ${{ github.event.repository.default_branch }} steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: - ref: litellm_internal_staging + ref: ${{ env.BASE_BRANCH }} persist-credentials: false - name: Set up uv uses: ./.github/actions/setup-uv-with-retries @@ -63,6 +65,6 @@ jobs: gh pr create --title "feat(models): sync together_ai model registry" \ --body-file "$RUNNER_TEMP/pr_body.md" \ --head "$branch" \ - --base litellm_internal_staging + --base "$BASE_BRANCH" env: GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 36fd656c231..e6d2264fbf0 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -74,6 +74,12 @@ jobs: - name: check_workflow_startup_safety run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py + - name: check_workflow_job_name_collisions + run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py + + - name: test_workflow_job_name_collisions + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py + - name: test_e2e_changed_gate run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py diff --git a/.github/workflows/test-litellm-ui-unit.yml b/.github/workflows/test-litellm-ui-unit.yml index 314efcc49d5..cd58f861a87 100644 --- a/.github/workflows/test-litellm-ui-unit.yml +++ b/.github/workflows/test-litellm-ui-unit.yml @@ -12,6 +12,7 @@ on: - "litellm_**" push: branches: + - main - litellm_internal_staging concurrency: diff --git a/.github/workflows/test-model-map.yml b/.github/workflows/test-model-map.yml deleted file mode 100644 index c2770e5da4c..00000000000 --- a/.github/workflows/test-model-map.yml +++ /dev/null @@ -1,37 +0,0 @@ -name: Validate model_prices_and_context_window.json - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - litellm_oss_staging - - "litellm_**" - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} - cancel-in-progress: ${{ github.event_name == 'pull_request' }} - -jobs: - validate-model-prices-json: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Validate model_prices_and_context_window.json - run: | - jq empty model_prices_and_context_window.json - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Check model_prices_and_context_window.schema.json is in sync - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py --check diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 9b8b132df62..4f56e78ddee 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -7,6 +7,7 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" + - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" @@ -22,6 +23,7 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" + - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" @@ -34,102 +36,89 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +env: + CARGO_TERM_COLOR: always + jobs: - rust-checks: - name: rustfmt, clippy, test + rust-lint: runs-on: ubuntu-latest timeout-minutes: 10 defaults: run: working-directory: litellm-rust - env: - CARGO_TERM_COLOR: always - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false - - name: Set up Rust - run: rustup toolchain install + - run: rustup toolchain install --no-self-update - - name: Cache Cargo registry and target - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - run: cargo fmt --check + + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: | ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-cargo- + ${{ runner.os }}-cargo-${{ github.job }}- - - name: Check Rust formatting - run: cargo fmt --check + - run: cargo clippy --workspace --all-targets --locked -- -D warnings - - name: Run Clippy - run: cargo clippy --workspace --all-targets --locked -- -D warnings + - run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings - - name: Run Clippy with Bedrock auth - run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings + - run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings - - name: Run Clippy with all gateway features - run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings - - - name: Run Rust tests - run: cargo test --workspace --locked - - - name: Run core tests with Bedrock auth - run: cargo test -p litellm-core --features bedrock-auth --locked - - # Not --all-features: python-config links libpython, which this job does not install. - - name: Run gateway tests with the server feature - run: cargo test -p litellm-ai-gateway --features server --locked - - release-wheel: - name: release wheel + rust-test: runs-on: ubuntu-latest - timeout-minutes: 20 - permissions: - contents: read - env: - CARGO_TERM_COLOR: always + timeout-minutes: 30 steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries + - uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - - name: Set up Rust - run: rustup toolchain install + - run: rustup toolchain install --no-self-update - - name: Build release wheel - run: uv build --wheel --out-dir dist + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-${{ github.job }}- - - name: Build panic contract wheel - run: >- + - run: cargo test --workspace --locked + working-directory: litellm-rust + + - run: cargo test -p litellm-core --features bedrock-auth --locked + working-directory: litellm-rust + + - run: cargo test -p litellm-ai-gateway --features server --locked + working-directory: litellm-rust + + - run: uv build --wheel --out-dir dist + + - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + + - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + + - run: >- uv build --wheel --out-dir panic-dist --config-setting "maturin.build-args=--features panic-test,extension-module" - - name: Smoke-test native panic unwinding - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - - name: Verify stripped native extension - env: - RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - - - name: Test native route wheel - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 33245ec5b5f..cc606339a20 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -116,6 +116,7 @@ jobs: tests/test_litellm/rerank_api tests/test_litellm/rust_bridge tests/test_litellm/sandbox + tests/test_litellm/skills tests/test_litellm/test_router tests/test_litellm/vector_stores tests/test_litellm/videos diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3d2fa3e51c8..b04e004aa1a 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -149,7 +149,7 @@ graph TD | `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user | | `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation | | `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation | -| `litellm_skills` | `proxy/hooks/skills_injection.py` | Skills injection | +| `litellm_skills` | `proxy/hooks/litellm_skills/main.py` | Skills injection | To add a new proxy hook, implement `CustomLogger` and register in `PROXY_HOOKS`. @@ -220,20 +220,20 @@ graph LR | Job | Interval | Purpose | Key Files | |-----|----------|---------|-----------| | `update_spend` | 60s | Batch write spend logs to PostgreSQL | `proxy/db/db_spend_update_writer.py` | -| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/management_helpers/budget_reset_job.py` | +| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/common_utils/reset_budget_job.py` | | `add_deployment` | 10s | Sync new model deployments from DB | `proxy/proxy_server.py` (`ProxyConfig`) | -| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/management_helpers/spend_log_cleanup.py` | -| `check_batch_cost` | 30min | Calculate costs for batch jobs | `proxy/management_helpers/check_batch_cost_job.py` | -| `check_responses_cost` | 30min | Calculate costs for responses API | `proxy/management_helpers/check_responses_cost_job.py` | -| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/management_helpers/key_rotation_manager.py` | +| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/db/db_transaction_queue/spend_log_cleanup.py` | +| `check_batch_cost` | 30min | Calculate costs for batch jobs | `enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py` | +| `check_responses_cost` | 30min | Calculate costs for responses API | `enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py` | +| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/common_utils/key_rotation_manager.py` | | `_run_background_health_check` | continuous | Health check model deployments | `proxy/proxy_server.py` | | `send_weekly_spend_report` | weekly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) | | `send_monthly_spend_report` | monthly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) | **Cost Attribution Flow:** 1. LLM response returns to `utils.py` wrapper after `litellm.acompletion()` completes -2. `update_response_metadata()` (`llm_response_utils/response_metadata.py`) is called -3. `logging_obj._response_cost_calculator()` (`litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`) +2. `update_response_metadata()` (`litellm_core_utils/llm_response_utils/response_metadata.py`) is called +3. `logging_obj._response_cost_calculator()` (`litellm_core_utils/litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`) 4. Cost is stored in `response._hidden_params["response_cost"]` 5. `proxy/common_request_processing.py` extracts cost from `hidden_params` and adds to response headers (`x-litellm-response-cost`) 6. `logging_obj.async_success_handler()` triggers callbacks including `_ProxyDBLogger.async_log_success_event()` diff --git a/CLAUDE.md b/CLAUDE.md index d9e9e8f1586..2bc39332817 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank -Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR +Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis @@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a Python max line length is 120, not 88 -When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing +Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on `litellm_internal_staging` in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. If your branch already carries a budget edit, drop it before opening the PR `make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..57ca267e504 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1808 + "limit": 1804 }, "reportRedeclaration": { "limit": 8 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 138 + "limit": 136 }, "reportUnusedImport": { "limit": 542 diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py new file mode 100644 index 00000000000..50aa40ba220 --- /dev/null +++ b/ci_cd/cost_map_guard.py @@ -0,0 +1,146 @@ +"""Guard the cost map on pull requests. + +Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file, +and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named +litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final + +from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors + +COST_MAP_PATH: Final = "model_prices_and_context_window.json" +BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" +SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" +GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH) +BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" + +CostMap = dict[str, object] + + +@dataclass(frozen=True, slots=True) +class Snapshot: + cost_map: str + backup: str + schema: str + + +def _parse_object(text: str, path: str) -> CostMap | str: + try: + parsed: Final = json.loads(text) + except json.JSONDecodeError as error: + return f"{path} is not valid JSON: {error}" + return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root" + + +def _rendered_schema(cost_map: CostMap) -> str: + try: + return render(build_schema(cost_map)) + except SystemExit as error: + return str(error) + + +def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]: + schema_text: Final = _rendered_schema(head_map) + if not schema_text.startswith("{"): + return (schema_text,) + backup_failure: Final = ( + () + if head.backup == head.cost_map + else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",) + ) + schema_failure: Final = ( + () + if head.schema == schema_text + else ( + f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; " + "run `python ci_cd/generate_model_prices_schema.py` and commit the result", + ) + ) + return ( + *backup_failure, + *schema_failure, + *( + f"{COST_MAP_PATH} does not validate against its schema: {error}" + for error in validation_errors(head_map, json.loads(schema_text))[:20] + ), + ) + + +def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]: + return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)} + + +def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]: + base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH) + if isinstance(base_map, str): + return (f"merge base: {base_map}",) + base_entries: Final = _entries(base_map) + head_entries: Final = _entries(head_map) + removed_fields: Final = tuple( + f"{key}.{field}" + for key, entry in base_entries.items() + if key in head_entries + for field in entry + if field not in head_entries[key] + ) + return ( + *( + f"bot PRs may only change the cost map files, not {path}" + for path in changed_files + if path not in GUARDED_PATHS + ), + *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map), + *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields), + *( + f"bot PRs may not change {key}" + for key in sorted(SPECIAL_ROOT_KEYS) + if base_map.get(key) != head_map.get(key) + ), + ) + + +def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]: + head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH) + if isinstance(head_map, str): + return (head_map,) + return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ())) + + +def _git(*args: str) -> str: + result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True) + return result.stdout if result.returncode == 0 else "" + + +def snapshot(revision: str) -> Snapshot: + return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS)) + + +def main(argv: Sequence[str]) -> int: + parser: Final = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", required=True, help="merge base of the pull request") + parser.add_argument("--head", required=True, help="head commit of the pull request") + parser.add_argument("--head-ref", required=True, help="head branch name of the pull request") + args: Final = parser.parse_args(argv) + bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX) + changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines()) + failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot) + contract: Final = "bot contract enforced" if bot else "human PR, file checks only" + if failures: + print(f"cost map guard failed ({contract}):") + print("\n".join(f"- {failure}" for failure in failures)) + return 1 + print(f"cost map guard passed ({contract})") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 62e943d0f42..43b9ec1aac2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1415,6 +1415,8 @@ dependencies = [ "litellm-config", "litellm-core", "reqwest", + "rustls 0.23.42", + "rustls-native-certs", "serde", "serde_json", "sha2 0.10.9", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 720c4545181..82de7f40069 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -28,6 +28,8 @@ pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" +rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } sha2 = "0.10" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 10369fa3bfd..74cf66e88a2 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -20,6 +20,10 @@ litellm-config.workspace = true # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true +# rustls and its root store are direct dependencies so `io::tls` can build the +# one TLS config the outbound dials use; see that module for why it has to. +rustls.workspace = true +rustls-native-certs.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 diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index cce56dd2121..7098d67993f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -3,3 +3,4 @@ pub mod ocr; pub mod realtime; pub mod realtime_pool; pub mod responses_ws; +pub(crate) mod tls; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 662f7328982..207c31dffa0 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; +use crate::io::tls::connect_upstream; + /// Environment variable holding the OpenAI API key (last-resort fallback). const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; @@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream( .map_err(|err| Error::Auth(err.to_string()))?, ); - let (upstream, _response) = connect_async(request) + let (upstream, _response) = connect_upstream(request) .await .map_err(|err| Error::Network(err.to_string()))?; Ok(upstream) @@ -284,6 +286,33 @@ mod tests { serde_json::from_str(raw).expect("valid event json") } + /// The realtime dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = dial_upstream( + "gpt-realtime", + "sk-test", + Some(&format!("wss://127.0.0.1:{port}")), + ) + .await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + #[test] fn resolve_api_key_prefers_param_then_blank_falls_through() { assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test"); diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 0b01747b1a5..9df3d0c6cc5 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use crate::io::tls::connect_upstream; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, @@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection { .map_err(|error| Error::InvalidRequest(error.to_string()))?; request.headers_mut().insert(header_name, header_value); } - let connect = connect_async(request); + let connect = connect_upstream(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { Error::Network("Responses WebSocket connection timed out".to_string()) })?, None => connect.await, }; - let (socket, _) = result.map_err(|error| match error { + let (socket, _) = result.map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), @@ -138,13 +140,13 @@ async fn dial_upstream( ); let result = tokio::time::timeout( Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS), - connect_async(request), + connect_upstream(request), ) .await .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?; result .map(|(socket, _)| socket) - .map_err(|error| match error { + .map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), @@ -324,6 +326,29 @@ mod tests { use tokio::net::TcpListener; use tokio_tungstenite::accept_async; + /// The Responses dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = + dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); let address = listener.local_addr().expect("local address"); diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs new file mode 100644 index 00000000000..a2562f60345 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -0,0 +1,80 @@ +//! Outbound WebSocket dials over a TLS config this crate builds once and owns. +//! +//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth` +//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that +//! `tokio-tungstenite` uses when handed no connector panics rather than guess +//! between them. Naming ring on a connector of our own settles that for these +//! dials without touching the process-wide default, and building the config +//! once keeps the platform trust store, which `tokio-tungstenite` would +//! otherwise re-read on every dial, off the dial path. + +use std::io; +use std::sync::{Arc, OnceLock}; + +use rustls::{ClientConfig, RootCertStore}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::Error; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::TlsError; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, +}; + +static TLS_CONFIG: OnceLock> = OnceLock::new(); + +fn build_config() -> Result> { + let native = rustls_native_certs::load_native_certs(); + let roots = { + let mut store = RootCertStore::empty(); + let (added, _ignored) = store.add_parsable_certificates(native.certs); + if added == 0 { + return Err(Box::new(Error::Io(io::Error::other(format!( + "no usable native root certificates: {:?}", + native.errors + ))))); + } + store + }; + + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map(|builder| builder.with_root_certificates(roots).with_no_client_auth()) + .map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error)))) +} + +fn tls_config() -> Result, Box> { + if let Some(config) = TLS_CONFIG.get() { + return Ok(Arc::clone(config)); + } + let built = Arc::new(build_config()?); + Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) +} + +pub(crate) async fn connect_upstream( + request: R, +) -> Result<(WebSocketStream>, Response), Box> +where + R: IntoClientRequest + Unpin, +{ + let request = request.into_client_request().map_err(Box::new)?; + let connector = match request.uri().scheme_str() { + Some("wss") => Some(Connector::Rustls(tls_config()?)), + _ => None, + }; + connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(Box::new) +} + +#[cfg(test)] +mod tests { + use super::build_config; + + #[test] + fn builds_a_usable_config_with_both_provider_features_enabled() { + let config = build_config().expect("a client config"); + + assert!(!config.crypto_provider().cipher_suites.is_empty()); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 446b323db3a..ed41f1ff9e7 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -265,6 +265,12 @@ impl CallLifecycleHooks for OcrLi Box::pin(async move { Ok(request) }) } + #[tracing::instrument( + name = "success_callback", + target = "litellm::function_trace", + level = "trace", + skip_all + )] fn async_log_success_event<'a>( &'a self, context: &'a CallLifecycleContext, @@ -288,6 +294,12 @@ impl CallLifecycleHooks for OcrLi }) } + #[tracing::instrument( + name = "failure_callback", + target = "litellm::function_trace", + level = "trace", + skip_all + )] fn async_log_failure_event<'a>( &'a self, context: &'a CallLifecycleContext, diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs new file mode 100644 index 00000000000..05f7d9610d5 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -0,0 +1,48 @@ +//! Guards the wiring, not just the helper: a `wss://` dial through the public +//! API has to resolve its own crypto provider, in a test binary where nothing +//! has installed a process-wide one, and has to leave it uninstalled. + +use std::collections::HashMap; +use std::time::Duration; + +use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; +use tokio::net::TcpListener; + +async fn dead_tls_server() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + port +} + +#[tokio::test] +async fn dialing_wss_returns_an_error_instead_of_panicking() { + let port = dead_tls_server().await; + + let result = ResponsesWebSocketConnection::connect_url( + &format!("wss://127.0.0.1:{port}/"), + &HashMap::new(), + Some(Duration::from_secs(10)), + ) + .await; + + assert!( + result.is_err(), + "a plain TCP server cannot finish a TLS handshake" + ); + assert!( + rustls::crypto::CryptoProvider::get_default().is_none(), + "the dial settles its provider on its own connector, not process-wide" + ); +} diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index c3a89f4394d..60e90ed2a7c 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -11,9 +11,13 @@ use litellm_ai_gateway::integrations::custom_logger::{ use litellm_ai_gateway::integrations::types::RequestMetadata; use litellm_ai_gateway::ocr::{OcrRequest, ocr}; use litellm_core::error::Error; +#[cfg(feature = "trace-parity")] +use litellm_core::observability::FunctionTrace; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; +#[cfg(feature = "trace-parity")] +use tracing::instrument::WithSubscriber; async fn read_http_headers(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -320,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall, ])); - let response = ocr(OcrRequest { + #[cfg(feature = "trace-parity")] + let trace = FunctionTrace::default(); + let api_base = format!("http://{addr}"); + let call = 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}")), + api_base: Some(&api_base), custom_llm_provider: Some("mistral"), extra_headers: None, optional_params: Map::new(), @@ -339,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { ..Default::default() }, litellm_call_id: Some("ocr-call-1"), - }) - .await - .expect("ocr request succeeds"); + }); + #[cfg(feature = "trace-parity")] + let call = call.with_subscriber(trace.dispatcher()); + let response = call.await.expect("ocr request succeeds"); assert_eq!(response["pages"][0]["markdown"], "ok"); assert_eq!( @@ -359,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { error_kind: None, }] ); + #[cfg(feature = "trace-parity")] + assert_eq!( + trace + .events() + .iter() + .filter(|event| event.function.ends_with("_callback")) + .map(|event| event.function) + .collect::>(), + vec!["success_callback"] + ); let request = server.await.expect("server task completes"); assert!(request.contains(r#""guarded_pre":true"#), "{request}"); @@ -388,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { }); let logger = Arc::new(RecordingOcrLogger::default()); - let err = ocr(OcrRequest { + #[cfg(feature = "trace-parity")] + let trace = FunctionTrace::default(); + let api_base = format!("http://{addr}"); + let call = 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}")), + api_base: Some(&api_base), custom_llm_provider: Some("mistral"), extra_headers: None, optional_params: Map::new(), @@ -404,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { guardrails: Vec::new(), request_metadata: RequestMetadata::default(), litellm_call_id: Some("ocr-call-2"), - }) - .await - .expect_err("provider error propagates"); + }); + #[cfg(feature = "trace-parity")] + let call = call.with_subscriber(trace.dispatcher()); + let err = call.await.expect_err("provider error propagates"); assert!(matches!(err, Error::Http { status: 500, .. })); server.await.expect("server task completes"); @@ -421,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { error_kind: Some("HttpError".to_string()), }] ); + #[cfg(feature = "trace-parity")] + assert_eq!( + trace + .events() + .iter() + .filter(|event| event.function.ends_with("_callback")) + .map(|event| event.function) + .collect::>(), + vec!["failure_callback"] + ); } #[tokio::test] diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs index ea7d9f4993e..bc3c962f7a3 100644 --- a/litellm-rust/crates/python-bridge/src/function_trace.rs +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -1,3 +1,4 @@ +use std::fmt::Display; use std::future::Future; use litellm_core::observability::{FunctionTrace, FunctionTraceEvent}; @@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber; #[derive(Serialize)] pub(crate) struct TracedResponse { - response: T, + #[serde(skip_serializing_if = "Option::is_none")] + response: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, trace: Vec, } pub(crate) async fn capture( future: impl Future>, -) -> Result, E> { +) -> Result, E> +where + E: Display, +{ let trace = FunctionTrace::default(); - let response = future.with_subscriber(trace.dispatcher()).await?; - Ok(TracedResponse { - response, - trace: trace.events(), + let result = future.with_subscriber(trace.dispatcher()).await; + let events = trace.events(); + Ok(match result { + Ok(response) => TracedResponse { + response: Some(response), + error: None, + trace: events, + }, + Err(error) => TracedResponse { + response: None, + error: Some(error.to_string()), + trace: events, + }, }) } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 3285da14d5f..bc51647cbad 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -486,10 +486,11 @@ asyncio.run(exercise()) let code = CString::new( r#" result = routes.echo("traced") -assert result == { - "response": "traced", - "trace": [{"function": "execute_echo", "depth": 0}], -} +assert result["response"] == "traced", result +assert [event["function"] for event in result["trace"]] == ["execute_echo"], result +failure = routes.echo("error") +assert failure["error"] == "invalid request: synthetic error", failure +assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure "#, ) .expect("Python source should not contain null bytes"); diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..dc2f40af46e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -495,6 +495,7 @@ public_model_groups: Optional[List[str]] = None public_agent_groups: Optional[List[str]] = None agent_search_embedding_model: Optional[str] = None mcp_tool_search: Optional[Mapping[str, object]] = None +skill_search_embedding_model: Optional[str] = None # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 97be5f77d79..959c7498479 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -25,6 +25,27 @@ class BatchCostUsageResult: failed_requests: int +_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) +_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) + + +def batch_cost_is_final(batch: Batch) -> bool: + """Whether this retrieve of the batch is the one to account its cost from. + + A batch still in flight has nothing to price, and a "completed" batch can report + no output_file_id for a moment before the output populates; pricing either records + $0 under the batch's single spend row and pins it there. Final means a completed + batch whose output file has arrived or whose counts prove no line succeeded, or + any other terminal status (failed, cancelled, expired). + """ + if batch.status not in _TERMINAL_BATCH_STATUSES: + return False + if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: + return True + request_counts: Final = batch.request_counts + return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 + + async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index d70f947469a..87350b5479c 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args +from openai.types.chat import ChatCompletion +from openai.types.responses import Response from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( FunctionCallOutput, @@ -33,7 +36,7 @@ from litellm.responses.sse_output_recovery import ( record_output_item_chunk, record_output_text_chunk, ) -from litellm.responses.utils import normalize_responses_api_stream_options +from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options from litellm.types.llms.openai import ( REASONING_EFFORT, ChatCompletionAnnotation, @@ -43,6 +46,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, Reasoning, ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, ResponsesAPIStreamEvents, ) from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -54,7 +58,7 @@ if TYPE_CHECKING: ) from pydantic import BaseModel - from litellm import LiteLLMLoggingObj, ModelResponse + from litellm import LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import ( ALL_RESPONSES_API_TOOL_PARAMS, @@ -69,6 +73,28 @@ if TYPE_CHECKING: from litellm.types.utils import Choices +_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) +_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset( + ChatCompletion.model_fields +) + + +def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]: + return MappingProxyType( + { + key: value + for key, value in (response_fields.items() if response_fields else ()) + if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS + } + ) + + +def _upstream_response_id(response_id: str | None) -> str | None: + if response_id is None: + return None + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id) + + class _ReasoningSummaryText(TypedDict): type: str text: str @@ -904,6 +930,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage), ) + model_response.id = _upstream_response_id(raw_response.id) or raw_response.id + for key, value in _provider_metadata(raw_response.model_extra).items(): + setattr(model_response, key, value) + # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {}) @@ -1359,14 +1389,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if event_type == "response.created": # Initial response creation event verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk) + created_response: Final = parsed_chunk.get("response") return ModelResponseStream( + id=_upstream_response_id(created_response.get("id")) if created_response else None, choices=[ StreamingChoices( index=0, delta=Delta(content=""), finish_reason=None, ) - ] + ], ) elif event_type == "response.output_item.added": # New output item added @@ -1534,6 +1566,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): from litellm.responses.utils import ResponseAPILoggingUtils usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage")) + provider_metadata: Final = _provider_metadata(response_data) return ModelResponseStream( choices=[ StreamingChoices( @@ -1546,6 +1579,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ], usage=usage, + provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict ) else: pass diff --git a/litellm/constants.py b/litellm/constants.py index ce744e9c58a..d53686e5e5b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1901,6 +1901,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset( } ) +PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = ( + "x-amzn-requestid", + "x-request-id", + "request-id", + "x-ms-request-id", + "apim-request-id", + "x-goog-request-id", + "cf-ray", +) + # Browser-facing security headers that a malicious or misconfigured upstream # provider must not be able to set on the proxy's own response. BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( diff --git a/litellm/litellm_core_utils/api_route_to_call_types.py b/litellm/litellm_core_utils/api_route_to_call_types.py index e3562095d7f..428d7563d4a 100644 --- a/litellm/litellm_core_utils/api_route_to_call_types.py +++ b/litellm/litellm_core_utils/api_route_to_call_types.py @@ -14,21 +14,48 @@ from typing import Final from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes +def _segment_matches(route_segment: str, pattern_segment: str) -> bool: + """ + Match one concrete path segment against one pattern segment. + A bare placeholder ({param}) matches any segment; a placeholder with a + literal suffix ({model}:generateContent) requires the segment to end with + that suffix and have a non-empty value before it. + """ + if not pattern_segment.startswith("{"): + return route_segment == pattern_segment + placeholder_end: Final = pattern_segment.find("}") + if placeholder_end == -1: + return route_segment == pattern_segment + literal_suffix: Final = pattern_segment[placeholder_end + 1 :] + if not literal_suffix: + return True + return route_segment.endswith(literal_suffix) and len(route_segment) > len(literal_suffix) + + +def _pattern_tail_spans_segments(pattern_tail: str) -> bool: + """ + Whether the pattern's last segment is a suffixed placeholder + ({model}:generateContent) that may absorb extra route segments, mirroring + FastAPI's {model_name:path} converter for slash-containing model names. + """ + return pattern_tail.startswith("{") and "}" in pattern_tail and not pattern_tail.endswith("}") + + def _route_matches_pattern(route: str, pattern: str) -> bool: """ Return True if the concrete route matches the pattern. - Pattern segments like {param} match any single path segment. + Pattern segments like {param} match any single path segment, and a + suffixed placeholder in the last segment may span multiple segments. """ route_parts: Final = route.strip("/").split("/") pattern_parts: Final = pattern.strip("/").split("/") - if len(route_parts) != len(pattern_parts): + if len(route_parts) < len(pattern_parts): return False - for r, p in zip(route_parts, pattern_parts): - if p.startswith("{") and p.endswith("}"): - continue - if r != p: - return False - return True + if len(route_parts) > len(pattern_parts) and not _pattern_tail_spans_segments(pattern_parts[-1]): + return False + head_count: Final = len(pattern_parts) - 1 + merged_parts: Final = (*route_parts[:head_count], "/".join(route_parts[head_count:])) + return all(_segment_matches(r, p) for r, p in zip(merged_parts, pattern_parts)) def get_call_types_for_route(route: str) -> Sequence[CallTypes] | None: diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 389e6f7f501..1fd79db15a6 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -21,13 +21,10 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( } ) -# The per-deployment Rust opt-in. -RUST_KWARG_KEY: Final = "rust" - # Keys `completion()` forwards from its own kwargs into `get_litellm_params`, # which are otherwise invisible to it because that call site passes explicit # named arguments rather than `**kwargs`. -FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY}) +FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls @@ -58,10 +55,6 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", - # The per-deployment Rust opt-in. `all_litellm_params` keeps it out - # of the provider body; this keeps it *in* litellm_params, which is - # where the chat completions handlers read it from. - RUST_KWARG_KEY, } ) | AWS_CREDENTIAL_KWARGS_KEYS diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 9cba5db8ab7..ba8738c8de0 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -53,12 +53,14 @@ class GetModelCostMap: _backup_model_count: int = -1 # -1 = not yet loaded + @staticmethod + def read_local_model_cost_map_text() -> str: + return files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8") + @staticmethod def load_local_model_cost_map() -> dict: """Load the local backup model cost map bundled with the package.""" - content: Final = json.loads( - files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8") - ) + content: Final = json.loads(GetModelCostMap.read_local_model_cost_map_text()) return content @classmethod diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 989777e9412..ca2cca5360f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -36,12 +36,13 @@ from litellm._logging import ( verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import _handle_completed_batch +from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final from litellm.caching.caching import DualCache, InMemoryCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + PROVIDER_REQUEST_ID_HEADERS, SENTRY_DENYLIST, SENTRY_PII_DENYLIST, ) @@ -255,6 +256,30 @@ _in_memory_loggers: Final[list[CustomLogger]] = [] _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys()) + +def _get_provider_request_id(original_exception: Exception) -> str | None: + try: + error_response: Final = getattr(original_exception, "response", None) + header_sources: Final = ( + _get_response_headers(original_exception), + getattr(error_response, "headers", None), + getattr(original_exception, "litellm_response_headers", None), + ) + return next( + ( + str(value) + for expected_header_name in PROVIDER_REQUEST_ID_HEADERS + for headers in header_sources + if isinstance(headers, Mapping) + for header_name, value in headers.items() + if isinstance(header_name, str) and header_name.lower() == expected_header_name and value + ), + None, + ) + except Exception: + return None + + ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys @@ -2899,13 +2924,6 @@ class Logging(LiteLLMLoggingBaseClass): ): # polling job will query these frequently, don't spam db logs return - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - - # check if file id is a unified file id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(result.id) - batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) @@ -2913,9 +2931,7 @@ class Logging(LiteLLMLoggingBaseClass): batch_failed_requests: Final = kwargs.get("batch_failed_requests", None) has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models)) - should_compute_batch_data: Final = ( - not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed" - ) + should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result) if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models @@ -3918,11 +3934,12 @@ class Logging(LiteLLMLoggingBaseClass): LiteLLMResponsesTransformationHandler, ) + served_id: Final = _provider_response_id(result) try: - return LiteLLMResponsesTransformationHandler().transform_response( + translated: Final = LiteLLMResponsesTransformationHandler().transform_response( model=self.model, raw_response=result, - model_response=litellm.ModelResponse(id=_provider_response_id(result)), + model_response=litellm.ModelResponse(id=served_id), logging_obj=self, request_data={}, messages=[], @@ -3930,6 +3947,8 @@ class Logging(LiteLLMLoggingBaseClass): litellm_params={}, encoding=litellm.encoding, ) + translated.id = served_id or translated.id + return translated except Exception as e: verbose_logger.debug( "Responses API -> ModelResponse translation failed for " @@ -3937,7 +3956,7 @@ class Logging(LiteLLMLoggingBaseClass): "usage-only ModelResponse to keep the spend_logs row.", str(e), ) - model_response: Final = litellm.ModelResponse(id=_provider_response_id(result)) + model_response: Final = litellm.ModelResponse(id=served_id) model_response.model = self.model usage: Final = getattr(result, "usage", None) if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage): @@ -5670,6 +5689,7 @@ class StandardLoggingPayloadSetup: rate_limit_category: Final = validate_rate_limit_category(getattr(original_exception, "category", None)) rate_limit_type: Final = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None)) budget_error: Final = original_exception if isinstance(original_exception, BudgetExceededError) else None + provider_request_id: Final = _get_provider_request_id(original_exception) if original_exception else None return StandardLoggingPayloadErrorInformation( error_code=error_status, @@ -5677,6 +5697,7 @@ class StandardLoggingPayloadSetup: llm_provider=_llm_provider_in_exception, traceback=_redact_string(traceback_info), error_message=_redact_string(error_message), + error_provider_request_id=provider_request_id, error_rate_limit_category=rate_limit_category, error_rate_limit_type=rate_limit_type, error_budget_entity_type=budget_error.entity_type if budget_error else None, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8e24302b440..68dc27ec25e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -26,6 +26,7 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, ServiceTier, Usage, + text_tokens_without_nested_reasoning, ) from litellm.utils import get_model_info @@ -513,6 +514,7 @@ def _get_token_base_cost( current_time: datetime | None = None, *, threshold_is_inclusive: bool = False, + missing_cache_read_uses_input: bool = False, ) -> tuple[float, float, float, float, float]: """ Return prompt cost, completion cost, and cache costs for a given model and usage. @@ -523,6 +525,9 @@ def _get_token_base_cost( `threshold_is_inclusive` switches that comparison to >=, for providers such as xAI that bill the higher tier once the prompt reaches the threshold. + `missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved + input rate instead of 0.0; an explicit 0.0 rate stays a real price either way. + Returns: Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) """ @@ -550,29 +555,16 @@ def _get_token_base_cost( float, _get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"), ) - cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key)) + cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None) ## CHECK IF ABOVE THRESHOLD # Optimization: collect threshold keys first to avoid sorting all model_info keys. - # Most models don't have threshold pricing, so we can return early. # Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority) # so that the threshold detection loop only processes standard keys. The # service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key. threshold_keys: Final = [ k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] - if not threshold_keys: - return _apply_off_peak_to_base_costs( - model_info, - current_time, - ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - cache_read_cost, - ), - ) # Only sort the threshold keys (typically 1-2 keys instead of 66+) threshold: float | None = None @@ -661,10 +653,7 @@ def _get_token_base_cost( ), ) - cache_read_cost = cast( - float, - _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost), - ) + cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) break except (IndexError, ValueError): @@ -672,6 +661,17 @@ def _get_token_base_cost( except Exception: continue + if cache_read_cost is None: + cache_read_cost = ( + _off_peak_rate( + _open_off_peak_block(model_info, current_time) or MappingProxyType({}), + "input_cost_per_token", + prompt_base_cost, + ) + if missing_cache_read_uses_input + else 0.0 + ) + return _apply_off_peak_to_base_costs( model_info, current_time, @@ -860,7 +860,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu ) or 0 ) - text_tokens: Final = ( + reported_text_tokens: Final = ( cast( int | None, getattr(usage.completion_tokens_details, "text_tokens", None), @@ -882,6 +882,12 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu or 0 ) video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0)) + text_tokens: Final = text_tokens_without_nested_reasoning( + completion_tokens=usage.completion_tokens, + text_tokens=reported_text_tokens, + reasoning_tokens=reasoning_tokens, + other_modality_tokens=audio_tokens + image_tokens + video_tokens, + ) return CompletionTokensDetailsResult( audio_tokens=audio_tokens, @@ -1409,6 +1415,57 @@ def get_token_type_cost_breakdown( ) +def calculate_prompt_caching_savings( + model_info: ModelInfo, + usage: Usage, + custom_llm_provider: str | None, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + billed_at: datetime | None = None, +) -> float: + """Read discount minus write premium, using the biller's rate and TTL resolution. + + Missing reads and unpublished (missing/zero) writes claim no saving or premium; + explicit zero reads remain free. An unpublished 1h price uses the ordinary write rate. + ``billed_at`` is the request's completion time, so off-peak windows resolve as the + biller saw them rather than at the later spend write. + """ + prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier=service_tier, + current_time=billed_at, + threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), + missing_cache_read_uses_input=True, + ) + write_rate: Final = cache_creation_cost or prompt_base_cost + write_rate_1h: Final = cache_creation_cost_above_1hr or write_rate + prompt_tokens_details: Final = parse_prompt_tokens_details(usage) + cache_read_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"], 0) + cache_creation_tokens: Final = max(prompt_tokens_details["cache_creation_tokens"], 0) + details: Final = prompt_tokens_details["cache_creation_token_details"] + cache_creation_details: Final = ( + CacheCreationTokenDetails( + ephemeral_5m_input_tokens=max(details.ephemeral_5m_input_tokens or 0, 0), + ephemeral_1h_input_tokens=max(details.ephemeral_1h_input_tokens or 0, 0), + ) + if details is not None + else None + ) + read_discount: Final = cache_read_tokens * max(prompt_base_cost - cache_read_cost, 0.0) + write_premium: Final = calculate_cache_writing_cost( + cache_creation_tokens=cache_creation_tokens, + cache_creation_token_details=cache_creation_details, + cache_creation_cost_above_1hr=write_rate_1h - prompt_base_cost, + cache_creation_cost=write_rate - prompt_base_cost, + ) + uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift( + model_info, vertex_location + ) + return (read_discount - write_premium) * uplift + + def calculate_image_response_cost_from_usage( model: str, image_response: ImageResponse, diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 81132fa89a5..24f3b8bca7f 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -2,7 +2,11 @@ Helper functions to handle images passed in messages """ +import asyncio import base64 +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from httpx import Response @@ -11,9 +15,11 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get +from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 +MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) @@ -72,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str: return result +def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) + return litellm.ImageFetchError( + "Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; " + f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}" + ) + + async def async_convert_url_to_base64(url: str) -> str: if url.startswith("data:") and ";base64," in url: return url @@ -93,6 +107,8 @@ async def async_convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise _rejected_image_fetch(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -119,8 +135,192 @@ def convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise _rejected_image_fetch(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( f"Error: Unable to fetch image from URL after 3 attempts. url={url}", ) + + +_REMOTE_URL_PREFIXES: Final = ("http://", "https://") + + +@dataclass(frozen=True, slots=True) +class _RemoteImage: + part: Mapping[str, object] + image_url: Mapping[str, object] | None + url: str + + +@dataclass(frozen=True, slots=True) +class _RemoteFile: + part: Mapping[str, object] + file: Mapping[str, object] + url: str + + +def _as_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one + + +def _remote_url(candidate: object) -> str | None: + return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None + + +_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"}) + + +@dataclass(frozen=True, slots=True) +class _RemoteSource: + part: Mapping[str, object] + source: Mapping[str, object] + url: str + + +@dataclass(frozen=True, slots=True) +class RemoteMedia: + url: str + fields: Mapping[str, object] + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def inline_every_remote_url(_media: RemoteMedia) -> bool: + return True + + +def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: + if fields.get("type") != "image_url": + return None + image_url: Final = fields.get("image_url") + image_url_fields: Final = _as_mapping(image_url) + url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url) + return _RemoteImage(fields, image_url_fields, url) if url is not None else None + + +def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None: + file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None + url: Final = _remote_url(file.get("file_id")) if file is not None else None + return _RemoteFile(fields, file, url) if file is not None and url is not None else None + + +def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None: + source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None + url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None + return _RemoteSource(fields, source, url) if source is not None and url is not None else None + + +def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None: + fields: Final = _as_mapping(part) + if fields is None: + return None + return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields) + + +def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia: + match remote: + case _RemoteImage(_, image_url, url): + return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS) + case _RemoteFile(_, file, url): + return RemoteMedia(url, file) + case _RemoteSource(_, source, url): + return RemoteMedia(url, source) + + +_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) + + +def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]: + return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({}) + + +def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str: + return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part + + +def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]: + kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part + return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part + + +def _base64_source(url: str, data_url: str) -> Mapping[str, str]: + fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1) + media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type + return {"type": "base64", "media_type": media_type, "data": data} # mutable-ok: json-serialized message part + + +def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]: + match remote: + case _RemoteImage(part, image_url, _): + return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part + case _RemoteFile(part, file, url): + return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part + case _RemoteSource(part, _, url): + return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part + + +def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: + content: Final = message.get("content") + return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one + + +def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object: + remote: Final = _parse_remote_part(part) + if remote is None or not should_inline(_remote_media(remote)): + return part + data_url: Final = data_urls.get(remote.url) + return _inline(remote, data_url) if data_url is not None else part + + +def _inline_message( + message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] +) -> AllMessageValues: + parts: Final = _content_parts(message) + if not parts: + return message + inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks + _inline_part(part, data_urls, should_inline) for part in parts + ] + inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message + return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined + + +async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str: + async with in_flight: + return await async_convert_url_to_base64(url) + + +async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]: + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls) + try: + return tuple(await asyncio.gather(*fetches)) + except BaseException: + for fetch in fetches: + fetch.cancel() + await asyncio.gather(*fetches, return_exceptions=True) + raise + + +async def async_inline_remote_media( + messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] + should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url, +) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] + remote_urls: Final = tuple( + dict.fromkeys( + remote.url + for message in messages + for part in _content_parts(message) + if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote)) + ) + ) + if not remote_urls: + return messages + data_urls: Final = await _fetch_data_urls(remote_urls) + inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) + return [ # mutable-ok: transform_request takes a list + _inline_message(message, inlined, should_inline) for message in messages + ] diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 1e43117933d..fa070a648f5 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): check but still resolve DNS and still rewrite HTTP to the resolved IP. """ +import asyncio import socket from ipaddress import ip_address, ip_network from typing import Any, Final, Protocol @@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response kwargs.pop("follow_redirects", None) headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} for _ in range(_MAX_REDIRECTS): - validated_url, original_host = validate_url(url) + validated_url, original_host = await asyncio.to_thread(validate_url, url) response = await fetcher.get( validated_url, headers={**headers_view["headers"], "Host": original_host}, diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index f2d405e9a17..039c462b38a 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig @@ -42,12 +43,37 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) +class AzureAIGPT5Config(OpenAIGPT5Config): + @classmethod + def _model_map_lookup_name(cls, model: str) -> str: + """Normalise a Foundry routing name to its cost-map key, when the map has one. + + A Foundry deployment and its OpenAI-hosted namesake are different products with + different capabilities, so ``azure_ai/`` is the entry to read whenever the map + carries it. Most gpt-5-family names have no ``azure_ai/`` row, though, and prefixing + those anyway costs them every flag: ``get_llm_provider`` re-resolves an ``azure_ai/`` + name to the azure provider when a global AZURE_AI_API_BASE points at an + openai.azure.com host, ``azure/`` is not a key either, so the lookup lands + nowhere and every effort answer degrades to False. A missing key defers to the base + resolver instead. + """ + prefixed: Final = model if model.startswith("azure_ai/") else f"azure_ai/{model}" + return prefixed if prefixed in litellm.model_cost else super()._model_map_lookup_name(model) + + +azureAIGPT5Config: Final = AzureAIGPT5Config() + + class AzureAIStudioConfig(OpenAIConfig): def get_supported_openai_params(self, model: str) -> list: model_supports_tool_choice = True # azure ai supports this by default if not supports_tool_choice(model=f"azure_ai/{model}"): model_supports_tool_choice = False - supported_params = super().get_supported_openai_params(model) + supported_params = ( + azureAIGPT5Config.get_supported_openai_params(model) + if azureAIGPT5Config.is_model_gpt_5_model(model) + else super().get_supported_openai_params(model) + ) if not model_supports_tool_choice: filtered_supported_params: Final = [] for param in supported_params: @@ -61,6 +87,27 @@ class AzureAIStudioConfig(OpenAIConfig): return supported_params + def map_openai_params( + self, + non_default_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature + optional_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: OpenAIConfig.map_openai_params signature + if not azureAIGPT5Config.is_model_gpt_5_model(model): + return super().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + return azureAIGPT5Config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + def _supports_stop_reason(self, model: str) -> bool: """ Check if the model supports stop tokens. diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index bbe1cc85df1..7bfc87a30d6 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -411,6 +411,10 @@ class BaseConfig(ABC): def has_custom_stream_wrapper(self) -> bool: return False + @property + def uses_async_transform_request(self) -> bool: + return False + @property def supports_stream_param_in_request_body(self) -> bool: """ diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 9a25d3294e0..4faf0aaaf30 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -1,5 +1,6 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC): ) -> tuple[dict, RequestFiles]: pass + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image, + image_edit_optional_request_params=dict(image_edit_optional_request_params), + litellm_params=litellm_params, + headers=dict(headers), + ) + def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict: """ Last pass on the request dict after ``transform_image_edit_request``, using the diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 67720451c00..38f280eef03 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( - async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): return _anthropic_request + @property + def uses_async_transform_request(self) -> bool: + return True + async def async_transform_request( self, model: str, @@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: - _anthropic_request: Final = self._build_bedrock_anthropic_request_base( + return self.transform_request( model=model, - messages=messages, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(_anthropic_request) - beta_list: Final = self._compute_bedrock_invoke_beta_headers( - model=model, - messages=messages, - optional_params=optional_params, - headers=headers, - ) - if beta_list: - _anthropic_request["anthropic_beta"] = beta_list - - return _anthropic_request - def _build_bedrock_anthropic_request_base( self, model: str, @@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None: - """ - Async version of document URL conversion for async completion paths. - """ - messages: Final = anthropic_request.get("messages") - if not isinstance(messages, list): - return - - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - - for block in content: - if not isinstance(block, dict) or block.get("type") != "document": - continue - source = block.get("source") - if not isinstance(source, dict) or source.get("type") != "url": - continue - source_url = source.get("url") - if not isinstance(source_url, str): - continue - - inferred_format: str | None = None - if source_url.lower().endswith(".pdf"): - inferred_format = "application/pdf" - base64_url = await async_convert_url_to_base64(url=source_url) - image_chunk = convert_to_anthropic_image_obj( - openai_image_url=base64_url, - format=inferred_format, - ) - block["source"] = { - "type": "base64", - "media_type": image_chunk["media_type"], - "data": image_chunk["data"], - } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: """ Convert tool search entries to the format supported by the Bedrock Invoke API. diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index 31fc079c0c9..7e2037c33f1 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth. from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): litellm_params: dict, headers: dict, ) -> dict: - model_id: Final = model.replace("mantle/", "", 1) - - request: Final = self._build_bedrock_anthropic_request_base( - model=model_id, - messages=messages, + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(request) - return self._restore_mantle_body_fields( - request=request, - model_id=model_id, - optional_params=optional_params, - ) @staticmethod def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict: diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 62b631a7671..ddab4f54d57 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/ import base64 import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -16,7 +17,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.litellm_core_utils.url_utils import safe_get +from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -37,6 +38,22 @@ else: LiteLLMLoggingObj = Any +_BFL_REQUEST_PARAMS: Final = ( + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + "aspect_ratio", + "steps", + "guidance", + "grow_mask", + "top", + "bottom", + "left", + "right", +) + + class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Configuration for Black Forest Labs image editing. @@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): BFL-specific params are passed through directly. """ optional_params: Final[dict[str, object]] = {} - - # Pass through BFL-specific params - bfl_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - # Kontext-specific - "aspect_ratio", - # Fill/Inpaint-specific - "steps", - "guidance", - "grow_mask", - # Expand-specific - "top", - "bottom", - "left", - "right", - ] - - # Convert TypedDict to regular dict for access - params_dict: Final = dict(image_edit_optional_params) - - for param in bfl_params: - if param in params_dict: - value = params_dict[param] - if value is not None: - optional_params[param] = value + params: Final[Mapping[str, object]] = image_edit_optional_params + for param in _BFL_REQUEST_PARAMS: + if (value := params.get(param)) is not None: + optional_params[param] = value # Set default output format if "output_format" not in optional_params: @@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): "input_image": b64_image, } - # Add optional params (only BFL-recognized parameters) - bfl_request_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - "aspect_ratio", - "steps", - "guidance", - "grow_mask", - "top", - "bottom", - "left", - "right", - ] for key, value in image_edit_optional_request_params.items(): - if key in bfl_request_params and value is not None: + if key in _BFL_REQUEST_PARAMS and value is not None: request_body[key] = value # Handle mask if provided (for inpainting) @@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") # BFL uses JSON, not multipart - return empty files - return request_body, [] + return request_body, () + + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + downloaded_image: Final = await self._fetch_remote_image(image) + downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask")) + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image if downloaded_image is None else downloaded_image, + image_edit_optional_request_params=( + dict(image_edit_optional_request_params) + if downloaded_mask is None + else {**image_edit_optional_request_params, "mask": downloaded_mask} + ), + litellm_params=litellm_params, + headers=dict(headers), + ) + + async def _fetch_remote_image(self, image: object) -> bytes | None: + candidate: Final = image[0] if isinstance(image, list) and image else image + if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")): + return None + response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0) + response.raise_for_status() + return response.content def transform_image_edit_response( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..2f561809940 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import ( Skill, ) from litellm.types.llms.openai import ( + AllMessageValues, CreateBatchRequest, CreateFileRequest, FileContentRequest, @@ -163,13 +164,10 @@ from litellm.utils import ( def _rust_responses_websocket_enabled( custom_llm_provider: str | None, - litellm_params: GenericLiteLLMParams, ) -> bool: from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return custom_llm_provider == "openai" and rust_enabled(request_override=request_override) + return custom_llm_provider == "openai" and rust_enabled() from .http_handler import get_shared_realtime_ssl_context @@ -488,7 +486,7 @@ class BaseLLMHTTPHandler: def completion( self, model: str, - messages: list, + messages: list[AllMessageValues], api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, @@ -507,7 +505,7 @@ class BaseLLMHTTPHandler: shared_session: Optional["ClientSession"] = None, ): json_mode: Final[bool] = optional_params.pop("json_mode", False) - extra_body: Final[dict | None] = optional_params.pop("extra_body", None) + extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None) provider_config = provider_config or ProviderConfigManager.get_provider_chat_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -522,14 +520,17 @@ class BaseLLMHTTPHandler: ) # get config from model, custom llm provider - headers = provider_config.validate_environment( - api_key=api_key, - headers=headers or {}, - model=model, - messages=messages, - optional_params=optional_params, - api_base=api_base, - litellm_params=litellm_params, + request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict + "dict[str, object]", + provider_config.validate_environment( + api_key=api_key, + headers=headers or {}, + model=model, + messages=messages, + optional_params=optional_params, + api_base=api_base, + litellm_params=litellm_params, + ), ) api_base = provider_config.get_complete_url( @@ -541,93 +542,117 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - data: dict[str, object] = provider_config.transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - - if extra_body is not None: - data = {**data, **extra_body} - - headers, signed_json_body = provider_config.sign_request( - headers=headers, - optional_params={ - **optional_params, - **_aws_signing_overrides(optional_params, litellm_params), - }, - request_data=data, - api_base=api_base, - api_key=api_key, - stream=stream, - fake_stream=fake_stream, - model=model, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key=api_key, - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) - - # Check if stream was converted for WebSearch interception - # This is set by the async_pre_request_hook in WebSearchInterceptionLogger - if litellm_params.get("_websearch_interception_converted_stream", False): - logging_obj.model_call_details["websearch_interception_converted_stream"] = True - - if acompletion is True: - if stream is True: - data = self._add_stream_param_to_request_body( - data=data, - provider_config=provider_config, + def sign_and_log( + transformed: dict[str, object], # mutable-ok: async_completion takes dict + ) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict + data: Final = {**transformed, **extra_body} if extra_body is not None else transformed + signed: Final = cast( # cast-ok: sign_request is declared as a bare dict + "tuple[dict[str, object], bytes | None]", + provider_config.sign_request( + headers=request_headers, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, + request_data=data, + api_base=api_base, + api_key=api_key, + stream=stream, fake_stream=fake_stream, - ) + model=model, + ), + ) + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": signed[0], + }, + ) + if litellm_params.get("_websearch_interception_converted_stream", False): + logging_obj.model_call_details["websearch_interception_converted_stream"] = True + return data, signed[0], signed[1] + + def dispatch_async( + data: dict[str, object], # mutable-ok: async_completion takes dict + signed_headers: dict[str, object], # mutable-ok: async_completion takes dict + signed_json_body: bytes | None, + ): + async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None + if stream is True: return self.acompletion_stream_function( model=model, messages=messages, api_base=api_base, - headers=headers, + headers=signed_headers, custom_llm_provider=custom_llm_provider, provider_config=provider_config, timeout=timeout, logging_obj=logging_obj, - data=data, + data=self._add_stream_param_to_request_body( + data=data, + provider_config=provider_config, + fake_stream=fake_stream, + ), fake_stream=fake_stream, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + client=async_client, litellm_params=litellm_params, json_mode=json_mode, optional_params=optional_params, signed_json_body=signed_json_body, ) + return self.async_completion( + custom_llm_provider=custom_llm_provider, + provider_config=provider_config, + api_base=api_base, + headers=signed_headers, + data=data, + timeout=timeout, + model=model, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + client=async_client, + json_mode=json_mode, + signed_json_body=signed_json_body, + shared_session=shared_session, + ) - else: - return self.async_completion( - custom_llm_provider=custom_llm_provider, - provider_config=provider_config, - api_base=api_base, - headers=headers, - data=data, - timeout=timeout, - model=model, - model_response=model_response, - logging_obj=logging_obj, - api_key=api_key, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - json_mode=json_mode, - signed_json_body=signed_json_body, - shared_session=shared_session, + if acompletion is True and provider_config.uses_async_transform_request: + + async def transform_then_dispatch(): + transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict + "dict[str, object]", + await provider_config.async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ), ) + return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) + + return transform_then_dispatch() + + data, signed_headers, signed_json_body = sign_and_log( + provider_config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ) + ) + + if acompletion is True: + return dispatch_async(data, signed_headers, signed_json_body) if stream is True: data = self._add_stream_param_to_request_body( @@ -641,7 +666,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, messages=messages, @@ -651,7 +676,7 @@ class BaseLLMHTTPHandler: completion_stream, headers = self.make_sync_call( provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, original_data=data, @@ -684,7 +709,7 @@ class BaseLLMHTTPHandler: sync_httpx_client=sync_httpx_client, provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, timeout=timeout, @@ -2403,9 +2428,7 @@ class BaseLLMHTTPHandler: return None from litellm.rust_bridge.configuration import rust_enabled - raw_request_override: Final = litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - if not rust_enabled(request_override=request_override): + if not rust_enabled(): return None if has_agentic_hook: return None @@ -6514,7 +6537,7 @@ class BaseLLMHTTPHandler: @asynccontextmanager async def _backend_connection(): - if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params): + if _rust_responses_websocket_enabled(custom_llm_provider): from litellm.rust_bridge import responses_websocket as rust_responses_websocket rust_backend: Final = await rust_responses_websocket.connect( @@ -6759,7 +6782,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - data, files = image_edit_provider_config.transform_image_edit_request( + data, files = await image_edit_provider_config.async_transform_image_edit_request( model=model, image=image, prompt=prompt, diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 1a67b33665b..42c9ef13730 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject from litellm.types.llms.vertex_ai import ContentType, PartType from litellm.utils import supports_reasoning -from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history +from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): if element.get("type") == "image_url": img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked _image_url, format, detail = _image_url_fields(img_element) - if _image_url and "https://" in _image_url: + if ( + _image_url + and "https://" in _image_url + and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX) + ): image_obj = convert_to_anthropic_image_obj(_image_url, format=format) converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj) if detail is not None: @@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): llm_provider="gemini", ) file_id = _file_field.get("file_id") - if file_id and ("http://" in file_id or "https://" in file_id): + if ( + file_id + and ("http://" in file_id or "https://" in file_id) + and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX) + ): # Convert HTTP/HTTPS file URL to base64 data try: base64_data = convert_url_to_base64(file_id) diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..494a72d6999 --- /dev/null +++ b/litellm/llms/gemini/google_genai/guardrail_translation/__init__.py @@ -0,0 +1,20 @@ +"""Google GenAI generateContent guardrail translation handler.""" + +from typing import Final + +from litellm.llms.gemini.google_genai.guardrail_translation.handler import ( + GoogleGenAIGenerateContentHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings: Final = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict) + CallTypes.generate_content: GoogleGenAIGenerateContentHandler, + CallTypes.agenerate_content: GoogleGenAIGenerateContentHandler, + CallTypes.generate_content_stream: GoogleGenAIGenerateContentHandler, + CallTypes.agenerate_content_stream: GoogleGenAIGenerateContentHandler, +} + +__all__ = ( + "GoogleGenAIGenerateContentHandler", + "guardrail_translation_mappings", +) diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py new file mode 100644 index 00000000000..e13e1e63cbb --- /dev/null +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -0,0 +1,255 @@ +""" +Google GenAI generateContent handler for Unified Guardrails. + +Extracts text from generateContent requests (systemInstruction.parts[].text +and contents[].parts[].text) and responses (candidates[].content.parts[].text), +applies the guardrail, and +writes the guardrailed text back in place. Requests and responses may be +dicts (wire format) or google-genai SDK objects; streaming chunks may +additionally be raw SSE frames, which are scanned for detection (a blocking +guardrail raises) without rewriting the frames. +""" + +import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamTransformSink, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth + +_EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _field(container: object, name: str) -> object | None: + if isinstance(container, dict): + return container.get(name) + return getattr(container, name, None) + + +def _part_text(part: object) -> str | None: + text: Final = _field(part, "text") + if isinstance(text, str) and text: + return text + return None + + +def _write_part_text(part: object, text: str) -> None: + if isinstance(part, dict): + part["text"] = text # rebind-ok: guardrail write-back rewrites the caller's part in place by handler contract + return + setattr(part, "text", text) # noqa: B010 # SDK parts are typed as object here; direct assignment cannot type-check + + +def _content_text_parts(content: object) -> tuple[object, ...]: + parts: Final = _field(content, "parts") + if not isinstance(parts, (list, tuple)): + return () + return tuple(part for part in parts if _part_text(part) is not None) + + +def _system_instruction(data: Mapping[str, object]) -> object | None: + return next( + ( + value + for container in (data, data.get("config")) + if container is not None + for key in ("systemInstruction", "system_instruction") + for value in (_field(container, key),) + if value is not None + ), + None, + ) + + +def _request_text_parts(data: Mapping[str, object]) -> tuple[object, ...]: + contents: Final = data.get("contents") + content_list: Final = ( + (contents,) if isinstance(contents, dict) else tuple(contents) if isinstance(contents, list) else () + ) + return ( + *_content_text_parts(_system_instruction(data)), + *(part for content in content_list for part in _content_text_parts(content)), + ) + + +def _response_text_parts(response: object) -> tuple[object, ...]: + candidates: Final = _field(response, "candidates") + if not isinstance(candidates, (list, tuple)): + return () + return tuple(part for candidate in candidates for part in _content_text_parts(_field(candidate, "content"))) + + +def _part_texts(text_parts: Sequence[object]) -> tuple[str, ...]: + return tuple(text for part in text_parts for text in (_part_text(part),) if text is not None) + + +def _texts_payload( + texts: Sequence[str], +) -> list[str]: # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] + return list(texts) # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] + + +def _write_back_texts(text_parts: Sequence[object], guardrailed_texts: Sequence[str] | None) -> None: + if not guardrailed_texts or len(guardrailed_texts) != len(text_parts): + return + for part, text in zip(text_parts, guardrailed_texts): + _write_part_text(part, text) + + +def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: + try: + parsed: Final = json.loads(payload) + except json.JSONDecodeError: + return None + if isinstance(parsed, dict): + return parsed + return None + + +def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: + return tuple( + text + for line in sse_text.splitlines() + if line.startswith("data:") + for payload in (line[len("data:") :].strip(),) + if payload and payload != "[DONE]" + for parsed in (_parse_json_dict_or_none(payload),) + if parsed is not None + for text in _part_texts(_response_text_parts(parsed)) + ) + + +def _chunk_sse_text(chunk: object) -> str | None: + if isinstance(chunk, bytes): + return chunk.decode("utf-8", errors="replace") + if isinstance(chunk, str): + return chunk + return None + + +def _accumulated_stream_text(responses_so_far: Sequence[object]) -> str: + object_texts: Final = tuple( + text + for chunk in responses_so_far + if _chunk_sse_text(chunk) is None + for text in _part_texts(_response_text_parts(chunk)) + ) + sse_text: Final = "".join(sse for chunk in responses_so_far for sse in (_chunk_sse_text(chunk),) if sse is not None) + return "".join(object_texts) + "".join(_sse_payload_texts(sse_text)) + + +class GoogleGenAIGenerateContentHandler(BaseTranslation): + """ + Guardrail translation for the google genai generateContent surface + (/models/{model}:generateContent, :streamGenerateContent, and the + litellm SDK generate_content call types). + """ + + async def process_input_messages( + self, + data: dict, # mutable-ok: base handler contract passes the proxy's request dict through to apply_guardrail + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> object: + text_parts: Final = _request_text_parts(data) + if not text_parts: + verbose_proxy_logger.debug("Google GenAI guardrail: no request text found, skipping") + return data + model: Final = data.get("model") + inputs: Final = ( + GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model) + if isinstance(model, str) + else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts))) + ) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + _write_back_texts(text_parts, guardrailed_inputs.get("texts")) + return data + + async def process_output_response( + self, + response: object, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Mapping[str, object] | None = None, + ) -> object: + text_parts: Final = _response_text_parts(response) + if not text_parts: + verbose_proxy_logger.debug("Google GenAI guardrail: no response text found, skipping") + return response + guardrail_request_data: Final = self._merged_request_data( + request_data=request_data, + user_api_key_dict=user_api_key_dict, + context_key="response", + context_value=response, + ) + model: Final = guardrail_request_data.get("model") + inputs: Final = ( + GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model) + if isinstance(model, str) + else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts))) + ) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=guardrail_request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + _write_back_texts(text_parts, guardrailed_inputs.get("texts")) + return response + + async def process_output_streaming_response( + self, + responses_so_far: Sequence[object], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Mapping[str, object] | None = None, + stream_transform_sink: StreamTransformSink | None = None, + ) -> object: + accumulated_text: Final = _accumulated_stream_text(responses_so_far) + if not accumulated_text: + return responses_so_far + guardrail_request_data: Final = self._merged_request_data( + request_data=request_data, + user_api_key_dict=user_api_key_dict, + context_key="responses_so_far", + context_value=responses_so_far, + ) + _guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=_texts_payload((accumulated_text,))), + request_data=guardrail_request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + def _merged_request_data( + self, + request_data: Mapping[str, object] | None, + user_api_key_dict: Optional["UserAPIKeyAuth"], + context_key: str, + context_value: object, + ) -> dict: # mutable-ok: CustomGuardrail.apply_guardrail requires a plain dict request payload + base: Final = request_data if request_data is not None else _EMPTY_REQUEST_DATA + user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + context_pairs: Final = ((context_key, context_value),) if context_key not in base else () + metadata_pairs: Final = ( + (("litellm_metadata", user_metadata),) if user_metadata and "litellm_metadata" not in base else () + ) + return dict((*base.items(), *context_pairs, *metadata_pairs)) # mutable-ok: apply_guardrail takes a plain dict diff --git a/litellm/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py index a6c88718f11..04a3a7dbc91 100644 --- a/litellm/llms/litellm_proxy/skills/constants.py +++ b/litellm/llms/litellm_proxy/skills/constants.py @@ -16,3 +16,8 @@ DEFAULT_MAX_ITERATIONS: Final[int] = 10 DEFAULT_SANDBOX_TIMEOUT: Final[int] = 120 """Default timeout in seconds for sandbox code execution.""" + +MAX_SKILLS_PER_SEARCH: Final[int] = 5000 +"""Upper bound on how many of the caller's accessible skills a single semantic +search embeds. Ranking runs in memory over this candidate set (no tsvector/DB-side +filtering yet), so this caps worst-case embedding cost per search request.""" diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 73f6ed23092..9b625cb0571 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -6,11 +6,15 @@ Used by the transformation layer and skills injection hook. """ import uuid +from collections.abc import Sequence from typing import Final from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache -from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX +from litellm.llms.litellm_proxy.skills.constants import ( + LITELLM_SKILL_ID_PREFIX, + MAX_SKILLS_PER_SEARCH, +) from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -131,6 +135,19 @@ class LiteLLMSkillsHandler: ) return [_prisma_skill_to_litellm(s) for s in skills] + @staticmethod + async def list_skills_for_search( + user_api_key_dict: UserAPIKeyAuth | None = None, + ) -> Sequence[LiteLLM_SkillsTable]: + """Every skill the caller can access, for ranking. Same owner-scope filter as + ``list_skills``, but unpaginated (up to ``MAX_SKILLS_PER_SEARCH``) since a query + must be scored against the whole accessible set, not one page of it.""" + return await LiteLLMSkillsHandler.list_skills( + limit=MAX_SKILLS_PER_SEARCH, + offset=0, + user_api_key_dict=user_api_key_dict, + ) + @staticmethod async def _load_skill(skill_id: str) -> object | None: """Cache-first read of the Prisma skill row. Owner-scope filtering diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py new file mode 100644 index 00000000000..f975c6c4cab --- /dev/null +++ b/litellm/llms/litellm_proxy/skills/skill_search.py @@ -0,0 +1,161 @@ +"""Semantic ranking over the LiteLLM-hosted skill registry, shared by GET /v1/skills?query= and the skill_search MCP tool.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import BaseModel, ConfigDict + +from litellm.llms.litellm_proxy.skills.constants import MAX_SKILLS_PER_SEARCH +from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler +from litellm.proxy.common_utils.semantic_text_index import ( + Embedder, + EmbeddingFailed, + SemanticTextIndex, + router_embedder, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + +DEFAULT_SKILL_SEARCH_TOP_K: Final = 5 +MAX_SKILL_SEARCH_TOP_K: Final = 100 +"""Matches the ``le=100`` bound GET /v1/skills?query= enforces via FastAPI's Query +validation, so the MCP tool can't return a larger payload than the REST endpoint allows.""" +MAX_SKILL_SEARCH_TEXT_CHARS: Final = 4000 +"""Per-skill cap on the title + description + instructions text that gets embedded, so one +search embeds at most ``MAX_SKILLS_PER_SEARCH * MAX_SKILL_SEARCH_TEXT_CHARS`` characters no +matter how long the stored instructions are.""" + + +@dataclass(frozen=True, slots=True) +class SkillSearchHit: + skill: LiteLLM_SkillsTable + score: float + + +@dataclass(frozen=True, slots=True) +class SkillSearchHits: + hits: tuple[SkillSearchHit, ...] + + +@dataclass(frozen=True, slots=True) +class SkillSearchNotConfigured: + reason: str + + +@dataclass(frozen=True, slots=True) +class SkillSearchEmbeddingFailed: + reason: str + + +@dataclass(frozen=True, slots=True) +class SkillSearchUnsupportedProvider: + reason: str + + +SkillSearchOutcome: TypeAlias = SkillSearchHits | SkillSearchNotConfigured | SkillSearchEmbeddingFailed +HostedSkillSearchOutcome: TypeAlias = SkillSearchOutcome | SkillSearchUnsupportedProvider + + +class SkillSearchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + skill_id: str + display_title: str | None + description: str | None + score: float + + +def skill_search_text(skill: LiteLLM_SkillsTable) -> str: + joined: Final = "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part) + return joined[:MAX_SKILL_SEARCH_TEXT_CHARS] + + +def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult: + return SkillSearchResult( + skill_id=hit.skill.skill_id, + display_title=hit.skill.display_title, + description=hit.skill.description, + score=hit.score, + ) + + +class SkillSearchIndex: + """Caches one vector per distinct skill text per embedding model, so repeat searches only embed the query.""" + + def __init__(self, max_entries: int = MAX_SKILLS_PER_SEARCH) -> None: + self._index: Final = SemanticTextIndex(max_entries=max_entries) + + async def search( + self, + query: str, + skills: Sequence[LiteLLM_SkillsTable], + top_k: int, + embed: Embedder, + embedding_model: str, + ) -> SkillSearchHits | SkillSearchEmbeddingFailed: + texts: Final = tuple(skill_search_text(skill) for skill in skills) + scores: Final = await self._index.scores(query, texts, embed, embedding_model) + if isinstance(scores, EmbeddingFailed): + return SkillSearchEmbeddingFailed(reason=scores.reason) + ranked: Final = sorted( + (SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)), + key=lambda hit: hit.score, + reverse=True, + ) + return SkillSearchHits(hits=tuple(ranked[:top_k])) + + +global_skill_search_index: Final = SkillSearchIndex() + + +async def search_skills( + query: str, + skills: Sequence[LiteLLM_SkillsTable], + top_k: int, + router: Router | None, + embedding_model: str | None, + index: SkillSearchIndex, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> SkillSearchOutcome: + if embedding_model is None: + return SkillSearchNotConfigured( + reason="skill search needs litellm_settings.skill_search_embedding_model set to an embedding model from model_list" + ) + if router is None: + return SkillSearchNotConfigured(reason="skill search needs a model_list so the embedding model can be called") + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, skills, top_k, embed, embedding_model) + + +async def search_hosted_skills( + custom_llm_provider: str | None, + query: str, + top_k: int, + router: Router | None, + embedding_model: str | None, + index: SkillSearchIndex, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> HostedSkillSearchOutcome: + """GET /v1/skills?query= for the skills LiteLLM hosts itself: only ``litellm_proxy`` has a registry to rank.""" + if custom_llm_provider != LlmProviders.LITELLM_PROXY.value: + return SkillSearchUnsupportedProvider(reason="query is only supported for custom_llm_provider=litellm_proxy") + skills: Final = await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict=user_api_key_dict) + return await search_skills( + query=query, + skills=skills, + top_k=top_k, + router=router, + embedding_model=embedding_model, + index=index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index c972dc349c9..9fc2d2cbb45 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -154,7 +154,7 @@ class LiteLLMSkillsTransformationHandler: user_api_key_dict=user_api_key_dict, ) - return self._db_skill_to_response(db_skill) + return self.db_skill_to_response(db_skill) def list_skills_handler( self, @@ -222,7 +222,9 @@ class LiteLLMSkillsTransformationHandler: user_api_key_dict=user_api_key_dict, ) - skills: Final = [self._db_skill_to_response(s) for s in db_skills] + skills: Final = [ # mutable-ok: ListSkillsResponse.data needs list[Skill]; never mutated after + self.db_skill_to_response(s) for s in db_skills + ] return ListSkillsResponse( data=skills, has_more=len(skills) >= limit, @@ -288,7 +290,7 @@ class LiteLLMSkillsTransformationHandler: skill_id=skill_id, user_api_key_dict=user_api_key_dict, ) - return self._db_skill_to_response(db_skill) + return self.db_skill_to_response(db_skill) def delete_skill_handler( self, @@ -354,7 +356,7 @@ class LiteLLMSkillsTransformationHandler: type=result.get("type", "skill_deleted"), ) - def _db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: + def db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: """ Convert a database skill record to Anthropic-compatible Skill response. @@ -375,4 +377,5 @@ class LiteLLMSkillsTransformationHandler: latest_version=db_skill.latest_version, source=db_skill.source or "custom", type="skill", + description=db_skill.description, ) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py new file mode 100644 index 00000000000..2b3264dc756 --- /dev/null +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -0,0 +1,210 @@ +""" +Support for Mistral Voxtral text-to-speech via ``/v1/audio/speech``. + +API reference: https://docs.mistral.ai/api/#tag/audio/operation/audio_speech_v1_audio_speech_post +""" + +import base64 +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig, + TextToSpeechRequestData, +) +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class MistralTextToSpeechException(BaseLLMException): + pass + + +class MistralTextToSpeechConfig(BaseTextToSpeechConfig): + TTS_BASE_URL: Final[str] = "https://api.mistral.ai/v1" + AUDIO_CONTENT_TYPES: Final[MappingProxyType[str, str]] = MappingProxyType( + { + "mp3": "audio/mpeg", + "wav": "audio/wav", + "pcm": "audio/pcm", + "flac": "audio/flac", + "opus": "audio/ogg", + } + ) + DROPPED_RESPONSE_HEADERS: Final[frozenset[str]] = frozenset( + {"content-encoding", "transfer-encoding", "content-length", "content-type"} + ) + OPENAI_VOICE_ALIASES: Final[MappingProxyType[str, str]] = MappingProxyType( + { + "alloy": "en_paul_neutral", + "echo": "gb_oliver_neutral", + "fable": "en_paul_cheerful", + "onyx": "en_paul_confident", + "nova": "gb_jane_sarcasm", + "shimmer": "gb_jane_sarcasm", + } + ) + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a plain list + return ["voice", "response_format"] # mutable-ok: base class contract returns a plain list + + def _map_openai_voice(self, voice_id: str) -> str: + return self.OPENAI_VOICE_ALIASES.get(voice_id.lower(), voice_id) + + def _resolve_voice_id(self, voice: object) -> str | None: + if isinstance(voice, str) and voice.strip(): + return self._map_openai_voice(voice.strip()) + if isinstance(voice, Mapping): + candidates: Final = (voice.get(key) for key in ("voice_id", "id", "name")) + resolved: Final = next( + (candidate.strip() for candidate in candidates if isinstance(candidate, str) and candidate.strip()), + None, + ) + return self._map_openai_voice(resolved) if resolved else None + return None + + def map_openai_params( + self, + model: str, + optional_params: Mapping[str, object], + voice: object = None, + drop_params: bool = False, + kwargs: Mapping[str, object] | None = None, + ) -> tuple[str | None, dict]: # mutable-ok: base class contract returns a plain dict + response_format: Final = optional_params.get("response_format") + ref_audio: Final = kwargs.get("ref_audio") if kwargs else None + voice_id_kwarg: Final = kwargs.get("voice_id") if kwargs else None + mapped_voice: Final = self._resolve_voice_id(voice) or self._resolve_voice_id(voice_id_kwarg) + mapped_params: Final = { # mutable-ok: base class contract returns a plain dict + key: value + for key, value in (("response_format", response_format), ("ref_audio", ref_audio)) + if isinstance(value, str) + } + return mapped_voice, mapped_params + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: base class contract returns a plain dict + resolved_key: Final = api_key or get_secret_str("MISTRAL_API_KEY") + if resolved_key is None: + raise MistralTextToSpeechException( + status_code=401, + message="Mistral API key is required. Set MISTRAL_API_KEY or pass api_key.", + ) + return { # mutable-ok: base class contract returns a plain dict + **headers, + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + } + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: Mapping[str, object], + ) -> str: + configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/") + versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1" + return f"{versioned_base}/audio/speech" + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: str | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> TextToSpeechRequestData: + response_format: Final = optional_params.get("response_format") + ref_audio: Final = optional_params.get("ref_audio") + request_data: Final[TextToSpeechRequestData] = { + "dict_body": { + "model": model, + "input": input, + **({"voice_id": voice} if voice else {}), + **({"response_format": response_format} if isinstance(response_format, str) else {}), + **({"ref_audio": ref_audio} if isinstance(ref_audio, str) else {}), + }, + "headers": {"Content-Type": "application/json"}, + } + return request_data + + def _requested_content_type(self, request: httpx.Request) -> str: + request_body: Final = json.loads(request.content or b"{}") + requested_format: Final = request_body.get("response_format") + if not isinstance(requested_format, str): + return "audio/mpeg" + return self.AUDIO_CONTENT_TYPES.get(requested_format, "audio/mpeg") + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + from litellm.types.llms.openai import HttpxBinaryResponseContent + + try: + response_json: Final = raw_response.json() + except (json.JSONDecodeError, ValueError): + raise MistralTextToSpeechException( + status_code=raw_response.status_code, + message=f"Non-JSON response from Mistral speech API: {raw_response.text[:500]}", + headers=raw_response.headers, + ) + audio_b64: Final = response_json.get("audio_data") + if not isinstance(audio_b64, str) or not audio_b64: + raise MistralTextToSpeechException( + status_code=500, + message=f"No audio_data in Mistral speech response. Response keys: {tuple(response_json.keys())}", + headers=raw_response.headers, + ) + try: + audio_bytes: Final = base64.b64decode(audio_b64, validate=True) + except ValueError: + raise MistralTextToSpeechException( + status_code=500, + message="Invalid base64 audio_data in Mistral speech response.", + headers=raw_response.headers, + ) + retained_headers: Final = tuple( + (key, value) + for key, value in raw_response.headers.items() + if key.lower() not in self.DROPPED_RESPONSE_HEADERS + ) + response_headers: Final = retained_headers + ( + ("content-length", str(len(audio_bytes))), + ("content-type", self._requested_content_type(raw_response.request)), + ) + binary_response: Final = httpx.Response( + status_code=200, + headers=response_headers, + content=audio_bytes, + request=raw_response.request, + ) + return HttpxBinaryResponseContent(binary_response) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict | httpx.Headers, # mutable-ok: BaseLLMException takes a plain dict or httpx.Headers + ) -> BaseLLMException: + return MistralTextToSpeechException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 6e9bb83b0a0..1b494ebad47 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -290,19 +290,14 @@ def handle_cohere_stream_chunk( ) -> ModelResponseStream: """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream. - ``prior_tool_calls_emitted`` lets the caller signal whether tool calls - were already emitted in earlier chunks of the same stream. When set, the - terminal consolidation chunk's tool calls are suppressed (they would - duplicate prior deltas); otherwise they are passed through so a stream - that delivers tool calls only on the terminal chunk doesn't silently - drop them. - - ``prior_text_emitted`` plays the analogous role for the ``text`` field: - when set, the terminal consolidation chunk's ``text`` is suppressed - (it would re-emit the full assembled response on top of prior deltas); - when unset (e.g. a degenerate stream that delivers the entire response - in a single SSE event carrying both ``chatHistory`` and ``finishReason``), - the text is passed through so the response content isn't silently lost. + OCI Cohere streams the answer as single-token ``text`` deltas, then restates + the whole assembled ``text`` on every chunk that carries ``toolCalls`` or + ``chatHistory`` (the tool-calls event and the terminal event). Once the + caller reports that earlier chunks already emitted text + (``prior_text_emitted``), those restatements are dropped so the client does + not see the answer twice; a stream whose only text lives on such a chunk + keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool + calls the terminal ``chatHistory`` chunk repeats. """ try: typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk) @@ -315,33 +310,10 @@ def handle_cohere_stream_chunk( if typed_chunk.index is None: typed_chunk.index = 0 - # OCI Cohere's terminal SSE event re-sends the full assembled response in - # `text` alongside a populated `chatHistory` and a non-null `finishReason`. - # Emitting that text would concatenate the whole response onto the - # already-streamed deltas. We require both signals to be present so that a - # future API change which adds `chatHistory` to intermediate chunks (or a - # rare early-populated case) doesn't silently drop legitimate token deltas. - is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None - # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive - # chunks) emit ``content=None`` rather than ``content=""`` so downstream - # stream-mergers that distinguish "no text in this delta" from "an - # explicitly empty text delta" behave correctly. - # - # We only suppress the terminal chunk's ``text`` when the caller has - # confirmed that text deltas were already emitted earlier — otherwise - # (e.g. a degenerate stream that delivers the whole response in a - # single SSE event), passing it through is the only chance to surface it. - text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text - - # Tool calls on the terminal consolidation chunk (whether from - # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what - # was already streamed in intermediate chunks. Re-emitting them would - # mint fresh `uuid4` IDs and cause downstream consumers to execute each - # tool call twice. We only suppress when the caller has confirmed that - # tool calls were already emitted earlier — otherwise (e.g. a short - # response that delivers tool calls exclusively on the terminal chunk), - # passing them through is the only chance to surface them. - cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls + restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None + restates_tool_calls: Final = typed_chunk.chatHistory is not None + text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text + cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls tool_calls: list[dict[str, object]] | None = None if cohere_tool_calls: diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index d7c2fcace09..926de3e8854 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,15 +1,17 @@ from collections.abc import Mapping, Sequence +from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints import httpx from openai.types.responses import ResponseReasoningItem -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -42,6 +44,30 @@ _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders. _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +class _ReasoningSupportEntry(BaseModel): + litellm_provider: str | None = None + supports_reasoning: bool | None = None + + +_BUNDLED_COST_MAP: Final = TypeAdapter(dict[str, _ReasoningSupportEntry]) + + +@lru_cache(maxsize=1) +def _bundled_openai_reasoning_models() -> frozenset[str]: + """OpenAI models the cost map shipped with this release flags as reasoning models. + + The live map can lag this release (a pinned mirror, or a proxy on newer code than the + map it fetches), and a lagging entry must never strip `reasoning` from a model this + release knows accepts it. + """ + bundled: Final = _BUNDLED_COST_MAP.validate_json(GetModelCostMap.read_local_model_cost_map_text()) + return frozenset( + name + for name, entry in bundled.items() + if entry.litellm_provider == LlmProviders.OPENAI.value and entry.supports_reasoning is True + ) + + class _DeleteResponseBody(TypedDict): """Decoded body of the Responses API delete call.""" @@ -95,13 +121,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: """Return True if the model supports reasoning.effort='none'.""" - from litellm.utils import _supports_factory + from litellm.utils import supports_none_reasoning_effort - return _supports_factory( - model=model, - custom_llm_provider=None, - key="supports_none_reasoning_effort", - ) + return supports_none_reasoning_effort(model=model, custom_llm_provider=None) @staticmethod def _effort_resolves_to_none(model: str, effort: str | None) -> bool: @@ -117,6 +139,28 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return OpenAIGPT5Config.effort_resolves_to_none(model, effort) + @staticmethod + def _supports_reasoning_param(model: str) -> bool: + from litellm.utils import _get_model_info_helper + + try: + info: Final = _get_model_info_helper( + model=model.split("/")[-1], custom_llm_provider=LlmProviders.OPENAI.value + ) + except Exception: + return True + declared: Final = info.get("supports_reasoning") + if declared is not None: + return declared + return info["key"] in _bundled_openai_reasoning_models() + + @staticmethod + def _requests_reasoning_effort(reasoning: object) -> bool: + effort: Final = ( + reasoning.get("effort") if isinstance(reasoning, Mapping) else getattr(reasoning, "effort", None) + ) + return effort is not None + @staticmethod def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None": """Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum. @@ -166,6 +210,23 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): if "max_output_tokens" in params: params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens")) + if ( + self.custom_llm_provider == LlmProviders.OPENAI + and self._requests_reasoning_effort(params.get("reasoning")) + and not self._supports_reasoning_param(model=model) + ): + if drop_params or litellm.drop_params: + params.pop("reasoning", None) + else: + raise litellm.UnsupportedParamsError( + message=( + f"{model} doesn't support `reasoning.effort` " + "(its model cost map entry lacks `supports_reasoning`). " + "To drop unsupported params set `litellm.drop_params = True`" + ), + status_code=400, + ) + if self._is_gpt_5_model(model=model): temperature: Final = params.get("temperature") if temperature is not None and temperature != 1: @@ -478,7 +539,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): processed_headers: Final = process_response_headers(raw_response_headers) try: response = ResponsesAPIResponse.model_validate(raw_response_json) - except Exception: + except ValidationError: verbose_logger.debug( "Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json ) @@ -870,7 +931,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): try: response = ResponsesAPIResponse.model_validate(raw_response_json) - except Exception: + except ValidationError: verbose_logger.debug( "Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json ) diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index c64fc583edc..f65b0876202 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( create_anthropic_image_param, select_anthropic_content_block_type_for_file, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload @@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) return self._transform_request_openai(model, messages, optional_params, stream, extra_body) + @property + def uses_async_transform_request(self) -> bool: + return True + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, object], # mutable-ok: BaseConfig signature + litellm_params: dict[str, object], # mutable-ok: BaseConfig signature + headers: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages + return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers) + def _transform_request_openai( self, model: str, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e2d62be6a69..13e2238fdf6 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works import json import os import re +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -27,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, response_schema_prompt, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( @@ -68,6 +70,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None # Shared sync client for GCS JSON API metadata reads so proxy/SSL settings # from litellm's HTTP stack apply (see Greptile review on PR #27278). _GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None +GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" _GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = { "image/jpg": "image/jpeg", } @@ -556,7 +559,7 @@ def _process_gemini_media( file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) - elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"): + elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX): # Gemini Files API URIs — the file is already uploaded to Google's # servers; pass the URI through as file_data without fetching it. # These URLs return 403 when accessed directly, so we must not try @@ -1307,6 +1310,23 @@ def sync_transform_request_body( ) +def _explicit_mime_type(fields: Mapping[str, object]) -> str | None: + hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type") + return hint if isinstance(hint, str) else None + + +def _ai_studio_inlines(media: RemoteMedia) -> bool: + return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX) + + +def _vertex_inlines(media: RemoteMedia) -> bool: + if media.url.startswith(GEMINI_FILES_API_URI_PREFIX): + return False + return media.url.startswith("http://") or ( + _explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None + ) + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1348,13 +1368,17 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(messages): + inlined_messages: Final = await async_inline_remote_media( + messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + ) + + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): # _transform_request_body may issue a sync httpx.get (up to 5s timeout) # via _get_gcs_object_content_type to fetch GCS object metadata. Run the # whole sync transformation on a worker thread so it does not block the # async event loop. return await asyncify(_transform_request_body)( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, @@ -1363,7 +1387,7 @@ async def async_transform_request_body( ) return _transform_request_body( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, diff --git a/litellm/main.py b/litellm/main.py index 253c9381337..75b7f7f10a5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client = _dispatch_client_http(ctx) + injected_client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR shared_session: Final = ctx.shared_session stream: Final = ctx.stream timeout: Final = ctx.timeout + client: Final = ( + injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None) + ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible try: - client = ( - HTTPHandler(timeout=timeout) if stream is False else None - ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible response: Final = base_llm_http_handler.completion( model=model, messages=messages, @@ -8389,6 +8389,34 @@ def speech( client=client, _is_async=aspeech or False, ) + elif custom_llm_provider == "mistral": + from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + ) + + mistral_tts_config: Final = text_to_speech_provider_config or MistralTextToSpeechConfig() + + if api_base is not None: + litellm_params_dict["api_base"] = api_base + if api_key is not None: + litellm_params_dict["api_key"] = api_key + + mistral_voice: Final[str | None] = voice if isinstance(voice, str) else None + + response = base_llm_http_handler.text_to_speech_handler( + model=model, + input=input, + voice=mistral_voice, + text_to_speech_provider_config=mistral_tts_config, + text_to_speech_optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params_dict, + logging_obj=logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=client, + _is_async=aspeech or False, + ) elif custom_llm_provider == "aws_polly": from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2c55828c1e2..b1ffc1583e4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3485,6 +3485,55 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://ai.azure.com/catalog/models/gpt-6-astra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, @@ -7189,7 +7238,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -7455,7 +7504,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -30742,6 +30791,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], "source": "https://developers.openai.com/api/docs/models/chat-latest", "supported_endpoints": [ "/v1/chat/completions", @@ -30759,6 +30811,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -34898,9 +34951,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-latest": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -37176,6 +37229,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37216,6 +37270,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37427,6 +37482,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37467,6 +37523,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -56467,9 +56524,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 6c68971f8d5..df3f9d2096b 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,6 +29,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.bindings import native_exception_types +from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -196,12 +198,6 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS -def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: - raw_request_override: Final = prepared_request.litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override) - - def _rust_bridge_optional_params( prepared_request: _PreparedOCRRequest, resolve_secret: Callable[[str], str | None], @@ -286,6 +282,33 @@ def _prepare_rust_ocr_call( ) +def _map_rust_ocr_error( + error: Exception, + prepared_request: _PreparedOCRRequest, + exception_types: tuple[type[BaseException], type[BaseException]] | None, +) -> Exception: + if exception_types is None: + return error + _, upstream_error = exception_types + if not isinstance(error, upstream_error): + return error + error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs + tuple[object, ...], error.args + ) + status_value: Final = error_args[0] if error_args else 0 + message_value: Final = error_args[1] if len(error_args) > 1 else str(error) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped + Callable[..., Exception], prepared_request.provider_config.get_error_class + ) + return error_factory( + error_message=message, + status_code=status or 500, + headers={}, # mutable-ok: provider error factories require a concrete header dict + ) + + def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], @@ -296,16 +319,19 @@ def _run_rust_ocr( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - rust_response: Final = 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, - ) + try: + rust_response: Final = 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, + ) + except Exception as error: + raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error if rust_response is None: return None return OCRResponse.model_validate(rust_response) @@ -321,16 +347,19 @@ async def _run_rust_aocr( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - rust_response: Final = 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, - ) + try: + rust_response: Final = 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, + ) + except Exception as error: + raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error if rust_response is None: return None return OCRResponse.model_validate(rust_response) @@ -430,7 +459,7 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): + if _rust_ocr_supported(prepared) and rust_enabled(): from litellm.secret_managers.main import get_secret_str rust_response: Final = await _run_rust_aocr( @@ -702,7 +731,7 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): + if _rust_ocr_supported(prepared) and rust_enabled(): from litellm.secret_managers.main import get_secret_str rust_response: Final = _run_rust_ocr( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 612596bc803..e7fd650a324 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -163,7 +163,10 @@ from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( id_jag_assertion_capture_gap_at_startup, ) -from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path +from litellm.proxy.middleware.per_request_root_path_middleware import ( + get_request_root_path, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import ( @@ -3858,7 +3861,7 @@ class MCPServerManager: if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig): # authorization_code's missing per-user token -> the per-server browser-OAuth # challenge, built here where the full MCPServer is in hand. - raise_user_oauth_challenge(server, root_path=get_server_root_path()) + raise_user_oauth_challenge(server, root_path=get_request_root_path()) if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): # token_exchange (OBO): a missing/rejected subject token -> the RFC 9728 challenge # pointing at the IdP the client must SSO with to obtain one, rather than an opaque @@ -3866,7 +3869,7 @@ class MCPServerManager: # Access) threads its claims blob into the challenge for the client to satisfy. raise_token_exchange_challenge( server, - root_path=get_server_root_path(), + root_path=get_request_root_path(), claims=err.unauthorized.claims, ) raise_public(err) @@ -3914,7 +3917,7 @@ class MCPServerManager: if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): return if subject_token is None and isinstance(spec.config, TokenExchangeConfig): - raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path()) + raise_token_exchange_challenge(resolved_server, root_path=get_request_root_path()) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return @@ -3922,7 +3925,7 @@ class MCPServerManager: if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): raise_token_exchange_challenge( resolved_server, - root_path=get_server_root_path(), + root_path=get_request_root_path(), claims=err.unauthorized.claims, ) raise_public(err) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index ea2318bd6f1..77979a15199 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -12,6 +12,7 @@ every other mode so the caller defers to v1 (parity-safe); it grows one branch p from __future__ import annotations import base64 +import os from typing import TYPE_CHECKING, Final, Literal, NoReturn from fastapi import HTTPException @@ -310,13 +311,30 @@ def raise_public(error: CredError) -> NoReturn: def oauth_protected_resource_path(root_path: str, server: MCPServer) -> str: """The server's RFC 9728 Protected Resource Metadata path, the shared anchor of both challenges. - ``root_path`` is the proxy's ``SERVER_ROOT_PATH``, resolved by the caller (the imperative shell) - so this stays a pure function of its inputs; ``"/"`` and ``""`` both mean no prefix. The path is - relative, so it resolves against the caller's own host (correct even behind a reverse proxy). + ``root_path`` is the prefix the request was routed under, resolved by the caller (the imperative + shell); ``"/"`` and ``""`` both mean no prefix. The path is relative, so it resolves against the + caller's own host (correct even behind a reverse proxy). + + URL structure depends on how the prefix is served: + + - The scalar ``SERVER_ROOT_PATH`` deployment registers the well-known routes with the prefix + *inserted* into the path (via :func:`well_known_root_suffix` at import time), matching RFC 8414 + §3 well-known path insertion. When ``root_path`` equals ``SERVER_ROOT_PATH`` the URL must use + the same insertion or a client fetching it 404s. + - The per-request ``SERVER_ROOT_PATHS`` deployment can't register routes per prefix (the prefix + set is dynamic and could contain many entries); the middleware strips the prefix from + ``scope["path"]`` and the router matches the un-inserted well-known route. The URL must place + the prefix *before* ``.well-known`` so the strip leaves a matching path. + + Picking the wrong form 404s the client's discovery fetch — the discovery document and the 401 + challenge would then disagree on where the resource metadata lives. """ prefix: Final = "" if root_path == "/" else root_path name: Final = server.alias or server.server_name or server.name or server.server_id - return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}" + scalar_env: Final = os.getenv("SERVER_ROOT_PATH", "").rstrip("/") + if not prefix or (scalar_env and prefix == scalar_env): + return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}" + return f"{prefix}/.well-known/oauth-protected-resource/mcp/{name}" def raise_user_oauth_challenge(server: MCPServer, *, root_path: str) -> NoReturn: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 5fbfad54a39..329dddbdf05 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,7 @@ import asyncio import importlib from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal @@ -8,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status +from starlette.datastructures import Headers from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT @@ -104,6 +106,9 @@ if MCP_AVAILABLE: from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.client import MCPClient + from litellm.llms.litellm_proxy.skills.skill_search import ( + DEFAULT_SKILL_SEARCH_TOP_K, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -188,10 +193,12 @@ if MCP_AVAILABLE: AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, coerce_top_k, handle_agent_search, handle_mcp_tool_call, handle_mcp_tool_search, + handle_skill_search, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj @@ -210,6 +217,14 @@ if MCP_AVAILABLE: ), user_api_key_dict=user_api_key_dict, ) + if tool_name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(tool_arguments.get("query", "")), + top_k=coerce_top_k( + tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K + ), + user_api_key_dict=user_api_key_dict, + ) rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) ( virtual_mcp_auth_header, @@ -1153,6 +1168,45 @@ if MCP_AVAILABLE: scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes + _STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization) + ) + + @dataclass(frozen=True, slots=True) + class _StagedServerTest: + request: NewMCPServerRequest + mcp_auth_header: str | None + oauth2_headers: dict[str, str] | None + + def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest: + """ + Resolve the credentials a not-yet-saved server config carries for a preview call. + + Both preview endpoints (``/test/connection`` and ``/test/tools/list``) must hand the + temporary client the same credentials, or a server that the saved connection reaches + fine fails one of them. + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request) + mcp_auth_header: Final = ( + request.credentials.get("auth_value") + if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict) + else None + ) + # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY): + # when the primary x-litellm-api-key header is absent, the Authorization value is the + # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream. + oauth2_headers: Final = ( + MCPRequestHandler._get_oauth2_headers_from_headers(headers) + if request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) + else None + ) + return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers) + async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: with anyio.move_on_after(deadline): return await client.list_tools(raise_on_error=True) @@ -1374,6 +1428,8 @@ if MCP_AVAILABLE: }, ) + staged: Final = _stage_server_test(new_mcp_server_request, request.headers) + async def _test_connection_operation(client): async def _noop(session): return "ok" @@ -1382,8 +1438,10 @@ if MCP_AVAILABLE: return {"status": "ok"} return await _execute_with_mcp_client( - new_mcp_server_request, + staged.request, _test_connection_operation, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, raw_headers=_safe_get_request_headers(request), ) @@ -1404,37 +1462,11 @@ if MCP_AVAILABLE: }, ) - new_mcp_server_request = _inherit_credentials_from_existing_server(new_mcp_server_request) + staged: Final = _stage_server_test(new_mcp_server_request, request.headers) # For OpenAPI spec servers, generate tools from the spec directly - if new_mcp_server_request.spec_path: - return await _preview_openapi_tools(new_mcp_server_request.spec_path) - - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, - ) - - headers: Final = request.headers - - mcp_auth_header: str | None = None - if new_mcp_server_request.auth_type in { - MCPAuth.api_key, - MCPAuth.bearer_token, - MCPAuth.basic, - MCPAuth.authorization, - }: - credentials: Final = getattr(new_mcp_server_request, "credentials", None) - if isinstance(credentials, dict): - mcp_auth_header = credentials.get("auth_value") - - # Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY): - # when the primary x-litellm-api-key header is absent, the Authorization value is the - # caller's LiteLLM key, not an upstream token, and must never be forwarded upstream. - oauth2_headers: dict[str, str] | None = None - if new_mcp_server_request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and headers.get( - MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY - ): - oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) + if staged.request.spec_path: + return await _preview_openapi_tools(staged.request.spec_path) async def _list_tools_operation(client): # Bound the whole pagination walk: without this the preview is limited only by the @@ -1465,9 +1497,9 @@ if MCP_AVAILABLE: } return await _execute_with_mcp_client( - new_mcp_server_request, + staged.request, _list_tools_operation, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, raw_headers=_safe_get_request_headers(request), ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 26f5d6e7c8c..0b424c31c4b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -911,15 +911,18 @@ if MCP_AVAILABLE: Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so the caller falls through to normal tool routing. """ + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, DEFAULT_AGENT_SEARCH_TOP_K, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, VIRTUAL_TOOL_NAMES, coerce_top_k, handle_agent_search, handle_mcp_tool_call, handle_mcp_tool_search, + handle_skill_search, ) if name not in VIRTUAL_TOOL_NAMES: @@ -961,6 +964,12 @@ if MCP_AVAILABLE: top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K), user_api_key_dict=user_api_key_auth, ) + if name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(args.get("query", "")), + top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K), + user_api_key_dict=user_api_key_auth, + ) virtual_logging_obj: Final = await _build_virtual_call_logging_obj( name=name, arguments=args, @@ -3848,12 +3857,14 @@ if MCP_AVAILABLE: # then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the # header lost, so the discovery flow needs this pre-emptive challenge. if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers: - from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph raise_token_exchange_challenge, ) - from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 + from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports proxy utils + get_request_root_path, + ) - raise_token_exchange_challenge(server, root_path=get_server_root_path()) + raise_token_exchange_challenge(server, root_path=get_request_root_path()) # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run # the exchange here at the transport edge, so a rejected subject raises the RFC 9728 diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index af02c11ad86..f19340d30cb 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -11,6 +11,7 @@ from pydantic import ValidationError from typing_extensions import ReadOnly, Required, assert_never import litellm +from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K from litellm.proxy.common_utils.semantic_text_index import ( Embedder, @@ -30,7 +31,10 @@ MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search" MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call" AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search" -VIRTUAL_TOOL_NAMES: Final = frozenset((MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME)) +SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search" +VIRTUAL_TOOL_NAMES: Final = frozenset( + (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, SKILL_SEARCH_TOOL_NAME) +) def coerce_top_k(value: Any, default: int = 5) -> int: @@ -199,8 +203,28 @@ _AGENT_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { } +_SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { + "name": SKILL_SEARCH_TOOL_NAME, + "description": "Find registered skills by describing what you need in natural language. Returns the best " + "matching skills you can access, ranked by semantic similarity, each with its skill_id, display_title, " + "description, and score.", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "What you need the skill to do, in natural language."}, + "top_k": { + "type": "integer", + "description": "Maximum number of skills to return.", + "default": DEFAULT_SKILL_SEARCH_TOP_K, + }, + }, + "required": _json_array("query"), + }, +} + + def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]: - return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION) + return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION) def _text_tool_result(text: str, is_error: bool) -> CallToolResult: @@ -223,7 +247,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI ) from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj await check_feature_access_for_user(user_api_key_dict, "agents") outcome: Final = await search_agents( @@ -234,6 +258,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): @@ -245,6 +270,39 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI assert_never(outcome) +async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth) -> CallToolResult: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TOP_K, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchNotConfigured, + global_skill_search_index, + search_skills, + skill_search_result, + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + outcome: Final = await search_skills( + query=query, + skills=await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict), + top_k=min(max(top_k, 1), MAX_SKILL_SEARCH_TOP_K), + router=llm_router, + embedding_model=litellm.skill_search_embedding_model, + index=global_skill_search_index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + match outcome: + case SkillSearchHits(hits): + results: Final = tuple(skill_search_result(hit).model_dump() for hit in hits) + return _text_tool_result(json.dumps(results), is_error=False) + case SkillSearchNotConfigured(reason) | SkillSearchEmbeddingFailed(reason): + return _text_tool_result(reason, is_error=True) + case _: + assert_never(outcome) + + async def handle_mcp_tool_search( query: str, top_k: int, @@ -257,7 +315,7 @@ async def handle_mcp_tool_search( raw_headers: dict[str, str] | None = None, ) -> CallToolResult: from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj settings: Final = mcp_tool_search_settings() if isinstance(settings, ValidationError): @@ -271,7 +329,7 @@ async def handle_mcp_tool_search( ) ranker: Final = ( SemanticToolRanker( - embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict), + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj), embedding_model=settings.embedding_model, index=global_mcp_tool_search_index, ) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 3f90e6c0a7a..50e0a961a49 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -295,15 +295,18 @@ class LazyFeatureMiddleware: # Short-circuit once every feature has loaded. if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features): path = scope.get("path", "") - # Strip SERVER_ROOT_PATH so prefix matching works under a server - # root path. Without this, requests like /api/v1/policies/... never - # match the registered prefixes (/policies/...) and lazy features - # stay unloaded — every endpoint under them returns 404. The + # Strip the request's root_path so prefix matching works under a + # server root path. Without this, requests like /api/v1/policies/... + # never match the registered prefixes (/policies/...) and lazy + # features stay unloaded — every endpoint under them returns 404. + # scope["root_path"] wins over the cached env scalar: FastAPI + # stamps SERVER_ROOT_PATH there, and PerRequestRootPathMiddleware + # resolves SERVER_ROOT_PATHS prefixes there per request. The # `+ "/"` boundary prevents false-positive matches (e.g. /apiv2 - # against root /api). If the path doesn't start with the prefix - # (e.g. a reverse proxy already stripped it), we leave it alone. - if self._root_path and path.startswith(self._root_path + "/"): - path = path[len(self._root_path) :] + # against root /api); a pre-stripped path is left alone. + root_path: Final = str(scope.get("root_path", "")).rstrip("/") or self._root_path + if root_path and path.startswith(root_path + "/"): + path = path[len(root_path) :] # rebind-ok: local strip after the boundary check above for feat in self._features: if feat.module_path in self._loaded: continue diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 91a97ad6544..c71761a7adb 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -4687,6 +4687,17 @@ "title": "Created At", "type": "string" }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, "display_title": { "anyOf": [ { @@ -4713,6 +4724,17 @@ ], "title": "Latest Version" }, + "search_score": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Search Score" + }, "source": { "title": "Source", "type": "string" @@ -4781,7 +4803,7 @@ "paths": { "/v1/skills": { "get": { - "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nReturns: ListSkillsResponse with list of skills", + "description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nPass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can\naccess by semantic similarity instead of paging through the whole registry:\n```bash\ncurl \"http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5\" -H \"Authorization: Bearer your-key\"\n```\n\nReturns: ListSkillsResponse with list of skills", "operationId": "list_skills_v1_skills_get", "parameters": [ { @@ -4849,6 +4871,39 @@ "default": "anthropic", "title": "Custom Llm Provider" } + }, + { + "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.", + "in": "query", + "name": "query", + "required": false, + "schema": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.", + "title": "Query" + } + }, + { + "description": "With query: the maximum number of ranked skills to return.", + "in": "query", + "name": "top_k", + "required": false, + "schema": { + "default": 5, + "description": "With query: the maximum number of ranked skills to return.", + "maximum": 100, + "minimum": 1, + "title": "Top K", + "type": "integer" + } } ], "responses": { @@ -9597,7 +9652,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -9605,7 +9662,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -11914,7 +11971,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -11922,7 +11981,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c28ac8848ba..abce10690e5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -848,6 +848,7 @@ class LiteLLMRoutes(enum.Enum): "/model/{model_id}/update", "/prompt/list", "/prompt/info", + "/vector_store/info", # Project read routes - endpoint scopes results to caller's teams (non-admin) "/project/list", "/project/info", @@ -1276,6 +1277,7 @@ class UpdateKeyRequest(KeyRequestBase): # else they will get overwritten duration: str | None = None spend: float | None = None + soft_budget: float | None = None metadata: dict | None = None temp_budget_increase: float | None = None temp_budget_expiry: datetime | None = None diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 76e3fe6c5ad..65a89bb2c7a 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -18,6 +18,7 @@ from litellm.types.agents import AgentResponse if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 @@ -132,6 +133,7 @@ async def search_agents( embedding_model: str | None, index: AgentSearchIndex, user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, ) -> AgentSearchOutcome: if embedding_model is None: return AgentSearchNotConfigured( @@ -139,6 +141,5 @@ async def search_agents( ) if router is None: return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called") - return await index.search( - query, agents, top_k, router_embedder(router, embedding_model, user_api_key_dict), embedding_model - ) + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, agents, top_k, embed, embedding_model) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index cc17672553b..aa8979a73c6 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -249,7 +249,7 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep async def _rank_agents_by_query( query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth ) -> tuple[AgentResponse, ...]: - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj outcome: Final = await search_agents( query=query, @@ -259,6 +259,7 @@ async def _rank_agents_by_query( embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py index 9390bf4c537..4426c0b547a 100644 --- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py @@ -2,11 +2,23 @@ Anthropic Skills API endpoints - /v1/skills """ -from typing import Final +from types import MappingProxyType +from typing import Annotated, Final import orjson -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from typing_extensions import ReadOnly, TypedDict, assert_never +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import ( + DEFAULT_SKILL_SEARCH_TOP_K, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchNotConfigured, + SkillSearchUnsupportedProvider, + global_skill_search_index, + search_hosted_skills, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -23,6 +35,51 @@ from litellm.types.llms.anthropic_skills import ( router: Final = APIRouter() +class _SkillSearchErrorDetail(TypedDict): + error: ReadOnly[str] + message: ReadOnly[str] + + +def _skill_search_error(status_code: int, error: str, message: str) -> HTTPException: + detail: Final[_SkillSearchErrorDetail] = {"error": error, "message": message} + return HTTPException(status_code=status_code, detail=detail) + + +async def _search_skills( + custom_llm_provider: str | None, query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth +) -> ListSkillsResponse: + from litellm.llms.litellm_proxy.skills.transformation import ( + LiteLLMSkillsTransformationHandler, + ) + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + + outcome: Final = await search_hosted_skills( + custom_llm_provider=custom_llm_provider, + query=query, + top_k=top_k, + router=llm_router, + embedding_model=litellm.skill_search_embedding_model, + index=global_skill_search_index, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + to_response: Final = LiteLLMSkillsTransformationHandler().db_skill_to_response + match outcome: + case SkillSearchHits(hits): + skills: Final = [ # mutable-ok: ListSkillsResponse.data requires list[Skill]; never mutated after + to_response(hit.skill).model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits + ] + return ListSkillsResponse(data=skills, has_more=False, next_page=None) + case SkillSearchUnsupportedProvider(reason): + raise _skill_search_error(400, "skill_search_unsupported_provider", reason) + case SkillSearchNotConfigured(reason): + raise _skill_search_error(400, "skill_search_not_configured", reason) + case SkillSearchEmbeddingFailed(reason): + raise _skill_search_error(503, "skill_search_unavailable", reason) + case _: + assert_never(outcome) + + @router.post( "/v1/skills", tags=["[beta] Anthropic Skills API"], @@ -134,32 +191,58 @@ async def list_skills( after_id: str | None = None, before_id: str | None = None, custom_llm_provider: str | None = "anthropic", + query: Annotated[ + str | None, + Query( + min_length=1, + description="Describe what you need in natural language to rank the skills you can access by " + "semantic similarity over their title and description. Each result carries a search_score. " + "Only supported for custom_llm_provider=litellm_proxy. Requires " + "litellm_settings.skill_search_embedding_model.", + ), + ] = None, + top_k: Annotated[ + int, + Query(ge=1, le=100, description="With query: the maximum number of ranked skills to return."), + ] = DEFAULT_SKILL_SEARCH_TOP_K, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ List skills on Anthropic. - + Requires `?beta=true` query parameter. - + Model-based routing (for multi-account support): - Pass model via header: `x-litellm-model: claude-account-1` - Pass model via query: `?model=claude-account-1` - Pass model via body: `{"model": "claude-account-1"}` - + Example usage: ```bash # Basic usage curl "http://localhost:4000/v1/skills?beta=true&limit=10" \ -H "Authorization: Bearer your-key" - + # With model-based routing curl "http://localhost:4000/v1/skills?beta=true&limit=10" \ -H "Authorization: Bearer your-key" \ -H "x-litellm-model: claude-account-1" ``` - + + Pass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can + access by semantic similarity instead of paging through the whole registry: + ```bash + curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" \ + -H "Authorization: Bearer your-key" + ``` + Returns: ListSkillsResponse with list of skills """ + if query is not None: + return await _search_skills( + custom_llm_provider=custom_llm_provider, query=query, top_k=top_k, user_api_key_dict=user_api_key_dict + ) + from litellm.proxy.proxy_server import ( general_settings, llm_router, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d6007a2d56e..1e4836654a1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1981,9 +1981,28 @@ def get_model_from_request( if vertex_match: model = vertex_match.group(1) + if route.lower().startswith("/bedrock"): + bedrock_model: Final = _model_from_bedrock_route(route) + return model if bedrock_model is None else bedrock_model + return model +def _model_from_bedrock_route(route: str) -> str | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + is_bedrock_count_tokens_endpoint, + ) + + bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) + if is_bedrock_count_tokens_endpoint(bedrock_endpoint): + return None + try: + return _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + return None + + def abbreviate_api_key(api_key: str) -> str: if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH: return "sk-..." diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5e6c9b34332..9720e4b1cf8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -756,7 +756,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse): content: AsyncGenerator[str, None], *, media_type: str | None = None, - headers: dict | None = None, + headers: Mapping[str, str] | None = None, status_code: int = status.HTTP_200_OK, upstream_generator: AsyncGenerator[str, None] | None = None, ) -> None: @@ -888,25 +888,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n" +def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + """`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384).""" + return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS}) + + +async def _resolve_stream_headers( + headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None +) -> Mapping[str, str]: + if refresh_headers is None: + return headers + try: + return await refresh_headers() + except Exception as e: # noqa: BLE001 # a stream whose first chunk is already paid for must not fail over its headers + verbose_proxy_logger.exception("Error refreshing streaming response headers: %s", e) + return headers + + async def create_response( generator: AsyncGenerator[str, None], media_type: str, - headers: dict, + headers: Mapping[str, str], default_status_code: int = status.HTTP_200_OK, request: Request | None = None, + refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None, ) -> StreamingResponse | JSONResponse: """ Create streaming response, checking if the first chunk is an error. If the first chunk is an error, return a standard JSON error response. Otherwise, return StreamingResponse and stream all content. + + ``refresh_headers`` is consulted once the first chunk has been buffered, for + callers whose headers can only be known then. """ - # Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE - # immediately instead of releasing the whole stream in one batch (issue #28384). - streaming_headers: Final = { - **headers, - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - } first_chunk_value: str | None = None final_status_code = default_status_code @@ -917,6 +931,7 @@ async def create_response( # Now get the first chunk from the actual generator first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) + resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers) if first_chunk_value is not None: try: @@ -943,7 +958,7 @@ async def create_response( return JSONResponse( status_code=final_status_code, content={"error": error_dict}, - headers=headers, + headers=resolved_headers, ) except Exception as e: verbose_proxy_logger.debug("Error parsing first chunk value: %s", e) @@ -972,7 +987,7 @@ async def create_response( return StreamingResponse( empty_gen(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)), status_code=default_status_code, ) except Exception as e: @@ -988,7 +1003,7 @@ async def create_response( return StreamingResponse( error_gen_message(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)), status_code=error_status, ) @@ -1010,7 +1025,7 @@ async def create_response( return _UpstreamClosingStreamingResponse( combined_generator(), media_type=media_type, - headers=streaming_headers, + headers=_sse_stream_headers(resolved_headers), status_code=final_status_code, upstream_generator=generator, ) @@ -1535,7 +1550,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _merge_passthrough_streaming_headers( response_headers: httpx.Headers | dict | None, - custom_headers: dict, + custom_headers: Mapping[str, str], ) -> dict: """ Merge upstream passthrough headers with proxy/custom headers. @@ -2143,14 +2158,45 @@ class ProxyBaseLLMRequestProcessing: return fallback_model_group @staticmethod - def _get_model_id_from_response(hidden_params: dict, data: dict) -> str: + def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str: """Extract model_id from hidden_params with fallback to litellm_metadata.""" model_id = hidden_params.get("model_id", None) or "" if not model_id: - litellm_metadata: Final = data.get("litellm_metadata", {}) or {} - model_info: Final = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", "") or "" - return model_id + litellm_metadata: Final = data.get("litellm_metadata") + model_info: Final = litellm_metadata.get("model_info") if isinstance(litellm_metadata, Mapping) else None + model_id = (model_info.get("id") or "") if isinstance(model_info, Mapping) else "" + return str(model_id) if model_id else "" + + def _stream_response_headers( + self, + *, + hidden_params: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + logging_obj: LiteLLMLoggingObj, + version: str | None, + callback_headers: Mapping[str, str], + ) -> Mapping[str, str]: + """The streaming response headers describing `hidden_params`' deployment.""" + return MappingProxyType( + { + **ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=self._get_model_id_from_response(hidden_params, self.data), + cache_key=hidden_params.get("cache_key") or "", + api_base=hidden_params.get("api_base") or "", + version=version, + response_cost=hidden_params.get("response_cost") or "", + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"), + request_data=self.data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **(hidden_params.get("additional_headers") or MappingProxyType({})), + ), + **callback_headers, + } + ) @staticmethod def _get_deployment_model_name( @@ -2419,31 +2465,32 @@ class ProxyBaseLLMRequestProcessing: if self._is_streaming_request( data=self.data, is_streaming_request=is_streaming_request ) or self._is_streaming_response(response): # use generate_responses to stream responses - custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=logging_obj.litellm_call_id, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=fastest_response_batch_completion, - request_data=self.data, - hidden_params=hidden_params, - litellm_logging_obj=logging_obj, - **additional_headers, - ) - # Call response headers hook for streaming success - callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + stream_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( data=self.data, user_api_key_dict=user_api_key_dict, response=response, request_headers=dict(request.headers), ) - if callback_headers: - custom_headers.update(callback_headers) + custom_headers: Final = self._stream_response_headers( + hidden_params=hidden_params, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + callback_headers=stream_callback_headers or MappingProxyType({}), + ) + + async def refresh_stream_headers() -> Mapping[str, str]: + """`custom_headers` rebuilt for whichever deployment served the stream.""" + if not getattr(response, "fallback_headers_adopted", False): + return custom_headers + return self._stream_response_headers( + hidden_params=get_hidden_params_dict(response), + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + callback_headers=stream_callback_headers or MappingProxyType({}), + ) # Preserve the original client-requested model (pre-alias mapping) for downstream # streaming generators. Pre-call processing can rewrite `self.data["model"]` for @@ -2581,6 +2628,7 @@ class ProxyBaseLLMRequestProcessing: media_type="text/event-stream", headers=custom_headers, request=request, + refresh_headers=refresh_stream_headers, ) ### CALL HOOKS ### - modify outgoing data @@ -3032,7 +3080,7 @@ class ProxyBaseLLMRequestProcessing: response: Any, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", - custom_headers: dict, + custom_headers: Mapping[str, str], request_headers: dict[str, str], ) -> Response | None: if not self._has_post_call_guardrails_for_passthrough(): diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 0820459af49..b8d3595163e 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -5,10 +5,11 @@ from __future__ import annotations import math from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass -from itertools import chain +from itertools import chain, islice from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias +from fastapi import HTTPException from openai import OpenAIError from pydantic import BaseModel, ConfigDict @@ -16,10 +17,14 @@ from litellm.exceptions import BudgetExceededError if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router Vector: TypeAlias = tuple[float, ...] +DEFAULT_MAX_CACHED_VECTORS: Final = 5000 +"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first.""" + class Embedder(Protocol): def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ... @@ -42,6 +47,16 @@ class _EmbeddingData(BaseModel): data: tuple[_EmbeddingItem, ...] +class _EmbeddingRequest(BaseModel): + """The /embeddings-shaped request as the pre-call hooks (rate limits, budgets, guardrails) hand it back.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + model: str + input: tuple[str, ...] + metadata: dict[str, object] # mutable-ok: the router mutates the metadata dict it is handed + + def cosine_similarity(left: Vector, right: Vector) -> float: dot: Final = sum(a * b for a, b in zip(left, right, strict=True)) norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)) @@ -57,23 +72,40 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj } -def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: +def router_embedder( + router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging +) -> Embedder: + """Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs.""" + async def embed(texts: Sequence[str]) -> Sequence[Vector]: - batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input + request: Final = { # mutable-ok: pre_call_hook mutates the request dict in place + "model": embedding_model, + "input": list(texts), # mutable-ok: Router.aembedding accepts only str | list input + "metadata": embedding_spend_metadata(user_api_key_dict), + } + processed: Final = _EmbeddingRequest.model_validate( + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding" + ) + ) response: Final = await router.aembedding( - model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) + model=processed.model, + input=list(processed.input), # mutable-ok: Router.aembedding accepts only str | list input + metadata=processed.metadata, ) return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) return embed -_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({}) +_CacheKey: TypeAlias = tuple[str, str] async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed: try: vectors: Final = tuple(await embed(texts)) + except HTTPException: + raise except (OpenAIError, ValueError, BudgetExceededError) as exc: return EmbeddingFailed(reason=f"embedding the search query failed: {exc}") if len(vectors) != len(texts): @@ -111,20 +143,34 @@ async def _embed_query_and_texts( class SemanticTextIndex: - """Caches one vector per distinct text per embedding model, so repeat searches only embed the query.""" + """Caches one vector per distinct text per embedding model, so repeat searches only embed the query. - def __init__(self) -> None: - self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({}) + Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first.""" - def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]: - kept: Final = MappingProxyType( + def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None: + self._max_entries: Final = max_entries + self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({}) + + def _cached(self, embedding_model: str) -> Mapping[str, Vector]: + return MappingProxyType( + {text: vector for (model, text), vector in self._vectors.items() if model == embedding_model} + ) + + def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]: + dimension: Final = len(embedded.query_vector) + touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts}) + untouched: Final = MappingProxyType( { - text: vector - for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items() - if len(vector) == len(embedded.query_vector) + key: vector + for key, vector in chain( + self._vectors.items(), + (((embedding_model, text), vector) for text, vector in embedded.vectors.items()), + ) + if key not in touched and (key[0] != embedding_model or len(vector) == dimension) } ) - return MappingProxyType({**kept, **embedded.vectors}) + ordered: Final = MappingProxyType({**untouched, **touched}) + return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None))) async def scores( self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str @@ -132,11 +178,10 @@ class SemanticTextIndex: """Cosine similarity of `query` to each entry of `texts`, in the same order.""" if not texts: return () - cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) - embedded: Final = await _embed_query_and_texts(embed, query, texts, cached) + embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model)) if isinstance(embedded, EmbeddingFailed): return embedded if not _same_dimension(embedded.query_vector, embedded.vectors, texts): return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions") - self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)}) + self._vectors = self._merged(embedding_model, embedded, texts) return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts) diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 9c637a62dc1..b33bffbaaa1 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -37,7 +37,9 @@ CACHE_TTL_1H_SECONDS: Final = 3600 AUTOROUTER_BENCHMARKS_SQL: Final = """ WITH windowed AS ( SELECT * FROM "LiteLLM_AutoRouterSession" - WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp + WHERE last_turn_at >= $1::timestamp + AND first_turn_at < $2::timestamp + AND ($3::text IS NULL OR api_key = $3::text) ), tier_maps AS ( SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e6880d521f1..914c961b145 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -14,6 +14,7 @@ import time import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload import litellm @@ -84,6 +85,25 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) +def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: + return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" + + +_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) + + +def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: + """Reduce a batch's cost row to what tells the retrieves apart when logging is off. + + A proxy run with spend logs disabled still needs one row per batch to charge it once, + so the row is written either way, but it carries no request of its own: no metadata, + no requester IP, no key, model, or token counts (LIT-7048). + """ + if disable_spend_logs is False: + return payload + return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -215,7 +235,12 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ) -> None: + ) -> bool: + """Record the request's spend, answering whether its cost still needs charging. + + False only for a batch retrieve whose cost row another retrieve already wrote, + so the caller leaves the key, team, and user counters alone (LIT-7048). + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -232,7 +257,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return True if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -262,11 +287,12 @@ class DBSpendUpdateWriter: if team_id is not None and team_id != "": payload["team_id"] = team_id + if not await self._record_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ): + return False + if disable_spend_logs is False: - await self._insert_spend_log_to_db( - payload=payload, - prisma_client=prisma_client, - ) await self._enqueue_tool_usage_transaction( payload=payload, completion_response=completion_response, @@ -306,6 +332,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return True except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -318,7 +345,102 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) - return + return True + + async def _record_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool + ) -> bool: + if prisma_client is not None and _is_batch_cost_row(payload): + return await self._claim_batch_cost_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) + if disable_spend_logs is False: + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + + async def _claim_batch_cost_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> bool: + """Write the batch's cost row now, or learn that another retrieve already did. + + Every retrieve of one batch shares this row, so the insert that lands first owns + the charge and every later one finds the row and charges nothing (LIT-7048). Only + a row that recorded a charge counts: a failed retrieve, a request whose client + picked the batch id as its call id, and the $0 row an older proxy left behind + while the batch was still running all leave the charge to be made. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + row: Final = _batch_cost_row_to_write(payload, disable_spend_logs) + spend_logs: Final = SpendLogsRepository(prisma_client).table + try: + claimed: Final = await spend_logs.create_many( + data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list + skip_duplicates=True, + ) + if claimed == 1: + return True + existing: Final = await spend_logs.find_unique( + where={"request_id": request_id} # mutable-ok: prisma where clause + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log + verbose_proxy_logger.warning( + "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e + ) + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) + return True + if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": + verbose_proxy_logger.warning( + "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + request_id, + getattr(existing, "call_type", None), + ) + return True + if existing.spend > 0: + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False + return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) + + async def _take_over_uncharged_batch_cost_row( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] + ) -> bool: + """Take the batch's cost row over from the poll that left it charging nothing. + + A pre-upgrade proxy wrote that row every time it polled the batch while it was still + running, so the charge is still to be made and the row still has to end up carrying + it. The row stops matching the moment it carries a charge, so it is one retrieve that + takes it over and charges, and every later one reads the charge and charges nothing. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + if payload["spend"] <= 0: + verbose_proxy_logger.debug( + "Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id + ) + return False + try: + taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many( + data=prisma_client.jsonify_object( + MappingProxyType({field: value for field, value in row.items() if field != "request_id"}) + ), + where={ # mutable-ok: prisma where clause + "request_id": request_id, + "call_type": CallTypes.aretrieve_batch.value, + "status": "success", + "spend": 0.0, + }, + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over + verbose_proxy_logger.warning( + "Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e + ) + return False + if taken_over == 0: + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False + return True async def _enqueue_tool_usage_transaction( self, @@ -380,6 +502,7 @@ class DBSpendUpdateWriter: llm_router=get_llm_router, cost_breakdown=metadata.get("cost_breakdown"), recorded_autorouter_savings=metadata.get("autorouter_savings"), + billed_at=payload.get("endTime"), ) transaction: Final = build_autorouter_turn_transaction( payload=payload, @@ -2066,6 +2189,7 @@ class DBSpendUpdateWriter: usage_object=usage_obj, cost_breakdown=_metadata.get("cost_breakdown"), recorded_autorouter_savings=_metadata.get("autorouter_savings"), + billed_at=payload.get("endTime"), ) daily_transaction: Final = BaseDailySpendTransaction( diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index 221c142d9d3..b756bfeb6f6 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -16,7 +16,12 @@ keeps the batched-DELETE path, so existing deployments are untouched. import re from collections.abc import Callable from datetime import date, datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, TypeAlias +from typing import ( + TYPE_CHECKING, + Final, + TypeAlias, + cast, # noqa: TID251 # db.tx is reached through untyped __getattr__ delegation +) from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -25,6 +30,8 @@ from litellm.constants import ( ) if TYPE_CHECKING: + from prisma.client import TransactionManager + from litellm.proxy.utils import PrismaClient SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs" @@ -116,6 +123,21 @@ def select_partitions_to_drop(partitions: list[tuple[str, datetime | None]], cut return [name for name, upper in partitions if upper is not None and upper <= cutoff] +_TX_COMMIT_SLACK: Final = timedelta(seconds=5) + + +def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager": + """ + Open an interactive transaction that outlives the statement bound it + carries. prisma's default 5s transaction timeout would close it mid + lock-wait, after which the engine answers the next call with a 422. + """ + return cast( # cast-ok: PrismaWrapper delegates tx via __getattr__ (untyped) + "TransactionManager", + prisma_client.db.tx(timeout=timedelta(milliseconds=timeout_ms) + _TX_COMMIT_SLACK), + ) + + class SpendLogsPartitionManager: def __init__( self, @@ -137,7 +159,7 @@ class SpendLogsPartitionManager: if budget_ms is None: return False try: - async with prisma_client.db.tx() as tx: + async with _bounded_tx(prisma_client, budget_ms) as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}") rows: Final = await tx.query_raw( """ @@ -172,7 +194,7 @@ class SpendLogsPartitionManager: wait for the lock and statement_timeout bounds the work itself, so a partition this run cannot get is simply left for the next one. """ - async with prisma_client.db.tx() as tx: + async with _bounded_tx(prisma_client, timeout_ms) as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}") await tx.execute_raw(statement) @@ -209,7 +231,7 @@ class SpendLogsPartitionManager: async def _list_partitions( self, prisma_client: "PrismaClient", timeout_ms: int ) -> list[tuple[str, datetime | None]]: - async with prisma_client.db.tx() as tx: + async with _bounded_tx(prisma_client, timeout_ms) as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") rows: Final = await tx.query_raw( """ diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 744b2959c73..afb9997f2e6 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -1202,7 +1202,13 @@ async def patch_guardrail( litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) - litellm_params = LitellmParams(**merged_litellm_params) + try: + litellm_params = LitellmParams(**merged_litellm_params) + except ValidationError as validation_error: + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {validation_error}", + ) from validation_error # Update guardrail_info if provided guardrail_info: Final = ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py index d53a4157e0e..1607dfff63e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional import litellm from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( @@ -20,10 +20,7 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("MCP Security: guardrail_name is required") - on_violation: Final[Literal["block", "alert"]] = cast( - Literal["block", "alert"], - getattr(litellm_params, "on_violation", "block"), - ) + on_violation: Final[Literal["block", "alert"]] = "block" if litellm_params.on_violation == "block" else "alert" mcp_security_guardrail: Final = MCPSecurityGuardrail( guardrail_name=guardrail_name, diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py index 0aaba4016cd..88cf92a4a8c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py @@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" event_hook=litellm_params.mode, default_on=litellm_params.default_on, file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None), + block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None), ) litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 84c4f118b00..0954fe1698a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -93,6 +93,7 @@ class PromptSecurityGuardrail(CustomGuardrail): check_tool_results: bool | None = None, file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS, file_sanitization_fail_open: bool | None = None, + block_on_file_modify: bool | None = None, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -124,6 +125,7 @@ class PromptSecurityGuardrail(CustomGuardrail): self.poll_interval = 2 # Seconds between polling attempts self.file_sanitization_timeout = file_sanitization_timeout self.file_sanitization_fail_open = file_sanitization_fail_open is not False + self.block_on_file_modify = block_on_file_modify is not False super().__init__(**kwargs) @@ -372,13 +374,7 @@ class PromptSecurityGuardrail(CustomGuardrail): result = await self.sanitize_file_content( file_data, filename, user_api_key_alias=user_api_key_alias ) - - if result.get("action") == "block": - violations = result.get("violations", []) - raise HTTPException( - status_code=400, - detail=f"Image blocked by Prompt Security. Violations: {', '.join(violations)}", - ) + self._raise_if_file_blocked(result, "Image") except HTTPException: raise except Exception as e: @@ -408,7 +404,7 @@ class PromptSecurityGuardrail(CustomGuardrail): file_data: bytes, filename: str, user_api_key_alias: str | None = None, - ) -> dict: + ) -> _SanitizeResult: """ Sanitize file content using Prompt Security API. Returns: dict with keys 'action', 'content', 'metadata' @@ -528,6 +524,17 @@ class PromptSecurityGuardrail(CustomGuardrail): raise HTTPException(status_code=408, detail="File sanitization timeout") + def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None: + action: Final = sanitization_result.get("action") + if action != "block" and not (action == "modify" and self.block_on_file_modify): + return + + violations: Final = sanitization_result.get("violations", ()) + raise HTTPException( + status_code=400, + detail=f"{resource_name} blocked by Prompt Security. Violations: {', '.join(violations)}", + ) + async def _process_image_url_item(self, item: dict, user_api_key_alias: str | None) -> dict: """Process and sanitize image_url items.""" image_url_data: Final = item.get("image_url", {}) @@ -547,13 +554,7 @@ class PromptSecurityGuardrail(CustomGuardrail): file_data, filename, user_api_key_alias=user_api_key_alias ) action: Final = sanitization_result.get("action") - - if action == "block": - violations: Final = sanitization_result.get("violations", []) - raise HTTPException( - status_code=400, - detail=f"File blocked by Prompt Security. Violations: {', '.join(violations)}", - ) + self._raise_if_file_blocked(sanitization_result, "File") if action == "modify": sanitized_content: Final = sanitization_result.get("content", "") @@ -615,13 +616,7 @@ class PromptSecurityGuardrail(CustomGuardrail): file_data, filename, user_api_key_alias=user_api_key_alias ) action: Final = sanitization_result.get("action") - - if action == "block": - violations: Final = sanitization_result.get("violations", []) - raise HTTPException( - status_code=400, - detail=f"Document blocked by Prompt Security. Violations: {', '.join(violations)}", - ) + self._raise_if_file_blocked(sanitization_result, "Document") if action == "modify": sanitized_content: Final = sanitization_result.get("content", "") diff --git a/litellm/proxy/hooks/litellm_skills/__init__.py b/litellm/proxy/hooks/litellm_skills/__init__.py index 751122ac51c..d24cdd37161 100644 --- a/litellm/proxy/hooks/litellm_skills/__init__.py +++ b/litellm/proxy/hooks/litellm_skills/__init__.py @@ -21,10 +21,7 @@ from litellm.llms.litellm_proxy.skills import ( code_execution_handler, get_litellm_code_execution_tool, ) -from litellm.proxy.hooks.litellm_skills.main import ( - SkillsInjectionHook, - skills_injection_hook, -) +from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook __all__ = [ "LITELLM_CODE_EXECUTION_TOOL", @@ -35,5 +32,4 @@ __all__ = [ "SkillsSandboxExecutor", "code_execution_handler", "get_litellm_code_execution_tool", - "skills_injection_hook", ] diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 9edbc6dbf1c..8848878d15b 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -29,6 +29,7 @@ import json from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol +import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger @@ -493,7 +494,6 @@ class SkillsInjectionHook(CustomLogger): Returns the final response with generated files inline. """ - import litellm from litellm.llms.litellm_proxy.skills.code_execution import ( LiteLLMInternalTools, ) @@ -723,7 +723,6 @@ print('No executable skill module found') Returns the final response with generated files inline. """ - import litellm from litellm.llms.litellm_proxy.skills.code_execution import ( LiteLLMInternalTools, ) @@ -913,11 +912,3 @@ print('No executable skill module found') verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to response", len(generated_files)) return response - - -# Global instance for registration -skills_injection_hook: Final = SkillsInjectionHook() - -import litellm - -litellm.logging_callback_manager.add_litellm_callback(skills_injection_hook) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index d6276c03155..c4fba8ecf9e 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger +from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -37,6 +38,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( CallTypes, + LiteLLMBatch, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) @@ -248,6 +250,18 @@ class _ProxyDBLogger(CustomLogger): ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + if ( + isinstance(completion_response, LiteLLMBatch) + and kwargs.get("call_type") == CallTypes.aretrieve_batch.value + and not batch_cost_is_final(completion_response) + ): + verbose_proxy_logger.debug( + "Cost tracking deferred for batch %s still in status %s", + completion_response.id, + completion_response.status, + ) + await _release_budget_reservation(budget_reservation=budget_reservation) + return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -285,7 +299,7 @@ class _ProxyDBLogger(CustomLogger): call_type=call_type, ): ## UPDATE DATABASE - await _update_database_and_spend_counters( + charged: Final = await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, increment_spend_counters=increment_spend_counters, user_api_key=user_api_key, @@ -302,6 +316,8 @@ class _ProxyDBLogger(CustomLogger): request_tags=tags, model_access_groups=model_access_groups, ) + if not charged: + return # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) @@ -578,9 +594,9 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, -) -> None: +) -> bool: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -605,6 +621,9 @@ async def _update_database_and_spend_counters( "Failed to invalidate budget reservation counters after release failed" ) raise + if not charged: + await _release_budget_reservation(budget_reservation=budget_reservation) + return False try: await increment_spend_counters( @@ -630,6 +649,7 @@ async def _update_database_and_spend_counters( finally: budget_reservation["finalized"] = True raise + return True async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d026c5510e6..56512570448 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2882,6 +2882,28 @@ def _add_guardrails_from_policies_in_metadata( ) +def add_guardrails_from_auth_metadata( + user_api_key_dict: UserAPIKeyAuth, + data: dict, # mutable-ok: writes guardrails into the live request dict, same contract as the helpers it wraps + metadata_variable_name: str, +) -> None: + """Resolve key, team, and project guardrails, direct and via policies, onto the request metadata.""" + _add_guardrails_from_key_or_team_metadata( + key_metadata=user_api_key_dict.metadata, + team_metadata=user_api_key_dict.team_metadata, + project_metadata=user_api_key_dict.project_metadata, + data=data, + metadata_variable_name=metadata_variable_name, + ) + _add_guardrails_from_policies_in_metadata( + key_metadata=user_api_key_dict.metadata, + team_metadata=user_api_key_dict.team_metadata, + project_metadata=user_api_key_dict.project_metadata, + data=data, + metadata_variable_name=metadata_variable_name, + ) + + async def move_guardrails_to_metadata( data: dict, _metadata_variable_name: str, @@ -2914,22 +2936,8 @@ async def move_guardrails_to_metadata( data.pop("policies", None) return - # Check key/team/project-level guardrails - _add_guardrails_from_key_or_team_metadata( - key_metadata=user_api_key_dict.metadata, - team_metadata=user_api_key_dict.team_metadata, - project_metadata=project_metadata, - data=data, - metadata_variable_name=_metadata_variable_name, - ) - - ######################################################################################### - # Add guardrails from policies attached to key/team/project metadata - ######################################################################################### - _add_guardrails_from_policies_in_metadata( - key_metadata=user_api_key_dict.metadata, - team_metadata=user_api_key_dict.team_metadata, - project_metadata=project_metadata, + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name=_metadata_variable_name, ) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 2a7813bc140..0f0323b45f8 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -645,6 +645,7 @@ async def get_auto_router_benchmarks( str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to 30 days before end_date)") ] = None, end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None, + api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None, ) -> AutoRouterBenchmarksResponse: """ Benchmarks for the auto-router dashboard: session shape, savings against the configured @@ -681,6 +682,7 @@ async def get_auto_router_benchmarks( AUTOROUTER_BENCHMARKS_SQL, start_day.isoformat(), (end_day + timedelta(days=1)).isoformat(), + api_key, ) rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ()) groups: Final = ( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ebcfab090b5..f46c4170071 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -196,6 +196,38 @@ class _ModelRowWhere(TypedDict): model_id: ReadOnly[str] +class _KeyUpdateResult(TypedDict): + token: ReadOnly[str] + data: ReadOnly[Mapping[str, object]] + + +class _KeyRowWhere(TypedDict): + token: ReadOnly[str] + + +class _BudgetRowWhere(TypedDict): + budget_id: ReadOnly[str] + + +class _BudgetRowSoftBudgetUpdate(TypedDict): + soft_budget: ReadOnly[float | None] + updated_by: ReadOnly[str] + + +class _BudgetRowSoftBudgetCreate(TypedDict): + soft_budget: ReadOnly[float] + created_by: ReadOnly[str] + updated_by: ReadOnly[str] + + +class _KeyUpdateTx(Protocol): + @property + def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... + + @property + def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ... + + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -1812,11 +1844,7 @@ async def generate_key_fn( status_code=400, detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) - if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): - raise HTTPException( - status_code=400, - detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, - ) + _validate_soft_budget_value(data.soft_budget) custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( _custom_key_generate_hook(proxy_server) @@ -2121,6 +2149,88 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ return non_default_values +def _validate_soft_budget_value(soft_budget: float | None) -> None: + if soft_budget is not None and (not math.isfinite(soft_budget) or soft_budget < 0): + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {soft_budget}"}, + ) + + +async def _update_key_soft_budget( + db: _KeyUpdateTx, + existing_key_row: LiteLLM_VerificationToken, + soft_budget: float | None, + changed_by: str, +) -> str | None: + existing_budget_id: Final = existing_key_row.budget_id + if existing_budget_id is not None: + budget_update: Final[_BudgetRowSoftBudgetUpdate] = {"soft_budget": soft_budget, "updated_by": changed_by} + budget_where: Final[_BudgetRowWhere] = {"budget_id": existing_budget_id} + await db.litellm_budgettable.update(where=budget_where, data=budget_update) + return existing_budget_id + if soft_budget is None: + return None + budget_create: Final[_BudgetRowSoftBudgetCreate] = { + "soft_budget": soft_budget, + "created_by": changed_by, + "updated_by": changed_by, + } + created_budget: Final = await db.litellm_budgettable.create(data=budget_create) + return created_budget.budget_id + + +async def _apply_soft_budget_update( + data: UpdateKeyRequest, + non_default_values: Mapping[str, object], + db: _KeyUpdateTx, + existing_key_row: LiteLLM_VerificationToken, + changed_by: str, +) -> Mapping[str, object]: + remaining: Final = MappingProxyType({k: v for k, v in non_default_values.items() if k != "soft_budget"}) + updated_budget_id: Final = await _update_key_soft_budget( + db=db, + existing_key_row=existing_key_row, + soft_budget=data.soft_budget, + changed_by=changed_by, + ) + if updated_budget_id is not None and existing_key_row.budget_id is None: + return MappingProxyType({**remaining, "budget_id": updated_budget_id}) + return remaining + + +async def _update_key_row_with_soft_budget( + prisma_client: PrismaClient, + key: str, + data: UpdateKeyRequest, + non_default_values: Mapping[str, object], + existing_key_row: LiteLLM_VerificationToken, + changed_by: str, +) -> _KeyUpdateResult: + hashed_token: Final = _hash_token_if_needed(key) + key_where: Final[_KeyRowWhere] = {"token": hashed_token} + tx: _KeyUpdateTx + async with prisma_client.tx() as tx: + update_values: Final = await _apply_soft_budget_update( + data=data, + non_default_values=non_default_values, + db=tx, + existing_key_row=existing_key_row, + changed_by=changed_by, + ) + updated_row: Final = await tx.litellm_verificationtoken.update( + where=key_where, + data=with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ), + ) + updated_data: Final[Mapping[str, object]] = ( + updated_row.model_dump() if updated_row is not None else MappingProxyType({}) + ) + result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data} + return result + + async def prepare_key_update_data( data: UpdateKeyRequest | RegenerateKeyRequest, existing_key_row: LiteLLM_VerificationToken, @@ -2659,6 +2769,7 @@ async def _validate_update_key_data( (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) or data.spend is not None or "budget_limits" in data.model_fields_set + or "soft_budget" in data.model_fields_set ) _existing_metadata: Final = getattr(existing_key_row, "metadata", None) @@ -2840,6 +2951,11 @@ async def update_key_fn( """ Update an existing API key's parameters. + The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + an explicit null clears it. The metadata-backed fields below are the exception, merging into the + stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + replaces the stored metadata wholesale. + Parameters: - key: Optional[str] - The key to update. Either key or key_alias must be provided. - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases) @@ -2857,7 +2973,7 @@ async def update_key_fn( - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. + - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget. - max_parallel_requests: Optional[int] - Rate limit for parallel requests - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} - tpm_limit: Optional[int] - Tokens per minute limit @@ -2913,6 +3029,7 @@ async def update_key_fn( """ from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, llm_router, premium_user, prisma_client, @@ -2928,6 +3045,8 @@ async def update_key_fn( detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) + _validate_soft_budget_value(data.soft_budget) + # get the row from db existing_key_row: Final = await _get_and_validate_existing_key( token=data.key, @@ -2984,10 +3103,22 @@ async def update_key_fn( existing_key_alias=existing_key_row.key_alias, ) - _data: Final = {**non_default_values, "token": key} if prisma_client is None: raise Exception("Not connected to DB!") - response: Final = await prisma_client.update_data(token=key, data=_data) + + changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name + response: Final = ( + await _update_key_row_with_soft_budget( + prisma_client=prisma_client, + key=key, + data=data, + non_default_values=non_default_values, + existing_key_row=existing_key_row, + changed_by=changed_by, + ) + if "soft_budget" in data.model_fields_set + else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key})) + ) # Delete - key from cache, since it's been updated! # key updated - a new model could have been added to this key. it should not block requests after this is done @@ -6427,7 +6558,7 @@ async def _list_key_helper( {"token": "desc"}, # fallback sort ] ), - include={"object_permission": True}, + include={"object_permission": True, "litellm_budget_table": True}, ) verbose_proxy_logger.debug("Fetched %s keys", len(keys)) diff --git a/litellm/proxy/middleware/per_request_root_path_middleware.py b/litellm/proxy/middleware/per_request_root_path_middleware.py new file mode 100644 index 00000000000..d5df91458d2 --- /dev/null +++ b/litellm/proxy/middleware/per_request_root_path_middleware.py @@ -0,0 +1,122 @@ +"""Per-request ``root_path`` resolution from ``SERVER_ROOT_PATHS``. + +``SERVER_ROOT_PATH`` is a startup scalar, so one deployment serves exactly one +client-visible URL path prefix; a request under any other prefix 404s before a +handler runs. When the ingress preserves several prefixes into one pod (e.g. +``/tenant-a/*`` and ``/tenant-b/*``), the matched prefix becomes that +request's ``scope["root_path"]`` instead: Starlette strips it during route +matching and rebuilds it into ``request.base_url``, so every emitted URL — +the MCP OAuth discovery ``resource`` (RFC 9728 §3) and the 401 challenges' +``resource_metadata`` among them — lands under the prefix the client called. +Opt-in: with ``SERVER_ROOT_PATHS`` unset the middleware is not added at all. +""" + +import os +from collections.abc import Sequence +from contextvars import ContextVar +from typing import Final + +from starlette.types import ASGIApp, Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger + +SERVER_ROOT_PATHS_ENV: Final = "SERVER_ROOT_PATHS" + +# The effective ``root_path`` for the currently-handled request. Populated by +# ``PerRequestRootPathMiddleware`` from the (possibly-mutated) scope so code +# that emits URLs off the request path — the 401 challenges' resource_metadata +# and ``get_custom_url``'s SSO callbacks among them — can pick up the prefix +# the client actually called without threading scope through every call site. +# ``None`` means "middleware did not run" (the ``SERVER_ROOT_PATHS`` env is +# unset, so no per-request prefix exists); readers fall back to the scalar +# ``SERVER_ROOT_PATH`` in that case, which matches the pre-middleware behavior. +_request_root_path_var: Final[ContextVar[str | None]] = ContextVar("_request_root_path_var", default=None) + + +def get_request_root_path() -> str: + """Return the effective ``root_path`` for the current request. + + Reads the value ``PerRequestRootPathMiddleware`` stashed for this request; + falls back through :func:`~litellm.proxy.utils.get_server_root_path` (i.e. + the ``SERVER_ROOT_PATH`` env) when the middleware did not run — the + scalar-only deployment. Delegating to the existing helper keeps every + existing ``monkeypatch.setattr("litellm.proxy.utils.get_server_root_path"`` + test override working, and keeps a single source of truth for the scalar. + """ + value: Final = _request_root_path_var.get() + if value is not None: + return value + # Lazy import: utils.py imports this module (via the lazy import inside + # get_custom_url), so a top-level import would build a cycle at load time. + from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 # lazy import breaks a two-way dep + + return get_server_root_path() + + +def normalize_root_paths(raw_paths: Sequence[str]) -> tuple[str, ...]: + """Strip whitespace and trailing slashes, dedupe, order longest-first; + warn and drop entries missing a leading ``/`` and the bare root.""" + kept: Final[list[str]] = [] # mutable-ok: local accumulator; escapes only as a tuple + for entry in raw_paths: + candidate = entry.strip() + if not candidate: + continue + if not candidate.startswith("/"): + verbose_proxy_logger.warning( + "%s entry %r does not start with '/' and will be ignored.", + SERVER_ROOT_PATHS_ENV, + entry, + ) + continue + candidate = candidate.rstrip("/") + if not candidate: + verbose_proxy_logger.warning( + "%s entry %r is the bare root and will be ignored; a root-mounted deployment needs no entry.", + SERVER_ROOT_PATHS_ENV, + entry, + ) + continue + if candidate not in kept: + kept.append(candidate) + return tuple(sorted(kept, key=len, reverse=True)) + + +def get_server_root_paths() -> tuple[str, ...]: + """The normalized ``SERVER_ROOT_PATHS`` prefixes, empty when unset.""" + configured: Final = os.getenv(SERVER_ROOT_PATHS_ENV, "") + if not configured.strip(): + return () + return normalize_root_paths(configured.split(",")) + + +class PerRequestRootPathMiddleware: + """Sets ``scope["root_path"]`` to the configured prefix matching the + request path on a whole-segment boundary. ``scope["path"]`` is left + untouched (Starlette strips ``root_path`` at route-match time). Must be + the outermost middleware so inner middlewares and the router see the + resolved value; a matched prefix overrides a scalar ``SERVER_ROOT_PATH`` + for that request. + """ + + def __init__(self, app: ASGIApp, root_paths: Sequence[str]) -> None: + self.app = app + self.root_paths: Final = normalize_root_paths(root_paths) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] in ("http", "websocket"): + path: Final = scope.get("path", "") + for prefix in self.root_paths: + if path == prefix or path.startswith(prefix + "/"): + scope["root_path"] = prefix # rebind-ok: ASGI middleware contract; Router and base_url read it + break + # Stash the effective root_path (matched prefix, or the scope's + # existing value when nothing matched — i.e. FastAPI's scalar + # SERVER_ROOT_PATH) so code that emits URLs off the request path + # picks the same prefix the router will resolve the request under. + token: Final = _request_root_path_var.set(str(scope.get("root_path", ""))) + try: + await self.app(scope, receive, send) + finally: + _request_root_path_var.reset(token) + return + await self.app(scope, receive, send) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 15eeddbc489..b1f282a0978 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -15,6 +15,7 @@ from typing import ( runtime_checkable, ) +from litellm.batches.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: enumerated the batch and none succeeded. A zero or unknown total means counts are unreported, so stay eligible and let the next poller pass revisit it. (#37713) """ - if response.output_file_id is not None: - return True - request_counts = response.request_counts - if request_counts is None: - return False - return request_counts.total > 0 and request_counts.completed == 0 + return batch_cost_is_final(response) async def update_batch_in_database( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 32da2658b99..b95547e2b54 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -712,6 +712,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = { BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"} +def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool: + return "count_tokens" in endpoint or "count-tokens" in endpoint + + def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: """ Extract model name from Bedrock endpoint path. @@ -986,8 +990,7 @@ async def bedrock_llm_proxy_route( request_body: Final = await _read_request_body(request=request) - # Special handling for count_tokens endpoints - if "count_tokens" in endpoint or "count-tokens" in endpoint: + if is_bedrock_count_tokens_endpoint(endpoint): return await handle_bedrock_count_tokens( endpoint=endpoint, request=request, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3cb9acc6110..f1f823e59b5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -322,7 +322,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): def get_response_headers( headers: httpx.Headers, litellm_call_id: str | None = None, - custom_headers: dict | None = None, + custom_headers: Mapping[str, str] | None = None, ) -> dict: # Exclude headers that uvicorn writes itself (server, date) and # encoding/length headers that don't survive re-serialization. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ba5714fe950..0915b8dd1b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -598,6 +598,10 @@ from litellm.proxy.middleware.admission_control_middleware import ( from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) +from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + get_server_root_paths, +) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.middleware.request_size_limit_middleware import ( RequestSizeLimitMiddleware, @@ -18384,6 +18388,22 @@ app.add_middleware( get_settings=lambda: get_admission_control_settings(general_settings), state=admission_control_state, ) +# Added last on purpose - last-added is outermost, and the client-visible URL +# prefix must be resolved into scope["root_path"] before any inner middleware +# or the router inspects the path. Only added when SERVER_ROOT_PATHS is +# configured, so the default deployment's middleware stack is unchanged. +_server_root_paths: Final = get_server_root_paths() +if _server_root_paths: + if server_root_path and server_root_path != "/": + verbose_proxy_logger.warning( + "Both SERVER_ROOT_PATH=%r and SERVER_ROOT_PATHS=%r are set. A request " + "matching a SERVER_ROOT_PATHS prefix overrides the scalar root_path for " + "that request; unmatched requests keep SERVER_ROOT_PATH. Configure one " + "mechanism or the other.", + server_root_path, + _server_root_paths, + ) + app.add_middleware(PerRequestRootPathMiddleware, root_paths=_server_root_paths) async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse": diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 66f8c2ea36f..cd781abee26 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -688,6 +688,13 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "CHATGPT", + "provider_display_name": "ChatGPT Subscription", + "litellm_provider": "chatgpt", + "credential_fields": [], + "default_model_placeholder": "chatgpt/gpt-5.4" + }, { "provider": "CLARIFAI", "provider_display_name": "Clarifai", diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ed2bc87597c..152f3befa15 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1389,12 +1389,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: def _count_input_tokens(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: - return litellm.token_counter( - model=model, - messages=request_body.get("messages") or [], - tools=request_body.get("tools"), - tool_choice=request_body.get("tool_choice"), - ) + try: + return litellm.token_counter( + model=model, + messages=request_body.get("messages") or (), + tools=request_body.get("tools"), + tool_choice=request_body.get("tool_choice"), + ) + except ValueError: + return _count_text_tokens(model=model, text=request_body.get("messages")) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -1442,11 +1445,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: int | None = None - for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): - requested = _to_int(request_body.get(key)) - if requested is not None: - break + requested: Final = _requested_output_tokens(request_body) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1457,9 +1456,19 @@ def _estimate_output_tokens( # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - if requested is None: - requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - return min(requested, model_ceiling) + return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling) + + +_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") + + +def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: + inference_config: Final = request_body.get("inferenceConfig") + candidates: Final = ( + *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), + inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, + ) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) def _count_text_tokens(model: str, text: object) -> int: diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 1d0eb12da75..c541b9b40e5 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -9,12 +9,17 @@ have been aggregated across models. """ from collections.abc import Callable, Mapping +from datetime import datetime from typing import TYPE_CHECKING, Final, NamedTuple import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY -from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit, generic_cost_per_token +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + _get_cost_per_unit, + calculate_prompt_caching_savings, + generic_cost_per_token, +) from litellm.types.integrations.anthropic_cache_control_hook import ( GATEWAY_INJECTED_CACHE_METADATA_KEY, GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT, @@ -32,42 +37,13 @@ class SavingsSpend(NamedTuple): gateway_injected_caching: float = 0.0 -def _input_cache_read_and_write_cost(info: ModelInfo | None) -> tuple[float, float, float]: - """ - Return ``(input_cost, cache_read_cost, cache_write_cost)`` per token. - - ``info`` is whatever pricing the caller resolved -- deployment rates when the - request came through a router deployment, public rates otherwise -- so a - negotiated price is honoured here rather than silently replaced by the list rate. - ``None`` falls open to ``(0.0, 0.0, 0.0)`` so savings degrade to zero rather than - raising inside the spend writer. - - Prices are read through ``_get_cost_per_unit``, the same accessor the cost - calculator uses, which coerces the string prices a ``config.yaml`` can produce - (``"3e-7"``) and resolves service-tier suffixes. - - An absent cache price mirrors the input cost, which yields a zero discount on the - read leg and a zero premium on the write leg. Mirroring rather than taking - ``_get_cost_per_unit``'s 0.0 default is load-bearing on the write leg: a zero write - price would make the premium ``0 - input_cost``, turning a model that simply has no - write pricing into a spurious extra saving. - - The two legs then differ on an explicit ``0.0``, and the asymmetry is deliberate. A - free cache *write* does not exist -- entries carrying a literal zero (``deepseek-chat`` - does) mean "no separate price", so a falsy write price also mirrors input. A free - cache *read* is real: 15 models charge for input and serve reads for nothing, which - is the largest discount available, so the read leg keeps its literal zero. - """ - if info is None: - return 0.0, 0.0, 0.0 - input_cost: Final = _get_cost_per_unit(info, "input_cost_per_token") or 0.0 - cache_read_cost: Final = _get_cost_per_unit(info, "cache_read_input_token_cost", default_value=None) - cache_write_cost: Final = _get_cost_per_unit(info, "cache_creation_input_token_cost", default_value=None) - return ( - input_cost, - input_cost if cache_read_cost is None else cache_read_cost, - cache_write_cost if cache_write_cost else input_cost, - ) +def _coerce_billed_at(value: datetime | str | None) -> datetime | None: + if isinstance(value, datetime) or value is None: + return value + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None class _ModelIdentity(NamedTuple): @@ -586,6 +562,7 @@ def compute_savings_spend( llm_router: "Callable[[], Router | None] | None" = None, cost_breakdown: Mapping[str, object] | None = None, recorded_autorouter_savings: object = None, + billed_at: datetime | str | None = None, ) -> SavingsSpend: """ Dollar savings for one request, split by optimization driver. @@ -595,24 +572,10 @@ def compute_savings_spend( premium paid to write those entries, both derived here from ``usage_object`` so no caller can hand in a count that disagrees with the usage record. - The net form follows from what the request would have cost with caching off. The - provider reports ``prompt_tokens`` as the inclusive total of three disjoint - partitions (uncached text, cache reads, cache writes), so an uncached counterfactual - bills every one of those tokens at the flat input rate:: - - would_have_cost = (text + reads + writes) * input - actually_cost = text * input + reads * read_rate + writes * write_rate - savings = reads * (input - read_rate) - writes * (write_rate - input) - - So the write leg subtracts the write PREMIUM, not the whole write cost: those tokens - had to be sent either way, and the counterfactual already pays the input rate for - them. The premium stays signed, because a handful of models price writes below their - input rate and there the write is a genuine extra saving. - - A request that only writes cache and gets no hits therefore reports negative savings, - which is accurate: it really did cost more than the uncached call would have. The - daily rollup increments arithmetically, so those rows offset positive ones in the - same bucket. + The uncached counterfactual pays the ordinary input rate for the same prompt size + and tier. Cache writes subtract only the premium over that rate, split by TTL. + Savings stay signed: a write-only request can lose money, and daily rollups net + those losses against read savings. Caching is reported twice. ``prompt_caching`` is every net dollar caching saved, whoever caused it, which is what a customer means by "what did caching save me". @@ -638,12 +601,9 @@ def compute_savings_spend( calls this and only auto-routed ones need one, so looking it up eagerly at the call site would fetch and discard it on the rest. - ``cost_breakdown`` is what the cost calculator recorded for this request, and it - carries both what the request really cost and the tier and region it was priced on. - Only the auto-router driver reads it. Compression and prompt caching price a - hypothetical token delta off flat rate keys, so they are blind to tiered pricing in - the same way; that is pre-existing behaviour on two shipped drivers rather than - something introduced here, and moving those numbers is its own change. + ``cost_breakdown`` supplies the biller's tier and region to caching and auto-router + savings. Caching also uses the logged prompt size and TTL split. Compression retains + its flat input-rate estimate; changing that counterfactual is a separate concern. ``recorded_autorouter_savings`` is the figure the logging path stamped on the spend log's metadata, honoured over recomputation so the rollup, the turn table and the @@ -658,13 +618,24 @@ def compute_savings_spend( pricing: Final = _effective_model_info(router_instance, model_id, model or "") or ( _model_info(identity) if identity else None ) - input_cost, cache_read_cost, cache_write_cost = _input_cache_read_and_write_cost(pricing) + input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0 compression: Final = max(compression_saved_tokens, 0) * input_cost - cache_read_input_tokens: Final = extract_cache_read_tokens(usage_object) - cache_creation_input_tokens: Final = extract_cache_creation_tokens(usage_object) - read_discount: Final = max(cache_read_input_tokens, 0) * max(input_cost - cache_read_cost, 0.0) - write_premium: Final = max(cache_creation_input_tokens, 0) * (cache_write_cost - input_cost) - prompt_caching: Final = read_discount - write_premium + usage: Final = _usage_from_spend_log(usage_object) + basis: Final = _pricing_basis(cost_breakdown) + billed_at_datetime: Final = _coerce_billed_at(billed_at) + prompt_caching: Final = ( + calculate_prompt_caching_savings( + model_info=pricing, + usage=usage, + custom_llm_provider=identity.provider if identity else custom_llm_provider, + service_tier=basis.service_tier, + data_residency=basis.data_residency, + vertex_location=basis.vertex_location, + billed_at=billed_at_datetime, + ) + if pricing is not None and usage is not None + else 0.0 + ) gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0 # The figure the logging path recorded wins, before the usage gate on purpose: a row diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b4e4dfeae67..ddf31cb1d8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -152,7 +152,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.repositories.budget_repository import BudgetRepository @@ -924,7 +924,13 @@ class ProxyLogging: "incoming_bearer_token": kwargs.get("incoming_bearer_token"), "metadata": {"headers": kwargs.get("headers") or {}}, } - + user_api_key_auth: Final = kwargs.get("user_api_key_auth") + if isinstance(user_api_key_auth, UserAPIKeyAuth): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_auth, + data=synthetic_data, + metadata_variable_name="metadata", + ) return synthetic_data def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None: @@ -7206,7 +7212,19 @@ def get_custom_url(request_base_url: str, route: str | None = None) -> str: else: base_url = request_base_url - server_root_path: Final = get_server_root_path() + # get_request_root_path() returns the prefix the router is actually + # resolving this request under: the matched SERVER_ROOT_PATHS entry when + # PerRequestRootPathMiddleware ran, otherwise the SERVER_ROOT_PATH scalar. + # This keeps the emitted URL under one prefix — the one the client called — + # instead of stacking the scalar onto a request already living under a + # dynamic prefix (which would produce /tenant-a/legacy/... — a path that + # doesn't exist). join_paths()'s tail-dedup then collapses the append when + # base_url (i.e. request.base_url) already ends in the same prefix. + from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports utils + get_request_root_path, + ) + + server_root_path: Final = get_request_root_path() if route is not None: if server_root_path != "": # First join base_url with server_root_path, then with route diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index f2070e6604c..8363aaee99a 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -161,6 +161,8 @@ async def can_user_access_vector_store( this vector store id. 5. The caller's team_id matches the vector store's team_id. + A dashboard session credential is evaluated against the same effective + contexts as listing (its own grants plus each real team of the user). Otherwise access is denied. """ if _is_proxy_admin(user_api_key_dict): @@ -169,7 +171,8 @@ async def can_user_access_vector_store( if vector_store.get("team_id") is None: return True - return await _is_vector_store_granted(vector_store, user_api_key_dict) + auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict) + return await _is_vector_store_granted_to_any(vector_store, auth_contexts) async def _is_vector_store_granted( @@ -219,7 +222,7 @@ async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> ) -async def _vector_store_listing_auth_contexts( +async def _vector_store_auth_contexts( user_api_key_dict: UserAPIKeyAuth, ) -> tuple[UserAPIKeyAuth, ...]: if not is_ui_session_credential(user_api_key_dict): @@ -250,7 +253,7 @@ async def filter_listable_vector_stores( if _is_proxy_admin(user_api_key_dict): return tuple(vector_stores) - auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict) + auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict) return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)]) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 540f197044f..d24eb8ffc62 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -4,7 +4,6 @@ Model repository for database operations on LiteLLM_ProxyModelTable. import json from collections.abc import Mapping, Sequence -from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.model import LiteLLM_ProxyModelTable @@ -109,7 +108,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]: """Find every model except the row currently being updated.""" records: Final = await self.table.find_many( - where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})}) + where={"model_id": {"not": model_id}} # mutable-ok: Prisma requires plain dicts for query serialization ) return tuple(self._to_model_list(records)) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 3ca7b0503bf..540d492beec 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -25,9 +25,15 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, SpecialEnums, Usage, + text_tokens_without_nested_reasoning, ) +def _output_token_detail(details: object, field: str) -> int | None: + value: Final = getattr(details, field, None) + return value if isinstance(value, int) else None + + def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything return isinstance(value, list) @@ -1137,11 +1143,22 @@ class ResponseAPILoggingUtils: response_api_usage, "output_tokens_details", None ) if output_tokens_details: + reasoning_tokens: Final = _output_token_detail(output_tokens_details, "reasoning_tokens") + image_tokens: Final = _output_token_detail(output_tokens_details, "image_tokens") + audio_tokens: Final = _output_token_detail(output_tokens_details, "audio_tokens") + reported_text_tokens: Final = _output_token_detail(output_tokens_details, "text_tokens") completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None), - image_tokens=getattr(output_tokens_details, "image_tokens", None), - text_tokens=getattr(output_tokens_details, "text_tokens", None), - audio_tokens=getattr(output_tokens_details, "audio_tokens", None), + reasoning_tokens=reasoning_tokens, + image_tokens=image_tokens, + text_tokens=None + if reported_text_tokens is None + else text_tokens_without_nested_reasoning( + completion_tokens=completion_tokens, + text_tokens=reported_text_tokens, + reasoning_tokens=reasoning_tokens or 0, + other_modality_tokens=(audio_tokens or 0) + (image_tokens or 0), + ), + audio_tokens=audio_tokens, ) extra_usage_fields: Final = { diff --git a/litellm/router.py b/litellm/router.py index 3f450661946..95cabfad4bd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -403,6 +403,10 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: + return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else () + + def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: """ Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still @@ -612,6 +616,38 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +class FallbackAwareStreamWrapper(CustomStreamWrapper): + """Base for the Router's chat-completion stream wrappers, which are built around the + attempt the Router picked first and have to repoint themselves when a fallback takes over.""" + + fallback_headers_adopted: bool = False + + def adopt_fallback_response_headers( + self, + fallback_response: object, + prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], + ) -> None: + """Repoint this wrapper at the deployment that served the stream. + + Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit + counters, `model_id` and `api_base` cannot reach the proxy's response headers or + its callbacks. + """ + self._response_headers = getattr(fallback_response, "_response_headers", None) + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + if fallback_hidden_params: + self._hidden_params = { # mutable-ok: the rest of litellm writes into _hidden_params + **fallback_hidden_params, + # dict() because add_retry_fallback_headers mutates additional_headers in place + "additional_headers": dict(fallback_headers), # mutable-ok: see above + } + self._base_hidden_params = { # mutable-ok: CustomStreamWrapper keeps this snapshot as a dict + **self._hidden_params, + "response_cost": None, + } + self.fallback_headers_adopted = True + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -2572,6 +2608,18 @@ class Router: return fallback_hidden_params, {} return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + @staticmethod + def _adopt_fallback_response_headers( + wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", + fallback_response: object, + ) -> tuple[dict[str, object], dict[str, object]]: + """Repoint the wrapper at `fallback_response`, returning its prepared hidden params.""" + prepared: Final = Router._prepare_fallback_hidden_params(fallback_response) + adopting_wrapper: Final = wrapper_ref() + if adopting_wrapper is not None: + adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) + return prepared + @staticmethod def _apply_fallback_hidden_params_to_item( fallback_item: object, @@ -2611,7 +2659,7 @@ class Router: held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() - class FallbackStreamWrapper(CustomStreamWrapper): + class FallbackStreamWrapper(FallbackAwareStreamWrapper): def __init__(self, async_generator: AsyncGenerator): # Copy attributes from the original model_response super().__init__( @@ -2619,6 +2667,7 @@ class Router: model=model_response.model, custom_llm_provider=model_response.custom_llm_provider, logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), ) self._async_generator = async_generator inner_chunks: Final[object] = getattr(model_response, "chunks", None) @@ -2695,8 +2744,17 @@ class Router: # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( + wrapper_ref, fallback_response + ) + fallback_headers_are_settled = False async for fallback_item in fallback_response: + if not fallback_headers_are_settled: + fallback_headers_are_settled = True # rebind-ok: one-shot latch + # a fallback that failed over again only repoints itself once it yields + prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields + Router._adopt_fallback_response_headers(wrapper_ref, fallback_response) + ) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -2738,7 +2796,11 @@ class Router: e, ) - return FallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) + # weak, so the generator closing over it does not keep the wrapper out of + # refcount teardown and delay the `finally` that releases the deployment slot + wrapper_ref: Final = weakref.ref(wrapped_response) + return wrapped_response @staticmethod def _extract_partial_responses_usage( @@ -3167,13 +3229,14 @@ class Router: """ from litellm.exceptions import MidStreamFallbackError - class SyncFallbackStreamWrapper(CustomStreamWrapper): + class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper): def __init__(self, sync_generator: Generator): super().__init__( completion_stream=sync_generator, model=model_response.model, custom_llm_provider=model_response.custom_llm_provider, logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), ) self._sync_generator = sync_generator if hasattr(model_response, "_hidden_params"): @@ -3229,8 +3292,17 @@ class Router: ) if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) + prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( + wrapper_ref, fallback_response + ) + fallback_headers_are_settled = False for fallback_item in fallback_response: + if not fallback_headers_are_settled: + fallback_headers_are_settled = True # rebind-ok: one-shot latch + # a fallback that failed over again only repoints itself once it yields + prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields + Router._adopt_fallback_response_headers(wrapper_ref, fallback_response) + ) Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item @@ -3268,7 +3340,10 @@ class Router: close_err, ) - return SyncFallbackStreamWrapper(stream_with_fallbacks()) + wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) + # weak, for the same reason as the async twin + wrapper_ref: Final = weakref.ref(wrapped_response) + return wrapped_response async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): """ @@ -4449,7 +4524,7 @@ class Router: self.fail_calls[model_name] += 1 raise e - async def aspeech(self, model: str, input: str, voice: str, **kwargs): + async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): """ Example Usage: @@ -4501,7 +4576,7 @@ class Router: ) raise e - async def _aspeech(self, model: str, input: str, voice: str, **kwargs): + async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): model_name: Final = model try: verbose_router_logger.debug("Inside _aspeech()- model: %s; kwargs: %s", model, kwargs) @@ -4525,7 +4600,7 @@ class Router: **{ **data, "input": input, - "voice": voice, + "voice": data.get("voice") if voice is None else voice, "client": model_client, **kwargs, } @@ -7458,6 +7533,21 @@ class Router: Context_Policy_Fallbacks={content_policy_fallbacks}", ) + @staticmethod + def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]: + failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) + status_code: Final = getattr(exception, "status_code", None) + if not failed_deployment_id or not isinstance(status_code, int): + return () + if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error + return () + already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped) + skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) + verbose_router_logger.debug( + "Retry skips deployments that already answered %s to this request: %s", status_code, skipped + ) + return skipped + @tracer.wrap() async def async_function_with_retries(self, *args, **kwargs): verbose_router_logger.debug("Inside async function with retries.") @@ -7553,6 +7643,12 @@ class Router: ## LOGGING if num_retries > 0: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) + first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry( + exception=original_exception, + already_skipped=kwargs.get("_retry_skipped_deployment_ids"), + ) + if first_skipped_ids: + kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it else: raise @@ -7622,6 +7718,12 @@ class Router: except Exception: raise e + skipped_ids = self._deployment_ids_to_skip_on_retry( + exception=e, + already_skipped=kwargs.get("_retry_skipped_deployment_ids"), + ) + if skipped_ids: + kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, @@ -12454,7 +12556,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - healthy_deployments = litellm.utils._get_order_filtered_deployments( + healthy_deployments = litellm.utils.get_order_filtered_deployments( cast(list[dict], healthy_deployments), target_order=_target_order ) @@ -12462,11 +12564,24 @@ class Router: ## this request via weighted-failover. Always honored, regardless of the ## router-level flag, so a stale exclusion key on kwargs cannot escape. _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( cast(list[dict], healthy_deployments), excluded_deployment_ids=_excluded_deployment_ids, ) + ## RETRY SKIP ## -> drop deployments that already refused this request with a + ## non-retryable status, unless that leaves nothing, so the caller still gets + ## the provider's own error instead of a no-deployments error. + _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None + ) + healthy_deployments = ( + litellm.utils.get_excluded_filtered_deployments( + healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids + ) + or healthy_deployments + ) + if len(healthy_deployments) == 0: exception: Final = await async_raise_no_deployment_exception( litellm_router_instance=self, @@ -13359,7 +13474,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - healthy_deployments = litellm.utils._get_order_filtered_deployments( + healthy_deployments = litellm.utils.get_order_filtered_deployments( healthy_deployments, target_order=_target_order ) @@ -13367,11 +13482,22 @@ class Router: ## this request via weighted-failover. See async counterpart in ## async_get_healthy_deployments for details. _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments = litellm.utils.get_excluded_filtered_deployments( healthy_deployments, excluded_deployment_ids=_excluded_deployment_ids, ) + ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. + _retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids( + request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None + ) + healthy_deployments = ( + litellm.utils.get_excluded_filtered_deployments( + healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids + ) + or healthy_deployments + ) + if len(healthy_deployments) == 0: model_ids = self.get_model_ids(model_name=model) _cooldown_time = self.cooldown_cache.get_min_cooldown( diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index d8b9b5ea8b2..7d4497fb6f7 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2835,8 +2835,9 @@ class ComplexityRouter(CustomLogger): where the prompt never arrives as messages. Probed on a COPY of request_kwargs because the owner pops routing bookkeeping off the - dict it is handed (`_target_order`, `_excluded_deployment_ids`), and this is a - speculative question about a model that may never be picked. + dict it is handed (`_target_order`, `_excluded_deployment_ids`, + `_retry_skipped_deployment_ids`), and this is a speculative question about a model + that may never be picked. Every way the owner says "nothing here can serve this" is a negative verdict: no healthy deployment for the group at all (BadRequestError, which ContextWindowExceededError diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py index 1e6412856a6..b82269d5824 100644 --- a/litellm/router_utils/auto_router_tuning_baseline.py +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -10,7 +10,7 @@ from pydantic import ValidationError from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig -TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline" +TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline_v2" HEURISTIC_V1_TUNING_FIELDS: Final = ( "tiers", @@ -29,6 +29,8 @@ HEURISTIC_V1_TUNING_FIELDS: Final = ( "keyword_tier_rules", ) +_TUNING_FIELD_SET: Final = frozenset(HEURISTIC_V1_TUNING_FIELDS) + _V1_SCORING_CLASSIFIER_TYPES: Final = frozenset({"heuristic", "heuristic_first", "hybrid"}) _AUTO_ROUTER_COMPLEXITY_PREFIX: Final = "auto_router/complexity_router" _EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) @@ -40,12 +42,16 @@ def _mapping(value: object) -> Mapping[str, object]: def tuning_fingerprint(complexity_router_config: object) -> str | None: - """Digest of normalized heuristic-v1 tuning fields, or None when the config is invalid.""" + """Digest of explicitly supplied heuristic-v1 tuning fields, normalized, or None when invalid.""" + raw: Final = _mapping(complexity_router_config) try: - validated: Final = ComplexityRouterConfig.model_validate(_mapping(complexity_router_config)) + validated: Final = ComplexityRouterConfig.model_validate(raw) except ValidationError: return None - payload: Final = validated.model_dump(mode="json", include=frozenset(HEURISTIC_V1_TUNING_FIELDS)) + supplied: Final = ((_TUNING_FIELD_SET - frozenset(("tier_model_configs",))) & frozenset(raw)) | ( + frozenset(("tier_model_configs",)) if validated.tier_model_configs else frozenset() + ) + payload: Final = validated.model_dump(mode="json", include=supplied) return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 9185d901a28..7b145c15a07 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -10,8 +10,8 @@ opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config r UnsupportedParamsError without an explicit true. xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at -all: every entry carrying supports_max_reasoning_effort is Claude-family, and -anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort +all: outside the gpt-6-astra rows every entry carrying supports_max_reasoning_effort is Claude-family, +and anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an explicit flag is the only signal that the tier is a real one rather than litellm rounding the level to a budget, and a missing flag costs advisory metadata rather than a rejected request. diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index c599667ab17..674bd8847f7 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -247,8 +247,7 @@ def rust_chat_completions_accepts( return False if stream: return False - request_override: Final = litellm_params.get("rust") if litellm_params is not None else None - if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None): + if not rust_enabled(): return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 515ab6edef1..5582027bb5d 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -1,13 +1,11 @@ from __future__ import annotations import os -import warnings from typing import Final DEFAULT_RUST_ENABLED: Final = False _TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) _GLOBAL_ENV_NAME: Final = "LITELLM_RUST" -_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR" class _RustConfiguration: @@ -26,49 +24,24 @@ def _parse_env_bool(value: str | None) -> bool | None: def resolve_rust_enabled( *, - request_override: bool | None, process_override: bool | None, environment_override: bool | None, - legacy_environment_override: bool | None = None, release_default: bool = DEFAULT_RUST_ENABLED, ) -> bool: - if request_override is not None: - return request_override if process_override is not None: return process_override if environment_override is not None: return environment_override - if legacy_environment_override is not None: - return legacy_environment_override return release_default -def rust_enabled(*, request_override: bool | None = None) -> bool: - if request_override is not None: - return request_override - process_override: Final = _CONFIGURATION.override - if process_override is not None: - return process_override - global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) - legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME)) - if legacy_override is not None: - warnings.warn( - f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead", - DeprecationWarning, - stacklevel=2, - ) +def rust_enabled() -> bool: return resolve_rust_enabled( - request_override=None, - process_override=None, - environment_override=global_override, - legacy_environment_override=legacy_override, + process_override=_CONFIGURATION.override, + environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), ) -def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: - return rust_enabled(request_override=request_override) - - def reset_rust_configuration() -> None: _CONFIGURATION.override = None diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 86038438f57..b7fdb5a98ef 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -7,12 +7,9 @@ from typing import Final, Protocol, cast # noqa: TID251 # native extension exp import httpx -from litellm.rust_bridge import configuration as _configuration +from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds -rust_ocr_enabled = _configuration.rust_ocr_enabled -rust = _configuration.rust - class RustOcr(Protocol): def __call__( @@ -44,49 +41,24 @@ class RustAocr(Protocol): raise NotImplementedError -class _Unset: - pass +def _as_ocr(value: object) -> RustOcr | None: + return cast(RustOcr, value) if callable(value) else None -_UNSET: Final[_Unset] = _Unset() +def _as_aocr(value: object) -> RustAocr | None: + return cast(RustAocr, value) if callable(value) else None -_rust_ocr_impl: RustOcr | None = None -_rust_aocr_impl: RustAocr | None = None - - -def set_rust_ocr( - *, - ocr: RustOcr | None | _Unset = _UNSET, - aocr: RustAocr | None | _Unset = _UNSET, -) -> None: - global _rust_ocr_impl, _rust_aocr_impl - if not isinstance(ocr, _Unset): - _rust_ocr_impl = ocr - if not isinstance(aocr, _Unset): - _rust_aocr_impl = aocr +_OCR: Final = NativeBinding("ocr", validate=_as_ocr) +_AOCR: Final = NativeBinding("aocr", validate=_as_aocr) 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: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustOcr, native_bridge.ocr) + return _OCR.load() 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: Final = get_native_bridge() - if native_bridge is None: - return None - return cast(RustAocr, getattr(native_bridge, "aocr", None)) + return _AOCR.load() def ocr( diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 3d71f6f8a50..6c81786accd 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -56,7 +56,6 @@ _STATE: Final = _RustTranscriptionState() def configure_rust_transcription( - enabled: bool = True, *, transcription: RustTranscription | None | _Unset = _UNSET, atranscription: RustAtranscription | None | _Unset = _UNSET, diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 002419dbad4..71fd78f11a3 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -182,10 +182,14 @@ def create_skill( if extra_body: create_request.update(extra_body) - # Route to LiteLLM DB if custom_llm_provider="litellm_proxy" + # Route to LiteLLM DB if custom_llm_provider="litellm_proxy". description/instructions + # arrive as top-level kwargs from the REST form endpoint, or nested in extra_body from + # the SDK convention used by other providers' create_request above. if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: return _get_litellm_skills_handler().create_skill_handler( display_title=display_title, + description=kwargs.get("description") or (extra_body.get("description") if extra_body else None), + instructions=kwargs.get("instructions") or (extra_body.get("instructions") if extra_body else None), files=files, metadata=_get_skill_request_metadata(kwargs, extra_body), user_id=kwargs.get("user_id"), diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ef28181eba5..02dee40f2a3 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -778,6 +778,9 @@ class ContentFilterConfigModel(BaseModel): ) +MCP_SECURITY_ON_VIOLATION: Final = frozenset({"block", "alert"}) + + class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails api_key: str | None = Field(default=None, description="API key for the guardrail service") api_base: str | None = Field(default=None, description="Base URL for the guardrail service API") @@ -886,9 +889,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up default=None, description="For /v1/realtime sessions: automatically close the session after this many guardrail violations.", ) - on_violation: Literal["warn", "end_session"] | None = Field( + on_violation: Literal["warn", "end_session", "block", "alert"] | None = Field( default=None, - description="For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + description=( + "For /v1/realtime sessions: 'warn' speaks the violation message and continues; " + "'end_session' speaks the message and closes the connection. " + "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning." + ), ) realtime_violation_message: str | None = Field( default=None, @@ -1093,6 +1100,15 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o except (TypeError, ValueError) as e: raise ValueError(f"timeout must be numeric, got {v!r}") from e + @model_validator(mode="after") + def validate_on_violation_for_guardrail(self) -> "LitellmParams": + if ( + self.on_violation in MCP_SECURITY_ON_VIOLATION + and self.guardrail != SupportedGuardrailIntegrations.MCP_SECURITY.value + ): + raise ValueError(f"on_violation={self.on_violation!r} is only supported by guardrail='mcp_security'") + return self + def __init__(self, **kwargs) -> None: default_on: Final = kwargs.pop("default_on", None) if default_on is not None: diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py index 51eefe7154f..4b27f9b17ef 100644 --- a/litellm/types/llms/anthropic_skills.py +++ b/litellm/types/llms/anthropic_skills.py @@ -57,6 +57,15 @@ class Skill(BaseModel): updated_at: str """ISO 8601 timestamp of when the skill was last updated""" + description: str | None = None + """Description of the skill. Populated for the LiteLLM-hosted registry + (custom_llm_provider="litellm_proxy"); Anthropic's list endpoint does not + return a description, so this is None there.""" + + search_score: float | None = None + """Semantic similarity to the ``query`` passed to ``GET /v1/skills``. None + unless a query was given.""" + class ListSkillsResponse(BaseModel): """Response from listing skills""" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py index 94f8161f44e..29f1b4bdcd6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py @@ -16,6 +16,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel): default=True, description="Whether file sanitization timeouts allow the original file through instead of blocking the request.", ) + block_on_file_modify: bool = Field( + default=True, + description="Whether a file sanitization `modify` verdict blocks the request instead of replacing the file content.", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/litellm/types/router.py b/litellm/types/router.py index 2dad22751de..0db482d8a58 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -307,7 +307,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): """ custom_llm_provider: str | None = None - rust: bool | None = None tpm: int | None = None rpm: int | None = None itpm: int | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 78ef6edfb19..9b5fb08a45f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -957,6 +957,14 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { CallTypes.agenerate_content_stream, CallTypes.generate_content_stream, ], + "/v1beta/models/{model}:generateContent": ( + CallTypes.agenerate_content, + CallTypes.generate_content, + ), + "/v1beta/models/{model}:streamGenerateContent": ( + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ), # MCP (Model Context Protocol) "/mcp/call_tool": [CallTypes.call_mcp_tool], # A2A (Agent-to-Agent) @@ -964,12 +972,12 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { "/a2a/{agent_id}/message/send": [CallTypes.asend_message, CallTypes.send_message], # Passthrough endpoints "/llm_passthrough": [ - CallTypes.llm_passthrough_route, CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, ], "/v1/llm_passthrough": [ - CallTypes.llm_passthrough_route, CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, ], "/v1/messages": [CallTypes.anthropic_messages], # OCR @@ -1625,6 +1633,17 @@ class Choices(SafeAttributeModel, OpenAIObject): setattr(self, key, value) +def text_tokens_without_nested_reasoning( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, +) -> int: + reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens + nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0)) + return text_tokens - nested_reasoning_tokens + + class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for older openai versions text_tokens: int | None = None """Text tokens generated by the model.""" @@ -3053,6 +3072,7 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): llm_provider: str | None traceback: str | None error_message: str | None + error_provider_request_id: ReadOnly[str | None] # error_rate_limit_category: # For 429 / rate-limit errors, the source of the rate limit. One of the # string values defined by `litellm.exceptions.RateLimitErrorCategory` diff --git a/litellm/utils.py b/litellm/utils.py index bc2f4a86f12..b2477d344dc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2805,6 +2805,13 @@ def supports_reasoning(model: str, custom_llm_provider: str | None = None) -> bo return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_reasoning") +def supports_none_reasoning_effort(model: str, custom_llm_provider: str | None = None) -> bool: + """ + Check if the given model accepts reasoning effort "none" and return a boolean value. + """ + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_none_reasoning_effort") + + def supports_native_structured_output(model: str, custom_llm_provider: str | None = None) -> bool: """ Check if the given model supports native structured outputs and return a boolean value. @@ -4889,7 +4896,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None: return order -def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: +def get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: if target_order is not None: return [d for d in healthy_deployments if _get_deployment_order(d) == target_order] @@ -4908,7 +4915,7 @@ def _get_order_filtered_deployments(healthy_deployments: list[dict], target_orde return healthy_deployments -def _get_excluded_filtered_deployments( +def get_excluded_filtered_deployments( healthy_deployments: list[dict], excluded_deployment_ids: Iterable[str] | None = None, ) -> list: @@ -4919,10 +4926,12 @@ def _get_excluded_filtered_deployments( across the remaining deployments in the same model group after one of them has failed. - If the filter would leave no deployments, an empty list is returned so the - caller raises its usual no-deployments error and the weighted-failover - helper falls through to the cross-group fallback path. Returning the - original unfiltered list here would re-include the just-failed deployment. + If the filter would leave no deployments, an empty list is returned and the + caller decides what that means. Weighted failover lets it raise the usual + no-deployments error and fall through to the cross-group fallback path; the + retry skip in `async_get_healthy_deployments` deliberately falls back to the + unfiltered list, so a request every deployment refused still comes back with + the provider's own error rather than a no-deployments one. """ if not excluded_deployment_ids: return healthy_deployments @@ -5210,13 +5219,16 @@ def _strip_openai_finetune_model_name(model_name: str) -> str: input: ft:gpt-3.5-turbo:my-org:custom_suffix:id output: ft:gpt-3.5-turbo + input: ft:gpt-4o-2024-08-06:my-org::id (OpenAI leaves the suffix empty when none was set) + output: ft:gpt-4o-2024-08-06 + Args: model_name (str): The full model name Returns: str: The stripped model name """ - return re.sub(r"(:[^:]+){3}$", "", model_name) + return re.sub(r"(:[^:]*){3}$", "", model_name) def _strip_model_name(model: str, custom_llm_provider: str | None) -> str: @@ -9451,6 +9463,12 @@ class ProviderConfigManager: ) return MinimaxTextToSpeechConfig() + elif litellm.LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + ) + + return MistralTextToSpeechConfig() elif litellm.LlmProviders.AWS_POLLY == provider: from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2c55828c1e2..b1ffc1583e4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3485,6 +3485,55 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://ai.azure.com/catalog/models/gpt-6-astra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, @@ -7189,7 +7238,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -7455,7 +7504,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -30742,6 +30791,9 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], "source": "https://developers.openai.com/api/docs/models/chat-latest", "supported_endpoints": [ "/v1/chat/completions", @@ -30759,6 +30811,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -34898,9 +34951,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-latest": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" @@ -37176,6 +37229,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37216,6 +37270,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37427,6 +37482,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -37467,6 +37523,7 @@ "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -56467,9 +56524,9 @@ "supports_audio_input": true }, "mistral/voxtral-mini-tts-2603": { + "input_cost_per_character": 1.6e-05, "litellm_provider": "mistral", "mode": "audio_speech", - "output_cost_per_character": 1.6e-05, "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", "supported_endpoints": [ "/v1/audio/speech" diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index d38a0eee3de..f245803408c 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -12,6 +12,10 @@ # - litellm/ Python -> `make lint` (test-linting.yml's lint job) # - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) # + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) +# - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py, +# scripts/test_quality_gate.py +# -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's +# test-tree ruff and test-quality budget steps) # - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) # - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml) # @@ -88,15 +92,14 @@ existing_files() { litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e/.*\.py$' +test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/(check_test_quality|test_quality_gate)\.py)$' spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' -# CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or -# scripts-only commit can't turn it red; scope the trigger there to skip the slow -# make lint when it couldn't catch anything. litellm_py_files=$(scope_match "$litellm_py_pattern") e2e_py_files=$(scope_match "$e2e_py_pattern") +test_tree_files=$(scope_match "$test_tree_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types @@ -136,6 +139,7 @@ if [ -n "$staged" ]; then } warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" + warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_pattern" "$test_tree_files" warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi @@ -288,6 +292,15 @@ if [ -n "$spec_files" ]; then set +m fi +if [ -n "$test_tree_files" ] && [ -z "$litellm_py_files" ]; then + echo "check: linting the test tree (ruff check --config ruff-tests.toml tests)" + uv run --no-sync ruff check --config ruff-tests.toml tests \ + || { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; } + echo "check: checking the test-quality budget (make lint-test-quality)" + make lint-test-quality \ + || { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; } +fi + if [ -n "${python_pid:-}" ]; then wait "$python_pid" || status=1 cat "$python_log"; rm -f "$python_log" @@ -313,10 +326,12 @@ summary_item() { echo "check: summary" summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope" summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope" +summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_files" \ + "no tests/ Python files or test-tree lint inputs in scope" summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope" summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope" -if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then +if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2 printf '%s\n' "$scope" | sed 's/^/ /' >&2 echo " A pass here is a no-op, not a lint verdict." >&2 diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py index 7d34b194f1c..4f4eeb17ec1 100644 --- a/scripts/test_quality_gate.py +++ b/scripts/test_quality_gate.py @@ -10,17 +10,12 @@ base. Every rule is seeded at exactly its count on the day the gate landed, so the suite's existing debt is grandfathered and any net-new violation trips the gate -immediately. ``--update`` ratchets a limit down by the violations this branch -fixed relative to its branch point (the merge-base), so the ceilings only ever -fall. Base counts are measured with the *current* checker, so a rule introduced -on this branch is counted at the base too and ratchets like every other one. - -Only ever falling is not the same as always falling, so the gate enforces the -second half: a branch that clears violations and leaves the ceiling above its -new count fails, naming the rules and telling the author to run -``make lint-budget-update``. Without that, a removed violation could come back -later under a ceiling nobody lowered. Drift already in the base is never -blamed, so this fires only on the branch that did the clearing. +immediately. ``--update`` ratchets a limit down by the violations fixed relative +to ``--base``, so the ceilings only ever fall. Base counts are measured with the +*current* checker, so a rule introduced on this branch is counted at the base too +and ratchets like every other one. The ratchet runs as a scheduled automation +against litellm_internal_staging, not on PR branches, so concurrent PRs never +race to edit the same limit. The deliberate difference from its sibling: this gate has no headroom anywhere. Type discipline seeded LIT010/LIT011 at 1.5x to leave room for an in-flight @@ -34,13 +29,14 @@ import argparse import json import re import shutil +import signal import subprocess import sys import tempfile from collections import Counter from collections.abc import Mapping, Sequence from pathlib import Path -from types import MappingProxyType +from types import FrameType, MappingProxyType from typing import Final, NamedTuple REPO_ROOT: Final = Path(__file__).resolve().parent.parent @@ -48,6 +44,7 @@ CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py" BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json" TARGET: Final = "tests" DEFAULT_BASE: Final = "origin/litellm_internal_staging" +TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP) _HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE) _FILE_HEADER: Final = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE) @@ -116,22 +113,33 @@ def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]: return MappingProxyType(dict(Counter(v.code for v in violations))) -def base_counts(ref: str) -> Mapping[str, int]: +def _exit_on_termination(signum: int, _frame: FrameType | None) -> None: + raise SystemExit(128 + signum) + + +def _install_termination_handlers() -> None: + for termination in TERMINATION_SIGNALS: + if signal.getsignal(termination) == signal.SIG_DFL: + signal.signal(termination, _exit_on_termination) + + +def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]: """Rule counts at `ref`, measured with the *current* rule logic rather than whatever the checker looked like at that commit.""" + _install_termination_handlers() parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_")) worktree: Final = parent / "wt" try: - _run(["git", "worktree", "add", "--detach", str(worktree), ref]) + _run(["git", "worktree", "add", "--detach", str(worktree), ref], cwd=repo_root) (worktree / "scripts").mkdir(parents=True, exist_ok=True) - checker: Final = worktree / "scripts" / "check_test_quality.py" - shutil.copy(CHECKER, checker) - return count_by_rule(_check(worktree, checker)) + base_checker: Final = worktree / "scripts" / "check_test_quality.py" + shutil.copy(checker, base_checker) + return count_by_rule(_check(worktree, base_checker)) finally: # Teardown must never raise, or it masks the real error when the body failed. subprocess.run( ["git", "worktree", "remove", "--force", str(worktree)], - cwd=REPO_ROOT, capture_output=True, text=True, + cwd=repo_root, capture_output=True, text=True, ) shutil.rmtree(parent, ignore_errors=True) @@ -144,21 +152,6 @@ def over_ceiling(head: Mapping[str, int], budget: Mapping[str, Mapping[str, int] ) -def unratcheted( - head: Mapping[str, int], - base: Mapping[str, int], - budget: Mapping[str, Mapping[str, int]], -) -> tuple[Breach, ...]: - """Rules this branch cleared without lowering the ceiling behind them. Requires - both `head < base`, so drift already in the base is never blamed on this change, - and `head < limit`, so a ceiling already at the count is left alone.""" - return tuple(sorted( - Breach(rule, head.get(rule, 0), spec["limit"], head.get(rule, 0) - base.get(rule, 0)) - for rule, spec in budget.items() - if head.get(rule, 0) < base.get(rule, 0) and head.get(rule, 0) < spec["limit"] - )) - - def evaluate( head: Mapping[str, int], base: Mapping[str, int], @@ -198,38 +191,15 @@ def introduced( return tuple(v for v in violations if v.line in changed.get(v.file, frozenset())) -def touches_measured_tree(base_point: str) -> bool: - """Whether this branch changed anything that can move a count. A branch that - touches neither the test tree nor the checker cannot have cleared a violation, - so the base scan is skipped and the gate stays cheap on the common change.""" - changed: Final = _run( - ["git", "diff", "--name-only", base_point, "--", TARGET, str(CHECKER.relative_to(REPO_ROOT))] - ) - return bool(changed.strip()) - - def cmd_check(base: str) -> None: budget: Final = json.loads(BUDGET_PATH.read_text()) head: Final = head_violations() head_counts: Final = count_by_rule(head) - base_point: Final = resolve_base_point(base) - if not over_ceiling(head_counts, budget) and not touches_measured_tree(base_point): + if not over_ceiling(head_counts, budget): print(f"OK: every TQ rule is within its test-suite ceiling (base {base})") return + base_point: Final = resolve_base_point(base) base_at_point: Final = base_counts(base_point) - stale: Final = unratcheted(head_counts, base_at_point, budget) - if stale: - print(f"FAIL: TQ-rule limits were left above the count this branch reached (base {base}):") - for breach in stale: - print( - f" {breach.rule}: this branch cleared {-breach.added} down to {breach.total}, " - f"but the limit is still {breach.cap}" - ) - print( - "Run `make lint-budget-update` and commit the lowered limits, so the " - "violations you cleared cannot come back under a ceiling nobody moved." - ) - raise SystemExit(1) breaches: Final = evaluate(head_counts, base_at_point, budget) if not breaches: print(f"OK: every TQ rule is within its test-suite ceiling (base {base})") diff --git a/tests/code_coverage_tests/check_workflow_job_name_collisions.py b/tests/code_coverage_tests/check_workflow_job_name_collisions.py new file mode 100644 index 00000000000..ae2c1d80c8f --- /dev/null +++ b/tests/code_coverage_tests/check_workflow_job_name_collisions.py @@ -0,0 +1,521 @@ +"""Catch workflow jobs that publish check runs under the same name. + +A ruleset's required status check names a check run and GitHub matches it by that +name alone. When two jobs publish the same name the required context stops +mapping to the job that proves it: the commit carries two check runs under one +name and nothing says which one the ruleset required. Both being green hides the +clash completely, so the context quietly stops meaning what the ruleset intended. +One job lands in the same place when its `name:` holds no matrix value, since +every combination it runs then reports under that one name. + +`.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while +`.github/workflows/test-mcp.yml` already published the required `test` context, +and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs. +This invariant has to be enforced here because CI cannot enforce it on itself. + +A job publishes its `name:` when it sets one, and otherwise its job id plus the +values of the combination it runs, the way GitHub writes `build (3.12)`. A name +carrying `${{ ... }}` publishes one check run per combination the matrix +produces: `exclude` rows drop combinations before `include` rows fold into the +survivors, and each `include` row's values stay together rather than crossing +with the other rows', so two shard lists that overlap collide even though their +templates read differently. Each expression is evaluated per combination over the +pieces a job name can hold: string literals, `matrix.`, `format()`, `==` and +`!=`, and the ` && || ` idiom, which is how the shards reach their +real ` / Run tests` names rather than staying opaque. + +Whatever the sweep cannot work out is left out of the comparison and reported +instead of guessed, because a guess that lands wrong fails a workflow GitHub +would have published perfectly well. A name still holding an expression once the +combination is filled in is usually one GitHub resolves per job, so it is one of +those: guessing that two jobs sharing such a template clash would fail workflows +over a context this sweep cannot read. The exception is a name whose leftover +expressions all read a `github.` property other than `github.job`, which one run +fills in the same way for every job in it, so those are compared against the +other jobs of their own workflow and stay out of the comparison across files, +where two workflows can run on different events. A matrix that is itself an +expression or that lists values which are not scalars, an `include` or `exclude` +row shaped the same way, a whole `strategy:` that comes from an expression, and a +call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind +one of them goes unseen, which leaves a merge no worse off than before this check +existed, where the opposite direction would block work that was fine. + +A job calling a local reusable workflow publishes one check run per job of the +callee, named ` / ` and chained through however many levels of +local calls it takes, which is why a caller's name never collides with a plain +job that happens to match it. A file under `.github/workflows/` that does not +read as one workflow at all is reported rather than skipped, since skipping it +silently would hide every job it holds. +""" + +import itertools +import operator +import re +import sys +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import yaml +from pydantic import BaseModel, Field, ValidationError + +REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent +WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" +EXPRESSION: Final = re.compile(r"\$\{\{(?P.*?)\}\}", re.DOTALL) +MATRIX_REF: Final = re.compile(r"^matrix\.(?P[\w-]+)$") +LITERAL: Final = re.compile(r"^'(?P[^']*)'$") +FORMAT_CALL: Final = re.compile(r"^format\((?P.*)\)$", re.DOTALL) +COMPARISON: Final = re.compile(r"^(?P.+?)\s*(?P==|!=)\s*(?P.+)$", re.DOTALL) +RUN_WIDE: Final = re.compile(r"^github\.(?!job\b)[\w.]+$") +GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}") +NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({}) +NO_CALLERS: Final[frozenset[str]] = frozenset() +SCALAR: Final = (str, int, float) +MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"}) +LOCAL_CALL_PREFIX: Final = "./" + + +@dataclass(frozen=True, slots=True) +class Unreadable: + reason: str + + +@dataclass(frozen=True, slots=True) +class Opaque: + reason: str + + +@dataclass(frozen=True, slots=True) +class Names: + """The check-run names a job publishes, beside the reasons the rest of them stay unknown.""" + + known: tuple[str, ...] = () + unknown: tuple[str, ...] = () + local: tuple[str, ...] = () + + +class Job(BaseModel): + name: object = None + uses: str | None = None + strategy: object = Field(default_factory=dict) + + +class Workflow(BaseModel): + jobs: Mapping[str, Job] = Field(default_factory=dict) + + +def scalar_text(value: object) -> str: + """A YAML scalar the way GitHub renders it, so `true` never reaches a name as `True`.""" + return str(value).lower() if isinstance(value, bool) else str(value) + + +def parse(source: str) -> tuple[Workflow, object] | Unreadable: + """The workflow plus its raw `on:` value, or why the file does not read as one.""" + try: + parsed: Final = yaml.safe_load(source) + except yaml.YAMLError: + return Unreadable("it does not read as one YAML document") + if not isinstance(parsed, dict): + return Unreadable("its top level is not a mapping of workflow keys") + try: + return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on")) + except ValidationError as error: + return Unreadable(f"{error.error_count()} of its job definitions have a shape GitHub would reject") + + +def events(raw_on: object) -> frozenset[str]: + if isinstance(raw_on, Mapping): + return frozenset(str(key) for key in raw_on) + if isinstance(raw_on, str): + return frozenset({raw_on}) + if isinstance(raw_on, Sequence): + return frozenset(str(event) for event in raw_on) + return frozenset() + + +def publishes_check_runs(raw_on: object) -> bool: + """A `workflow_call`-only workflow posts its check runs through callers, never itself.""" + return events(raw_on) != frozenset({"workflow_call"}) + + +def scalar_list(value: object) -> tuple[str, ...] | Opaque: + """One matrix key's values, or why the combinations it produces cannot be worked out.""" + if not isinstance(value, Sequence) or isinstance(value, str): + return Opaque("a matrix key holds something other than a list of values") + if any(not isinstance(item, SCALAR) for item in value): + return Opaque("a matrix key lists values that are not plain scalars") + return tuple(scalar_text(item) for item in value) + + +def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...] | Opaque: + listed: Final = tuple( + (str(key), scalar_list(values)) for key, values in matrix.items() if str(key) not in MATRIX_DIRECTIVES + ) + opaque: Final = next((values for _, values in listed if isinstance(values, Opaque)), None) + if opaque is not None: + return opaque + return tuple((key, values) for key, values in listed if not isinstance(values, Opaque)) + + +def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque: + """One `include` or `exclude` row, or why the combinations they shape cannot be worked out.""" + rows: Final = matrix.get(directive) + if rows is None: + return () + if not isinstance(rows, Sequence) or isinstance(rows, str): + return Opaque(f"a matrix `{directive}` is itself an expression rather than a list of rows") + mappings: Final = tuple(row for row in rows if isinstance(row, Mapping)) + if len(mappings) != len(rows): + return Opaque(f"a matrix `{directive}` row is not a mapping of values") + if any(not isinstance(value, SCALAR) for row in mappings for value in row.values()): + return Opaque(f"a matrix `{directive}` row holds a value that is not a plain scalar") + return tuple(MappingProxyType({str(key): scalar_text(value) for key, value in row.items()}) for row in mappings) + + +def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: + """GitHub removes a combination that carries every value one `exclude` row names.""" + return all(combination.get(key) == value for key, value in row.items()) + + +def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool: + """GitHub folds an `include` row into a combination only where it overwrites no listed value.""" + return all(combination[key] == value for key, value in row.items() if key in combination) + + +def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]]) -> Mapping[str, str]: + additions: Final = {key: value for row in rows if extends(row, combination) for key, value in row.items()} + return MappingProxyType({**combination, **additions}) + + +def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mapping[str, str], ...]: + if not listed: + return () + return tuple( + MappingProxyType(dict(zip((key for key, _ in listed), values))) + for values in itertools.product(*(values for _, values in listed)) + ) + + +def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque: + """One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does.""" + if not isinstance(job.strategy, Mapping): + return Opaque("its whole `strategy` comes from an expression") + matrix: Final = job.strategy.get("matrix") + if matrix is None: + return () + if not isinstance(matrix, Mapping): + return Opaque("the matrix itself comes from an expression") + listed: Final = listed_values(matrix) + if isinstance(listed, Opaque): + return listed + rows: Final = directive_rows(matrix, "include") + if isinstance(rows, Opaque): + return rows + dropped: Final = directive_rows(matrix, "exclude") + if isinstance(dropped, Opaque): + return dropped + kept: Final = tuple( + combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped) + ) + standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept)) + return (*(extended(combination, rows) for combination in kept), *standalone) + + +def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]: + depth, quoted = state + if char == "'": + return depth, not quoted + if quoted: + return depth, quoted + return depth + int(char == "(") - int(char == ")"), quoted + + +def split_outside(text: str, token: str) -> tuple[str, ...]: + """`text` cut on every `token` that sits outside quotes and parentheses.""" + states: Final = tuple(itertools.accumulate(text, scanned, initial=(0, False))) + cuts: Final = tuple( + index + for index in range(len(text) - len(token) + 1) + if text.startswith(token, index) and states[index] == (0, False) + ) + starts: Final = (0, *(cut + len(token) for cut in cuts)) + return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text)))) + + +def formatted(template: str, arguments: Sequence[str]) -> str | None: + """GitHub's `format()` fills `{0}`-style holes and escapes braces, so anything richer resolves to nothing.""" + residue: Final = GITHUB_PLACEHOLDER.sub("", template) + if "{" in residue or "}" in residue: + return None + try: + return template.format(*arguments) + except (IndexError, KeyError, ValueError): + return None + + +def value_of(text: str, values: Mapping[str, str]) -> str | None: + expression: Final = text.strip() + literal: Final = LITERAL.match(expression) + if literal is not None: + return literal.group("text") + reference: Final = MATRIX_REF.match(expression) + if reference is not None: + return values.get(reference.group("key")) + call: Final = FORMAT_CALL.match(expression) + if call is None: + return None + arguments: Final = tuple(value_of(part, values) for part in split_outside(call.group("args"), ",")) + resolved: Final = tuple(argument for argument in arguments if argument is not None) + if not resolved or len(resolved) != len(arguments): + return None + return formatted(resolved[0], resolved[1:]) + + +def holds(condition: str, values: Mapping[str, str]) -> bool | None: + comparison: Final = COMPARISON.match(condition.strip()) + if comparison is None: + return None + left: Final = value_of(comparison.group("left"), values) + right: Final = value_of(comparison.group("right"), values) + if left is None or right is None: + return None + return (left == right) == (comparison.group("operator") == "==") + + +def evaluate(body: str, values: Mapping[str, str]) -> str | None: + """The single string this expression yields, or None when its shape is not understood.""" + branches: Final = tuple(split_outside(alternative, "&&") for alternative in split_outside(body, "||")) + outcomes: Final = tuple(tuple(holds(part, values) for part in branch[:-1]) for branch in branches) + if any(outcome is None for branch in outcomes for outcome in branch): + return None + taken: Final = next((branch[-1] for branch, outcome in zip(branches, outcomes) if all(outcome)), None) + return None if taken is None else value_of(taken, values) + + +def resolved_span(span: re.Match[str], values: Mapping[str, str]) -> str: + substitution: Final = evaluate(span.group("body"), values) + return span.group(0) if substitution is None else substitution + + +def rendered(template: str, values: Mapping[str, str]) -> str: + return EXPRESSION.sub(lambda span: resolved_span(span, values), template) + + +def comparable(name: str) -> bool: + """A name still holding an expression is one GitHub resolves per job, so it is nothing to compare.""" + return EXPRESSION.search(name) is None + + +def run_wide(name: str) -> bool: + """A name whose leftover expressions one workflow run fills in the same way for every job in it.""" + return all(RUN_WIDE.match(span.group("body").strip()) is not None for span in EXPRESSION.finditer(name)) + + +def settled(names: Sequence[str]) -> Names: + unresolved: Final = tuple(name for name in names if not comparable(name)) + return Names( + tuple(name for name in names if comparable(name)), + tuple(f"its name stays `{name}`" for name in unresolved if not run_wide(name)), + tuple(name for name in unresolved if run_wide(name)), + ) + + +def expand(template: str, job: Job) -> Names: + combinations: Final = matrix_combinations(job) + if isinstance(combinations, Opaque): + return Names((), (combinations.reason,)) + over: Final = combinations or (NO_MATRIX,) + return settled(tuple(rendered(template, values) for values in over)) + + +def suffixed(job_id: str, combination: Mapping[str, str]) -> str: + """The name GitHub gives a job with no `name:`, its id plus the combination it runs.""" + return f"{job_id} ({', '.join(combination.values())})" if combination else job_id + + +def published_names(job_id: str, job: Job) -> Names: + if job.name is not None: + return expand(scalar_text(job.name), job) + combinations: Final = matrix_combinations(job) + if isinstance(combinations, Opaque): + return Names((), (combinations.reason,)) + suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in combinations)) + return Names(suffixes or (job_id,)) + + +def callee_path(job: Job) -> str | None: + if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX): + return None + return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0] + + +def joined(groups: Sequence[Names]) -> Names: + return Names( + tuple(name for group in groups for name in group.known), + tuple(reason for group in groups for reason in group.unknown), + tuple(name for group in groups for name in group.local), + ) + + +def tagged(names: Names) -> tuple[tuple[str, bool], ...]: + """Each name a job publishes beside whether only its own workflow's run settles it.""" + return (*((name, False) for name in names.known), *((name, True) for name in names.local)) + + +def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None: + path: Final = callee_path(job) + if path is None: + return "it calls a reusable workflow outside this repository" + if path in callers: + return f"its call to {path} loops back on itself" + return None if path in workflows else f"it calls {path}, which this checkout does not hold" + + +def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS) -> Names: + prefixes: Final = published_names(job_id, job) + if job.uses is None: + return prefixes + blocker: Final = call_blocker(job, workflows, callers) + if blocker is not None: + return Names((), (*prefixes.unknown, blocker)) + path: Final = callee_path(job) or "" + suffixes: Final = joined( + tuple( + job_names(callee_id, callee_job, workflows, callers | {path}) + for callee_id, callee_job in workflows[path].jobs.items() + ) + ) + composed: Final = tuple( + (f"{prefix} / {suffix}", prefix_local or suffix_local) + for prefix, prefix_local in tagged(prefixes) + for suffix, suffix_local in tagged(suffixes) + ) + return Names( + tuple(name for name, is_local in composed if not is_local), + (*prefixes.unknown, *suffixes.unknown), + tuple(name for name, is_local in composed if is_local), + ) + + +def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]: + parsed: Final = {rel: parse(source) for rel, source in sources.items()} + return MappingProxyType({rel: entry for rel, entry in parsed.items() if not isinstance(entry, Unreadable)}) + + +def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]: + parsed: Final = {rel: parse(source) for rel, source in sources.items()} + return tuple( + f"{rel} sits in the workflows directory but {entry.reason}, so none of its jobs were checked." + for rel, entry in sorted(parsed.items()) + if isinstance(entry, Unreadable) + ) + + +def scanned_jobs(sources: Mapping[str, str]) -> Iterator[tuple[str, str, Names]]: + parsed: Final = readable(sources) + workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()} + for rel, (workflow, raw_on) in parsed.items(): + if not publishes_check_runs(raw_on): + continue + for job_id, job in workflow.jobs.items(): + yield rel, job_id, job_names(job_id, job, workflows) + + +def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]: + for rel, job_id, names in scanned_jobs(sources): + for name in names.known: + yield name, f"{rel} job `{job_id}`" + + +def blind_spots(sources: Mapping[str, str]) -> tuple[str, ...]: + """Jobs whose published names GitHub decides at run time, which no offline sweep can compare.""" + return tuple( + f"{rel} job `{job_id}` publishes a name this check cannot work out because {reason}." + for rel, job_id, names in scanned_jobs(sources) + for reason in sorted(names.unknown) + ) + + +def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]: + for name, pairs in itertools.groupby(sorted(published(sources)), key=operator.itemgetter(0)): + yield name, tuple(owner for _, owner in pairs) + + +def clash(name: str, owners: Sequence[str]) -> str | None: + """Why one name is ambiguous, whether two jobs carry it or one job repeats it over its matrix.""" + jobs: Final = tuple(dict.fromkeys(owners)) + if len(jobs) > 1: + return ( + f"`{name}` is published by {len(jobs)} jobs: {', '.join(jobs)}. A required status check matching " + f"that name cannot say which job proves it; give one of them a distinct `name:` or job id." + ) + if len(owners) > 1: + return ( + f"`{name}` is published {len(owners)} times by {jobs[0]}, once per matrix combination. A required " + f"status check matching that name cannot say which run proves it; put a matrix value in its `name:`." + ) + return None + + +def local_published(sources: Mapping[str, str]) -> Iterator[tuple[tuple[str, str], str]]: + """Names their own workflow's run settles, keyed by the file whose run settles them.""" + for rel, job_id, names in scanned_jobs(sources): + for name in names.local: + yield (rel, name), f"job `{job_id}`" + + +def local_clash(rel: str, name: str, owners: Sequence[str]) -> str | None: + """Why one workflow's own run lands several of its jobs on one check run.""" + if len(owners) < 2: + return None + jobs: Final = tuple(dict.fromkeys(owners)) + return ( + f"`{name}` is published {len(owners)} times inside {rel}, by {', '.join(jobs)}. One run fills that " + f"expression in the same way throughout, so they all land on one check run; make the names differ." + ) + + +def local_clashes(sources: Mapping[str, str]) -> tuple[str, ...]: + grouped: Final = itertools.groupby(sorted(local_published(sources)), key=operator.itemgetter(0)) + found: Final = tuple(local_clash(rel, name, tuple(owner for _, owner in pairs)) for (rel, name), pairs in grouped) + return tuple(message for message in found if message is not None) + + +def collisions(sources: Mapping[str, str]) -> tuple[str, ...]: + found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources)) + return (*(message for message in found if message is not None), *local_clashes(sources)) + + +def workflow_sources() -> Mapping[str, str]: + """Repo-relative posix paths to text, the keys `uses: ./...` resolves against.""" + return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))} + + +def report(header: str, problems: Sequence[str]) -> None: + if problems: + print(f"ERROR: {header}:\n - " + "\n - ".join(problems), file=sys.stderr) + + +def exit_code(sources: Mapping[str, str]) -> int: + unread: Final = unreadable(sources) + found: Final = collisions(sources) + blind: Final = blind_spots(sources) + if blind: + print("NOTE: names left out of the comparison:\n - " + "\n - ".join(blind)) + report("Some workflows could not be read", unread) + report("Check-run names are not unique", found) + if unread or found: + return 1 + + print(f"Check-run names are unique across {len(sources)} workflows") + return 0 + + +def main() -> int: + return exit_code(workflow_sources()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/code_coverage_tests/test_workflow_job_name_collisions.py b/tests/code_coverage_tests/test_workflow_job_name_collisions.py new file mode 100644 index 00000000000..0f5ba43bd7a --- /dev/null +++ b/tests/code_coverage_tests/test_workflow_job_name_collisions.py @@ -0,0 +1,880 @@ +from typing import Final + +from check_workflow_job_name_collisions import ( + Unreadable, + blind_spots, + callee_path, + collisions, + exit_code, + parse, + published, + unreadable, + workflow_sources, +) + +REUSABLE_BASE: Final = """on: + workflow_call: +jobs: + run: + name: >- + ${{ matrix.python-version == '3.12' && 'Run tests' + || format('Run tests (Python {0})', matrix.python-version) }} + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] +""" + +SHARD_CALLER: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.shard }} + uses: ./.github/workflows/base.yml + strategy: + matrix: + include: + - shard: core-utils +""" + + +CORRELATED_ROWS: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.shard }} on ${{ matrix.test-path }} + runs-on: ubuntu-latest + strategy: + matrix: + include: + - shard: core-utils + test-path: tests/core + - shard: proxy + test-path: tests/proxy +""" + +LISTED_PLUS_ROW: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.python-version }} ${{ matrix.label }} + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] + include: + - label: fast +""" + +NAMELESS_MATRIX: Final = """on: pull_request +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12", "3.13"] +""" + +EXCLUDED_PAIR: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.os }}-${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + matrix: + os: [ubuntu, macos] + python-version: ["3.12", "3.13"] + exclude: + - os: macos + python-version: "3.13" +""" + +EXCLUDED_KEY: Final = """on: pull_request +jobs: + unit: + name: ${{ matrix.os }}-${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + matrix: + os: [ubuntu, macos] + python-version: ["3.12", "3.13"] + exclude: + - os: macos +""" + +BOOLEAN_MATRIX: Final = """on: pull_request +jobs: + unit: + name: cache ${{ matrix.cached }} + runs-on: ubuntu-latest + strategy: + matrix: + cached: [true, false] +""" + +UNFILLABLE_FORMAT: Final = """on: pull_request +jobs: + unit: + name: ${{ format('{0} {1}', matrix.shard) }} + runs-on: ubuntu-latest + strategy: + matrix: + shard: [core] +""" + + +def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None: + assert collisions(workflow_sources()) == () + + +def test_every_workflow_in_the_repo_parses_into_jobs() -> None: + unparsed: Final = tuple( + rel + for rel, source in workflow_sources().items() + if isinstance(entry := parse(source), Unreadable) or not entry[0].jobs + ) + + assert unparsed == () + + +def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None: + sources: Final = workflow_sources() + parsed: Final = tuple(entry for text in sources.values() if not isinstance(entry := parse(text), Unreadable)) + unresolved: Final = tuple( + job.uses + for workflow, _ in parsed + for job in workflow.jobs.values() + if (callee := callee_path(job)) is not None and callee not in sources + ) + + assert unresolved == () + + +def test_two_jobs_falling_back_to_the_same_job_id_collide() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`test` is published by 2 jobs" in found[0] + assert "a.yml job `test`" in found[0] and "b.yml job `test`" in found[0] + + +def test_an_explicit_name_overrides_the_job_id_and_clears_the_collision() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n name: Sweep tests\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_an_explicit_name_matching_another_job_id_collides() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n sweep:\n name: test\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`test` is published by 2 jobs" in found[0] + + +def test_two_callers_of_one_reusable_workflow_collide_on_a_shared_matrix_value() -> None: + base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + caller: Final = ( + "on: pull_request\n" + "jobs:\n" + " {job}:\n" + " name: ${{{{ matrix.shard }}}}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: {shard}\n" + ) + sources: Final = { + ".github/workflows/base.yml": base, + "unit.yml": caller.format(job="unit", shard="proxy-auth"), + "proxy-db.yml": caller.format(job="proxy-db", shard="proxy-auth"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`proxy-auth / Run tests` is published by 2 jobs" in found[0] + + +def test_distinct_matrix_values_through_one_reusable_workflow_do_not_collide() -> None: + base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + caller: Final = ( + "on: pull_request\n" + "jobs:\n" + " {job}:\n" + " name: ${{{{ matrix.shard }}}}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: {shard}\n" + ) + sources: Final = { + ".github/workflows/base.yml": base, + "unit.yml": caller.format(job="unit", shard="proxy-auth"), + "proxy-db.yml": caller.format(job="proxy-db", shard="budgets"), + } + + assert collisions(sources) == () + + +def test_a_reusable_caller_does_not_collide_with_a_plain_job_of_the_same_name() -> None: + sources: Final = { + ".github/workflows/base.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n" + ), + "unit.yml": ( + "on: pull_request\n" + "jobs:\n" + " unit:\n" + " name: ${{ matrix.shard }}\n" + " uses: ./.github/workflows/base.yml\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: proxy-behavior\n" + ), + "postgres.yml": ( + "on: pull_request\n" + "jobs:\n" + " postgres:\n" + " name: ${{ matrix.shard }}\n" + " runs-on: ubuntu-latest\n" + " strategy:\n" + " matrix:\n" + " include:\n" + " - shard: proxy-behavior\n" + ), + } + + assert collisions(sources) == () + + +def test_a_workflow_call_only_workflow_publishes_nothing_of_its_own() -> None: + sources: Final = { + "base.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n", + "other.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_a_workflow_call_workflow_that_also_runs_on_pull_request_still_publishes() -> None: + sources: Final = { + "base.yml": "on:\n workflow_call:\n pull_request:\njobs:\n run:\n runs-on: ubuntu-latest\n", + "other.yml": "on: pull_request\njobs:\n run:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`run` is published by 2 jobs" in found[0] + + +def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\n" + "jobs:\n" + " build:\n" + " name: Analyze (${{ matrix.language }})\n" + " runs-on: ubuntu-latest\n" + " strategy:\n" + " matrix:\n" + " language: [python, go]\n" + ), + "b.yml": "on: pull_request\njobs:\n go:\n name: Analyze (go)\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`Analyze (go)` is published by 2 jobs" in found[0] + + +def test_two_workflows_sharing_a_run_wide_template_are_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.event_name }}}}-build\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert blind_spots(sources) == () + + +def test_two_jobs_of_one_workflow_sharing_a_run_wide_template_are_a_collision() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + " two:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published 2 times inside a.yml, by job `one`, job `two`" in found[0] + assert exit_code(sources) == 1 + + +def test_a_run_wide_template_carrying_a_matrix_value_does_not_collide_inside_one_workflow() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-${{ matrix.shard }}\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n shard: [core, extras]\n" + ), + } + + assert collisions(sources) == () + assert blind_spots(sources) == () + + +def test_a_run_wide_name_repeated_over_a_matrix_by_one_job_is_a_collision() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n shard: [core, extras]\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "is published 2 times inside a.yml, by job `one`" in found[0] + + +def test_a_name_reading_the_job_it_sits_in_stays_out_of_the_comparison() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n" + " two:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n" + ), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_run_wide_caller_name_collides_through_the_workflow_it_calls() -> None: + sources: Final = { + ".github/workflows/a.yml": ( + "on: pull_request\njobs:\n" + " one:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n" + " two:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n" + ), + ".github/workflows/c.yml": "on:\n workflow_call:\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "github.event_name }} / build` is published 2 times inside .github/workflows/a.yml" in found[0] + + +def test_a_run_wide_name_inside_a_called_workflow_collides_under_the_caller() -> None: + sources: Final = { + ".github/workflows/a.yml": ("on: pull_request\njobs:\n one:\n uses: ./.github/workflows/c.yml\n"), + ".github/workflows/c.yml": ( + "on:\n workflow_call:\njobs:\n" + " build:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n" + " lint:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n" + ), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`one / ${{ github.event_name }}` is published 2 times inside .github/workflows/a.yml" in found[0] + + +def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ github.workflow }}}} / build\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert exit_code(sources) == 0 + + +def test_a_format_call_python_accepts_but_github_does_not_publishes_nothing_to_compare() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: ${{ format('{0.real}', matrix.shard) }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n" + ), + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_a_format_call_padding_its_argument_publishes_nothing_to_compare() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: ${{ format('{0:>8}', matrix.shard) }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n" + ), + "b.yml": "on: pull_request\njobs:\n two:\n name: ' core'\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_an_exclude_row_that_is_not_a_mapping_is_reported_rather_than_skipped() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n v: [1, 2]\n exclude:\n - oops\n" + ), + "b.yml": "on: pull_request\njobs:\n other:\n name: build (1)\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_an_exclude_row_holding_a_non_scalar_never_drops_every_combination() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n v: [1, 2]\n exclude:\n - cfg: {k: 1}\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + assert blind_spots(sources) != () + + +def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None: + template: Final = ( + "on: pull_request\njobs:\n {job}:\n name: ${{{{ matrix.shard }}}}\n runs-on: ubuntu-latest\n" + ) + sources: Final = { + "a.yml": template.format(job="one"), + "b.yml": template.format(job="two"), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None: + sources: Final = { + "notes.yml": "just a string\n", + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + } + + found: Final = unreadable(sources) + + assert len(found) == 1 + assert "notes.yml" in found[0] + assert collisions(sources) == () + + +def test_a_workflow_holding_a_job_shape_github_would_reject_is_reported() -> None: + sources: Final = {"a.yml": "on: pull_request\njobs:\n test:\n uses: [not, a, string]\n"} + + found: Final = unreadable(sources) + + assert len(found) == 1 + assert "a.yml" in found[0] + + +def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None: + names: Final = frozenset( + name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) + ) + + assert names == frozenset({"core-utils / Run tests", "core-utils / Run tests (Python 3.13)"}) + + +def test_a_conditional_name_never_publishes_the_branch_its_condition_rules_out() -> None: + names: Final = frozenset( + name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER}) + ) + + assert "core-utils / Run tests (Python 3.12)" not in names + + +def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_same_name() -> None: + sources: Final = { + ".github/workflows/base.yml": REUSABLE_BASE, + "unit.yml": SHARD_CALLER, + "postgres.yml": ("on: pull_request\njobs:\n legacy:\n name: core-utils / Run tests\n"), + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`core-utils / Run tests` is published by 2 jobs" in found[0] + + +def test_a_name_reading_two_matrix_keys_publishes_only_the_pairs_each_include_row_holds() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS})) + + assert names == frozenset({"core-utils on tests/core", "proxy on tests/proxy"}) + + +def test_a_name_reading_two_matrix_keys_never_publishes_a_pair_no_include_row_holds() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS})) + + assert "core-utils on tests/proxy" not in names + assert "proxy on tests/core" not in names + + +def test_an_include_row_carrying_no_listed_key_extends_every_listed_combination() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": LISTED_PLUS_ROW})) + + assert names == frozenset({"3.12 fast", "3.13 fast"}) + + +def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None: + unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name) + + assert unresolved == () + + +def test_a_matrix_job_with_no_name_publishes_the_id_and_values_github_appends() -> None: + names: Final = frozenset(name for name, _ in published({"a.yml": NAMELESS_MATRIX})) + + assert names == frozenset({"build (3.12)", "build (3.13)"}) + + +def test_a_matrix_job_with_no_name_does_not_collide_with_a_plain_job_carrying_its_id() -> None: + sources: Final = { + "a.yml": NAMELESS_MATRIX, + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert collisions(sources) == () + + +def test_a_matrix_job_with_no_name_collides_with_the_suffixed_name_github_writes() -> None: + sources: Final = { + "a.yml": NAMELESS_MATRIX, + "b.yml": "on: pull_request\njobs:\n legacy:\n name: build (3.13)\n runs-on: ubuntu-latest\n", + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`build (3.13)` is published by 2 jobs" in found[0] + + +def test_an_excluded_combination_publishes_no_check_run() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_PAIR})) + + assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13", "macos-3.12"}) + + +def test_an_exclude_row_naming_one_key_drops_every_combination_carrying_it() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_KEY})) + + assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13"}) + + +def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None: + names: Final = frozenset(name for name, _ in published({"unit.yml": BOOLEAN_MATRIX})) + + assert names == frozenset({"cache true", "cache false"}) + + +def test_a_format_call_its_arguments_cannot_fill_publishes_nothing_to_compare() -> None: + sources: Final = {"unit.yml": UNFILLABLE_FORMAT} + + assert frozenset(name for name, _ in published(sources)) == frozenset() + assert "its name stays" in blind_spots(sources)[0] + + +def test_a_call_to_a_workflow_outside_the_repo_is_reported_rather_than_guessed() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", + "b.yml": "on: pull_request\njobs:\n unit:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"}) + assert collisions(sources) == () + assert "outside this repository" in blind_spots(sources)[0] + + +def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None: + sources: Final = { + ".github/workflows/leaf.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: Leaf\n runs-on: ubuntu-latest\n" + ), + ".github/workflows/mid.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: Mid\n uses: ./.github/workflows/leaf.yml\n" + ), + "top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/mid.yml\n", + } + + names: Final = frozenset(name for name, _ in published(sources)) + + assert names == frozenset({"Top / Mid / Leaf"}) + + +def test_a_job_name_that_is_not_a_string_still_publishes_the_value_github_renders() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n sweep:\n name: 2024\n runs-on: ubuntu-latest\n", + "b.yml": 'on: pull_request\njobs:\n other:\n name: "2024"\n runs-on: ubuntu-latest\n', + } + + found: Final = collisions(sources) + + assert len(found) == 1 + assert "`2024` is published by 2 jobs" in found[0] + + +def test_the_check_fails_when_a_file_in_the_workflows_directory_cannot_be_read() -> None: + assert exit_code({"notes.yml": "just a string\n"}) == 1 + + +def test_the_check_fails_when_two_jobs_publish_one_check_run_name() -> None: + plain: Final = "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n" + + assert exit_code({"a.yml": plain, "b.yml": plain}) == 1 + + +def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n", + "b.yml": "on: pull_request\njobs:\n sweep:\n runs-on: ubuntu-latest\n", + } + + assert exit_code(sources) == 0 + + +def test_two_callers_of_one_reusable_workflow_named_from_its_inputs_do_not_collide() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n run:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n" + ), + "caller.yml": ( + "on: pull_request\njobs:\n" + " alpha:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: alpha\n" + " beta:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: beta\n" + ), + } + + assert collisions(sources) == () + assert len(blind_spots(sources)) == 2 + + +def test_a_matrix_that_is_itself_an_expression_never_collapses_onto_the_bare_job_id() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n strategy:\n" + " matrix: ${{ fromJson(needs.plan.outputs.matrix) }}\n runs-on: ubuntu-latest\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"build"}) + assert collisions(sources) == () + assert "the matrix itself comes from an expression" in blind_spots(sources)[0] + + +def test_a_matrix_listing_objects_never_collapses_onto_the_bare_job_id() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n build:\n strategy:\n matrix:\n target:\n" + " - os: ubuntu\n - os: windows\n runs-on: ubuntu-latest\n" + ), + "b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"build"}) + assert collisions(sources) == () + assert "not plain scalars" in blind_spots(sources)[0] + + +def test_a_call_to_a_workflow_file_the_checkout_does_not_hold_is_reported() -> None: + sources: Final = {"a.yml": "on: pull_request\njobs:\n unit:\n uses: ./.github/workflows/gone.yml\n"} + + assert collisions(sources) == () + assert "which this checkout does not hold" in blind_spots(sources)[0] + + +def test_reusable_workflows_calling_each_other_in_a_loop_are_reported_not_followed() -> None: + sources: Final = { + ".github/workflows/a.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: A\n uses: ./.github/workflows/b.yml\n" + ), + ".github/workflows/b.yml": ( + "on:\n workflow_call:\njobs:\n call:\n name: B\n uses: ./.github/workflows/a.yml\n" + ), + "top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/a.yml\n", + } + + assert collisions(sources) == () + assert any("loops back on itself" in spot for spot in blind_spots(sources)) + + +def test_a_caller_still_publishes_the_callee_jobs_it_can_read() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n" + " lint:\n name: Lint\n runs-on: ubuntu-latest\n" + " suite:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n" + ), + "caller.yml": "on: pull_request\njobs:\n call:\n name: A\n uses: ./.github/workflows/callee.yml\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"A / Lint"}) + assert len(blind_spots(sources)) == 1 + + +def test_a_name_the_check_cannot_work_out_is_reported_without_failing_the_check() -> None: + sources: Final = { + "a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n", + } + + assert blind_spots(sources) != () + assert exit_code(sources) == 0 + + +def test_a_caller_whose_own_name_is_unreadable_publishes_none_of_its_callee_names() -> None: + sources: Final = { + ".github/workflows/callee.yml": ( + "on:\n workflow_call:\njobs:\n lint:\n name: Lint\n runs-on: ubuntu-latest\n" + ), + "caller.yml": ( + "on: pull_request\njobs:\n call:\n name: ${{ matrix.suite }}\n" + " uses: ./.github/workflows/callee.yml\n" + ), + "other.yml": "on: pull_request\njobs:\n plain:\n name: Lint\n runs-on: ubuntu-latest\n", + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"Lint"}) + assert collisions(sources) == () + assert "its name stays" in blind_spots(sources)[0] + + +def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_matches() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n unit:\n" + " name: ${{ matrix.python-version }} ${{ matrix.label }}\n" + " runs-on: ubuntu-latest\n strategy:\n matrix:\n" + ' python-version: ["3.12", "3.13"]\n' + " include:\n" + ' - python-version: "3.12"\n' + " label: fast\n" + ) + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"}) + assert len(blind_spots(sources)) == 1 + + +def test_a_job_whose_whole_strategy_is_an_expression_is_reported_rather_than_rejecting_the_file() -> None: + sources: Final = { + "plan.yml": ( + "on: pull_request\njobs:\n plan:\n name: Plan\n runs-on: ubuntu-latest\n" + " fan:\n strategy: ${{ fromJSON(needs.plan.outputs.strategy) }}\n runs-on: ubuntu-latest\n" + ) + } + + assert unreadable(sources) == () + assert frozenset(name for name, _ in published(sources)) == frozenset({"Plan"}) + assert "`strategy` comes from an expression" in blind_spots(sources)[0] + assert exit_code(sources) == 0 + + +def test_one_job_publishing_one_name_for_every_matrix_combination_is_a_collision() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n build:\n name: Run tests\n runs-on: ubuntu-latest\n" + ' strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n' + ) + } + + found: Final = collisions(sources) + assert len(found) == 1 + assert "`Run tests` is published 2 times by unit.yml job `build`" in found[0] + assert exit_code(sources) == 1 + + +def test_a_name_carrying_a_matrix_value_publishes_one_name_per_combination_without_colliding() -> None: + sources: Final = { + "unit.yml": ( + "on: pull_request\njobs:\n build:\n name: Run tests ${{ matrix.python-version }}\n" + ' runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n' + ) + } + + assert frozenset(name for name, _ in published(sources)) == frozenset({"Run tests 3.12", "Run tests 3.13"}) + assert collisions(sources) == () + assert exit_code(sources) == 0 + + +def test_a_file_that_is_not_valid_yaml_is_reported_rather_than_raising() -> None: + sources: Final = {"broken.yml": "jobs:\n build: [\n"} + + assert unreadable(sources) == ( + "broken.yml sits in the workflows directory but it does not read as one YAML " + "document, so none of its jobs were checked.", + ) + assert exit_code(sources) == 1 + + +def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() -> None: + sources: Final = {"two.yml": "on: pull_request\n---\non: push\n"} + + assert len(unreadable(sources)) == 1 + assert exit_code(sources) == 1 + + +def test_an_exclude_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: build\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n python: ['3.11', '3.12']\n" + " exclude: ${{ fromJson(vars.SKIP) }}\n" + ), + } + + found: Final = blind_spots(sources) + + assert collisions(sources) == () + assert len(found) == 1 + assert "a matrix `exclude` is itself an expression" in found[0] + + +def test_an_include_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None: + sources: Final = { + "a.yml": ( + "on: pull_request\njobs:\n one:\n name: build-${{ matrix.python }}\n runs-on: ubuntu-latest\n" + " strategy:\n matrix:\n python: ['3.11']\n" + " include: ${{ fromJson(vars.EXTRA) }}\n" + ), + } + + found: Final = blind_spots(sources) + + assert collisions(sources) == () + assert len(found) == 1 + assert "a matrix `include` is itself an expression" in found[0] diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index b83300cc8d5..1183096b81e 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -52,6 +52,15 @@ The suites run against a live proxy, so bring one up first by running the litell Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE` +The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it: + +```yaml +general_settings: + store_prompts_in_spend_logs: true +``` + +Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config .yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts + A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships ### The pull request check diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index c64fd6150af..81832bebf49 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -3,6 +3,7 @@ - {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"} - {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"} - {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"} +- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"} - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"} - {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"} - {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"} diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d571fb36546..860d96a50b4 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -5,6 +5,8 @@ - {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"} - {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"} - {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"} +- {id: mgmt.key.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "key_management_endpoints.py:2829", rationale: "A partial /key/update changes only the field it names; alias, models, limits, budget window, team and metadata read back unchanged on every gateway replica"} +- {id: mgmt.key.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "key_management_endpoints.py:2829", rationale: "An explicit null on /key/update clears max_budget and budget_duration, and the derived budget_reset_at with it, on every gateway replica"} - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"} - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"} - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index b50551ec105..6b69677d490 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -7,7 +7,7 @@ - {id: reliability.retry.timeout.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: timeout, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:44", rationale: "Timeout retried per policy"} - {id: reliability.retry.429.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: "429", assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:46", rationale: "429 retried per RateLimitErrorRetries policy"} - {id: reliability.retry.auth.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: auth, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:42", rationale: "Transient auth glitch retry"} -- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", rationale: "Multi-attempt on context error"} +- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", fail_before_fix: proven, rationale: "A context-window 400 under BadRequestErrorRetries retries onto a sibling deployment in the same model group, instead of coming straight back as the 400 the deployment that just refused it returned"} - {id: reliability.cooldown.5xx.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "5xx", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:40", rationale: "Deployment cools after repeated 5xx, recovers after cooldown_time"} - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"} - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"} diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py index c6d87473c21..49d698938ce 100644 --- a/tests/e2e/guardrails/test_presidio_masking_e2e.py +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -19,6 +19,12 @@ PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip. Each guardrail registers with an explicit presidio_filter_scope so only the configured hook's callback exists (the default "both" registers input masking AND a post_call output masker), and is deleted on teardown. + +The spend-log audit test requires general_settings.store_prompts_in_spend_logs: +true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy +process environment before startup. Setting it only on pytest has no effect. +Without this opt-in, redacting guardrail_response is expected proxy behavior; +this suite deliberately requires the detected-entity details to remain visible. """ from __future__ import annotations @@ -26,16 +32,22 @@ from __future__ import annotations import os import time from collections.abc import Callable -from typing import Literal +from typing import Final, Literal import pytest -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from e2e_config import unique_marker -from e2e_http import Result, Success +from e2e_http import Result, StreamingResponse, Success from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody from lifecycle import ResourceManager -from models import AnthropicMessagesResponse, ChatResponse +from models import ( + AnthropicMessagesResponse, + ChatResponse, + GuardrailEntityMatch, + GuardrailRunRecord, + SpendLogRow, +) pytestmark = pytest.mark.e2e @@ -276,3 +288,114 @@ class TestPresidioPostCallMasking: f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}" ) time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + + +_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"} + +_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch]) + + +def _applied_guardrails(outcome: StreamingResponse) -> str: + return outcome.headers.get("x-litellm-applied-guardrails", "") + + +def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]: + metadata = row.metadata + if metadata is None: + return () + return tuple(metadata.guardrail_information or ()) + + +def _poll_until_guardrail_applied( + client: GuardrailsClient, key: str, guardrail_name: str, prompt: str +) -> StreamingResponse: + deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128) + while time.monotonic() < deadline: + if last.ok and guardrail_name in _applied_guardrails(last): + return last + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128) + return last + + +class TestPresidioSpendLogRecord: + @pytest.mark.covers( + "guardrail.presidio.pre_call.logs_masked_entities", + exercised_on=["chat_completions"], + ) + def test_masking_run_is_recorded_on_the_spend_log( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-log-{unique_marker()}" + _register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES) + + email = _fake_email() + prompt = _pii_prompt(unique_marker(), email) + + outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt) + assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}" + assert name in _applied_guardrails(outcome), ( + "the response must carry x-litellm-applied-guardrails naming the guardrail; without it " + f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}" + ) + + request_id = ChatResponse.model_validate_json(outcome.body).id + assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}" + + rows = client.proxy.poll_logs_for_request_id( + request_id, + predicate=lambda logged: bool(_guardrail_records(logged[0])), + ) + assert rows, f"no spend log row ever appeared for request {request_id}" + + records = _guardrail_records(rows[0]) + assert records, ( + f"the spend log for {request_id} carries no guardrail_information, so the dashboard's " + "guardrail panel would render nothing for a request the guardrail demonstrably ran on" + ) + + record = next( + (entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"), + None, + ) + assert record is not None, ( + "guardrail_information carries no pre_call record for this guardrail, only " + f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}" + ) + assert record.guardrail_status == "success", ( + f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}" + ) + assert record.guardrail_provider == "presidio", ( + f"the record must attribute the run to presidio so the dashboard picks the right " + f"renderer, got {record.guardrail_provider!r}" + ) + + counts = record.masked_entity_count or {} + assert _LOGGED_ENTITIES.keys() <= counts.keys(), ( + f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}" + ) + assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), ( + f"each masked entity must be counted at least once, got {counts}" + ) + + assert not isinstance(record.guardrail_response, str), ( + "This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config " + "or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup " + "(not just the pytest environment). Default prompt redaction is valid proxy behavior, " + f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}" + ) + entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response) + assert entities, ( + "guardrail_response must carry the detected entities; the dashboard's Detected Entities " + "list and its per-entity scores are rendered from exactly this array" + ) + assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), ( + f"the detected entities must cover what was masked, got " + f"{sorted(entity.entity_type for entity in entities)}" + ) + for entity in entities: + assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}" + assert 0 <= entity.start < entity.end, ( + f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}" + ) diff --git a/tests/e2e/management/test_key_lifecycle_e2e.py b/tests/e2e/management/test_key_lifecycle_e2e.py new file mode 100644 index 00000000000..4c8effc4d24 --- /dev/null +++ b/tests/e2e/management/test_key_lifecycle_e2e.py @@ -0,0 +1,299 @@ +"""Live e2e: one virtual key walked through its whole lifecycle, read back on every +gateway replica. + +Create, read, partial update, clear, enforce, delete: one method per step, and every +step creates its own team and key (both deleted on teardown) so a step reruns or skips +on its own. Writes go through the control plane; read-backs poll every URL in +PROXY_REPLICA_URLS until each replica converges, because a write that is visible on the +gateway that took it and stale on its neighbour is exactly the failure this file exists +to catch. Revocation is the slowest of those: a deleted key stays usable on the other +replicas until their auth cache entry expires, so the delete step polls each of them +rather than asserting once. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import Result, StreamingResponse, Success, UnknownApiError, unwrap +from lifecycle import ResourceManager +from management_client import MODEL_ACCESS_DENIED_MARKER, ManagementClient +from models import ( + CLEAR, + ChatBody, + ChatMessage, + KeyGenerateBody, + KeyGenerateResponse, + KeyInfo, + KeyInfoParams, + KeyInfoResponse, + KeyMetadata, + KeyUpdateBody, + LiteLLMParamsBody, + TeamNewBody, +) +from proxy_client import Converged, NotConverged, Poller, await_converged, await_converged_everywhere +from transport import Transport + +pytestmark = pytest.mark.e2e + +BACKING_MODEL: Final = "gpt-4o-mini" +DENIED_MODEL: Final = "gpt-5.5" +MAX_BUDGET: Final = 25.0 +TPM_LIMIT: Final = 313131 +RPM_LIMIT: Final = 323232 +UPDATED_RPM_LIMIT: Final = 424242 +BUDGET_DURATION: Final = "30d" + + +@dataclass(frozen=True, slots=True) +class CreatedKey: + written: KeyGenerateBody + response: KeyGenerateResponse + + @property + def key(self) -> str: + return self.response.key + + +@pytest.fixture(scope="module") +def mock_deployment(client: ManagementClient) -> Iterator[str]: + """A deployment that answers from a canned response, so the enforcement step needs no + provider key. The alias carries a unique marker, like every other model this suite + registers, so concurrent runs never share one model group.""" + model_name: Final = f"e2e-key-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.create_model(model_name, LiteLLMParamsBody(model=BACKING_MODEL, mock_response="ok")) + try: + yield model_name + finally: + client.proxy.delete_model(model_id) + + +def _await[T](client: ManagementClient, poller: Poller[T], converged: Callable[[T], bool], failure: str) -> T: + outcome: Final = await_converged( + poller, + converged=converged, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Converged(result=result): + return result + case NotConverged(last_result=last): + pytest.fail(f"{failure}; last outcome: {last}") + + +def _chat_poller(transport: Transport, key: str, model: str) -> Poller[StreamingResponse]: + return lambda: transport.send( + "/chat/completions", + headers=transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"say hi {unique_marker()}")], + max_tokens=16, + ), + ) + + +def _create_key(client: ManagementClient, resources: ResourceManager, model: str) -> CreatedKey: + marker: Final = unique_marker() + team_id: Final = client.create_team(TeamNewBody(team_alias=f"e2e-key-lifecycle-team-{marker}")) + resources.defer(lambda: client.delete_team(team_id)) + written: Final = KeyGenerateBody( + key_alias=f"e2e-key-lifecycle-{marker}", + models=[model], + max_budget=MAX_BUDGET, + tpm_limit=TPM_LIMIT, + rpm_limit=RPM_LIMIT, + budget_duration=BUDGET_DURATION, + metadata=KeyMetadata(tag=marker), + team_id=team_id, + ) + response: Final = unwrap(client.generate_key(written)) + resources.defer(lambda: client.proxy.delete_key(response.key)) + return CreatedKey(written=written, response=response) + + +def _key_info_everywhere( + client: ManagementClient, key: str, settled: Callable[[KeyInfo], bool] +) -> Mapping[str, KeyInfo]: + def converged(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and settled(result.data.info) + + reads: Final = client.proxy.read_back_everywhere( + "/key/info", params=KeyInfoParams(key=key), response_type=KeyInfoResponse, converged=converged + ) + return MappingProxyType({replica: unwrap(read).info for replica, read in reads.items()}) + + +def _is_key_not_found(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _assert_reads_back(info: KeyInfo, expected: KeyGenerateBody, replica: str) -> None: + for field, observed, wanted in ( + ("key_alias", info.key_alias, expected.key_alias), + ("models", info.models, expected.models), + ("max_budget", info.max_budget, expected.max_budget), + ("tpm_limit", info.tpm_limit, expected.tpm_limit), + ("rpm_limit", info.rpm_limit, expected.rpm_limit), + ("budget_duration", info.budget_duration, expected.budget_duration), + ("team_id", info.team_id, expected.team_id), + ("metadata", info.metadata, expected.metadata), + ): + assert observed == wanted, f"{replica}: /key/info reports {field}={observed!r}, expected {wanted!r}" + + +def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> None: + _ = _await( + client, + _chat_poller(client.proxy.transport, key, model), + lambda outcome: outcome.ok, + f"chat on {model} never succeeded for the key before the deadline", + ) + + +def _warm_every_replica(client: ManagementClient, key: str, model: str) -> None: + """Serve one call from every replica, so each has the key in its auth cache. Without + this the revocation check below would only prove a replica rejects a key it never + knew, which is true of any random string.""" + for replica, transport in client.proxy.replicas.items(): + _ = _await( + client, + _chat_poller(transport, key, model), + lambda outcome: outcome.ok, + f"{replica}: chat on {model} never succeeded for the key before the deadline", + ) + + +def _assert_chat_rejected_everywhere(client: ManagementClient, key: str, model: str) -> None: + outcomes: Final = await_converged_everywhere( + {replica: _chat_poller(transport, key, model) for replica, transport in client.proxy.replicas.items()}, + converged=lambda outcome: outcome.status_code == 401, + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + for replica, outcome in outcomes.items(): + assert isinstance(outcome, Converged), ( + f"{replica}: the deleted key was still accepted on chat after " + f"{client.proxy.poll_timeout}s, last status {outcome.last_result.status_code}" + ) + + +class TestKeyLifecycle: + def test_create_echoes_every_field_written( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + response: Final = created.response + for field, observed, wanted in ( + ("key_alias", response.key_alias, created.written.key_alias), + ("models", response.models, created.written.models), + ("max_budget", response.max_budget, created.written.max_budget), + ("tpm_limit", response.tpm_limit, created.written.tpm_limit), + ("rpm_limit", response.rpm_limit, created.written.rpm_limit), + ("budget_duration", response.budget_duration, created.written.budget_duration), + ("team_id", response.team_id, created.written.team_id), + ("metadata", response.metadata, created.written.metadata), + ): + assert observed == wanted, f"/key/generate echoed {field}={observed!r}, sent {wanted!r}" + + def test_read_reflects_the_create_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + infos: Final = _key_info_everywhere( + client, created.key, lambda info: info.key_alias == created.written.key_alias + ) + for replica, info in infos.items(): + _assert_reads_back(info, created.written, replica) + assert info.budget_reset_at is not None, ( + f"{replica}: /key/info reports no budget_reset_at for budget_duration={BUDGET_DURATION!r}" + ) + + @pytest.mark.covers("mgmt.key.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_field( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + before: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == RPM_LIMIT) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, rpm_limit=UPDATED_RPM_LIMIT))) + + after: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == UPDATED_RPM_LIMIT) + for replica, info in after.items(): + _assert_reads_back(info, created.written.model_copy(update={"rpm_limit": UPDATED_RPM_LIMIT}), replica) + assert info.budget_reset_at == before[replica].budget_reset_at, ( + f"{replica}: budget_reset_at moved from {before[replica].budget_reset_at!r} to " + f"{info.budget_reset_at!r} on a /key/update that did not name budget_duration" + ) + + @pytest.mark.covers("mgmt.key.update.clear_persists") + def test_explicit_null_clears_the_budget_and_its_reset_time( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + _ = _key_info_everywhere(client, created.key, lambda info: info.max_budget == MAX_BUDGET) + + _ = unwrap(client.update_key(KeyUpdateBody(key=created.key, max_budget=CLEAR, budget_duration=CLEAR))) + + cleared: Final = _key_info_everywhere(client, created.key, lambda info: info.max_budget is None) + for replica, info in cleared.items(): + assert info.budget_duration is None, ( + f"{replica}: budget_duration={info.budget_duration!r} survived an explicit null" + ) + assert info.budget_reset_at is None, ( + f"{replica}: clearing budget_duration left budget_reset_at={info.budget_reset_at!r}" + ) + _assert_reads_back( + info, created.written.model_copy(update={"max_budget": None, "budget_duration": None}), replica + ) + + def test_key_serves_its_model_and_is_denied_others( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + created: Final = _create_key(client, resources, mock_deployment) + + _poll_chat_ok(client, created.key, mock_deployment) + + denied: Final = client.chat_status(created.key, DENIED_MODEL, f"say hi {unique_marker()}") + assert denied.status_code == 403, ( + f"chat on {DENIED_MODEL!r} outside the key's model list must be denied 403, got " + f"{denied.status_code}: {denied.body[:300]}" + ) + assert MODEL_ACCESS_DENIED_MARKER in denied.body, ( + f"403 body must be a model-access denial, got: {denied.body[:300]}" + ) + + def test_delete_revokes_info_and_chat_on_every_replica( + self, client: ManagementClient, resources: ResourceManager, mock_deployment: str + ) -> None: + """The teardown's deferred delete fires again on the already-deleted key by + design: the deferred cleanup must survive this test failing before the + in-body delete, and a repeat /key/delete is a cheap no-op the warn-only + teardown absorbs.""" + created: Final = _create_key(client, resources, mock_deployment) + _warm_every_replica(client, created.key, mock_deployment) + + client.delete_key_strict(created.key) + + _ = client.proxy.read_back_everywhere( + "/key/info", + params=KeyInfoParams(key=created.key), + response_type=KeyInfoResponse, + converged=_is_key_not_found, + ) + _assert_chat_rejected_everywhere(client, created.key, mock_deployment) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2c6c0e9bbd4..9f2654e0eec 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -8,9 +8,9 @@ from __future__ import annotations from collections.abc import Sequence from datetime import datetime -from typing import Literal +from typing import Final, Literal -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -49,6 +49,7 @@ class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None priority: str | None = None batch_enqueued_token_limit: int | None = None + tag: str | None = None class ObjectPermission(BaseModel): @@ -81,6 +82,14 @@ class KeyGenerateBody(BaseModel): class KeyGenerateResponse(BaseModel): key: str + key_alias: str | None = None + models: list[str] = [] + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + budget_duration: str | None = None + team_id: str | None = None + metadata: KeyMetadata | None = None class KeyRegenerateBody(BaseModel): @@ -122,6 +131,7 @@ class KeyInfo(BaseModel): blocked: bool | None = None spend: float | None = None max_budget: float | None = None + budget_duration: str | None = None budget_reset_at: str | None = None budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None @@ -291,6 +301,7 @@ class RouterSettingsOverride(BaseModel): context_window_fallbacks: list[dict[str, list[str]]] | None = None content_policy_fallbacks: list[dict[str, list[str]]] | None = None num_retries: int | None = None + model_group_retry_policy: dict[str, dict[str, int]] | None = None enable_tag_filtering: bool | None = None @@ -573,6 +584,27 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- +class GuardrailEntityMatch(BaseModel): + entity_type: str + score: float + start: int + end: int + + +class GuardrailRunRecord(BaseModel): + guardrail_name: str | None = None + guardrail_mode: str | None = None + guardrail_status: str | None = None + guardrail_provider: str | None = None + masked_entity_count: dict[str, int] | None = None + guardrail_response: object | None = None + + +class SpendLogMetadata(BaseModel): + applied_guardrails: list[str] | None = None + guardrail_information: list[GuardrailRunRecord] | None = None + + class SpendLogRow(BaseModel): request_id: str | None = None api_key: str | None = None @@ -589,6 +621,7 @@ class SpendLogRow(BaseModel): completion_tokens: int | None = None total_tokens: int | None = None request_tags: list[str] | None = None + metadata: SpendLogMetadata | None = None class SpendLogs(RootModel[list[SpendLogRow]]): @@ -918,12 +951,34 @@ class CredentialCreateResponse(BaseModel): # ---------- key / team / user / organization management ---------- +class Cleared(BaseModel): + """An explicit JSON null in a merge-patch body. The transport drops `None` fields + before sending (`exclude_none`), so `None` means "leave the stored value alone"; a + field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it.""" + + model_config = ConfigDict(frozen=True) + + @model_serializer + def _as_null(self) -> None: + return None + + +CLEAR: Final = Cleared() + + class KeyUpdateBody(BaseModel): + """POST /key/update is a merge patch: a field left `None` is dropped from the body and + keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration` + clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale.""" + key: str models: list[str] | None = None key_alias: str | None = None tpm_limit: int | None = None rpm_limit: int | None = None + max_budget: float | Cleared | None = None + budget_duration: str | Cleared | None = None + metadata: KeyMetadata | None = None class KeyBlockBody(BaseModel): diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 3b15e57dfcf..520cbfde5a9 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -16,6 +16,8 @@ from datetime import datetime from types import MappingProxyType from typing import Final +from pydantic import BaseModel + from e2e_http import ( AnthropicHeaders, AuthHeaders, @@ -149,9 +151,7 @@ def await_servable( last_result: Result[ModelsListResponse] | None = None while True: t = now() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds remaining = phase_deadline - t if remaining <= 0: if ( @@ -164,9 +164,7 @@ def await_servable( poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) - listed = isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ) + listed = isinstance(last_result, Success) and any(entry.id == model_name for entry in last_result.data.data) t = now() if not listed: first_seen_at = None @@ -179,9 +177,7 @@ def await_servable( elif t - first_seen_at >= db_sync_seconds: return Servable() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) + phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds wait = min(interval, phase_deadline - now()) if wait > 0: sleep(wait) @@ -239,6 +235,88 @@ def servable_timeout_message( ) +type Poller[T] = Callable[[], T] + + +@dataclass(frozen=True, slots=True) +class Converged[T]: + result: T + + +@dataclass(frozen=True, slots=True) +class NotConverged[T]: + """The deadline passed without a read satisfying the predicate; `last_result` is + the final read, so the caller can tell a stale body from a failed request.""" + + last_result: T + + +type ConvergeOutcome[T] = Converged[T] | NotConverged[T] + + +def await_converged[T]( + poll: Poller[T], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ConvergeOutcome[T]: + """Poll until a read satisfies `converged` or `timeout` elapses. + + Polls before testing the deadline, so a zero or already-spent budget still gets one + attempt, and sleeps only min(interval, time left), so the attempt that lands exactly + on the deadline is taken rather than skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + while True: + result = poll() + if converged(result): + return Converged(result=result) + remaining = deadline - now() + if remaining <= 0: + return NotConverged(last_result=result) + sleep(min(interval, remaining)) + + +def await_converged_everywhere[T]( + pollers: Mapping[str, Poller[T]], + *, + converged: Callable[[T], bool], + timeout: float, + interval: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Mapping[str, ConvergeOutcome[T]]: + """`await_converged` against every replica in turn, each with the full budget, so a + replica that lags behind the one a write landed on is polled until it catches up + rather than failing on its first stale read.""" + return MappingProxyType( + { + replica: await_converged( + poll, converged=converged, timeout=timeout, interval=interval, now=now, sleep=sleep + ) + for replica, poll in pollers.items() + } + ) + + +def first_lagging_replica[T]( + outcomes: Mapping[str, ConvergeOutcome[T]], +) -> tuple[str, NotConverged[T]] | None: + return next( + ((replica, outcome) for replica, outcome in outcomes.items() if isinstance(outcome, NotConverged)), + None, + ) + + +def converge_timeout_message(*, what: str, replica: str, timeout: float, last_result: object) -> str: + return ( + f"{what} on {replica} never converged within {timeout}s of the write " + f"(control/data-plane propagation issue); last read: {last_result}" + ) + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport @@ -290,6 +368,52 @@ class ProxyClient: ) ).info + def read_back_everywhere[R: BaseModel]( + self, + path: str, + *, + params: BaseModel, + response_type: type[R], + converged: Callable[[Result[R]], bool], + ) -> Mapping[str, Result[R]]: + """GET `path` under the master key on every replica in PROXY_REPLICA_URLS (the + data-plane URL alone when the stack exports no per-gateway addresses), polling + each to poll_timeout until its read satisfies `converged`. Returns that read per + replica, or fails naming the first replica that never converged and its last + read. Behind a load balancer the single address proves one replica converged, + not all of them; only per-gateway addresses make this a fleet-wide proof.""" + outcomes: Final = await_converged_everywhere( + { + url: self._body_poller(transport, path, params, response_type) + for url, transport in self.replicas.items() + }, + converged=converged, + timeout=self.poll_timeout, + interval=self.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + lagging: Final = first_lagging_replica(outcomes) + if lagging is not None: + replica, outcome = lagging + raise AssertionError( + converge_timeout_message( + what=f"GET {path}", + replica=replica, + timeout=self.poll_timeout, + last_result=outcome.last_result, + ) + ) + return MappingProxyType( + {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} + ) + + @staticmethod + def _body_poller[R: BaseModel]( + transport: Transport, path: str, params: BaseModel, response_type: type[R] + ) -> Poller[Result[R]]: + return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it (config override merged over cost-map defaults).""" @@ -320,9 +444,7 @@ class ProxyClient: response_type=FileListResponse, ) - def list_fine_tuning_jobs( - self, key: str, params: FineTuningJobsParams - ) -> Result[FineTuningJobsResponse]: + def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]: return self.transport.get( "/v1/fine_tuning/jobs", headers=self.transport.bearer(key), @@ -375,7 +497,11 @@ class ProxyClient: ) ).model_id written_at = time.monotonic() - self._await_model_servable(body.model_name, listed_for) + try: + self._await_model_servable(body.model_name, listed_for) + except BaseException: + self.delete_model(model_id) + raise settle_propagation(written_at) return model_id diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 5822058003c..1efcb1a045b 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -73,10 +73,25 @@ def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_always_picked_small_context_deployment(proxy: ProxyClient, name: str) -> str: + """The always-picked half of a retry pair on the smallest-context model OpenAI + still serves: it holds all of the model group's shuffle weight, so an oversized + prompt opens on it and earns a real context-window refusal, which never benches + a deployment, so only the retry itself can steer the request off it.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY, weight=1), + model_info=ModelInfoBody(), + ) + ) + + def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str: """The other half of a retry pair: healthy, but weight 0, so the weighted shuffle - never opens on it. It is reachable only once its sibling is benched and the - weighted pick falls through to a uniform one over what is left.""" + never opens on it. It is reachable only once its sibling is out of the running, + benched by a cooldown or skipped by the retry, and the weighted pick falls through + to a uniform one over what is left.""" return proxy.register_model( ModelNewBody( model_name=name, diff --git a/tests/e2e/router/test_reliability_retries_e2e.py b/tests/e2e/router/test_reliability_retries_e2e.py index 5441412935c..da45cb46a46 100644 --- a/tests/e2e/router/test_reliability_retries_e2e.py +++ b/tests/e2e/router/test_reliability_retries_e2e.py @@ -1,13 +1,17 @@ """Live e2e: a request that fails on its first deployment is retried inside its own model group and still comes back a completion. -The model group is a pair: an always-timing-out deployment that holds all of the -group's shuffle weight, and a healthy backup at weight 0. The weighted pick always -opens on the timing-out one, its first Timeout benches it (an -`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls -through to the only deployment left. So the customer sees a completion and the -proxy reports that it took a retry to get there, with no random first pick in the -middle of it. +Each model group is a pair: a deployment that always refuses and holds all of the +group's shuffle weight, plus a healthy backup at weight 0. The weighted pick always +opens on the refusing one, so the customer sees a completion only if the retry +lands on the backup, and the proxy reports that it took a retry to get there, with +no random first pick in the middle of it. + +The timeout pair relies on cooldown: the first Timeout benches the timing-out +deployment (an `allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`) and the +retry falls through to the only deployment left. The context-window pair cannot: +a 400 never benches a deployment, so the retry policy's `BadRequestErrorRetries` +has to steer the retry off the deployment that just refused the prompt. """ from __future__ import annotations @@ -16,20 +20,48 @@ import pytest from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker +from e2e_http import StreamingResponse from lifecycle import ResourceManager from models import RouterSettingsOverride from reliability_support import ( chat_override, completion_tokens_of, content_of, + create_always_picked_small_context_deployment, create_always_timing_out_deployment, create_zero_weight_backup_deployment, finish_reason_of, + oversized_prompt, ) pytestmark = pytest.mark.e2e +def assert_retry_landed_on_backup(resp: StreamingResponse) -> None: + assert resp.status_code == 200, ( + f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + ) + + attempted = resp.headers.get("x-litellm-attempted-retries") + assert attempted is not None, "response is missing the x-litellm-attempted-retries header" + assert int(attempted) >= 1, ( + f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " + "opened on the refusing deployment, so this proves nothing about retries" + ) + + content = content_of(resp) + finish_reason = finish_reason_of(resp) + completion_tokens = completion_tokens_of(resp) or 0 + assert isinstance(content, str), ( + f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" + ) + assert content or (finish_reason == "length" and completion_tokens > 0), ( + f"the retry returned empty content with finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " + f"was spent on non-visible reasoning (body={resp.body[:300]})" + ) + + class TestReliabilityRetries: @pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries") def test_timeout_on_first_deployment_succeeds_on_retry( @@ -49,25 +81,27 @@ class TestReliabilityRetries: override=RouterSettingsOverride(num_retries=2), ) - assert resp.status_code == 200, ( - f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + assert_retry_landed_on_backup(resp) + + @pytest.mark.covers("reliability.retry.context_window.succeeds_within_retries") + def test_context_window_refusal_on_first_deployment_succeeds_on_retry( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-retry-{unique_marker()}" + small_context = create_always_picked_small_context_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(small_context)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + resp = chat_override( + client.proxy, + scoped_key, + group, + oversized_prompt(unique_marker()), + override=RouterSettingsOverride( + num_retries=2, + model_group_retry_policy={group: {"BadRequestErrorRetries": 2}}, + ), ) - attempted = resp.headers.get("x-litellm-attempted-retries") - assert attempted is not None, "response is missing the x-litellm-attempted-retries header" - assert int(attempted) >= 1, ( - f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " - "opened on the timing-out deployment, so this proves nothing about retries" - ) - - content = content_of(resp) - finish_reason = finish_reason_of(resp) - completion_tokens = completion_tokens_of(resp) or 0 - assert isinstance(content, str), ( - f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" - ) - assert content or (finish_reason == "length" and completion_tokens > 0), ( - f"the retry returned empty content with finish_reason={finish_reason!r}, " - f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " - f"was spent on non-visible reasoning (body={resp.body[:300]})" - ) + assert_retry_landed_on_backup(resp) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index a508c97b9fb..2caac58333f 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -1,9 +1,11 @@ -"""Harness coverage for the model barrier that gates on every replica. +"""Harness coverage for the barriers that gate on every replica. No proxy needed and no ``e2e`` marker: this pins that a model registered through the control plane only counts as servable once every configured replica lists it -on /v1/models, which is what keeps a two-gateway stack from handing a test a -model that one gateway has not reloaded yet. The fakes are plain pollers and an +on /v1/models, and that a management write only counts as read back once every +replica's read satisfies the caller's predicate, which is what keeps a two-gateway +stack from handing a test a model or a key that one gateway has not caught up on +yet. The fakes are plain pollers standing in for each replica's transport plus an injected clock, so nothing here monkeypatches anything. """ @@ -12,18 +14,33 @@ from __future__ import annotations from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain, repeat +from types import MappingProxyType from typing import Final import pytest from e2e_config import parse_replica_urls -from e2e_http import Success -from models import ModelListEntry, ModelsListResponse -from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere +from e2e_http import Result, Success +from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from proxy_client import ( + Poller, + ConvergeOutcome, + Converged, + ModelsPoller, + NotConverged, + NotServableOn, + Servable, + await_converged_everywhere, + await_servable_everywhere, + first_lagging_replica, + converge_timeout_message, +) MODEL: Final = "gpt-under-test" TIMEOUT: Final = 10.0 INTERVAL: Final = 2.0 +RPM_BEFORE_UPDATE: Final = 100 +RPM_AFTER_UPDATE: Final = 200 @dataclass @@ -78,6 +95,91 @@ class TestAwaitServableEverywhere: assert _await(pollers) == Servable() +def _key_info(rpm_limit: int) -> Success[KeyInfoResponse]: + return Success(status_code=200, data=KeyInfoResponse(info=KeyInfo(rpm_limit=rpm_limit))) + + +def _reads(results: Iterable[Result[KeyInfoResponse]]) -> Poller[Result[KeyInfoResponse]]: + it: Final = iter(results) + return lambda: next(it) + + +def _updated(result: Result[KeyInfoResponse]) -> bool: + return isinstance(result, Success) and result.data.info.rpm_limit == RPM_AFTER_UPDATE + + +def _converge( + pollers: Mapping[str, Poller[Result[KeyInfoResponse]]], clock: FakeClock +) -> Mapping[str, ConvergeOutcome[Result[KeyInfoResponse]]]: + return await_converged_everywhere( + pollers, + converged=_updated, + timeout=TIMEOUT, + interval=INTERVAL, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitConvergedEverywhere: + def test_waits_for_the_replica_that_lags_behind_the_write(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 2), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert outcomes == { + "gateway-1": Converged(result=_key_info(RPM_AFTER_UPDATE)), + "gateway-2": Converged(result=_key_info(RPM_AFTER_UPDATE)), + } + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * INTERVAL + + def test_names_the_replica_that_never_converges_with_its_last_read(self) -> None: + clock: Final = FakeClock() + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))), + "gateway-2": _reads(repeat(_key_info(RPM_BEFORE_UPDATE))), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) == ( + "gateway-2", + NotConverged(last_result=_key_info(RPM_BEFORE_UPDATE)), + ) + assert clock.elapsed == TIMEOUT + message: Final = converge_timeout_message( + what="GET /key/info", + replica="gateway-2", + timeout=TIMEOUT, + last_result=_key_info(RPM_BEFORE_UPDATE), + ) + assert "gateway-2" in message and "/key/info" in message and str(RPM_BEFORE_UPDATE) in message + + def test_each_replica_gets_its_own_full_budget(self) -> None: + """A replica that converges late must not eat into the next replica's budget: both + need most of the timeout here, so one shared deadline would starve the second.""" + clock: Final = FakeClock() + slow: Final = chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + pollers: Final = MappingProxyType( + { + "gateway-1": _reads(slow), + "gateway-2": _reads( + chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE))) + ), + } + ) + outcomes: Final = _converge(pollers, clock) + assert first_lagging_replica(outcomes) is None + assert clock.elapsed == 2 * 3 * INTERVAL + + class TestParseReplicaUrls: def test_splits_and_trims_the_gateway_addresses(self) -> None: raw: Final = " http://127.0.0.1:4010/, http://127.0.0.1:4011 " diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts index 3a62252915c..71774c95d24 100644 --- a/tests/e2e/ui/constants.ts +++ b/tests/e2e/ui/constants.ts @@ -15,6 +15,8 @@ export const UI_BASE_URL = ( // writable path (the image runner already exports TMPDIR) to relocate them. export const ARTIFACT_DIR = process.env.E2E_UI_ARTIFACT_DIR || "."; +export const MOCK_PRESIDIO_URL = (process.env.E2E_MOCK_PRESIDIO_URL || "http://127.0.0.1:8091").replace(/\/+$/, ""); + const storagePath = (name: string): string => path.join(ARTIFACT_DIR, name); // Storage state paths for each role diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 58939ca2b9a..bce09b49e10 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -1,50 +1,142 @@ -/** - * Source of truth for the App Router migration E2E suites. - * - * Add an entry (legacy sidebar page id -> route segment) once a page's migration - * has MERGED to the branch under test. Consumers pick it up automatically: - * - migration smoke (tests/migration/migratedPages.spec.ts), via MIGRATED_E2E_SEGMENTS: - * default mount: npm run e2e:migration - * server-root-path mount: SERVER_ROOT_PATH=/ npm run e2e:migration:root - * - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts) - * - * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. - */ -export const MIGRATED_E2E_PAGES: Record = { - "api-keys": "api-keys", - models: "models-and-endpoints", - api_ref: "api-reference", - "llm-playground": "playground", - projects: "projects", - "access-groups": "access-groups", - budgets: "budgets", - workflows: "workflows", - "guardrails-monitor": "guardrails-monitor", - "mcp-servers": "mcp-servers", - "search-tools": "search-tools", - "tag-management": "tag-management", - "vector-stores": "vector-stores", - memory: "memory", - policies: "policies", - guardrails: "guardrails", - prompts: "prompts", - "tool-policies": "tool-policies", - skills: "skills", - caching: "caching", - "cost-tracking": "cost-tracking", - "transform-request": "transform-request", - "ui-theme": "ui-theme", - logs: "logs", - "admin-panel": "admin-panel", - "logging-and-alerts": "logging-and-alerts", - "model-hub-table": "model-hub-table", - new_usage: "usage", - usage: "old-usage", - agents: "agents", - "router-settings": "router-settings", - users: "users", - teams: "teams", - organizations: "organizations", -}; +export type MigratedPage = Readonly<{ + segment: string; + linkName: string | RegExp; + group?: string; + content: Readonly<{ role: "heading" | "tab" | "button"; name: string }> | Readonly<{ text: string }>; + unlicensedText?: string; +}>; -export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; +export const MIGRATED_E2E_PAGES: Readonly> = { + "api-keys": { segment: "api-keys", linkName: "Virtual Keys", content: { role: "heading", name: "Virtual Keys" } }, + models: { + segment: "models-and-endpoints", + linkName: "Models + Endpoints", + content: { role: "heading", name: "Model Management" }, + }, + api_ref: { + segment: "api-reference", + linkName: "API Reference", + content: { role: "heading", name: "OpenAI Compatible Proxy: API Reference" }, + }, + "llm-playground": { segment: "playground", linkName: "Playground", content: { role: "tab", name: "Chat" } }, + projects: { + segment: "projects", + linkName: /^Projects(?: Beta)?$/, + content: { role: "heading", name: "Projects" }, + }, + "access-groups": { + segment: "access-groups", + linkName: "Access Groups", + content: { role: "heading", name: "Access Groups" }, + }, + budgets: { segment: "budgets", linkName: "Budgets", content: { role: "heading", name: "Budgets" } }, + workflows: { + segment: "workflows", + linkName: "Workflow Runs", + group: "Agentic", + content: { text: "Workflow Runs" }, + }, + "guardrails-monitor": { + segment: "guardrails-monitor", + linkName: "Guardrails Monitor", + content: { role: "heading", name: "Guardrails Monitor" }, + }, + "mcp-servers": { + segment: "mcp-servers", + linkName: "MCP Servers", + content: { role: "heading", name: "MCP Servers" }, + }, + "search-tools": { + segment: "search-tools", + linkName: "Search Tools", + group: "Tools", + content: { role: "heading", name: "Search Tools" }, + }, + "tag-management": { + segment: "tag-management", + linkName: "Tag Management", + group: "Experimental", + content: { role: "heading", name: "Tag Management" }, + }, + "vector-stores": { + segment: "vector-stores", + linkName: "Vector Stores", + group: "Tools", + content: { role: "heading", name: "Vector Store Management" }, + }, + memory: { segment: "memory", linkName: "Memory", group: "Agentic", content: { role: "heading", name: "Memory" } }, + policies: { segment: "policies", linkName: "Policies", content: { role: "tab", name: "Policy Simulator" } }, + guardrails: { segment: "guardrails", linkName: "Guardrails", content: { role: "tab", name: "Guardrails" } }, + prompts: { + segment: "prompts", + linkName: "Prompts", + group: "Experimental", + content: { role: "button", name: "Add New Prompt" }, + }, + "tool-policies": { + segment: "tool-policies", + linkName: "Tool Policies", + group: "Tools", + content: { role: "heading", name: "Tool Policies" }, + }, + skills: { segment: "skills", linkName: "Skills", content: { role: "heading", name: "Skills" } }, + caching: { segment: "caching", linkName: "Response Cache", content: { role: "tab", name: "Cache Settings" } }, + "cost-tracking": { + segment: "cost-tracking", + linkName: "Cost Tracking", + group: "Settings", + content: { text: "Cost Tracking Settings" }, + }, + "transform-request": { + segment: "transform-request", + linkName: "API Playground", + group: "Experimental", + content: { role: "heading", name: "Playground" }, + }, + "ui-theme": { + segment: "ui-theme", + linkName: "UI Theme", + group: "Settings", + content: { role: "heading", name: "UI Theme Customization" }, + }, + logs: { segment: "logs", linkName: "Logs", content: { role: "heading", name: "Request Logs" } }, + "admin-panel": { + segment: "admin-panel", + linkName: "Admin Settings", + group: "Settings", + content: { role: "heading", name: "Admin Access" }, + }, + "logging-and-alerts": { + segment: "logging-and-alerts", + linkName: "Logging & Alerts", + group: "Settings", + content: { role: "tab", name: "Logging Callbacks" }, + }, + "model-hub-table": { + segment: "model-hub-table", + linkName: "AI Hub", + content: { role: "heading", name: "AI Hub" }, + }, + new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage View" } }, + usage: { + segment: "old-usage", + linkName: "Old Usage", + group: "Experimental", + content: { role: "tab", name: "All Up" }, + }, + agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } }, + "router-settings": { + segment: "router-settings", + linkName: "Router Settings", + group: "Settings", + content: { role: "heading", name: "Routing Settings" }, + }, + users: { segment: "users", linkName: "Internal Users", content: { role: "tab", name: "Users" } }, + teams: { segment: "teams", linkName: "Teams", content: { role: "heading", name: "Teams" } }, + organizations: { + segment: "organizations", + linkName: "Organizations", + content: { text: "Click on an organization ID to view its details." }, + unlicensedText: "This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.", + }, +}; diff --git a/tests/e2e/ui/fixtures/mock_presidio_server/server.py b/tests/e2e/ui/fixtures/mock_presidio_server/server.py new file mode 100644 index 00000000000..0fc97c4a065 --- /dev/null +++ b/tests/e2e/ui/fixtures/mock_presidio_server/server.py @@ -0,0 +1,107 @@ +""" +Mock Presidio analyzer + anonymizer for UI e2e tests. +Serves POST /analyze and POST /anonymize over a fixed set of regex recognizers. +""" + +import os +import re + +import uvicorn +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware + + +def _luhn_ok(candidate): + digits = [int(c) for c in candidate if c.isdigit()] + doubled = [d * 2 - 9 if d * 2 > 9 else d * 2 for d in digits[-2::-2]] + return len(digits) >= 13 and (sum(digits[::-1][::2]) + sum(doubled)) % 10 == 0 + + +RECOGNIZERS = ( + ("EMAIL_ADDRESS", re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), 1.0, None), + ("CREDIT_CARD", re.compile(r"\b(?:\d[ -]?){13,19}\b"), 1.0, _luhn_ok), + ("US_SSN", re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), 0.85, None), + ("PHONE_NUMBER", re.compile(r"\b(?:\+?\d{1,2}[ .-]?)?(?:\(\d{3}\)|\d{3})[ .-]?\d{3}[ .-]?\d{4}\b"), 0.75, None), + ("IP_ADDRESS", re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), 0.6, None), + ("URL", re.compile(r"\bhttps?://[^\s]+"), 0.5, None), +) + +app = FastAPI(title="Mock Presidio Server") +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["*"], + allow_headers=["*"], +) + + +def _detect(text, wanted): + found = [ + {"entity_type": name, "start": m.start(), "end": m.end(), "score": score} + for name, pattern, score, validate in RECOGNIZERS + if wanted is None or name in wanted + for m in pattern.finditer(text) + if validate is None or validate(m.group()) + ] + ranked = sorted(found, key=lambda r: (-r["score"], r["start"])) + kept = [] + for candidate in ranked: + overlaps = any(candidate["start"] < k["end"] and k["start"] < candidate["end"] for k in kept) + if not overlaps: + kept.append(candidate) + return sorted(kept, key=lambda r: r["start"]) + + +@app.get("/health") +async def health(): + return {"status": "ok"} + + +@app.post("/analyze") +async def analyze(request: Request): + body = await request.json() + text = body.get("text", "") or "" + entities = body.get("entities") + wanted = set(entities) if entities else None + threshold = body.get("score_threshold") or 0.0 + return [ + {**result, "analysis_explanation": None, "recognition_metadata": {"recognizer_name": "MockRecognizer"}} + for result in _detect(text, wanted) + if result["score"] >= threshold + ] + + +@app.post("/anonymize") +async def anonymize(request: Request): + body = await request.json() + text = body.get("text", "") or "" + results = sorted(body.get("analyzer_results") or [], key=lambda r: r["start"]) + + pieces = [] + items = [] + cursor = 0 + for result in results: + start, end = result["start"], result["end"] + if start < cursor: + continue + placeholder = f"<{result['entity_type']}>" + pieces.append(text[cursor:start]) + masked_start = sum(len(p) for p in pieces) + pieces.append(placeholder) + items.append( + { + "operator": "replace", + "entity_type": result["entity_type"], + "start": masked_start, + "end": masked_start + len(placeholder), + "text": placeholder, + } + ) + cursor = end + pieces.append(text[cursor:]) + + return {"text": "".join(pieces), "items": list(reversed(items))} + + +if __name__ == "__main__": + uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("MOCK_PRESIDIO_PORT", "8091"))) diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 5f87d40b2fc..e0e7b4da396 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -1,5 +1,29 @@ import { Page } from "../fixtures/pages"; import { Page as PlaywrightPage, expect } from "@playwright/test"; +import { UI_BASE_URL } from "../constants"; + +export const sidebarLink = (page: PlaywrightPage, name: string | RegExp) => + page.getByRole("complementary").getByRole("link", { name, exact: true }); + +export async function clickSidebarLink(page: PlaywrightPage, name: string | RegExp, groupName?: string): Promise { + const link = sidebarLink(page, name); + if (groupName && !(await link.isVisible())) { + const group = page.getByRole("complementary").getByRole("button", { name: groupName, exact: true }); + await expect(group).toBeVisible(); + if ((await group.getAttribute("aria-expanded")) === "false") { + await group.click(); + } + } + await link.click(); +} + +export async function expectUiRoute(page: PlaywrightPage, segment: string): Promise { + const root = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); + const expected = new URL(`${root}/ui/${segment}`, UI_BASE_URL); + await expect(page, `navigate to ${expected.pathname}`).toHaveURL( + (url) => url.origin === expected.origin && url.pathname.replace(/\/+$/, "") === expected.pathname, + ); +} /** * Navigates to a specific page using the page query parameter. diff --git a/tests/e2e/ui/run_e2e.sh b/tests/e2e/ui/run_e2e.sh index 67e3225f668..beb1bc8bf3b 100755 --- a/tests/e2e/ui/run_e2e.sh +++ b/tests/e2e/ui/run_e2e.sh @@ -14,7 +14,7 @@ set -euo pipefail # # Ports default to 4000 / 5432 / 8090 and can be moved when another checkout # already holds them: -# PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 ./run_e2e.sh +# PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 MOCK_PRESIDIO_PORT=8191 ./run_e2e.sh # # In CI (CI=true), expects: # - PostgreSQL already running on 127.0.0.1:5432 @@ -29,6 +29,7 @@ DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard" IS_CI="${CI:-false}" CONTAINER_NAME="litellm-e2e-postgres-$$" MOCK_PID="" +MOCK_PRESIDIO_PID="" PROXY_PID="" PROXY_LOG="" @@ -40,7 +41,8 @@ PROXY_LOG="" PROXY_PORT="${PROXY_PORT:-4000}" POSTGRES_PORT="${POSTGRES_PORT:-5432}" MOCK_LLM_PORT="${MOCK_LLM_PORT:-8090}" -export MOCK_LLM_PORT +MOCK_PRESIDIO_PORT="${MOCK_PRESIDIO_PORT:-8091}" +export MOCK_LLM_PORT MOCK_PRESIDIO_PORT # --- Ensure common tool paths are available (local dev only) --- if [ "$IS_CI" = "false" ]; then @@ -82,6 +84,7 @@ fi cleanup() { echo "Cleaning up..." [ -n "$MOCK_PID" ] && kill "$MOCK_PID" 2>/dev/null || true + [ -n "$MOCK_PRESIDIO_PID" ] && kill "$MOCK_PRESIDIO_PID" 2>/dev/null || true [ -n "$PROXY_PID" ] && kill "$PROXY_PID" 2>/dev/null || true [ -n "$PROXY_LOG" ] && rm -f "$PROXY_LOG" || true if [ "$IS_CI" = "false" ]; then @@ -110,9 +113,9 @@ if [ "$IS_CI" = "false" ]; then # to someone else's :5432 (a psql session, a running app, a Prisma engine # talking to a remote database) aborts the run with "port 5432 is in use" # while nothing is actually bound locally. - for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT"; do + for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT" "$MOCK_PRESIDIO_PORT"; do if lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; then - echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT)" + echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT / MOCK_PRESIDIO_PORT)" exit 1 fi done @@ -143,6 +146,7 @@ fi # --- Credentials --- export LITELLM_MASTER_KEY="sk-1234" export MOCK_LLM_URL="http://127.0.0.1:${MOCK_LLM_PORT}/v1" +export E2E_MOCK_PRESIDIO_URL="http://127.0.0.1:${MOCK_PRESIDIO_PORT}" export DISABLE_SCHEMA_UPDATE="true" # The suite resolves its target from E2E_UI_BASE_URL (constants.ts), which # otherwise defaults to :4000 -- so without this a relocated stack would be @@ -151,6 +155,7 @@ export DISABLE_SCHEMA_UPDATE="true" export E2E_UI_BASE_URL="${E2E_UI_BASE_URL:-http://127.0.0.1:${PROXY_PORT}}" # Ensure the proxy serves UI at /ui (not behind a subpath) export SERVER_ROOT_PATH="" +export PROXY_BASE_URL="$E2E_UI_BASE_URL" # Boot with an external logout URL so proxyLogoutUrl.spec.ts can assert the # redirect. This same value is exported to the Playwright process below (the # spec's skip guard reads it). Safe for the rest of the suite — nothing else @@ -203,6 +208,20 @@ for i in $(seq 1 15); do sleep 1 done +echo "=== Starting mock Presidio server ===" +uv run --no-sync python "$SCRIPT_DIR/fixtures/mock_presidio_server/server.py" & +MOCK_PRESIDIO_PID=$! + +PRESIDIO_READY=0 +for i in $(seq 1 15); do + if curl -sf http://127.0.0.1:${MOCK_PRESIDIO_PORT}/health >/dev/null 2>&1; then PRESIDIO_READY=1; break; fi + sleep 1 +done +if [ "$PRESIDIO_READY" -ne 1 ]; then + echo "Mock Presidio server never answered /health on port ${MOCK_PRESIDIO_PORT}" >&2 + exit 1 +fi + # --- LiteLLM proxy --- echo "=== Starting LiteLLM proxy ===" cd "$REPO_ROOT" diff --git a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts index 1e43c7a2b22..0267eeae909 100644 --- a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts +++ b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts @@ -7,6 +7,7 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traff interface StoredGuardrail { guardrail_id: string; guardrail_name: string | null; + litellm_params?: { pii_entities_config?: Record | null } | null; } async function listGuardrails(page: PlaywrightPage): Promise { @@ -249,6 +250,17 @@ test.describe("Guardrails", () => { const row = page.getByRole("row").filter({ hasText: guardrailName }); await expect(row).toHaveCount(1, { timeout: 15_000 }); + const stored = await findGuardrail(page, guardrailName); + const piiConfig = stored?.litellm_params?.pii_entities_config ?? {}; + expect( + Object.keys(piiConfig).length, + "Select All & Mask persisted no PII entities, so the guardrail would mask nothing", + ).toBeGreaterThan(0); + expect( + Object.entries(piiConfig).filter(([, action]) => action !== "MASK"), + "Select All & Mask must persist every entity with the MASK action", + ).toEqual([]); + await navigateToPage(page, Page.Teams); await dismissFeedbackPopup(page); await clickTeamId(page, E2E_TEAM_NO_ADMIN_ID); diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts new file mode 100644 index 00000000000..d4ed5308342 --- /dev/null +++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts @@ -0,0 +1,157 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, MOCK_PRESIDIO_URL } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { CHAT_MODEL_A, masterKey, rootPath, waitForSpendLogByPrompt } from "../../helpers/traffic"; +import { openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground"; + +const RAW_EMAIL = "jane.doe@example.com"; +const RAW_PHONE = "555-867-5309"; + +const visibleTestId = (page: PlaywrightPage, id: string) => page.getByTestId(id).filter({ visible: true }); + +const requestLogsRows = (page: PlaywrightPage) => + page.locator("table").filter({ visible: true }).first().locator("tbody tr"); + +async function createPresidioGuardrail(page: PlaywrightPage, guardrailName: string): Promise { + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: /Add New Guardrail/i }).click(); + await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click(); + + const dialog = page.getByRole("dialog", { name: "Create guardrail" }); + await expect(dialog).toBeVisible({ timeout: 10_000 }); + + await dialog.getByLabel("Guardrail Name").fill(guardrailName); + + const providerSelect = dialog.getByRole("combobox", { name: "Guardrail Provider" }); + await providerSelect.click(); + await providerSelect.fill("Presidio"); + await page.getByRole("option", { name: "Presidio PII" }).click(); + + await dialog.getByLabel("Mode", { exact: true }).click(); + await page.keyboard.type("pre_call"); + await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 }); + await page.keyboard.press("Enter"); + await expect(dialog.getByText("pre_call", { exact: true })).toBeVisible({ timeout: 5_000 }); + await dialog.getByText("Create guardrail", { exact: true }).click(); + + await dialog.getByLabel("presidio_analyzer_api_base").fill(MOCK_PRESIDIO_URL); + await dialog.getByLabel("presidio_anonymizer_api_base").fill(MOCK_PRESIDIO_URL); + + await dialog.getByRole("button", { name: "Next" }).click(); + await expect(dialog.getByText("Configure PII Protection")).toBeVisible({ timeout: 10_000 }); + await dialog.getByRole("button", { name: "Select All & Mask" }).click(); + + await dialog.getByRole("button", { name: "Create Guardrail" }).click(); + await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(1, { timeout: 15_000 }); +} + +async function deleteGuardrail(page: PlaywrightPage, guardrailName: string): Promise { + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + const row = page.getByRole("row").filter({ hasText: guardrailName }); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + await row.getByRole("button", { name: "Open guardrail actions" }).click(); + await page.getByRole("menuitem", { name: "Delete" }).click(); + const deleteModal = page.getByRole("dialog", { name: "Delete Guardrail" }); + await expect(deleteModal).toBeVisible({ timeout: 5_000 }); + await deleteModal.getByRole("button", { name: "Delete", exact: true }).click(); + await expect(page.getByText(`Guardrail "${guardrailName}" deleted successfully`)).toBeVisible({ timeout: 10_000 }); +} + +test.describe("Presidio PII guardrail, end to end from the dashboard", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("masks PII sent from the Playground and shows the run in Logs", async ({ page, request }) => { + const guardrailName = `e2e-presidio-story-${Date.now()}`; + const marker = `case-ref-${Math.random().toString(36).slice(2, 10)}`; + const prompt = `${marker}. Email me at ${RAW_EMAIL} or call ${RAW_PHONE}.`; + + await createPresidioGuardrail(page, guardrailName); + + await openPlayground(page); + await selectModel(page, CHAT_MODEL_A); + + const guardrailSelect = onlyVisible(page.getByPlaceholder("Select guardrails")); + await expect(guardrailSelect).toBeVisible({ timeout: 20_000 }); + await guardrailSelect.click(); + await guardrailSelect.fill(guardrailName); + await onlyVisible(page.getByRole("option", { name: guardrailName, exact: true })).click({ timeout: 20_000 }); + await page.keyboard.press("Escape"); + + const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false })); + await expect(input).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => { + await input.fill(prompt); + await sendButton(page).click(); + const res = await request.get(`${rootPath()}/spend/logs`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + if (!res.ok()) return false; + const rows: { metadata?: { applied_guardrails?: string[] } }[] = await res.json(); + return (Array.isArray(rows) ? rows : []).some((row) => + (row.metadata?.applied_guardrails ?? []).includes(guardrailName), + ); + }, + { + message: `the playground never produced a request that ran ${guardrailName}`, + timeout: 90_000, + intervals: [5_000], + }, + ) + .toBe(true); + + const requestId = await waitForSpendLogByPrompt(request, marker); + + const stored = await request.get(`${rootPath()}/spend/logs?request_id=${requestId}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(stored.ok(), `spend log read failed: ${stored.status()}`).toBe(true); + const storedBody = JSON.stringify(await stored.json()); + expect(storedBody, "the raw email reached the spend log, so the prompt was stored unmasked").not.toContain( + RAW_EMAIL, + ); + expect(storedBody, "the raw phone number reached the spend log, so the prompt was stored unmasked").not.toContain( + RAW_PHONE, + ); + expect(storedBody, "the stored prompt carries no EMAIL_ADDRESS placeholder, so nothing was masked").toContain( + "", + ); + expect(storedBody, "the stored prompt carries no PHONE_NUMBER placeholder, so nothing was masked").toContain( + "", + ); + + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + const search = visibleTestId(page, "datatable-search"); + await expect(search).toBeVisible({ timeout: 20_000 }); + await search.fill(requestId); + const row = requestLogsRows(page).filter({ hasText: requestId }); + await expect(row, `no logs row for request ${requestId}`).toHaveCount(1, { timeout: 30_000 }); + await row.click(); + + const drawer = page.getByRole("dialog").first(); + await expect(onlyVisible(drawer.getByText("Guardrails & Policy Compliance"))).toBeVisible({ timeout: 20_000 }); + await expect(onlyVisible(drawer.getByText(`Pre-call guardrail: ${guardrailName}`))).toBeVisible({ timeout: 20_000 }); + const maskedPrompt = drawer.getByText(`${marker}. Email me at or call .`); + await expect(onlyVisible(maskedPrompt)).toBeVisible({ timeout: 20_000 }); + + await onlyVisible(drawer.getByText("2 matched")).click(); + await expect(onlyVisible(drawer.getByText("Detected Entities (2)"))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("EMAIL_ADDRESS", { exact: true }))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("PHONE_NUMBER", { exact: true }))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("Score: 1.00", { exact: true }))).toBeVisible({ timeout: 10_000 }); + await expect(onlyVisible(drawer.getByText("Score: 0.75", { exact: true }))).toBeVisible({ timeout: 10_000 }); + + await expect(drawer.getByText(RAW_EMAIL)).toHaveCount(0); + await expect(drawer.getByText(RAW_PHONE)).toHaveCount(0); + + await deleteGuardrail(page, guardrailName); + }); +}); diff --git a/tests/e2e/ui/tests/migration/README.md b/tests/e2e/ui/tests/migration/README.md index d6b33598ec4..59933502463 100644 --- a/tests/e2e/ui/tests/migration/README.md +++ b/tests/e2e/ui/tests/migration/README.md @@ -1,17 +1,25 @@ # App Router migration smoke A growing E2E smoke for pages migrated from the legacy `?page=` switch to App -Router path routes. For each migrated page it clicks the page's sidebar link, checks -the URL is the path route and the page renders, reloads it, then clicks off to a -legacy page and back to confirm navigation still works. It runs in two situations: -the default mount and a non-root `SERVER_ROOT_PATH` mount. +Router path routes. For each page it clicks the sidebar link by its accessible +name, verifies the destination's content, reloads it, then visits Virtual Keys +and returns. It runs at the default mount and a non-root `SERVER_ROOT_PATH` mount + +Link selection does not depend on `href` formatting. URL assertions compare the +origin and pathname, allowing a trailing slash, query string, and fragment while +rejecting another route or mount. Reloads must return a successful document, +and each journey must finish without uncaught browser errors ## Adding a page -When a page's migration merges, add its route segment to -`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` -in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up -automatically. +Add an entry to `tests/e2e/ui/fixtures/migratedPages.ts`, keyed by the legacy page +ID. Specify its route segment, accessible link name, sidebar group if collapsed, +and distinctive visible content such as a heading or tab. Keep expectations +independent of the application's route table so an incorrect destination fails +the test. Both navigation suites use this fixture + +For a licensed-only page, `unlicensedText` describes the expected upgrade notice. +The authenticated session's license claim determines which content must render ## Running @@ -31,4 +39,9 @@ SERVER_ROOT_PATH=/litellm npm run e2e:migration:root ``` `globalSetup` logs in once per role; the admin storage state is reused for these -tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`. +tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login` + +`tests/navigation/sidebar.spec.ts` also checks the navigation helpers against +equivalent link formats on the live dashboard and a deep link containing a query +string and fragment. The link-format cases change only the rendered `href` +attribute to exercise the locator contract; destination pages and APIs remain live diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 547330190bd..f2dbd2dbc77 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -1,105 +1,73 @@ import { test, expect, type Page } from "@playwright/test"; -import { MIGRATED_E2E_SEGMENTS } from "../../fixtures/migratedPages"; +import { MIGRATED_E2E_PAGES, type MigratedPage } from "../../fixtures/migratedPages"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { dismissFeedbackPopup } from "../../helpers/navigation"; +import { clickSidebarLink, dismissFeedbackPopup, expectUiRoute, sidebarLink } from "../../helpers/navigation"; +import { proxyIsPremium } from "../../helpers/premium"; -/** - * App Router migration smoke as a user journey: start where the proxy lands you, - * click a migrated page in the sidebar, confirm it routed and rendered, reload it - * (the check a wrong server_root_path breaks), bounce to a legacy page and back, - * and, once two pages are migrated, navigate directly between two migrated pages. - * - * Driven by MIGRATED_E2E_SEGMENTS, so it grows as pages are migrated. Set - * SERVER_ROOT_PATH (e.g. "/litellm") to exercise the non-root mount; leave it - * unset for the default mount. Boot the proxy with the matching value first. - */ -const ROOT = process.env.SERVER_ROOT_PATH ?? ""; +const ROOT = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); +const apiKeys = MIGRATED_E2E_PAGES["api-keys"]; -const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); -const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`); -// Scope nav lookups to the sidebar (a `complementary` landmark). The top bar -// now renders a breadcrumb whose current-page item is also a "Virtual Keys" -// link, so an unscoped locator would match two elements. -const sidebar = (page: Page) => page.getByRole("complementary"); -const virtualKeysLink = (page: Page) => sidebar(page).getByRole("link", { name: "Virtual Keys", exact: true }); - -/** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */ -async function expectRendered(page: Page) { - await expect(virtualKeysLink(page)).toBeVisible({ timeout: 20_000 }); +async function expectContent(page: Page, destination: MigratedPage): Promise { + await expect(sidebarLink(page, apiKeys.linkName)).toBeVisible({ timeout: 20_000 }); + const main = page.getByRole("main"); + if (destination.unlicensedText && !proxyIsPremium()) { + await expect(main.getByText(destination.unlicensedText, { exact: true })).toBeVisible(); + return; + } + const content = destination.content; + const landmark = + "role" in content + ? main.getByRole(content.role, { name: content.name, exact: true }) + : main.getByText(content.text, { exact: true }); + await expect(landmark).toBeVisible(); } -/** - * Click a migrated page's sidebar link. Migrated items render as ; - * nested ones live under collapsible groups whose children only render while the - * group is open, so expand collapsed groups until the link is clickable. - */ -async function clickSidebar(page: Page, segment: string) { - const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); - const collapsedGroups = sidebar(page).getByRole("button", { expanded: false }); - for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - const stillCollapsed = await collapsedGroups.count(); - if (stillCollapsed === 0) break; - await collapsedGroups.first().click(); - await expect(collapsedGroups).toHaveCount(stillCollapsed - 1); - } - await link.click(); +async function navigateToDestination(page: Page, destination: MigratedPage): Promise { + await clickSidebarLink(page, destination.linkName, destination.group); + await expectUiRoute(page, destination.segment); + await dismissFeedbackPopup(page); + await expectContent(page, destination); } test.use({ storageState: ADMIN_STORAGE_PATH }); test.describe("App Router migrated pages", () => { - for (const segment of MIGRATED_E2E_SEGMENTS) { - test(`${segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => { + for (const destination of Object.values(MIGRATED_E2E_PAGES)) { + test(`${destination.segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => { const pageErrors: string[] = []; - page.on("pageerror", (e) => pageErrors.push(String(e))); + page.on("pageerror", (error) => pageErrors.push(String(error))); - // 1. Start where the proxy lands us. - await page.goto(`${ROOT}/ui/`); + const landing = await page.goto(`${ROOT}/ui/`); + expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); - await expectRendered(page); + await expectContent(page, apiKeys); - // 2. Click the migrated page in the sidebar -> path route + rendered. - await clickSidebar(page, segment); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - // 3. Reload the path route directly; a wrong server_root_path 404s here. - await page.reload(); + await navigateToDestination(page, destination); + + const reloaded = await page.reload(); + expect(reloaded?.ok(), `${destination.segment} document loads on reload`).toBe(true); + await expectUiRoute(page, destination.segment); await dismissFeedbackPopup(page); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - // 4. Click the Virtual Keys sidebar link to the api-keys landing (now a path route), then back. - await virtualKeysLink(page).click(); - await expect(page).toHaveURL(pathRe("api-keys")); - await dismissFeedbackPopup(page); - await expectRendered(page); - // 5. Click back to the migrated page. - await clickSidebar(page, segment); - await expect(page).toHaveURL(pathRe(segment)); - await expectRendered(page); - expect(pageErrors, `page errors during ${segment} journey`).toEqual([]); + await expectContent(page, destination); + + await navigateToDestination(page, apiKeys); + await navigateToDestination(page, destination); + expect(pageErrors, `page errors during ${destination.segment} journey`).toEqual([]); }); } test("navigates directly between two migrated pages", async ({ page }) => { - test.skip(MIGRATED_E2E_SEGMENTS.length < 2, "needs >= 2 migrated pages"); - const [first, second] = MIGRATED_E2E_SEGMENTS; const pageErrors: string[] = []; - page.on("pageerror", (e) => pageErrors.push(String(e))); + page.on("pageerror", (error) => pageErrors.push(String(error))); - await page.goto(`${ROOT}/ui/`); + const landing = await page.goto(`${ROOT}/ui/`); + expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); + await expectContent(page, apiKeys); - await clickSidebar(page, first); - await expect(page).toHaveURL(pathRe(first)); - await expectRendered(page); - await clickSidebar(page, second); - await expect(page).toHaveURL(pathRe(second)); - await expectRendered(page); - // Back to the first migrated page. - await clickSidebar(page, first); - await expect(page).toHaveURL(pathRe(first)); - await expectRendered(page); - - expect(pageErrors, "page errors during migrated -> migrated nav").toEqual([]); + for (const destination of [apiKeys, MIGRATED_E2E_PAGES.models, apiKeys]) { + await navigateToDestination(page, destination); + } + expect(pageErrors, "page errors during migrated page navigation").toEqual([]); }); }); diff --git a/tests/e2e/ui/tests/navigation/sidebar.spec.ts b/tests/e2e/ui/tests/navigation/sidebar.spec.ts index b220dc09ae2..eaa4985c659 100644 --- a/tests/e2e/ui/tests/navigation/sidebar.spec.ts +++ b/tests/e2e/ui/tests/navigation/sidebar.spec.ts @@ -3,7 +3,13 @@ import { Role } from "../../fixtures/roles"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { menuLabelToPage } from "../../fixtures/menuMappings"; -import { navigateToPage } from "../../helpers/navigation"; +import { + clickSidebarLink, + dismissFeedbackPopup, + expectUiRoute, + navigateToPage, + sidebarLink, +} from "../../helpers/navigation"; import { MIGRATED_E2E_PAGES } from "../../fixtures/migratedPages"; import type { Page as PlaywrightPage } from "@playwright/test"; @@ -11,7 +17,7 @@ const sidebarButtons = { [Role.ProxyAdmin]: [ "Virtual Keys", "Playground", - "Models", + "Models + Endpoints", "Usage", "Teams", "Internal Users", @@ -22,9 +28,9 @@ const sidebarButtons = { /** Migrated pages live at a path route; legacy pages keep the ?page= query param. */ async function expectPageUrl(page: PlaywrightPage, pageKey: string): Promise { - const migratedSegment = MIGRATED_E2E_PAGES[pageKey]; - if (migratedSegment) { - await expect(page).toHaveURL(new RegExp(`/ui/${migratedSegment}/?($|\\?)`)); + const migratedPage = MIGRATED_E2E_PAGES[pageKey]; + if (migratedPage) { + await expectUiRoute(page, migratedPage.segment); } else { await expect(page).toHaveURL(new RegExp(`[?&]page=${pageKey}(&|$)`)); } @@ -51,12 +57,7 @@ for (const { role, storage } of roles) { throw new Error(`No page mapping found for menu label: ${buttonLabel}`); } - // Sidebar items are links inside the `complementary` landmark; scoping - // there avoids the top-bar breadcrumb, which also links the page name. - const tab = page.getByRole("complementary").getByRole("link", { name: buttonLabel }); - await expect(tab).toBeVisible(); - - await tab.click(); + await clickSidebarLink(page, buttonLabel); await expectPageUrl(page, expectedPage); } @@ -81,5 +82,41 @@ for (const { role, storage } of roles) { await navigateToPage(page, Page.LlmPlayground); await expectPageUrl(page, Page.LlmPlayground); }); + + for (const format of ["without trailing slash", "absolute with query and fragment", "relative"] as const) { + test(`sidebar locator tolerates hrefs ${format}`, async ({ page }) => { + await page.goto("/ui/"); + await dismissFeedbackPopup(page); + const link = sidebarLink(page, "Models + Endpoints"); + await expect(link).toBeVisible(); + const destination = new URL("/ui/models-and-endpoints/", page.url()); + const href = + format === "without trailing slash" + ? destination.pathname.replace(/\/$/, "") + : format === "relative" + ? "./models-and-endpoints/" + : `${destination.href}?source=navigation-smoke#overview`; + + await link.evaluate((element, value) => element.setAttribute("href", value), href); + await expect(link).toHaveAttribute("href", href); + await clickSidebarLink(page, "Models + Endpoints"); + + await expectUiRoute(page, "models-and-endpoints"); + await expect( + page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }), + ).toBeVisible(); + }); + } + + test("route assertion tolerates a query string and fragment on a deep link", async ({ page }) => { + const response = await page.goto("/ui/models-and-endpoints/?source=navigation-smoke#overview"); + expect(response?.ok()).toBe(true); + await expectUiRoute(page, "models-and-endpoints"); + await expect( + page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }), + ).toBeVisible(); + expect(new URL(page.url()).search).toBe("?source=navigation-smoke"); + expect(new URL(page.url()).hash).toBe("#overview"); + }); }); } diff --git a/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts new file mode 100644 index 00000000000..f7930378b1e --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts @@ -0,0 +1,145 @@ +import { test, expect, type APIRequestContext, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { proxyIsPremium } from "../../helpers/premium"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function guardrailId(request: APIRequestContext, name: string): Promise { + const res = await request.get("/v2/guardrails/list", { headers: auth() }); + expect(res.ok(), `GET /v2/guardrails/list (${res.status()})`).toBe(true); + const rows = (await res.json()).guardrails as { guardrail_id: string; guardrail_name: string | null }[]; + return rows.find((row) => row.guardrail_name === name)?.guardrail_id; +} + +async function teamGuardrails(page: PlaywrightPage, teamId: string): Promise { + const body = await readBack<{ team_info: { metadata: { guardrails?: string[] } | null } }>( + page, + `/team/info?team_id=${encodeURIComponent(teamId)}`, + ); + return body.team_info.metadata?.guardrails ?? []; +} + +async function keywordPromptStatus(request: APIRequestContext, apiKey: string, keyword: string): Promise { + const res = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `please tell me about ${keyword}` }] }, + }); + return res.status(); +} + +test.describe("Proxy Admin - Team guardrail removal", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Clearing a team's only guardrail on the Settings tab lets blocked traffic through again", async ({ + page, + request, + }) => { + test.skip(!proxyIsPremium(), "proxy under test is unlicensed, so team guardrails are premium-gated"); + + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const guardrailName = `e2e-team-guardrail-${stamp}`; + const bannedKeyword = `e2eteamban${stamp}`; + const teamAlias = `e2e-guardrail-team-${stamp}`; + + let teamId = ""; + let teamKey = ""; + try { + const guardrailRes = await request.post("/guardrails", { + headers: auth(), + data: { + guardrail: { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "litellm_content_filter", + mode: "pre_call", + default_on: false, + blocked_words: [{ keyword: bannedKeyword, action: "BLOCK" }], + }, + }, + }, + }); + expect( + guardrailRes.ok(), + `POST /guardrails failed (${guardrailRes.status()}): ${await guardrailRes.text()}`, + ).toBe(true); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { team_alias: teamAlias, models: [CHAT_MODEL_A], metadata: { guardrails: [guardrailName] } }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + teamId = (await teamRes.json()).team_id as string; + + const keyRes = await request.post("/key/generate", { headers: auth(), data: { team_id: teamId } }); + expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); + teamKey = (await keyRes.json()).key as string; + + await expect + .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), { + message: "the team's guardrail never started refusing the banned keyword", + timeout: 60_000, + }) + .toBe(400); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + const chip = page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }); + await expect(chip).toBeVisible({ timeout: 10_000 }); + await chip.locator('[data-slot="combobox-chip-remove"]').click(); + await expect(chip).toHaveCount(0, { timeout: 10_000 }); + + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => teamGuardrails(page, teamId), { + message: "the team still carries a guardrail in /team/info after the save", + timeout: 20_000, + }) + .toEqual([]); + + await page.reload(); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await expect(page.getByRole("combobox", { name: "Select guardrails" })).toBeVisible({ timeout: 15_000 }); + await expect( + page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }), + "the removed guardrail is gone from the Settings tab after a reload", + ).toHaveCount(0); + + await expect + .poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), { + message: "the team key is still refused for a keyword whose guardrail was removed", + timeout: 60_000, + }) + .toBe(200); + + const served = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: `please tell me about ${bannedKeyword}` }], + }, + }); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + } finally { + if (teamKey) { + await request.post("/key/delete", { headers: auth(), data: { keys: [teamKey] } }); + } + if (teamId) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + const id = await guardrailId(request, guardrailName); + if (id) { + await request.delete(`/guardrails/${id}`, { headers: auth() }); + } + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts new file mode 100644 index 00000000000..c7325e78f75 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts @@ -0,0 +1,129 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; + +interface TeamInfoResponse { + team_info: { + models: string[]; + members_with_roles: { user_id?: string; role?: string }[]; + }; + team_memberships: { + user_id: string; + litellm_budget_table: { max_budget: number | null } | null; + }[]; +} + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function teamInfo(page: PlaywrightPage, teamId: string): Promise { + return readBack(page, `/team/info?team_id=${encodeURIComponent(teamId)}`); +} + +function roleOf(info: TeamInfoResponse, userId: string): string | undefined { + return info.team_info.members_with_roles.find((member) => member.user_id === userId)?.role; +} + +function budgetOf(info: TeamInfoResponse, userId: string): number | null | undefined { + return info.team_memberships.find((membership) => membership.user_id === userId)?.litellm_budget_table?.max_budget; +} + +function otherMembers(info: TeamInfoResponse, userId: string): string[] { + return info.team_info.members_with_roles + .filter((member) => member.user_id !== userId) + .map((member) => `${member.user_id}:${member.role}`) + .sort(); +} + +test.describe("Proxy Admin - Team member edit", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + const createdTeams: string[] = []; + const createdUsers: string[] = []; + + test.afterEach(async ({ request }) => { + for (const teamId of createdTeams.splice(0)) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + for (const userId of createdUsers.splice(0)) { + await request.post("/user/delete", { headers: auth(), data: { user_ids: [userId] } }); + } + }); + + test("Editing a member's role and per-member budget persists and survives a reload", async ({ page, request }) => { + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const memberId = `e2e-member-edit-${stamp}`; + const teamAlias = `e2e-member-edit-team-${stamp}`; + + createdUsers.push(memberId); + const created = await request.post("/user/new", { + headers: auth(), + data: { user_id: memberId, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { + team_alias: teamAlias, + models: [CHAT_MODEL_A], + members_with_roles: [{ user_id: memberId, role: "admin" }], + }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + const teamId = (await teamRes.json()).team_id as string; + createdTeams.push(teamId); + + const before = await teamInfo(page, teamId); + expect(roleOf(before, memberId), "the member starts out as a team admin").toBe("admin"); + expect(budgetOf(before, memberId) ?? null, "the member starts out with no per-member budget").toBeNull(); + expect( + otherMembers(before, memberId).length, + "the team has another member for the edit to leave alone", + ).toBeGreaterThan(0); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Members" }).click(); + + const memberRow = page.locator("tr", { hasText: memberId }).first(); + await expect(memberRow).toBeVisible({ timeout: 10_000 }); + await memberRow.getByTestId("edit-member").click(); + + const modal = page.getByRole("dialog", { name: "Edit Member" }); + await expect(modal).toBeVisible({ timeout: 10_000 }); + + await modal.getByLabel(/^Role/).click(); + await page.getByRole("option", { name: "User", exact: true }).click(); + await modal.getByLabel(/Team Member Budget \(USD\)/).fill("5"); + await modal.getByRole("button", { name: "Save Changes" }).click(); + + await expect(page.getByText("Team member updated successfully").first()).toBeVisible({ timeout: 10_000 }); + + await expect + .poll( + async () => { + const info = await teamInfo(page, teamId); + return [roleOf(info, memberId), budgetOf(info, memberId)]; + }, + { message: "the member's role and budget never landed in /team/info", timeout: 20_000 }, + ) + .toEqual(["user", 5]); + + await page.reload(); + await page.getByRole("tab", { name: "Members" }).click(); + const reloadedRow = page.locator("tr", { hasText: memberId }).first(); + await expect(reloadedRow).toBeVisible({ timeout: 15_000 }); + await expect(reloadedRow.getByText("user", { exact: true }), "role shown after a reload").toBeVisible(); + await expect(reloadedRow.getByText("$5.00"), "per-member budget shown after a reload").toBeVisible(); + + const after = await teamInfo(page, teamId); + expect(after.team_info.models, "model access untouched by a member edit").toEqual(before.team_info.models); + expect(otherMembers(after, memberId), "the rest of the roster untouched by a member edit").toEqual( + otherMembers(before, memberId), + ); + }); +}); diff --git a/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts new file mode 100644 index 00000000000..75fb3be9b64 --- /dev/null +++ b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts @@ -0,0 +1,177 @@ +import { test, expect, type Browser, type BrowserContext, type Page as PlaywrightPage } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +const PASSWORD = "E2e-Member-Perms-Pass-1!"; + +const auth = () => ({ Authorization: `Bearer ${masterKey()}` }); + +async function sessionKey(page: PlaywrightPage): Promise { + const cookie = (await page.context().cookies()).find((candidate) => candidate.name === "token"); + expect(cookie?.value, "logged-in session carries a token cookie").toBeTruthy(); + const payload = JSON.parse(Buffer.from(cookie!.value.split(".")[1], "base64url").toString("utf-8")) as { + key?: string; + }; + expect(payload.key, "session JWT carries the virtual key the dashboard calls with").toMatch(/^sk-/); + return payload.key!; +} + +async function signIn(browser: Browser, email: string): Promise { + const context = await browser.newContext({ storageState: { cookies: [], origins: [] } }); + const page = await context.newPage(); + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page.getByPlaceholder("Enter your password").fill(PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + return context; +} + +test.describe("Team Admin - Member permissions", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("Granting /key/generate lets a plain member create a team key that serves traffic", async ({ + browser, + request, + }) => { + const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`; + const adminId = `e2e-perm-admin-${stamp}`; + const memberId = `e2e-perm-member-${stamp}`; + const adminEmail = `${adminId}@test.local`; + const memberEmail = `${memberId}@test.local`; + const teamAlias = `e2e-perm-team-${stamp}`; + + const createUser = async (userId: string, email: string): Promise => { + const created = await request.post("/user/new", { + headers: auth(), + data: { user_id: userId, user_email: email, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new for ${userId} (${created.status()}): ${await created.text()}`).toBe(true); + const password = await request.post("/user/update", { + headers: auth(), + data: { user_id: userId, password: PASSWORD }, + }); + expect(password.ok(), `POST /user/update for ${userId} (${password.status()})`).toBe(true); + }; + + let teamId = ""; + const createdKeys: string[] = []; + const contexts: BrowserContext[] = []; + try { + await createUser(adminId, adminEmail); + await createUser(memberId, memberEmail); + + const teamRes = await request.post("/team/new", { + headers: auth(), + data: { + team_alias: teamAlias, + models: [CHAT_MODEL_A], + members_with_roles: [ + { user_id: adminId, role: "admin" }, + { user_id: memberId, role: "user" }, + ], + }, + }); + expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + teamId = (await teamRes.json()).team_id as string; + + const memberContext = await signIn(browser, memberEmail); + contexts.push(memberContext); + const memberPage = memberContext.pages()[0]; + const memberSessionKey = await sessionKey(memberPage); + + const refused = await memberPage.request.post("/key/generate", { + headers: { Authorization: `Bearer ${memberSessionKey}`, "Content-Type": "application/json" }, + data: { team_id: teamId, key_alias: `e2e-perm-denied-${stamp}` }, + }); + expect(refused.status(), "a plain member cannot mint a team key before the grant").toBe(401); + expect(await refused.text()).toContain("/key/generate"); + + const adminContext = await signIn(browser, adminEmail); + contexts.push(adminContext); + const adminPage = adminContext.pages()[0]; + await navigateToPage(adminPage, Page.Teams); + await clickTeamId(adminPage, teamId); + await adminPage.getByRole("tab", { name: "Member Permissions" }).click(); + + for (const route of ["/key/generate", "/key/update"]) { + await adminPage.getByRole("row").filter({ hasText: route }).getByRole("checkbox").check(); + } + await adminPage.getByRole("button", { name: "Save Changes" }).click(); + await expect(adminPage.getByText("Permissions updated successfully").first()).toBeVisible({ timeout: 10_000 }); + + await expect + .poll( + async () => { + const res = await request.get(`/team/permissions_list?team_id=${encodeURIComponent(teamId)}`, { + headers: auth(), + }); + if (!res.ok()) return []; + return ((await res.json()).team_member_permissions ?? []) as string[]; + }, + { message: "the granted permissions never landed in /team/permissions_list", timeout: 20_000 }, + ) + .toEqual(expect.arrayContaining(["/key/generate", "/key/update"])); + + const keyAlias = `e2e-perm-key-${stamp}`; + await navigateToPage(memberPage, Page.ApiKeys); + await memberPage.getByRole("button", { name: /Create New Key/i }).click(); + await expect(memberPage.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + await memberPage.getByLabel(/Key Name/).fill(keyAlias); + + const teamSelect = memberPage.getByTestId("team-dropdown").getByRole("combobox"); + await teamSelect.click(); + await memberPage.keyboard.type(teamAlias); + await memberPage.getByRole("option", { name: teamAlias }).first().click(); + + await memberPage.getByRole("combobox", { name: "Select models" }).click(); + await memberPage.getByRole("option", { name: "All Team Models", exact: true }).click(); + await memberPage.keyboard.press("Escape"); + + await memberPage.getByRole("button", { name: "Create Key", exact: true }).click(); + const saveDialog = memberPage.getByRole("dialog", { name: "Save your Key" }); + await expect(saveDialog).toBeVisible({ timeout: 15_000 }); + const apiKey = (await saveDialog.locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + createdKeys.push(apiKey); + await memberPage.keyboard.press("Escape"); + + await expect + .poll( + async () => { + const res = await request.get( + `/key/list?team_id=${encodeURIComponent(teamId)}&return_full_object=true&size=100`, + { headers: auth() }, + ); + if (!res.ok()) return null; + const row = ((await res.json()).keys as Record[]).find( + (candidate) => candidate.key_alias === keyAlias, + ); + return row ? [row.user_id, row.team_id] : null; + }, + { message: `key ${keyAlias} never appeared on the team with the member as its owner`, timeout: 20_000 }, + ) + .toEqual([memberId, teamId]); + + const served = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `member key ping ${stamp}` }] }, + }); + expect(served.status(), "the delegated key is a real key the gateway serves").toBe(200); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + } finally { + for (const context of contexts) { + await context.close(); + } + for (const key of createdKeys) { + await request.post("/key/delete", { headers: auth(), data: { keys: [key] } }); + } + if (teamId) { + await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } }); + } + await request.post("/user/delete", { headers: auth(), data: { user_ids: [adminId, memberId] } }); + } + }); +}); diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 7bc61c40ea0..c2272f3d20d 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -169,6 +169,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) matching = [row for row in rows if row["router_name"] == router] assert len(matching) == 1 @@ -181,6 +182,33 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): assert grouped["session_seconds"] == pytest.approx(60.0) +async def test_the_benchmarks_aggregate_can_filter_to_one_key(db): + router = f"r-{uuid.uuid4()}" + first_key = f"k-{uuid.uuid4()}" + second_key = f"k-{uuid.uuid4()}" + await _turn(db, first_key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, saved=0.5) + await _turn(db, second_key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, saved=9.0) + + rows = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + first_key, + ) + matching = [row for row in rows if row["router_name"] == router] + assert len(matching) == 1 + assert matching[0]["sessions"] == 1 + assert matching[0]["saved_spend"] == pytest.approx(0.5) + + unknown_key_rows = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + f"k-{uuid.uuid4()}", + ) + assert [row for row in unknown_key_rows if row["router_name"] == router] == [] + + async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db): key = f"k-{uuid.uuid4()}" router = f"r-{uuid.uuid4()}" @@ -191,6 +219,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) matching = sorted( (row for row in rows if row["router_name"] == router), @@ -253,6 +282,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {"simple": 2, "complex": 1} @@ -280,6 +310,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} @@ -294,6 +325,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {} diff --git a/tests/rust-python-harness/shared/parity/compare.py b/tests/rust-python-harness/shared/parity/compare.py index adf85e5c8d7..88239ed042a 100644 --- a/tests/rust-python-harness/shared/parity/compare.py +++ b/tests/rust-python-harness/shared/parity/compare.py @@ -78,3 +78,4 @@ def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent validate_harness(baseline, candidate, baseline_user_agent) assert_request_parity(baseline.requests, candidate.requests) assert_value_parity(baseline.report, candidate.report) + assert_value_parity(baseline.callbacks, candidate.callbacks, path="$.callbacks") diff --git a/tests/rust-python-harness/shared/parity/models.py b/tests/rust-python-harness/shared/parity/models.py index 898b58d23ee..5612e218824 100644 --- a/tests/rust-python-harness/shared/parity/models.py +++ b/tests/rust-python-harness/shared/parity/models.py @@ -37,6 +37,25 @@ class SDKError(BaseModel): llm_provider: str | None +class CallbackObservation(BaseModel): + model_config = ConfigDict(frozen=True) + + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ] + phase: Literal["success", "failure"] + model: str | None + call_type: str | None + litellm_call_id: str | None + metadata: JsonValue + kwargs: JsonValue + payload: JsonValue + error: SDKError | None + + class SDKJsonChunk(BaseModel): model_config = ConfigDict(frozen=True) @@ -119,6 +138,7 @@ class Execution(BaseModel): requests: tuple[CapturedRequest, ...] report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class SDKCommand(BaseModel): @@ -133,6 +153,7 @@ class WorkerSuccess(BaseModel): status: Literal["ok"] = "ok" report: SDKReport + callbacks: tuple[CallbackObservation, ...] | None = None class WorkerFailure(BaseModel): diff --git a/tests/rust-python-harness/shared/parity/runner.py b/tests/rust-python-harness/shared/parity/runner.py index 43a583382cb..feae0201aee 100644 --- a/tests/rust-python-harness/shared/parity/runner.py +++ b/tests/rust-python-harness/shared/parity/runner.py @@ -39,9 +39,7 @@ class SubprocessRunner: return ( sys.executable, "-m", - ".".join( - self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts - ), + ".".join(self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts), "--parity-worker", provider_url, ) @@ -113,7 +111,11 @@ class SubprocessWorker: ) assert isinstance(result, WorkerSuccess) try: - return Execution(requests=self.provider.take_requests(len(responses)), report=result.report) + return Execution( + requests=self.provider.take_requests(len(responses)), + report=result.report, + callbacks=result.callbacks, + ) except AssertionError: self.provider.reset() raise diff --git a/tests/rust-python-harness/shared/parity/test_parity.py b/tests/rust-python-harness/shared/parity/test_parity.py index 83daccdf8ba..7c8355fbfd8 100644 --- a/tests/rust-python-harness/shared/parity/test_parity.py +++ b/tests/rust-python-harness/shared/parity/test_parity.py @@ -7,7 +7,7 @@ import pytest from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr from .compare import assert_model_parity, assert_parity -from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report +from .models import CallbackObservation, CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report SENTINEL: Final = "python-parity-fallback" @@ -69,6 +69,56 @@ def test_parity_rejects_response_difference() -> None: assert_parity(python, rust, SENTINEL) +def test_parity_distinguishes_unobserved_callbacks_from_zero_events() -> None: + python: Final = _execution(user_agent=SENTINEL) + rust: Final = _execution(user_agent="litellm-rust").model_copy(update={"callbacks": ()}) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_payload_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"payload": {"model": "changed", "pages": []}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + +def test_parity_rejects_callback_kwargs_difference() -> None: + observation: Final = CallbackObservation( + hook="log_success_event", + phase="success", + model="test-model", + call_type="ocr", + litellm_call_id="test-call", + metadata={"profile": "success"}, + kwargs={"model": "test-model"}, + payload={"model": "test-model", "pages": []}, + error=None, + ) + python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)}) + rust: Final = _execution(user_agent="litellm-rust").model_copy( + update={"callbacks": (observation.model_copy(update={"kwargs": {"model": "changed"}}),)} + ) + + with pytest.raises(AssertionError, match=r"\$\.callbacks"): + assert_parity(python, rust, SENTINEL) + + def test_parity_rejects_error_difference() -> None: python: Final = Execution( requests=(), diff --git a/tests/rust-python-harness/shared/tracing/native.py b/tests/rust-python-harness/shared/tracing/native.py index 4f988f65294..688995cbc4b 100644 --- a/tests/rust-python-harness/shared/tracing/native.py +++ b/tests/rust-python-harness/shared/tracing/native.py @@ -19,7 +19,8 @@ class _TraceEventPayload(BaseModel): class TraceResponsePayload(BaseModel): model_config = ConfigDict(strict=True, extra="forbid") - response: object + response: object = None + error: str | None = None trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload] diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py index f668e178eef..95346bbd496 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py @@ -18,8 +18,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", note=( - "Recorded sync/async SDK parity; invalid-model provider errors differ, " - "and Reducto lacks a Rust contract." + "Recorded sync/async SDK parity with focused success/error callback profiles; " + "Reducto lacks a Rust contract, and known provider parity gaps remain." ), ), surface="sdk", diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index e72980752f2..5c4abc78081 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -1,23 +1,31 @@ from __future__ import annotations import asyncio +import datetime +import queue import sys import tempfile +import time import traceback -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator, Mapping from contextlib import contextmanager from enum import Enum from functools import partial from pathlib import Path from typing import Annotated, Final, Literal, cast +from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from .....shared.parity.compare import assert_parity from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures from .....shared.parity.models import ( + JSON_VALUE_ADAPTER, + CallbackObservation, + Execution, SDKCommand, SDKError, SDKReport, @@ -39,6 +47,9 @@ from .fixtures.config import configured_fixture_directory from .fixtures.models import OcrParityCase, OcrSdkInput API_KEY: Final = "test-key" +CALLBACK_DELAY_SECONDS: Final = 0.05 +CALLBACK_DRAIN_TIMEOUT_SECONDS: Final = 10.0 +CALLBACK_TERMINALS: Final[tuple[Literal["success", "failure"], ...]] = ("success", "failure") PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),)) RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),)) @@ -75,10 +86,148 @@ class InvalidOcrWorkerCase(BaseModel): case: InvalidOcrCase -OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")] +class CallbackOcrWorkerCase(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["callback"] = "callback" + case: OcrParityCase + terminal: Literal["success", "failure"] + + +OcrWorkerCase = Annotated[ + RecordedOcrWorkerCase | InvalidOcrWorkerCase | CallbackOcrWorkerCase, + Field(discriminator="kind"), +] OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase) +class RecordingCallback(CustomLogger): + def __init__(self) -> None: + self.message_logging: Final = True + self.turn_off_message_logging: Final = False + self._observations: Final[queue.SimpleQueue[CallbackObservation]] = queue.SimpleQueue() + + def _normalized_kwargs(self, value: object, key: str | None = None) -> JsonValue: + if value is None or isinstance(value, (bool, int, float)): + return value + if isinstance(value, str): + if key != "api_base": + return value + parsed: Final = urlsplit(value) + return urlunsplit(("", "", parsed.path, parsed.query, parsed.fragment)) + if isinstance(value, datetime.datetime): + return "datetime" + if isinstance(value, Exception): + return sdk_error_report(value).model_dump(mode="json") + if isinstance(value, BaseModel): + return self._normalized_kwargs(value.model_dump(mode="json"), key) + if isinstance(value, Mapping): + if any(not isinstance(map_key, str) for map_key in value): + raise TypeError("callback kwarg mappings must use string keys") + return { + map_key: self._normalized_kwargs(map_value, map_key) + for map_key, map_value in value.items() + } + if isinstance(value, (list, tuple)): + return [self._normalized_kwargs(item) for item in value] + raise TypeError(f"unsupported callback kwarg type: {type(value)}") + + def _record( + self, + hook: Literal[ + "log_success_event", + "async_log_success_event", + "log_failure_event", + "async_log_failure_event", + ], + phase: Literal["success", "failure"], + kwargs: dict[str, object], + response_obj: object, + ) -> None: + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_litellm_params) if isinstance(raw_litellm_params, Mapping) else {} + ) + raw_metadata: Final = litellm_params.get("metadata") + metadata_mapping: Final[Mapping[str, object]] = ( + cast(Mapping[str, object], raw_metadata) if isinstance(raw_metadata, Mapping) else {} + ) + metadata: Final = JSON_VALUE_ADAPTER.validate_python( + {key: metadata_mapping[key] for key in ("callback_profile", "sdk_route") if key in metadata_mapping} + ) + raw_error: Final = kwargs.get("exception") + error: Final = sdk_error_report(raw_error) if isinstance(raw_error, Exception) else None + normalized_kwargs: Final = self._normalized_kwargs(kwargs) + payload_source: Final = ( + response_obj.model_dump(mode="json") if isinstance(response_obj, BaseModel) else response_obj + ) + payload: Final = JSON_VALUE_ADAPTER.validate_python(payload_source) + raw_model: Final = kwargs.get("model") + raw_call_type: Final = kwargs.get("call_type") + raw_call_id: Final = kwargs.get("litellm_call_id") + self._observations.put( + CallbackObservation( + hook=hook, + phase=phase, + model=raw_model if isinstance(raw_model, str) else None, + call_type=str(raw_call_type) if raw_call_type is not None else None, + litellm_call_id=raw_call_id if isinstance(raw_call_id, str) else None, + metadata=metadata, + kwargs=normalized_kwargs, + payload=payload, + error=error, + ) + ) + + def log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_success_event", "success", kwargs, response_obj) + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_success_event", "success", kwargs, response_obj) + + def log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + time.sleep(CALLBACK_DELAY_SECONDS) + self._record("log_failure_event", "failure", kwargs, response_obj) + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + del start_time, end_time + await asyncio.sleep(CALLBACK_DELAY_SECONDS) + self._record("async_log_failure_event", "failure", kwargs, response_obj) + + def observations(self) -> tuple[CallbackObservation, ...]: + observations: Final = tuple(self._observations.get_nowait() for _ in range(self._observations.qsize())) + return tuple(sorted(observations, key=lambda observation: observation.hook)) + + INVALID_OCR_CASES: Final = ( InvalidOcrCase( name="unsupported_provider", @@ -229,6 +378,106 @@ def _execute_invalid_sdk_case( return _execute_sdk_call(call_kwargs, route, event_loop) +def _callback_call_id(route: SDKRoute, terminal: Literal["success", "failure"]) -> str: + return f"ocr-callback-{route.value}-{terminal}" + + +def _callback_metadata(route: SDKRoute, terminal: Literal["success", "failure"]) -> dict[str, str]: + return {"callback_profile": terminal, "sdk_route": route.value} + + +def _drain_callback_delivery(route: SDKRoute, event_loop: asyncio.AbstractEventLoop) -> None: + if route is SDKRoute.AOCR: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def drain_async_callbacks() -> None: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=CALLBACK_DRAIN_TIMEOUT_SECONDS) + await GLOBAL_LOGGING_WORKER.stop() + + event_loop.run_until_complete(drain_async_callbacks()) + + from litellm.litellm_core_utils.thread_pool_executor import executor + + executor.shutdown(wait=True, cancel_futures=False) + + +def _execute_callback_sdk_case( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + mock_url: str, + event_loop: asyncio.AbstractEventLoop, +) -> WorkerSuccess: + callback: Final = RecordingCallback() + call_kwargs: Final = { + **_call_kwargs(case.litellm_input, mock_url, route), + "callbacks": [callback], + "litellm_call_id": _callback_call_id(route, terminal), + "litellm_trace_id": _callback_call_id(route, terminal), + "metadata": _callback_metadata(route, terminal), + } + report: Final = _execute_sdk_call(call_kwargs, route, event_loop) + _drain_callback_delivery(route, event_loop) + return WorkerSuccess(report=report, callbacks=callback.observations()) + + +def _assert_callback_lifecycle( + execution: Execution, + route: SDKRoute, + terminal: Literal["success", "failure"], +) -> None: + observations: Final = execution.callbacks + assert observations is not None, f"{route.value} {terminal} callbacks were not observed" + expected_hooks: Final = ( + (f"log_{terminal}_event",) + if route is SDKRoute.OCR + else ("async_log_success_event",) + if terminal == "success" + else ("async_log_failure_event", "log_failure_event") + ) + actual_hooks: Final = tuple(observation.hook for observation in observations) + assert actual_hooks == expected_hooks, ( + f"{route.value} {terminal} expected callback hooks {expected_hooks}, received {actual_hooks}" + ) + expected_call_id: Final = _callback_call_id(route, terminal) + expected_metadata: Final = _callback_metadata(route, terminal) + for observation in observations: + assert observation.phase == terminal + assert observation.call_type == route.value + assert observation.litellm_call_id == expected_call_id + assert observation.metadata == expected_metadata + assert observation.model + if terminal == "success": + assert isinstance(execution.report, SDKSuccess) + assert observation.payload == execution.report.response + assert observation.error is None + else: + assert isinstance(execution.report, SDKError) + assert observation.payload is None + assert observation.error is not None + assert observation.error.exception_type + assert observation.error.message + assert observation.error.status_code is not None + assert observation.error.status_code >= 400 + + +def _check_callback_ocr_sdk_parity( + case: OcrParityCase, + route: SDKRoute, + terminal: Literal["success", "failure"], + case_file: Path, + runner: SubprocessRunner, +) -> None: + with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: + python_worker, rust_worker = workers + python: Final = python_worker.execute(case_file, route.value, case.provider_responses) + rust: Final = rust_worker.execute(case_file, route.value, case.provider_responses) + + _assert_callback_lifecycle(python, route, terminal) + _assert_callback_lifecycle(rust, route, terminal) + assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + + def _check_recorded_ocr_sdk_parity( ocr_fixture: OcrParityCase, route: SDKRoute, @@ -276,6 +525,25 @@ def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path return case_file +def _callback_fixture( + fixtures: tuple[OcrParityCase, ...], + terminal: Literal["success", "failure"], +) -> OcrParityCase: + matching: Final = tuple( + fixture + for fixture in fixtures + if fixture.litellm_input.contract == "mistral" + and ( + all(response.status_code < 400 for response in fixture.provider_responses) + if terminal == "success" + else any(response.status_code >= 400 for response in fixture.provider_responses) + ) + ) + if not matching: + raise AssertionError(f"no recorded Mistral OCR {terminal} fixture is available for callback parity") + return min(matching, key=lambda fixture: fixture_id(fixture.litellm_input, fixture.litellm_input.model)) + + @contextmanager def parity_checks() -> Generator[tuple[E2ECheck, ...]]: fixtures: Final = tuple( @@ -298,6 +566,17 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: _write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case)) for index, case in enumerate(INVALID_OCR_CASES) ) + callback_cases: Final[tuple[tuple[Literal["success", "failure"], OcrParityCase], ...]] = tuple( + (terminal, _callback_fixture(fixtures, terminal)) for terminal in CALLBACK_TERMINALS + ) + callback_files: Final = tuple( + _write_worker_case( + directory, + len(recorded_files) + len(invalid_files) + index, + CallbackOcrWorkerCase(case=case, terminal=terminal), + ) + for index, (terminal, case) in enumerate(callback_cases) + ) with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: recorded: Final = tuple( E2ECheck( @@ -315,7 +594,15 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]: for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True) for route in SDKRoute ) - yield (*recorded, *invalid) + callbacks: Final = tuple( + E2ECheck( + f"callback:{route.value}:{terminal}", + partial(_check_callback_ocr_sdk_parity, case, route, terminal, case_file, runner), + ) + for (terminal, case), case_file in zip(callback_cases, callback_files, strict=True) + for route in SDKRoute + ) + yield (*recorded, *invalid, *callbacks) def _execute_worker_command( @@ -333,6 +620,8 @@ def _execute_worker_command( return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop)) case InvalidOcrWorkerCase(case=invalid): return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop)) + case CallbackOcrWorkerCase(case=callback_case, terminal=terminal): + return _execute_callback_sdk_case(callback_case, route, terminal, mock_url, event_loop) except Exception: return WorkerFailure(error=traceback.format_exc()) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index da560f99730..04659b25382 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -16,6 +16,7 @@ TraceFailureSource = Literal["python", "rust", "harness"] class RouteFixture: kwargs: dict[str, object] provider_responses: tuple[RecordedHttpResponse, ...] + expected_failure: bool = False @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index f8d7c55d4e2..eb6c9233565 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -2,15 +2,16 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable +from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast from ....shared.parity.replay import replay_server from ....shared.reporting.models import Surface -from ....shared.tracing.native import native_trace_events +from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, pipeline_projection -from ..models import RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario +from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario from ..reporting import TraceComparisonArtifact @@ -18,9 +19,22 @@ class SdkCall(Protocol): def __call__(self, **kwargs: object) -> object: ... +@dataclass(frozen=True, slots=True) +class _CollectedTrace: + events: tuple[FunctionTraceEvent, ...] + error: str | None = None + + def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: async def invoke_async() -> object: - return await cast(Awaitable[object], function(**kwargs)) + try: + return await cast(Awaitable[object], function(**kwargs)) + finally: + await asyncio.sleep(0) + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + await GLOBAL_LOGGING_WORKER.stop() if asynchronous: return asyncio.run(invoke_async()) @@ -48,16 +62,30 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) +def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None: + try: + _invoke(function, kwargs, asynchronous=asynchronous) + except Exception as error: + return f"{type(error).__name__}: {error}" + return None + + def _collect( - function: SdkCall, kwargs: dict[str, object], engine: Engine, *, asynchronous: bool -) -> tuple[FunctionTraceEvent, ...]: + function: SdkCall, + fixture: RouteFixture, + engine: Engine, + *, + asynchronous: bool, +) -> _CollectedTrace: + kwargs: Final = fixture.kwargs if engine == "rust": - return native_trace_events(_invoke(function, kwargs, asynchronous=asynchronous)) + payload: Final = TraceResponsePayload.model_validate(_invoke(function, kwargs, asynchronous=asynchronous)) + return _CollectedTrace(native_trace_events(payload), payload.error) import litellm with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: - _invoke(function, kwargs, asynchronous=asynchronous) - return tuple(profiler.events) + error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous) + return _CollectedTrace(tuple(profiler.events), error) def collect_trace( @@ -68,22 +96,30 @@ def collect_trace( return function try: with replay_server() as provider: - fixture: Final = spec.fixture(engine, provider.url) - for response in fixture.provider_responses: + base_fixture: Final = spec.fixture(engine, provider.url) + for response in base_fixture.provider_responses: provider.enqueue_response(response) - kwargs: Final = { - **fixture.kwargs, - "api_key": "test-key", - "api_base": provider.url, - **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), - } - events: Final = _collect(function, kwargs, engine, asynchronous=asynchronous) + fixture: Final = RouteFixture( + kwargs={ + **base_fixture.kwargs, + "api_key": "test-key", + "api_base": provider.url, + **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), + }, + provider_responses=base_fixture.provider_responses, + expected_failure=base_fixture.expected_failure, + ) + collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") - if not events: + if fixture.expected_failure and collected.error is None: + return TraceExecutionFailure(engine, "call succeeded but the scenario expects failure") + if not fixture.expected_failure and collected.error is not None: + return TraceExecutionFailure(engine, collected.error) + if not collected.events: return TraceExecutionFailure(engine, "trace is empty") - return events + return collected.events def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFailure) -> str | None: diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index fe214f45339..2a4a1b3a152 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -19,6 +19,20 @@ COMMON_MAPPINGS: Final = ( mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), ) +SUCCESS_CALLBACK_SYNC_MAPPING: Final = mapping( + rust_span="success_callback", + python_frame=r"BoundedLoggingThreadPoolExecutor\.submit$", +) +SUCCESS_CALLBACK_ASYNC_MAPPING: Final = mapping( + rust_span="success_callback", + python_frame=r"Logging\.async_success_handler$", +) +FAILURE_CALLBACK_MAPPING: Final = mapping( + rust_span="failure_callback", + python_frame=r"Logging\.(?:async_)?failure_handler$", +) +IGNORED_SUCCESS_CALLBACK_MAPPING: Final = mapping(rust_span="success_callback") + SYNC_MAPPINGS: Final = ( *COMMON_MAPPINGS, mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), @@ -39,6 +53,21 @@ ASYNC_MAPPINGS: Final = ( ), ) +CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING) +CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING) +CALLBACK_FAILURE_SYNC_MAPPINGS: Final = ( + *COMMON_MAPPINGS, + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + FAILURE_CALLBACK_MAPPING, +) +CALLBACK_FAILURE_ASYNC_MAPPINGS: Final = ( + *COMMON_MAPPINGS, + mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), + FAILURE_CALLBACK_MAPPING, +) + + AZURE_COMMON_MAPPINGS: Final = ( *COMMON_MAPPINGS[:7], mapping( @@ -99,6 +128,34 @@ def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "mistral/mistral-ocr-latest") +def _callback_fixture(engine: Engine, *, failure: bool) -> RouteFixture: + fixture: Final = _fixture(engine, "mistral/mistral-ocr-latest") + provider_responses: Final = ( + ( + RecordedHttpResponse.from_bytes( + 400, + (HttpHeader(name="content-type", value="application/json"),), + b'{"message":"trace callback provider failure"}', + ), + ) + if failure + else fixture.provider_responses + ) + return RouteFixture( + kwargs=fixture.kwargs, + provider_responses=provider_responses, + expected_failure=failure, + ) + + +def _mistral_callback_success_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _callback_fixture(engine, failure=False) + + +def _mistral_callback_failure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _callback_fixture(engine, failure=True) + + def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture( engine, @@ -304,36 +361,50 @@ TRACE_SUITE: Final = TraceSuite( name="mistral", fixture=_mistral_fixture, mappings=COMMON_MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + ), + TraceScenario( + name="mistral-callback-success", + fixture=_mistral_callback_success_fixture, + mappings=COMMON_MAPPINGS, + sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, + async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + ), + TraceScenario( + name="mistral-callback-failure", + fixture=_mistral_callback_failure_fixture, + mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING), + sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, + async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, ), TraceScenario( name="azure-ai", fixture=_azure_fixture, mappings=AZURE_COMMON_MAPPINGS, - sync_mappings=AZURE_SYNC_MAPPINGS, - async_mappings=AZURE_ASYNC_MAPPINGS, + sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="azure-document-intelligence", fixture=_azure_document_intelligence_fixture, mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS, - sync_mappings=DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, - async_mappings=DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, + sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="vertex-ai", fixture=_vertex_fixture, mappings=VERTEX_COMMON_MAPPINGS, - sync_mappings=VERTEX_SYNC_MAPPINGS, - async_mappings=VERTEX_ASYNC_MAPPINGS, + sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), TraceScenario( name="vertex-deepseek", fixture=_vertex_deepseek_fixture, mappings=DEEPSEEK_COMMON_MAPPINGS, - sync_mappings=DEEPSEEK_SYNC_MAPPINGS, - async_mappings=DEEPSEEK_ASYNC_MAPPINGS, + sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), ), ), ) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py index 3e6c4060134..0e771f0dc17 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -230,6 +230,10 @@ _HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple( "test_ocr_exception_type_uses_resolved_provider_context", "Python wraps bridge exceptions into public errors.", ), + ( + "test_rust_upstream_error_uses_ocr_provider_error_mapping", + "Python maps native upstream errors through the selected OCR provider config.", + ), ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."), ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."), ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."), diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index b2cf253d164..a30474245c6 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -126,16 +126,6 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge -def test_bare_rust_still_toggles_ocr(): - from litellm.rust_bridge.ocr import rust_ocr_enabled - - litellm.rust(True) - assert rust_ocr_enabled() is True - - litellm.rust(False) - assert rust_ocr_enabled() is False - - def test_load_rust_amessages_returns_injected_impl(): bridge = RecordingAsyncMessages() litellm.rust(True) @@ -214,7 +204,7 @@ async def test_amessages_wrapper_forwards_args(): def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", - "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), + "litellm_params": GenericLiteLLMParams(api_key="sk-azure"), "has_agentic_hook": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", @@ -282,18 +272,6 @@ async def test_gate_uses_process_enable_without_request_override(): assert bridge.calls[0]["custom_llm_provider"] == "azure_ai" -@pytest.mark.asyncio -async def test_gate_skips_rust_when_flag_false(): - bridge = ExplodingAsyncMessages() - litellm.rust(True) - rust_messages.set_rust_messages(amessages=bridge) - - response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) - - assert response is None - assert bridge.calls == 0 - - @pytest.mark.asyncio async def test_gate_invokes_rust_for_native_anthropic_provider(): bridge = RecordingAsyncMessages() @@ -302,7 +280,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): response = await _gate( custom_llm_provider="anthropic", - litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=True), + litellm_params=GenericLiteLLMParams(api_key="sk-ant"), api_key="sk-ant", api_base="https://api.anthropic.com", headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..976a96f2db1 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -21,11 +21,12 @@ from types import MappingProxyType import httpx import pytest import respx +from openai.types.batch import BatchRequestCounts import litellm import litellm.batches.batch_utils as bu -from litellm.types.utils import Usage +from litellm.types.utils import LiteLLMBatch, Usage # --------------------------------------------------------------------------- # # Builders for batch OUTPUT file rows. @@ -1718,3 +1719,57 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): assert usage.total_tokens == 0 assert "does not understand" in caplog.text assert "inputTextTokenCount" in caplog.text + + +# --------------------------------------------------------------------------- # +# batch_cost_is_final +# --------------------------------------------------------------------------- # + +def _retrieved_batch( + status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None +) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="validating", + output_file_id=output_file_id, + request_counts=counts, + ).model_copy(update={"status": status}) + + +class TestBatchCostIsFinal: + """Every retrieve of one batch writes the same spend row, so the first retrieve + that prices it decides the row for good. A poll before the output exists must + therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" + + @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) + def test_in_flight_batch_is_not_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status)) is False + + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_with_output_is_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True + + def test_completed_without_output_and_unknown_counts_is_not_final(self): + assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False + + def test_completed_without_output_and_zero_counts_is_not_final(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_but_successful_lines_is_not_final(self): + counts = BatchRequestCounts(total=2, completed=2, failed=0) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_without_output_and_every_line_failed_is_final(self, status): + counts = BatchRequestCounts(total=2, completed=0, failed=2) + assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True + + @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) + def test_other_terminal_statuses_are_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status)) is True diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 6590718878d..d4d47b145d1 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3952,3 +3952,198 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool function_call_output = next(item for item in response if item.get("type") == "function_call_output") assert function_call_output["output"] == [{"type": "input_text", "text": "1 tool found"}] + + +def _litellm_encoded_response_id(upstream_id: str) -> str: + from litellm.responses.utils import ResponsesAPIRequestUtils + + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure", model_id="deployment-1", response_id=upstream_id + ) + + +def test_transform_response_keeps_upstream_id_and_provider_extras(): + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + content_filters = [ + {"blocked": False, "source_type": "prompt", "content_filter_results": {"hate": {"filtered": False}}} + ] + raw_response = ResponsesAPIResponse.model_validate( + { + "id": _litellm_encoded_response_id("resp_azure_123"), + "created_at": 1734366691, + "object": "response", + "model": "gpt-5.6", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_weather", + "arguments": '{"city": "Seattle"}', + "status": "completed", + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "service_tier": "default", + "content_filters": content_filters, + "max_tool_calls": None, + "background": False, + "top_logprobs": 0, + "store": True, + } + ) + model_response = ModelResponse( + id="chatcmpl-local", + created=1734366691, + model=None, + object="chat.completion", + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + result = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.6", + raw_response=raw_response, + model_response=model_response, + logging_obj=Mock(), + request_data={"model": "gpt-5.6"}, + messages=[{"role": "user", "content": "What is the weather in Seattle?"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + dumped = result.model_dump() + + assert dumped["id"] == "resp_azure_123" + assert dumped["object"] == "chat.completion" + assert dumped["service_tier"] == "default" + assert dumped["content_filters"] == content_filters + assert "max_tool_calls" not in dumped, "a null provider field must not appear as a null top-level key" + assert "output" not in dumped and "status" not in dumped, ( + "Responses schema fields must not leak into the chat response" + ) + assert not {"background", "top_logprobs", "store"} & dumped.keys(), ( + "Responses API bookkeeping must not ride along as chat metadata" + ) + assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather" + + +def test_bridged_response_is_priced_by_the_reported_service_tier(): + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse + + raw_response = ResponsesAPIResponse.model_validate( + { + "id": "resp_flex", + "created_at": 1734366691, + "object": "response", + "model": "gpt-5.4", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + "service_tier": "flex", + } + ) + + result = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + pricing = litellm.model_cost["gpt-5.4"] + flex_cost = 1000 * pricing["input_cost_per_token_flex"] + 100 * pricing["output_cost_per_token_flex"] + standard_cost = 1000 * pricing["input_cost_per_token"] + 100 * pricing["output_cost_per_token"] + + cost = litellm.completion_cost(completion_response=result, custom_llm_provider="openai") + + assert cost == pytest.approx(flex_cost) + assert cost < standard_cost + + +def test_streaming_chunks_carry_the_upstream_response_id(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + encoded_id = _litellm_encoded_response_id("resp_azure_stream") + events = [ + {"type": "response.created", "response": {"id": encoded_id, "output": []}}, + {"type": "response.output_text.delta", "delta": "Hel"}, + {"type": "response.completed", "response": {"id": encoded_id, "output": [{"type": "message"}]}}, + ] + + ids = [iterator.chunk_parser(event).id for event in events] + + assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}" + + +def test_streaming_final_chunk_carries_provider_metadata(): + from unittest.mock import MagicMock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + content_filters = [{"blocked": False, "source_type": "completion", "content_filter_results": {}}] + events = [ + {"type": "response.created", "response": {"id": "resp_azure_stream", "output": []}}, + {"type": "response.output_text.delta", "delta": "Hello"}, + { + "type": "response.completed", + "response": { + "id": "resp_azure_stream", + "output": [{"type": "message"}], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + "service_tier": "default", + "content_filters": content_filters, + "background": False, + }, + }, + ] + stream = CustomStreamWrapper( + completion_stream=iter([iterator.chunk_parser(event) for event in events]), + model="gpt-5.6", + custom_llm_provider="azure", + logging_obj=MagicMock(), + ) + + chunks = [chunk.model_dump() for chunk in stream] + + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["service_tier"] == "default" + assert chunks[-1]["content_filters"] == content_filters + assert "background" not in chunks[-1] + assert all("service_tier" not in chunk for chunk in chunks[:-1]) diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index 62c95cb100b..a4f32df46ae 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -7,9 +7,12 @@ # 4. Added proper cleanup in fixtures # 5. Added worker-specific isolation for parallel execution +import base64 import importlib import os from pathlib import Path +from types import SimpleNamespace +import httpx import pytest import asyncio @@ -595,3 +598,43 @@ def pytest_sessionfinish(session, exitstatus): _close_handler_if_needed(getattr(litellm, "aclient", None)) _close_handler_if_needed(getattr(litellm, "client", None)) _run_coroutine_if_needed(close_litellm_async_clients()) + + +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch 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 df680b7cb0e..59f0938e338 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 @@ -49,6 +49,44 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +@pytest.mark.parametrize("prompt_tokens", [100, 200000, 200001]) +@pytest.mark.parametrize("read_rate", [None, 0.0, 0.25e-6]) +@pytest.mark.parametrize("service_tier", [None, "priority"]) +def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, service_tier): + info = { + "input_cost_per_token": 3e-6, + "input_cost_per_token_priority": 4e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "input_cost_per_token_above_200k_tokens_priority": 8e-6, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": read_rate, + } + usage = Usage(prompt_tokens=prompt_tokens, prompt_tokens_details={"cached_tokens": 100}) + billed = _get_token_base_cost(info, usage, service_tier=service_tier) + savings = _get_token_base_cost(info, usage, service_tier=service_tier, missing_cache_read_uses_input=True) + prompt_cost, _ = generic_cost_per_token("policy-fixture", usage, "openai", service_tier=service_tier, model_info=info) + assert billed[4] == pytest.approx(read_rate or 0.0) + assert savings[:4] == billed[:4] + assert savings[4] == pytest.approx(billed[0] if read_rate is None else read_rate) + assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4]) + + +def test_missing_cache_read_uses_off_peak_input_rate(): + from datetime import datetime, timezone + + info = { + "input_cost_per_token": 3e-6, + "off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 5e-6}, + } + when = datetime(2026, 9, 7, 12, tzinfo=timezone.utc) + billed = _get_token_base_cost(info, Usage(prompt_tokens=100), current_time=when) + savings = _get_token_base_cost( + info, Usage(prompt_tokens=100), current_time=when, missing_cache_read_uses_input=True + ) + assert billed[4] == 0.0 + assert savings[0] == savings[4] == 5e-6 + + def test_reasoning_tokens_no_price_set(_local_model_cost_map): # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token) @@ -2008,7 +2046,14 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) -@pytest.mark.parametrize("model,zone_multiplier", [("azure/gpt-6-astra", 1.0), ("azure/us/gpt-6-astra", 1.1)]) +@pytest.mark.parametrize( + "model,custom_llm_provider,zone_multiplier", + [ + ("azure/gpt-6-astra", "azure", 1.0), + ("azure/us/gpt-6-astra", "azure", 1.1), + ("azure_ai/gpt-6-astra", "azure_ai", 1.0), + ], +) @pytest.mark.parametrize( "prompt_tokens,input_side_multiplier,output_multiplier", [(100000, 1.0, 1.0), (300000, 2.0, 1.5)], @@ -2016,6 +2061,7 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( _local_model_cost_map, model, + custom_llm_provider, zone_multiplier, prompt_tokens, input_side_multiplier, @@ -2023,7 +2069,8 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( ): """Microsoft Foundry sells gpt-6-astra at the OpenAI rates: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens on Standard Global, with the input side doubling and output 1.5x above 272K - prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. + prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. A Foundry + deployment reached through the azure_ai route bills the same Standard Global sheet. """ cached_tokens = 50000 cache_write_tokens = 40000 @@ -2041,7 +2088,7 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( prompt_cost, completion_cost = generic_cost_per_token( model=model, usage=usage, - custom_llm_provider="azure", + custom_llm_provider=custom_llm_provider, ) input_side = zone_multiplier * input_side_multiplier @@ -2051,6 +2098,18 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( assert completion_cost == pytest.approx(zone_multiplier * output_multiplier * completion_tokens * 5e-5) +def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): + usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + + standard = generic_cost_per_token(model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai") + flex = generic_cost_per_token( + model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai", service_tier="flex" + ) + + assert flex == standard + assert standard == pytest.approx((1000 * 1e-05, 100 * 5e-05)) + + @pytest.mark.parametrize( "model,expected_none,expected_xhigh,expected_minimal", [ @@ -4766,3 +4825,100 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r ) assert cost == expected_cost + + +def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map: None) -> None: + """ + Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with + reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top. + """ + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=346, + completion_tokens=29, + total_tokens=375, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=152, image_tokens=194, audio_tokens=0, cached_tokens=128 + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=29, audio_tokens=0, reasoning_tokens=19 + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert completion_cost == pytest.approx(29 * info["output_cost_per_token"]) + assert completion_cost - breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"]) + assert prompt_cost == pytest.approx( + 24 * info["input_cost_per_token"] + + 128 * info["cache_read_input_token_cost"] + + 194 * info["input_cost_per_image_token"] + ) + + +def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens( + _local_model_cost_map: None, +) -> None: + """Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=100, + completion_tokens=44, + total_tokens=144, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=25, audio_tokens=0, reasoning_tokens=19 + ), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert completion_cost == pytest.approx(44 * info["output_cost_per_token"]) + + +def test_generic_cost_per_token_strips_only_the_reasoning_share_when_text_over_reports( + _local_model_cost_map: None, +) -> None: + """Text over-reported past the reasoning share keeps its extra tokens billed; only the nested reasoning is netted out.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=120, + completion_tokens=100, + total_tokens=220, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=100, audio_tokens=70, reasoning_tokens=10), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"]) + assert completion_cost == pytest.approx( + 100 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] + ) + + +def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None: + """Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it.""" + + model = "gpt-realtime-2.1-mini" + usage = Usage( + prompt_tokens=120, + completion_tokens=100, + total_tokens=220, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=30, audio_tokens=70, reasoning_tokens=20), + ) + + _, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage) + + info = litellm.get_model_info(model=model, custom_llm_provider="openai") + assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"]) + assert completion_cost == pytest.approx( + 30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] + ) diff --git a/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py b/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py new file mode 100644 index 00000000000..42ca91bfd8f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_api_route_to_call_types.py @@ -0,0 +1,112 @@ +""" +Tests for route -> CallTypes resolution (api_route_to_call_types). + +Regression coverage for the guardrail route table bugs: +- placeholder segments with a literal suffix ({model}:generateContent) never matched +- the /v1beta generateContent routes were missing from the table +- /llm_passthrough listed the sync call type first, resolving consumers that + take call_types[0] to a handler-less type +""" + +from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, +) +from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes + + +class TestGenerateContentRouteResolution: + def test_bare_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/models/gemini-2.5-flash:generateContent") + assert call_types is not None + assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content] + + def test_v1beta_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/v1beta/models/gemini-2.5-flash:generateContent") + assert call_types is not None + assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content] + + def test_bare_stream_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/models/gemini-2.5-flash:streamGenerateContent") + assert call_types is not None + assert list(call_types) == [ + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ] + + def test_v1beta_stream_generate_content_route_resolves(self): + call_types = get_call_types_for_route("/v1beta/models/gemini-2.5-flash:streamGenerateContent") + assert call_types is not None + assert list(call_types) == [ + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ] + + def test_slash_containing_model_name_resolves(self): + call_types = get_call_types_for_route("/v1beta/models/gemini/gemini-2.5-flash:generateContent") + assert call_types is not None + assert list(call_types) == [CallTypes.agenerate_content, CallTypes.generate_content] + + def test_empty_model_name_does_not_match(self): + assert get_call_types_for_route("/models/:generateContent") is None + + def test_unrelated_model_action_does_not_match(self): + assert get_call_types_for_route("/models/gemini-2.5-flash:countTokens") is None + + +class TestPassthroughRouteOrdering: + def test_llm_passthrough_lists_async_call_type_first(self): + call_types = get_call_types_for_route("/llm_passthrough") + assert call_types is not None + assert list(call_types) == [ + CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, + ] + + def test_v1_llm_passthrough_lists_async_call_type_first(self): + call_types = get_call_types_for_route("/v1/llm_passthrough") + assert call_types is not None + assert list(call_types) == [ + CallTypes.allm_passthrough_route, + CallTypes.llm_passthrough_route, + ] + + +class TestFirstCallTypeHasTranslationHandler: + def test_first_call_type_is_translatable_whenever_any_is(self): + """ + Consumers (unified guardrail post-call and streaming resolution) take + call_types[0]. A route whose first call type lacks a guardrail + translation handler while a later one has it silently skips guardrail + scanning, so the table must list a handler-backed call type first. + """ + from litellm.llms import load_guardrail_translation_mappings + + mappings = load_guardrail_translation_mappings() + misordered = { + route: [call_type.value for call_type in call_types] + for route, call_types in API_ROUTE_TO_CALL_TYPES.items() + if call_types + and call_types[0] not in mappings + and any(call_type in mappings for call_type in call_types) + } + assert misordered == {} + + +class TestExistingRouteResolutionUnchanged: + def test_exact_route_still_resolves(self): + call_types = get_call_types_for_route("/chat/completions") + assert call_types is not None + assert CallTypes.acompletion in call_types + + def test_single_segment_placeholder_still_resolves(self): + call_types = get_call_types_for_route("/a2a/my-agent/message/send") + assert call_types is not None + assert list(call_types) == [CallTypes.asend_message, CallTypes.send_message] + + def test_longer_route_does_not_collapse_into_bare_placeholder_pattern(self): + call_types = get_call_types_for_route("/responses/resp_123/input_items") + assert call_types is not None + assert list(call_types) == [CallTypes.alist_input_items] + + def test_unknown_route_returns_none(self): + assert get_call_types_for_route("/not/a/real/route") is None diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 956da571d43..fb4cb494bee 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -215,32 +215,3 @@ class TestMetadataFallsBackToLitellmMetadata: assert result["metadata"] is not litellm_metadata result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} - - -class TestRustOptIn: - """`rust: true` is a litellm param, so it has to reach `litellm_params`. - - `all_litellm_params` keeps it out of the provider body; without it also - being carried into `litellm_params` the chat completions handlers cannot - see the opt-in and the Rust path is silently never taken. - """ - - def test_rust_is_an_optional_kwargs_key(self): - assert "rust" in _OPTIONAL_KWARGS_KEYS - - def test_rust_is_forwarded_from_completion_kwargs(self): - from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS - - assert "rust" in FORWARDED_KWARGS_KEYS - - def test_rust_survives_into_litellm_params(self): - params = get_litellm_params(rust=True) - assert params["rust"] is True - - def test_rust_is_absent_when_the_deployment_did_not_set_it(self): - assert "rust" not in get_litellm_params() - - def test_rust_stays_out_of_the_provider_body(self): - from litellm.types.utils import all_litellm_params - - assert "rust" in all_litellm_params diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index 893472d63ae..8fa4bd6c14d 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -1,3 +1,7 @@ +import asyncio +import copy +import time +import uuid from unittest.mock import patch import pytest @@ -7,9 +11,13 @@ import litellm from litellm import constants from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( + MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, + RemoteMedia, async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import SSRFError @pytest.fixture(autouse=True) @@ -107,9 +115,7 @@ class StreamingLargeImageClient: request=Request("GET", url), ) # Mock the iter_bytes method to return our generator - response.iter_bytes = lambda chunk_size=8192: generate_chunks( - size_bytes, chunk_size - ) + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) return response @@ -207,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): """ # Simulate a 1 petabyte file (1,000,000 GB) # Without streaming protection, this would cause OOM or hang indefinitely - client = StreamingLargeImageClient( - size_mb=1_000_000_000, include_content_length=False - ) + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) monkeypatch.setattr(litellm, "module_level_client", client) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -268,3 +272,259 @@ def test_image_size_limit_disabled(monkeypatch): assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + messages = [ + {"role": "system", "content": "be terse"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": image_url, "detail": "low"}}, + {"type": "image_url", "image_url": image_url}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + {"type": "file", "file": {"file_id": pdf_url}}, + {"type": "file", "file": {"file_id": image_url, "format": "image/png"}}, + {"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc"}}, + ], + }, + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages) + + data_url = async_only_image_fetch.data_url + base64_png = async_only_image_fetch.base64_png + assert inlined[0] == {"role": "system", "content": "be terse"} + assert inlined[1]["content"] == [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": data_url, "detail": "low"}}, + {"type": "image_url", "image_url": data_url}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}}, + {"type": "file", "file": {"format": "image/png", "file_data": data_url}}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": base64_png}, + "title": "the doc", + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}}, + {"type": "document", "source": {"type": "file", "file_id": "file_abc"}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url]) + assert messages == snapshot + + +async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch): + files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" + files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}" + hinted_image = f"https://img.example/{uuid.uuid4()}.png" + plain_image = f"https://img.example/{uuid.uuid4()}.png" + hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf" + seen = [] + + def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool: + seen.append(media) + return not media.url.startswith(files_api_prefix) and "format" not in media.fields + + messages = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": plain_image}}, + {"type": "image_url", "image_url": plain_image}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, + ], + } + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api) + + assert inlined[0]["content"] == [ + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, + {"type": "image_url", "image_url": async_only_image_fetch.data_url}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, + ] + assert async_only_image_fetch.fetched == [plain_image] + assert [(media.url, dict(media.fields)) for media in seen[:5]] == [ + (files_api_pdf, {"file_id": files_api_pdf}), + (hinted_image, {"url": hinted_image, "format": "image/png"}), + (plain_image, {"url": plain_image}), + (plain_image, {}), + (hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}), + ] + assert messages == snapshot + + +async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it( + async_only_image_fetch, +): + shared = f"https://img.example/{uuid.uuid4()}.png" + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": shared}}, + ], + } + ] + + inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields) + + assert inlined[0]["content"] == [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, + ] + assert async_only_image_fetch.fetched == [shared] + + +async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch): + missing = f"http://img.example/{uuid.uuid4()}-missing.png" + slow = f"http://img.example/{uuid.uuid4()}-slow.png" + slow_fetch_outcomes = [] + + async def serve(client, url, **kwargs): + if url == missing: + return Response(404, request=Request("GET", url)) + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + slow_fetch_outcomes.append("cancelled") + raise + slow_fetch_outcomes.append("finished") + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": missing}}, + {"type": "image_url", "image_url": {"url": slow}}, + ], + } + ] + started = time.perf_counter() + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media(messages) + + assert slow_fetch_outcomes == ["cancelled"] + assert time.perf_counter() - started < 1 + + +_SSRF_VERDICTS = ( + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." + ), + SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"), + SSRFError("No addresses found for 'internal.example'"), +) + + +def _assert_verdict_free_messages(messages, url): + assert len(messages) == len(_SSRF_VERDICTS) + assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve" + message = messages[0] + assert "The proxy could not resolve this host or its URL policy rejected it" in message + assert "user_url_allowed_hosts" in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message + + +async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): + attempts = [] + messages = [] + url = f"http://internal.example/{uuid.uuid4()}.png" + + for verdict in _SSRF_VERDICTS: + + async def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise verdict + + monkeypatch.setattr(image_handling, "async_safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + await async_convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_verdict_free_messages(messages, url) + + +def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): + attempts = [] + messages = [] + url = f"http://internal.example/{uuid.uuid4()}.png" + + for verdict in _SSRF_VERDICTS: + + def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise verdict + + monkeypatch.setattr(image_handling, "safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_verdict_free_messages(messages, url) + + +async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): + in_flight = {"now": 0, "peak": 0} + + async def serve_png_slowly(client, url, **kwargs): + in_flight["now"] += 1 + in_flight["peak"] = max(in_flight["peak"], in_flight["now"]) + await asyncio.sleep(0.01) + in_flight["now"] -= 1 + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly) + urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)] + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}] + + inlined = await async_inline_remote_media(messages) + + assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"]) + + +async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}], + }, + ] + + assert await async_inline_remote_media(messages) is messages + assert async_only_image_fetch.fetched == [] + + +async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch): + async def serve_404(client, url, **kwargs): + return Response(404, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_404) + url = f"http://img.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}]) 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 0f44cbbb2a7..0fdca755685 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,7 +1,9 @@ +import asyncio import contextlib +import datetime import os import sys -import asyncio +from typing import Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -632,6 +634,86 @@ class TestRetrieveBatchCostPassesModelIdentity: assert captured["model_info"]["input_cost_per_token"] == 0.0 +class TestRetrieveBatchPricesOnlyFinalBatches: + """Regression (LIT-7048): retrieving a provider-id batch priced it on every poll. + + Every retrieve of one batch logs under the same spend row, so pricing a poll + that landed before the output existed wrote that row at $0 and pinned it there. + Only a final batch gets priced; an in-flight poll carries no cost at all. + """ + + @staticmethod + def _logging_obj() -> LitellmLogging: + obj = LitellmLogging( + model="gpt-5.6-luna", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="batch-call-2", + function_id="f", + ) + obj.custom_llm_provider = "openai" + return obj + + @staticmethod + def _batch(status: str, output_file_id: str | None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_6a9c99e185588190877d391f8b9d7f8a", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="validating", + output_file_id=output_file_id, + ).model_copy(update={"status": status}) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status", "output_file_id"), + [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)], + ) + async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None: + from litellm.litellm_core_utils import litellm_logging as logging_module + + handle_completed_batch = AsyncMock() + monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + batch = self._batch(status, output_file_id) + + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + + handle_completed_batch.assert_not_awaited() + assert "response_cost" not in batch._hidden_params + + @pytest.mark.asyncio + async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None: + from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.litellm_core_utils import litellm_logging as logging_module + from litellm.types.utils import Usage + + handle_completed_batch = AsyncMock( + return_value=BatchCostUsageResult( + cost=8e-06, + usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), + models=["gpt-5.6-luna"], + successful_requests=2, + failed_requests=0, + ) + ) + monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + batch = self._batch("completed", "file-out") + + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + + handle_completed_batch.assert_awaited_once() + assert batch._hidden_params["response_cost"] == 8e-06 + assert batch.usage is not None + assert batch.usage.total_tokens == 35 + + class TestAnthropicPassthroughCustomPricing: """Verify the Anthropic pass-through handler forwards custom pricing.""" @@ -4391,6 +4473,39 @@ def test_handle_anthropic_messages_response_logging_translates_bare_responses_ap assert result.usage.total_tokens == 18 # type: ignore[attr-defined] +def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id(): + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + served_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream" + ) + logging_obj = _anthropic_messages_logging_obj() + result = logging_obj._handle_anthropic_messages_response_logging( + result=ResponsesAPIResponse( + id=served_id, + created_at=1700000000, + output=[ + ResponseOutputMessage( + id="msg-1", + type="message", + role="assistant", + status="completed", + content=[ResponseOutputText(annotations=[], text="hi", type="output_text")], + ) + ], + usage=ResponseAPIUsage(input_tokens=2, output_tokens=1, total_tokens=3), + service_tier="flex", + ) + ) + + assert isinstance(result, ModelResponse) + assert result.id == served_id, "the spend log row must keep the id the caller was served" + assert result.service_tier == "flex" + + def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() @@ -6326,6 +6441,113 @@ def test_get_standard_logging_object_payload_survives_logging_obj_without_timing assert payload["hidden_params"]["litellm_overhead_time_ms"] is None +@pytest.mark.parametrize( + ("header_name", "header_source"), + ( + ("x-amzn-RequestId", "response"), + ("x-request-id", "response"), + ("request-id", "response"), + ("x-ms-request-id", "response"), + ("apim-request-id", "response"), + ("x-goog-request-id", "response"), + ("cf-ray", "response"), + ("X-Request-Id", "litellm_response_headers"), + ("X-MS-Request-ID", "headers"), + ), +) +def test_failure_standard_logging_payload_captures_provider_request_id( + logging_obj: LitellmLogging, + header_name: str, + header_source: Literal["response", "litellm_response_headers", "headers"], +): + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + request_id = "provider-request-123" + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response(429, headers={header_name: request_id}, request=request) + provider_error = httpx.HTTPStatusError("provider error", request=request, response=response) + if header_source == "litellm_response_headers": + response.headers.clear() + provider_error.litellm_response_headers = {header_name: request_id} + elif header_source == "headers": + response.headers.clear() + provider_error.headers = {header_name: request_id} + now = datetime.datetime.now() + + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "test-model", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="failure", + original_exception=provider_error, + ) + + assert payload is not None + assert payload["error_information"] is not None + assert payload["error_information"]["error_provider_request_id"] == request_id + + +def test_get_error_information_ignores_unsupported_headers() -> None: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response(429, headers={"retry-after": "3"}, request=request) + provider_error = httpx.HTTPStatusError("provider error", request=request, response=response) + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] is None + + +def test_get_error_information_uses_header_precedence_and_fallback() -> None: + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + request = httpx.Request("POST", "https://provider.example/v1/chat/completions") + response = httpx.Response( + 429, + headers={"x-request-id": "response-id", "x-amzn-requestid": "amazon-id"}, + request=request, + ) + provider_error = RateLimitError( + message="provider error", + llm_provider="test-provider", + model="test-model", + response=response, + headers={"retry-after": "3"}, + ) + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] == "amazon-id" + + +def test_get_error_information_ignores_malformed_headers() -> None: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + provider_error = Exception("provider error") + provider_error.headers = [("x-request-id", "provider-request-123")] + + error_information = StandardLoggingPayloadSetup.get_error_information(provider_error) + + assert error_information["error_provider_request_id"] is None + + +def test_get_provider_request_id_ignores_header_lookup_errors() -> None: + from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id + + class HeaderLookupError(Exception): + @property + def response(self) -> object: + raise RuntimeError("headers unavailable") + + assert _get_provider_request_id(HeaderLookupError("provider error")) is None + + def test_get_standard_logging_object_payload_failure_status_keeps_overhead_none(logging_obj): """A post_call guardrail can fail the request after the upstream call succeeded; the failure payload keeps litellm_overhead_time_ms None, matching responses that carry their own _hidden_params.""" diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index aaaa43a0dc4..fccdc1a2a0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -1,5 +1,9 @@ +import asyncio import socket +import threading +import time +import httpx import pytest import litellm @@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): detail = str(exc.value) assert "attacker.example.com" not in detail assert "api.internal-corp.example" not in detail + + +async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): + loop_thread = threading.current_thread() + resolver_threads = [] + + def slow_getaddrinfo(host, port, *args, **kwargs): + resolver_threads.append(threading.current_thread()) + time.sleep(0.4) + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo) + + class FakeClient: + async def get(self, url, **kwargs): + return httpx.Response(200, request=httpx.Request("GET", url)) + + ticks = [time.perf_counter()] + + async def heartbeat(): + while True: + await asyncio.sleep(0.01) + ticks.append(time.perf_counter()) + + beating = asyncio.create_task(heartbeat()) + response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png") + beating.cancel() + + assert response.status_code == 200 + assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) + assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2 diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 043537f8c1f..b4b173b20c3 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2256,7 +2256,7 @@ class TestRustChatCompletionsHook: def _reset_bridge(self, monkeypatch): from litellm.rust_bridge import chat_completions as bridge - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions( chat_completions=None, achat_completions=None, decline=None ) @@ -2282,7 +2282,7 @@ class TestRustChatCompletionsHook: "logging_obj": MagicMock(), "optional_params": {"max_tokens": 16}, "timeout": 30.0, - "litellm_params": {"rust": True}, + "litellm_params": {}, "acompletion": False, "headers": {}, "client": None, @@ -2366,7 +2366,8 @@ class TestRustChatCompletionsHook: ) assert seen["call"][0]["optional_params"]["max_tokens"] == 7 - def test_without_the_opt_in_the_core_is_never_consulted(self): + def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "0") from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -2587,6 +2588,7 @@ class TestRustChatCompletionsHook: def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch): """The suppression must not swallow the log on the ordinary path.""" + monkeypatch.setenv("LITELLM_RUST", "0") from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.anthropic.chat.transformation import AnthropicConfig diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py index 56b754c3476..af7befecc33 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -82,3 +82,22 @@ class TestTheNormalizedTierIsTheTierSent: self, local_model_cost_map, model, provider, effort, expected ): assert _reasoning_effort_sent(model, provider, effort) == expected + + @pytest.mark.parametrize( + "model, provider", + [ + ("gpt-6-astra", "azure_ai"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ("gpt-6-astra", "azure"), + ("us/gpt-6-astra", "azure"), + ], + ) + def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts( + self, local_model_cost_map, model, provider + ): + """The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows + say so and the adapter sends the tier below instead of the rejected one.""" + assert _reasoning_effort_sent(model, provider, "max") == "xhigh" + + def test_the_openai_hosted_twin_still_sends_max(self, local_model_cost_map): + assert _reasoning_effort_sent("gpt-6-astra", "openai", "max") == "max" diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 33fbb4e8fc7..f8cc0b5071e 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -3,6 +3,8 @@ from unittest.mock import MagicMock, patch import pytest +import litellm +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.azure_ai.azure_model_router.transformation import ( AzureModelRouterConfig, ) @@ -138,6 +140,46 @@ def test_azure_ai_validate_environment_with_azure_ad_token(): assert headers["Content-Type"] == "application/json" +@pytest.fixture +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) + + +def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map): + optional_params = AzureAIStudioConfig().map_openai_params( + non_default_params={"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9}, + optional_params={}, + model="gpt-6-astra", + drop_params=False, + ) + + assert optional_params == {"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9} + + +def test_a_gpt_5_name_without_a_foundry_row_keeps_reading_its_own_entry( + monkeypatch: pytest.MonkeyPatch, _local_model_cost_map +): + """Most gpt-5-family names have no azure_ai/ row. Reading an azure_ai/ key for those finds + nothing, and an openai.azure.com base sends the name down the azure provider, which has no key + for it either, so every effort answer would silently fall back to false and take temperature, + top_p and logprobs down with it.""" + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.openai.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "placeholder") + + optional_params = litellm.utils.get_optional_params( + model="gpt-5.1-chat-latest", + custom_llm_provider="azure_ai", + temperature=0.2, + top_p=0.9, + logprobs=True, + ) + + assert optional_params["temperature"] == 0.2 + assert optional_params["top_p"] == 0.9 + assert optional_params["logprobs"] is True + + def test_azure_ai_grok_stop_parameter_handling(): """ Test that Grok models properly handle stop parameter filtering in Azure AI Studio. diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 0d7573a2536..cbf160c451f 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,15 +1,19 @@ import asyncio import json +import uuid from unittest.mock import patch +import httpx import pytest # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. +import litellm from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler def test_get_supported_params_thinking(): @@ -714,3 +718,95 @@ def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking assert result["thinking"] == {"type": "adaptive"} assert result["output_config"] == {"effort": "high"} + + +async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch): + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "document", "source": {"type": "url", "url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert async_only_image_fetch.fetched == [pdf_url] + assert { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png}, + } in captured["body"]["messages"][0]["content"] diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py new file mode 100644 index 00000000000..a8448f5fa7a --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -0,0 +1,99 @@ +import json +import uuid + +import httpx + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/mantle/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch): + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "A lease"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="bedrock/mantle/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "document", "source": {"type": "url", "url": pdf_url}}, + ], + } + ], + aws_access_key_id="AKIAEXAMPLE", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + client=client, + ) + + assert response.choices[0].message.content == "A lease" + assert async_only_image_fetch.fetched == [pdf_url] + assert { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png}, + } in captured["body"]["messages"][0]["content"] diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index c4d6896b17b..f34b8eb1fb9 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -48,7 +48,7 @@ RESOLVED_CREDENTIALS = Credentials( @pytest.fixture(autouse=True) def reset_bridge(monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions( chat_completions=None, achat_completions=None, decline=None ) @@ -87,7 +87,7 @@ def _completion_kwargs(**overrides): "optional_params": {"maxTokens": 16}, "acompletion": False, "timeout": 30.0, - "litellm_params": {"rust": True}, + "litellm_params": {}, "extra_headers": None, "client": None, "api_key": None, @@ -157,7 +157,8 @@ def test_the_core_receives_the_untranslated_openai_messages(): ] -def test_without_the_opt_in_the_core_is_never_consulted(): +def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch): + monkeypatch.setenv("LITELLM_RUST", "0") seen = _inject() try: _run(litellm_params={}) @@ -401,9 +402,10 @@ def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines(): assert logging_obj.pre_call.call_count == 1 -def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(): +def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch): """The suppression must not swallow the log on a request the gate declined, so a deployment with no `rust` flag keeps exactly the log it always had.""" + monkeypatch.setenv("LITELLM_RUST", "0") logging_obj = MagicMock() response = _run( logging_obj=logging_obj, @@ -491,6 +493,7 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke """With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no credentials at all. Preparing the Rust handoff must not dereference that None: the bearer token signs the request on its own.""" + monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") client = _sync_client_returning_converse_response() @@ -520,6 +523,7 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co """The deployment's AWS profile does not exist, so resolving SigV4 credentials raises; a bearer-token deployment must still serve the request, since the bearer token alone signs it.""" + monkeypatch.setenv("LITELLM_RUST", "0") if configured_through == "env_var": monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") else: diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py index ec243b7058d..d9d6e813d86 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -16,6 +16,9 @@ import httpx import pytest +from litellm.llms.black_forest_labs.image_edit import ( + transformation as bfl_transformation, +) from litellm.llms.black_forest_labs.image_edit.transformation import ( BlackForestLabsImageEditConfig, ) @@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation: assert data["output_format"] == "jpeg" # BFL uses JSON, not multipart - files should be empty - assert files == [] + assert files == () def test_transform_image_edit_request_with_mask(self): """Test request transformation with mask for inpainting.""" @@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation: def test_use_multipart_form_data_returns_false(self): """Test that use_multipart_form_data returns False for BFL.""" assert self.config.use_multipart_form_data() is False + + +async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch): + served = b"png-bytes-from-cdn" + fetched = [] + + def forbid_sync_fetch(client, url, **kwargs): + raise AssertionError(f"sync image fetch ran on the event loop: {url}") + + async def serve(client, url, **kwargs): + fetched.append((url, kwargs.get("timeout"))) + return httpx.Response(200, content=served, request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image="https://cdn.example/photo.png", + image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == served + assert base64.b64decode(data["mask"]) == served + assert data["seed"] == 7 + assert files == () + assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)] + + +async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch): + def refuse(*args, **kwargs): + raise AssertionError("no network fetch expected for local image bytes") + + monkeypatch.setattr(bfl_transformation, "safe_get", refuse) + monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=[BytesIO(b"first"), BytesIO(b"other")], + image_edit_optional_request_params={"mask": b"mask-bytes"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == b"first" + assert base64.b64decode(data["mask"]) == b"mask-bytes" + + +async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch): + fetched = [] + + async def serve(client, url, **kwargs): + fetched.append(url) + return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran")) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=["https://cdn.example/a.png", "https://cdn.example/b.png"], + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert fetched == ["https://cdn.example/a.png"] + assert base64.b64decode(data["input_image"]) == b"first-bytes" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 023e2d8843f..e16855da8cb 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import threading import time from unittest.mock import AsyncMock, Mock, patch @@ -17,7 +18,8 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, ) -from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -30,7 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import TranscriptionResponse +from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -2689,20 +2691,18 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h @pytest.mark.parametrize( - "custom_llm_provider, litellm_params, expected", - [ - ("openai", GenericLiteLLMParams(rust=True), True), - ("openai", GenericLiteLLMParams(), False), - ("openai", GenericLiteLLMParams(rust=False), False), - ("azure", GenericLiteLLMParams(rust=True), False), - ("hosted_vllm", GenericLiteLLMParams(rust=True), False), - (None, GenericLiteLLMParams(rust=True), False), - ], + "custom_llm_provider, enabled, expected", + [("openai", True, True), ("openai", False, False), ("azure", True, False), + ("hosted_vllm", True, False), (None, True, False)], ) -def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag( - custom_llm_provider, litellm_params, expected +def test_the_rust_responses_websocket_needs_openai_and_process_enablement( + custom_llm_provider, enabled, expected, monkeypatch ): - assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected + from litellm.rust_bridge import configuration + + configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") + assert _rust_responses_websocket_enabled(custom_llm_provider) is expected def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch): @@ -3186,3 +3186,288 @@ async def test_async_container_list_handler_transforms_success_response(): assert [container.id for container in response.data] == ["cntr_a"] assert response.has_more is True + + +class _TransformRecordingConfig(BaseConfig): + def __init__(self, transform_async: bool): + self.transform_async = transform_async + self.transform_calls = [] + self.sign_threads = [] + + @property + def uses_async_transform_request(self) -> bool: + return self.transform_async + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, non_default_params, optional_params, model, drop_params): + return optional_params + + def validate_environment( + self, headers, model, messages, optional_params, litellm_params, api_key=None, api_base=None + ): + return {} + + def transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("sync") + return {"transformed_by": "sync"} + + async def async_transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("async") + return {"transformed_by": "async"} + + def sign_request( + self, headers, optional_params, request_data, api_base, api_key=None, model=None, stream=None, fake_stream=None + ): + self.sign_threads.append(threading.current_thread()) + return headers, None + + def transform_response( + self, + model, + raw_response, + model_response, + logging_obj, + request_data, + messages, + optional_params, + litellm_params, + encoding, + api_key=None, + json_mode=None, + ): + model_response.choices[0].message.content = raw_response.json()["transformed_by"] + return model_response + + def get_error_class(self, error_message, status_code, headers): + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + + def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False): + return litellm.OpenAIGPTConfig().get_model_response_iterator( + streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode + ) + + +def _start_async_completion(config, logging_obj=None): + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + pending = BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + client=client, + provider_config=config, + ) + return pending, captured + + +async def test_completion_awaits_async_transform_request_when_config_opts_in(): + config = _TransformRecordingConfig(transform_async=True) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == [] + + response = await pending + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.choices[0].message.content == "async" + + +async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_transform(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + + +async def test_completion_keeps_sync_transform_request_before_returning_by_default(): + config = _TransformRecordingConfig(transform_async=False) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == ["sync"] + + response = await pending + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.choices[0].message.content == "sync" + + +def _sse_echoing_transformed_by(request): + transformed_by = json.loads(request.content)["transformed_by"] + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "stub-model", + "choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}], + } + return httpx.Response( + 200, + content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + request=request, + ) + + +def _streaming_logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="async-transform-stream", + function_id="f", + ) + logging_obj.update_environment_variables( + model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + return logging_obj + + +async def test_completion_streams_after_the_async_transform_request(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by)) + + stream = await BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=_streaming_logging_obj(), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + stream=True, + client=client, + provider_config=config, + ) + collected = [chunk async for chunk in stream] + + assert config.transform_calls == ["async"] + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async" + + +class _ImageEditRecordingConfig(BaseImageEditConfig): + def __init__(self): + self.transform_calls = [] + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, image_edit_optional_params, model, drop_params): + return dict(image_edit_optional_params) + + def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None): + return {} + + def get_complete_url(self, model, api_base, litellm_params): + return "https://images.example/v1/edits" + + def use_multipart_form_data(self): + return False + + def transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("sync") + return {"transformed_by": "sync"}, [] + + async def async_transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("async") + return {"transformed_by": "async"}, [] + + def transform_image_edit_response(self, model, raw_response, logging_obj): + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])]) + + +def _echo_json_transport(captured): + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + return httpx.MockTransport(handle) + + +async def test_async_image_edit_handler_awaits_the_async_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_echo_json_transport(captured)) + + response = await BaseLLMHTTPHandler().async_image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.data[0].b64_json == "async" + + +def test_image_edit_handler_keeps_the_sync_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = HTTPHandler() + client.client = httpx.Client(transport=_echo_json_transport(captured)) + + response = BaseLLMHTTPHandler().image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.data[0].b64_json == "sync" diff --git a/tests/test_litellm/llms/gemini/google_genai/__init__.py b/tests/test_litellm/llms/gemini/google_genai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py new file mode 100644 index 00000000000..4119ce99423 --- /dev/null +++ b/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py @@ -0,0 +1,234 @@ +""" +Tests for the Google GenAI generateContent guardrail translation handler. +""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.llms.gemini.google_genai.guardrail_translation.handler import ( + GoogleGenAIGenerateContentHandler, +) +from litellm.types.utils import CallTypes + + +class GuardrailBlockedError(Exception): + pass + + +def _mock_guardrail(returned_texts): + guardrail = MagicMock() + guardrail.apply_guardrail = AsyncMock(return_value={"texts": returned_texts}) + return guardrail + + +@pytest.mark.asyncio +async def test_input_contents_text_is_guardrailed_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked question"]) + data = { + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "raw question"}]}], + } + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + call_kwargs = guardrail.apply_guardrail.call_args.kwargs + assert call_kwargs["inputs"]["texts"] == ["raw question"] + assert call_kwargs["inputs"]["model"] == "gemini-2.5-flash" + assert call_kwargs["input_type"] == "request" + assert result["contents"][0]["parts"][0]["text"] == "masked question" + + +@pytest.mark.asyncio +async def test_input_system_instruction_text_is_scanned_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked instruction", "masked question"]) + data = { + "model": "gemini-2.5-flash", + "systemInstruction": {"role": "system", "parts": [{"text": "prohibited instruction"}]}, + "contents": [{"role": "user", "parts": [{"text": "benign question"}]}], + } + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [ + "prohibited instruction", + "benign question", + ] + assert result["systemInstruction"]["parts"][0]["text"] == "masked instruction" + assert result["contents"][0]["parts"][0]["text"] == "masked question" + + +@pytest.mark.asyncio +async def test_input_config_nested_snake_case_system_instruction_is_scanned(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["clean"]) + instruction_part = SimpleNamespace(text="prohibited instruction") + data = { + "contents": [], + "config": SimpleNamespace(system_instruction=SimpleNamespace(parts=[instruction_part])), + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["prohibited instruction"] + assert instruction_part.text == "clean" + + +@pytest.mark.asyncio +async def test_input_without_text_skips_guardrail(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail([]) + data = {"model": "gemini-2.5-flash", "contents": [{"role": "user", "parts": [{"inlineData": {}}]}]} + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + guardrail.apply_guardrail.assert_not_called() + assert result is data + + +@pytest.mark.asyncio +async def test_output_dict_response_text_is_guardrailed_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked answer"]) + response = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "harmful answer"}]}, + "finishReason": "STOP", + } + ] + } + + result = await handler.process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data={"model": "gemini-2.5-flash"}, + ) + + call_kwargs = guardrail.apply_guardrail.call_args.kwargs + assert call_kwargs["inputs"]["texts"] == ["harmful answer"] + assert call_kwargs["input_type"] == "response" + assert call_kwargs["request_data"]["response"] is response + assert result["candidates"][0]["content"]["parts"][0]["text"] == "masked answer" + + +@pytest.mark.asyncio +async def test_output_object_response_text_is_guardrailed_and_written_back(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["masked answer"]) + part = SimpleNamespace(text="harmful answer") + response = SimpleNamespace( + candidates=[SimpleNamespace(content=SimpleNamespace(parts=[part]), finish_reason="STOP")] + ) + + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] + assert part.text == "masked answer" + + +@pytest.mark.asyncio +async def test_output_without_text_skips_guardrail(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail([]) + + result = await handler.process_output_response(response={"candidates": []}, guardrail_to_apply=guardrail) + + guardrail.apply_guardrail.assert_not_called() + assert result == {"candidates": []} + + +@pytest.mark.asyncio +async def test_output_blocking_guardrail_exception_propagates(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = MagicMock() + guardrail.apply_guardrail = AsyncMock(side_effect=GuardrailBlockedError("blocked")) + response = {"candidates": [{"content": {"parts": [{"text": "harmful answer"}]}}]} + + with pytest.raises(GuardrailBlockedError): + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + +@pytest.mark.asyncio +async def test_streaming_dict_chunks_accumulate_text(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["clean"]) + chunks = [ + {"candidates": [{"content": {"parts": [{"text": "harmful "}]}}]}, + {"candidates": [{"content": {"parts": [{"text": "answer"}]}}]}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + ) + + call_kwargs = guardrail.apply_guardrail.call_args.kwargs + assert call_kwargs["inputs"]["texts"] == ["harmful answer"] + assert call_kwargs["input_type"] == "response" + assert result is chunks + + +@pytest.mark.asyncio +async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail(["clean"]) + frame_one = "data: " + json.dumps({"candidates": [{"content": {"parts": [{"text": "harmful "}]}}]}) + "\r\n\r\n" + frame_two = "data: " + json.dumps({"candidates": [{"content": {"parts": [{"text": "answer"}]}}]}) + "\r\n\r\n" + split_at = len(frame_one) // 2 + chunks = [frame_one[:split_at], frame_one[split_at:] + frame_two[:5], frame_two[5:].encode("utf-8")] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + ) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] + + +@pytest.mark.asyncio +async def test_streaming_blocking_guardrail_exception_propagates(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = MagicMock() + guardrail.apply_guardrail = AsyncMock(side_effect=GuardrailBlockedError("blocked")) + chunks = [{"candidates": [{"content": {"parts": [{"text": "harmful"}]}}]}] + + with pytest.raises(GuardrailBlockedError): + await handler.process_output_streaming_response(responses_so_far=chunks, guardrail_to_apply=guardrail) + + +@pytest.mark.asyncio +async def test_streaming_without_text_skips_guardrail(): + handler = GoogleGenAIGenerateContentHandler() + guardrail = _mock_guardrail([]) + + result = await handler.process_output_streaming_response(responses_so_far=[], guardrail_to_apply=guardrail) + + guardrail.apply_guardrail.assert_not_called() + assert result == [] + + +def test_generate_content_call_types_are_registered(): + from litellm.llms.gemini.google_genai.guardrail_translation import ( + guardrail_translation_mappings, + ) + + for call_type in ( + CallTypes.generate_content, + CallTypes.agenerate_content, + CallTypes.generate_content_stream, + CallTypes.agenerate_content_stream, + ): + assert guardrail_translation_mappings[call_type] is GoogleGenAIGenerateContentHandler + + +def test_discovery_finds_generate_content_handler(): + from litellm.llms import load_guardrail_translation_mappings + + mappings = load_guardrail_translation_mappings() + assert mappings[CallTypes.agenerate_content] is GoogleGenAIGenerateContentHandler + assert mappings[CallTypes.agenerate_content_stream] is GoogleGenAIGenerateContentHandler diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py new file mode 100644 index 00000000000..a0f22a59f0c --- /dev/null +++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py @@ -0,0 +1,436 @@ +import asyncio +import json +from collections.abc import Sequence +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from openai import APIConnectionError + +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TEXT_CHARS, + SkillSearchEmbeddingFailed, + SkillSearchHits, + SkillSearchIndex, + SkillSearchNotConfigured, + search_skills, + skill_search_text, +) +from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.common_utils.semantic_text_index import Vector, cosine_similarity + +CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1") + +TRANSLATOR: Final = LiteLLM_SkillsTable( + skill_id="translate-file", + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", +) +SQL_ANALYST: Final = LiteLLM_SkillsTable( + skill_id="warehouse-sql-analyst", + display_title="Warehouse SQL Analyst", + description="Runs SQL against the inventory database", +) +TRIP_PLANNER: Final = LiteLLM_SkillsTable( + skill_id="trip-planner", + display_title="Trip Planner", + description="Books flights and hotels", +) +SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER) + +VECTORS: Final = MappingProxyType( + { + "language translation": (1.0, 0.0, 0.0), + skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0), + skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0), + skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0), + } +) + + +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + +def _embedding_router() -> MagicMock: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], + ) + ) + return router + + +class FakeEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + return tuple(VECTORS[text] for text in texts) + + +class FixedDimensionEmbedder: + def __init__(self, dimensions: int) -> None: + self.dimensions: Final = dimensions + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + await asyncio.sleep(0) + return tuple((1.0,) * self.dimensions for _ in texts) + + +class TestSkillSearchText: + def test_joins_title_description_and_instructions(self) -> None: + assert skill_search_text(TRANSLATOR) == ( + "Document Translator\n" + "Converts files from one language into another\n" + "Take an uploaded document and produce it in the target language" + ) + + def test_missing_fields_fall_back_to_whatever_is_present(self) -> None: + assert skill_search_text(LiteLLM_SkillsTable(skill_id="bare", display_title="bare")) == "bare" + + def test_all_fields_absent_is_an_empty_string(self) -> None: + assert skill_search_text(LiteLLM_SkillsTable(skill_id="empty")) == "" + + def test_oversized_instructions_are_cut_so_one_skill_cannot_blow_up_the_embedding_batch(self) -> None: + bloated = LiteLLM_SkillsTable( + skill_id="bloated", display_title="Bloated", instructions="x" * (MAX_SKILL_SEARCH_TEXT_CHARS * 3) + ) + text = skill_search_text(bloated) + assert len(text) == MAX_SKILL_SEARCH_TEXT_CHARS + assert text.startswith("Bloated\n") + + +class TestCosineSimilarity: + def test_identical_direction_scores_one(self) -> None: + assert cosine_similarity((2.0, 0.0), (1.0, 0.0)) == pytest.approx(1.0) + + def test_orthogonal_scores_zero(self) -> None: + assert cosine_similarity((1.0, 0.0), (0.0, 1.0)) == pytest.approx(0.0) + + def test_zero_vector_scores_zero_instead_of_dividing(self) -> None: + assert cosine_similarity((0.0, 0.0), (1.0, 0.0)) == 0.0 + + +class TestSkillSearchIndex: + @pytest.mark.asyncio + async def test_ranks_by_similarity_and_truncates_to_top_k(self) -> None: + outcome = await SkillSearchIndex().search( + "language translation", SKILLS, top_k=2, embed=FakeEmbedder(), embedding_model="m" + ) + assert isinstance(outcome, SkillSearchHits) + assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file", "trip-planner"] + assert outcome.hits[0].score > outcome.hits[1].score + + @pytest.mark.asyncio + async def test_second_search_only_embeds_the_query(self) -> None: + index = SkillSearchIndex() + embedder = FakeEmbedder() + await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("language translation", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert len(embedder.calls[0]) == 1 + len(SKILLS) + assert embedder.calls[1] == ("language translation",) + + @pytest.mark.asyncio + async def test_switching_embedding_models_does_not_reuse_cached_vectors(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="small") + wide = FixedDimensionEmbedder(2) + outcome = await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="wide") + assert isinstance(outcome, SkillSearchHits) + assert len(wide.calls[0]) == 1 + len(SKILLS) + + @pytest.mark.asyncio + async def test_cached_vectors_of_another_dimension_are_re_embedded(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m") + fallback = FixedDimensionEmbedder(2) + outcome = await index.search("language translation", SKILLS, top_k=5, embed=fallback, embedding_model="m") + assert isinstance(outcome, SkillSearchHits) + assert fallback.calls == [ + ("language translation",), + ("language translation", *(skill_search_text(skill) for skill in SKILLS)), + ] + + @pytest.mark.asyncio + async def test_re_embedding_a_subset_drops_the_other_skills_old_vectors(self) -> None: + index = SkillSearchIndex() + await index.search("language translation", SKILLS, top_k=5, embed=FakeEmbedder(), embedding_model="m") + wide = FixedDimensionEmbedder(2) + await index.search("language translation", SKILLS[:1], top_k=5, embed=wide, embedding_model="m") + await index.search("language translation", SKILLS, top_k=5, embed=wide, embedding_model="m") + assert wide.calls[-1] == ("language translation", *(skill_search_text(skill) for skill in SKILLS[1:])) + + @pytest.mark.asyncio + async def test_concurrent_searches_keep_each_others_vectors(self) -> None: + index = SkillSearchIndex() + embedder = FixedDimensionEmbedder(3) + await asyncio.gather( + index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m"), + index.search("q", SKILLS[1:], top_k=5, embed=embedder, embedding_model="m"), + ) + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q",) + + @pytest.mark.asyncio + async def test_least_recently_searched_skills_are_evicted_once_the_index_is_full(self) -> None: + index = SkillSearchIndex(max_entries=len(SKILLS)) + embedder = FixedDimensionEmbedder(3) + newcomer = LiteLLM_SkillsTable(skill_id="newcomer", display_title="Newcomer") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m") + await index.search("q", (newcomer,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[1])) + + @pytest.mark.asyncio + async def test_deleted_skills_stop_occupying_the_index_after_enough_new_ones(self) -> None: + index = SkillSearchIndex(max_entries=2) + embedder = FixedDimensionEmbedder(3) + for generation in range(50): + skill = LiteLLM_SkillsTable(skill_id=f"gen-{generation}", display_title=f"Generation {generation}") + await index.search("q", (skill,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[0])) + + @pytest.mark.asyncio + async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None: + async def mixed(texts: Sequence[str]) -> Sequence[Vector]: + return ((1.0, 0.0), *((1.0, 0.0, 0.0) for _ in texts[1:])) + + outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=mixed, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + assert "mixed dimensions" in outcome.reason + + @pytest.mark.asyncio + async def test_no_accessible_skills_returns_no_hits_without_embedding(self) -> None: + embedder = FakeEmbedder() + outcome = await SkillSearchIndex().search("anything", (), top_k=5, embed=embedder, embedding_model="m") + assert outcome == SkillSearchHits(hits=()) + assert embedder.calls == [] + + @pytest.mark.asyncio + async def test_provider_error_becomes_embedding_failed(self) -> None: + async def failing(texts: Sequence[str]) -> Sequence[Vector]: + raise APIConnectionError(request=MagicMock()) + + outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=failing, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + assert "embedding the search query failed" in outcome.reason + + @pytest.mark.asyncio + async def test_wrong_vector_count_becomes_embedding_failed(self) -> None: + async def short(texts: Sequence[str]) -> Sequence[Vector]: + return ((1.0, 0.0, 0.0),) + + outcome = await SkillSearchIndex().search("q", SKILLS, top_k=5, embed=short, embedding_model="m") + assert isinstance(outcome, SkillSearchEmbeddingFailed) + + +class TestSearchSkills: + @pytest.mark.asyncio + async def test_no_embedding_model_is_not_configured(self) -> None: + outcome = await search_skills( + "q", + SKILLS, + 5, + router=MagicMock(), + embedding_model=None, + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchNotConfigured) + assert "skill_search_embedding_model" in outcome.reason + + @pytest.mark.asyncio + async def test_no_router_is_not_configured(self) -> None: + outcome = await search_skills( + "q", + SKILLS, + 5, + router=None, + embedding_model="m", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchNotConfigured) + + @pytest.mark.asyncio + async def test_router_embeddings_are_read_from_the_response(self) -> None: + router = _embedding_router() + outcome = await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + assert isinstance(outcome, SkillSearchHits) + assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file"] + assert router.aembedding.await_args.kwargs["model"] == "text-embedding-3-small" + + @pytest.mark.asyncio + async def test_embedding_spend_is_attributed_to_the_calling_key(self) -> None: + router = _embedding_router() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), + ) + metadata = router.aembedding.await_args.kwargs["metadata"] + assert metadata["user_api_key"] == "hashed-caller-key" + assert metadata["user_api_key_team_id"] == "team-1" + assert metadata["user_api_key_user_id"] == "user-1" + + @pytest.mark.asyncio + async def test_key_limits_are_checked_against_the_real_embedding_call_before_it_runs(self) -> None: + router = _embedding_router() + key_limits = _pass_through_key_limits() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + checked = key_limits.pre_call_hook.await_args.kwargs + assert checked["user_api_key_dict"] is CALLER + assert checked["call_type"] == "aembedding" + assert checked["data"]["model"] == "text-embedding-3-small" + assert checked["data"]["input"] == router.aembedding.await_args.kwargs["input"] + assert checked["data"]["metadata"]["user_api_key"] == "hashed-caller-key" + + @pytest.mark.asyncio + async def test_the_embedding_model_sees_the_request_as_the_guardrails_rewrote_it(self) -> None: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": [1.0, 0.0, 0.0]} for i in range(len(input))], + ) + ) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type: { + **data, + "input": ["[MASKED]" for _ in data["input"]], + "metadata": {**data["metadata"], "guardrail": "masked"}, + } + ) + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + sent = tuple(call.kwargs for call in router.aembedding.await_args_list) + assert sent + assert all(set(call["input"]) == {"[MASKED]"} for call in sent) + assert all(call["metadata"]["guardrail"] == "masked" for call in sent) + + @pytest.mark.asyncio + async def test_a_key_over_its_limit_never_reaches_the_embedding_model(self) -> None: + router = _embedding_router() + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded")) + with pytest.raises(ProxyRateLimitError) as raised: + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + assert raised.value.status_code == 429 + router.aembedding.assert_not_awaited() + + +@pytest.fixture +def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + list_for_search = AsyncMock(return_value=list(SKILLS)) + monkeypatch.setattr( + "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search + ) + return list_for_search + + +@pytest.fixture +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = _pass_through_key_limits() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + router = _embedding_router() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") + return router + + +class TestHandleSkillSearchMCP: + @pytest.mark.asyncio + async def test_top_k_is_clamped_to_the_same_ceiling_as_rest( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + from litellm.llms.litellm_proxy.skills.skill_search import MAX_SKILL_SEARCH_TOP_K + from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search + + many_skills: Final = tuple( + LiteLLM_SkillsTable( + skill_id=f"skill-{i}", display_title=TRIP_PLANNER.display_title, description=TRIP_PLANNER.description + ) + for i in range(MAX_SKILL_SEARCH_TOP_K + 50) + ) + accessible_skills.return_value = list(many_skills) + + result = await handle_skill_search( + query="language translation", top_k=10_000, user_api_key_dict=UserAPIKeyAuth(user_id="u") + ) + assert result.isError is False + assert len(json.loads(result.content[0].text)) == MAX_SKILL_SEARCH_TOP_K + + @pytest.mark.asyncio + async def test_top_k_below_one_is_raised_to_one( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + from litellm.proxy._experimental.mcp_server.tool_search import handle_skill_search + + result = await handle_skill_search( + query="language translation", top_k=0, user_api_key_dict=UserAPIKeyAuth(user_id="u") + ) + assert result.isError is False + assert len(json.loads(result.content[0].text)) == 1 diff --git a/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py new file mode 100644 index 00000000000..d819a79cef1 --- /dev/null +++ b/tests/test_litellm/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -0,0 +1,197 @@ +import base64 +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig +from litellm.llms.mistral.audio_speech.transformation import ( + MistralTextToSpeechConfig, + MistralTextToSpeechException, +) +from litellm.utils import ProviderConfigManager + +SPEECH_URL: Final = "https://api.mistral.ai/v1/audio/speech" + + +def test_mistral_text_to_speech_config_installed(): + config: Final = ProviderConfigManager.get_provider_text_to_speech_config( + model="voxtral-mini-tts-2603", + provider=litellm.LlmProviders.MISTRAL, + ) + assert isinstance(config, BaseTextToSpeechConfig) + assert isinstance(config, MistralTextToSpeechConfig) + + +def test_map_openai_params_drops_speed_and_instructions(): + config: Final = MistralTextToSpeechConfig() + voice, params = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={"response_format": "wav", "speed": 1.5, "instructions": "sound cheerful"}, + voice="en_paul_neutral", + ) + assert voice == "en_paul_neutral" + assert params == {"response_format": "wav"} + + +def test_map_openai_params_accepts_voice_dict_and_ref_audio(): + config: Final = MistralTextToSpeechConfig() + voice, params = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice={"voice_id": "1f3a8b0c-voice-uuid"}, + kwargs={"ref_audio": "bXktdm9pY2Utc2FtcGxl"}, + ) + assert voice == "1f3a8b0c-voice-uuid" + assert params == {"ref_audio": "bXktdm9pY2Utc2FtcGxl"} + + +def test_transform_request_builds_mistral_body(): + config: Final = MistralTextToSpeechConfig() + data: Final = config.transform_text_to_speech_request( + model="voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + optional_params={"response_format": "wav"}, + litellm_params={}, + headers={}, + ) + assert data["dict_body"] == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert data["headers"] == {"Content-Type": "application/json"} + + +def test_transform_request_omits_voice_for_ref_audio_cloning(): + config: Final = MistralTextToSpeechConfig() + data: Final = config.transform_text_to_speech_request( + model="voxtral-mini-tts-2603", + input="clone me", + voice=None, + optional_params={"ref_audio": "bXktdm9pY2Utc2FtcGxl"}, + litellm_params={}, + headers={}, + ) + assert data["dict_body"] == { + "model": "voxtral-mini-tts-2603", + "input": "clone me", + "ref_audio": "bXktdm9pY2Utc2FtcGxl", + } + + +def test_get_complete_url_default_base(): + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=None, litellm_params={}) + assert url == SPEECH_URL + + +@pytest.mark.parametrize( + "api_base", + ["https://custom.api.example.com/v1/", "https://custom.api.example.com/v1", "https://custom.api.example.com"], +) +def test_get_complete_url_custom_base_always_versioned(api_base: str): + config: Final = MistralTextToSpeechConfig() + url: Final = config.get_complete_url(model="voxtral-mini-tts-2603", api_base=api_base, litellm_params={}) + assert url == "https://custom.api.example.com/v1/audio/speech" + + +def test_validate_environment_sets_bearer_header(): + config: Final = MistralTextToSpeechConfig() + headers: Final = config.validate_environment( + headers={"x-custom": "1"}, + model="voxtral-mini-tts-2603", + api_key="sk-mistral-test", + ) + assert headers == { + "x-custom": "1", + "Authorization": "Bearer sk-mistral-test", + "Content-Type": "application/json", + } + + +def test_validate_environment_requires_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + config: Final = MistralTextToSpeechConfig() + with pytest.raises(MistralTextToSpeechException, match="MISTRAL_API_KEY"): + config.validate_environment(headers={}, model="voxtral-mini-tts-2603") + + +def test_transform_response_decodes_base64_audio(): + config: Final = MistralTextToSpeechConfig() + audio_bytes: Final = b"RIFF-fake-wav-bytes" + raw_response: Final = httpx.Response( + 200, + json={"audio_data": base64.b64encode(audio_bytes).decode()}, + headers={"x-request-id": "req-123"}, + request=httpx.Request( + "POST", + SPEECH_URL, + json={"model": "voxtral-mini-tts-2603", "input": "hi", "response_format": "wav"}, + ), + ) + result: Final = config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert result.content == audio_bytes + assert result.response.headers["content-type"] == "audio/wav" + assert result.response.headers["content-length"] == str(len(audio_bytes)) + assert result.response.headers["x-request-id"] == "req-123" + + +def test_transform_response_missing_audio_data_raises(): + config: Final = MistralTextToSpeechConfig() + raw_response: Final = httpx.Response( + 200, + json={"detail": "unexpected"}, + request=httpx.Request("POST", SPEECH_URL, json={"model": "voxtral-mini-tts-2603", "input": "hi"}), + ) + with pytest.raises(MistralTextToSpeechException, match="audio_data"): + config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + +def test_map_openai_params_maps_openai_voice_aliases(): + config: Final = MistralTextToSpeechConfig() + alloy_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="alloy", + ) + nova_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="Nova", + ) + passthrough_voice, _ = config.map_openai_params( + model="voxtral-mini-tts-2603", + optional_params={}, + voice="en_paul_happy", + ) + assert alloy_voice == "en_paul_neutral" + assert nova_voice == "gb_jane_sarcasm" + assert passthrough_voice == "en_paul_happy" + + +def test_transform_response_invalid_base64_raises(): + config: Final = MistralTextToSpeechConfig() + raw_response: Final = httpx.Response( + status_code=200, + json={"audio_data": "QUJD!QUJD"}, + request=httpx.Request("POST", SPEECH_URL), + ) + with pytest.raises(MistralTextToSpeechException, match="base64"): + config.transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index 5dd44d72d68..729a2d25f41 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -445,6 +445,72 @@ class TestOCICohereToolCalls: assert result.choices[0].index == 0 assert result.choices[0].finish_reason == "stop" # COMPLETE is mapped to stop + _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris." + _TOOL_TURN_DELTAS = [ + "I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", ".", + ] + _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}] + _TOOL_TURN_HISTORY = [ + {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."}, + {"role": "CHATBOT", "message": _TOOL_TURN_TEXT, "toolCalls": _TOOL_TURN_CALLS}, + ] + _TOOL_TURN_TERMINAL_TOGETHER = [ + { + "apiFormat": "COHERE", + "text": _TOOL_TURN_TEXT, + "chatHistory": _TOOL_TURN_HISTORY, + "finishReason": "COMPLETE", + "toolCalls": _TOOL_TURN_CALLS, + }, + ] + _TOOL_TURN_TERMINAL_SPLIT = [ + { + "apiFormat": "COHERE", + "text": _TOOL_TURN_TEXT, + "chatHistory": _TOOL_TURN_HISTORY, + "toolCalls": _TOOL_TURN_CALLS, + }, + {"apiFormat": "COHERE", "finishReason": "COMPLETE"}, + ] + + @staticmethod + def _drain_cohere_stream(events): + wrapper = OCIStreamWrapper( + completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock() + ) + chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events] + content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks) + tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])] + finish_reasons = [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] + return content, tool_calls, finish_reasons + + @pytest.mark.parametrize("terminal_events", [_TOOL_TURN_TERMINAL_TOGETHER, _TOOL_TURN_TERMINAL_SPLIT]) + def test_cohere_tool_turn_streams_the_answer_once(self, terminal_events): + """OCI restates the whole answer on the tool-calls chunk and again on the terminal chunk; + the client must read it exactly once, with one tool call and one finish reason.""" + deltas = [{"apiFormat": "COHERE", "text": token} for token in self._TOOL_TURN_DELTAS] + tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS} + + content, tool_calls, finish_reasons = self._drain_cohere_stream([*deltas, tool_calls_event, *terminal_events]) + + assert content == self._TOOL_TURN_TEXT + assert [(call["function"]["name"], call["function"]["arguments"]) for call in tool_calls] == [ + ("get_weather", '{"city": "Paris"}') + ] + assert finish_reasons == ["stop"] + + def test_cohere_tool_turn_without_preamble_deltas_keeps_the_only_text(self): + """When the tool-calls chunk carries the only copy of the text, dropping it would lose the answer.""" + tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS} + + content, tool_calls, finish_reasons = self._drain_cohere_stream( + [tool_calls_event, *self._TOOL_TURN_TERMINAL_TOGETHER] + ) + + assert content == self._TOOL_TURN_TEXT + assert len(tool_calls) == 1 + assert finish_reasons == ["stop"] + def test_cohere_parameter_mapping_excludes_tool_choice(self): """Test that tool_choice is excluded from Cohere parameter mapping""" config = OCIChatConfig() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index a5b748e391c..c5902b32a06 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -2020,3 +2020,108 @@ class TestFlattenToolSchemaCombinatorsWiring: assert result["tools"][0] is opaque_tool assert "anyOf" not in result["tools"][1]["parameters"] + + +class TestReasoningFollowsModelSupport: + """Responses API clients like Codex send `reasoning` on every request, and OpenAI 400s it + on non-reasoning models like gpt-4o. drop_params must strip it there, the same way the + chat completions surface already strips reasoning_effort for those models. + """ + + @pytest.mark.parametrize( + "model, reasoning_survives", + [ + ("gpt-4o", False), + ("gpt-4.1", False), + ("gpt-4o-mini", False), + ("ft:gpt-4o-2024-08-06:my-org::abc123", False), + ("chat-latest", True), + ("gpt-5.6", True), + ("o3", True), + ("o3-deep-research", True), + ("o4-mini-deep-research", True), + ("codex-mini-latest", True), + ("ft:o4-mini-2025-04-16:my-org::abc123", True), + ("computer-use-preview", True), + ], + ) + def test_drop_params_strips_reasoning_by_model(self, local_model_cost_map, model, reasoning_survives): + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "auto"}}, + model=model, + drop_params=True, + ) + assert ("reasoning" in mapped) is reasoning_survives + + @pytest.mark.parametrize( + "model, reasoning_survives", + [ + ("gpt-4o", False), + ("chat-latest", True), + ("o3", True), + ("o3-deep-research", True), + ], + ) + def test_a_cost_map_older_than_this_release_never_strips_a_known_reasoning_model( + self, local_model_cost_map, monkeypatch, model, reasoning_survives + ): + lagging = { + name: {field: value for field, value in entry.items() if field != "supports_reasoning"} + for name, entry in litellm.model_cost.items() + } + monkeypatch.setattr(litellm, "model_cost", lagging) + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model=model, + drop_params=True, + ) + assert ("reasoning" in mapped) is reasoning_survives + + def test_without_drop_params_the_error_is_litellms_400(self, local_model_cost_map, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as excinfo: + OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model="gpt-4o", + drop_params=False, + ) + assert excinfo.value.status_code == 400 + assert "reasoning.effort" in str(excinfo.value) + assert "cost map" in str(excinfo.value) + + @pytest.mark.parametrize("drop_params", [True, False]) + @pytest.mark.parametrize( + "reasoning", + [{"summary": "auto"}, {"effort": None, "summary": "auto"}, {}], + ) + def test_reasoning_without_an_effort_passes_through_on_non_reasoning_models( + self, local_model_cost_map, monkeypatch, drop_params, reasoning + ): + monkeypatch.setattr(litellm, "drop_params", drop_params) + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": dict(reasoning)}, + model="gpt-4o", + drop_params=drop_params, + ) + assert mapped["reasoning"] == reasoning + + def test_an_explicit_supports_reasoning_false_beats_the_bundled_floor(self, local_model_cost_map, monkeypatch): + overridden = { + name: ({**entry, "supports_reasoning": False} if name == "o3" else entry) + for name, entry in litellm.model_cost.items() + } + monkeypatch.setattr(litellm, "model_cost", overridden) + mapped = OpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model="o3", + drop_params=True, + ) + assert "reasoning" not in mapped + + def test_azure_deployments_keep_reasoning_even_on_a_non_reasoning_model_name(self, local_model_cost_map): + mapped = AzureOpenAIResponsesAPIConfig().map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium"}}, + model="gpt-4o", + drop_params=True, + ) + assert mapped["reasoning"] == {"effort": "medium"} diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 25a961c3413..5687a319f06 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -6,6 +6,7 @@ Tests tool calling request/response transformations and chat completions import asyncio import os import copy +import uuid import json from typing import Any, Dict, List @@ -945,3 +946,48 @@ class TestSnowflakeChatCompletion: assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + + +async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Green"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="snowflake/claude-sonnet-4-6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index fad310fc5c0..f135acd094f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,6 +1,13 @@ +import json +import uuid +from unittest.mock import Mock + +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.gemini import transformation from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -338,3 +345,133 @@ def test_map_function_enterprise_web_search_snake_case(): assert len(result) == 1 assert "enterpriseWebSearch" in result[0] + + +async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "Green"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "Green" + assert async_only_image_fetch.fetched == [image_url] + assert image_url not in captured["body"] + assert async_only_image_fetch.base64_png in captured["body"] + + +async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch): + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize these"}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "A report" + assert async_only_image_fetch.fetched == [] + file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part] + assert file_parts == [ + {"mime_type": "application/pdf", "file_uri": files_api_pdf}, + {"mime_type": "image/png", "file_uri": files_api_image}, + ] + + +async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fetch_itself(async_only_image_fetch): + plain_http_png = f"http://img.example/{uuid.uuid4()}.png" + extensionless_https = f"https://cdn.example/files/{uuid.uuid4().hex}" + https_png = f"https://img.example/{uuid.uuid4()}.png" + hinted_extensionless = f"https://cdn.example/files/{uuid.uuid4().hex}" + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe these"}, + {"type": "image_url", "image_url": {"url": plain_http_png}}, + {"type": "image_url", "image_url": {"url": extensionless_https}}, + {"type": "image_url", "image_url": {"url": https_png}}, + {"type": "image_url", "image_url": {"url": hinted_extensionless, "mime_type": "image/webp"}}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + ], + } + ] + + body = await transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-3.8-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=Mock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project="qa-project", + vertex_location="us-central1", + vertex_auth_header=None, + ) + + inlined = {"inline_data": {"mime_type": "image/png", "data": async_only_image_fetch.base64_png}} + assert body["contents"][0]["parts"] == [ + {"text": "Describe these"}, + inlined, + inlined, + {"file_data": {"mime_type": "image/png", "file_uri": https_png}}, + {"file_data": {"mime_type": "image/webp", "file_uri": hinted_extensionless}}, + {"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https]) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 1c2e07e0d24..c34833221cc 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -9,6 +9,7 @@ import httpx import pytest import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration @@ -17,6 +18,7 @@ from litellm.rust_bridge import configuration # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") +rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" @@ -38,6 +40,10 @@ class CapturedException(Exception): pass +class RustUpstreamError(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" @@ -182,6 +188,9 @@ class FakeOCRConfig: ) -> str: return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" + def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + def build_prepared_request( *, @@ -215,11 +224,13 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -229,7 +240,7 @@ def fake_bridge(): """Enable the Rust path with an injected recording bridge (no native wheel).""" bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) return bridge @@ -238,34 +249,14 @@ def fake_async_bridge(): """Enable the async Rust path with an injected recording bridge.""" bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) return bridge -def test_rust_toggles_flag(): - assert rust_bridge.rust_ocr_enabled() is False - litellm.rust(True) - assert rust_bridge.rust_ocr_enabled() is True - litellm.rust(False) - assert rust_bridge.rust_ocr_enabled() is False - - -def test_env_var_enables_rust_ocr(monkeypatch): - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert rust_bridge.rust_ocr_enabled() is True - - -def test_explicit_false_overrides_process_enable(): - litellm.rust(True) - - assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False - - def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) assert rust_bridge.load_rust_ocr() is bridge @@ -329,7 +320,7 @@ def test_native_bridge_available_reflects_loader(monkeypatch): def test_load_rust_aocr_returns_injected_impl(): bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) assert rust_bridge.load_rust_aocr() is bridge @@ -338,7 +329,8 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.override(bridge) + rust_bridge._AOCR.override(async_bridge) litellm.rust(False) assert rust_bridge.load_rust_ocr() is bridge @@ -350,16 +342,18 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): def test_explicit_ocr_none_clears_injected_impl(monkeypatch): monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: None, ) bridge = RecordingBridge() async_bridge = RecordingAsyncBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge, aocr=async_bridge) + rust_bridge._OCR.override(bridge) + rust_bridge._AOCR.override(async_bridge) - rust_bridge.set_rust_ocr(ocr=None, aocr=None) + rust_bridge._OCR.override(None) + rust_bridge._AOCR.override(None) assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None @@ -368,7 +362,7 @@ 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"), + rust_bridge_bindings, "get_native_bridge", lambda: None, ) @@ -385,7 +379,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] monkeypatch.setattr( - importlib.import_module("litellm.rust_bridge"), + rust_bridge_bindings, "get_native_bridge", lambda: fake_module, ) @@ -406,7 +400,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) response = rust_bridge.ocr( model="mistral-ocr-latest", document=DOCUMENT, @@ -441,7 +435,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=bridge) + rust_bridge._AOCR.override(bridge) response = await rust_bridge.aocr( model="mistral-ocr-maas", document=DOCUMENT, @@ -470,7 +464,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) response = ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -500,10 +494,24 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): } +def test_rust_upstream_error_uses_ocr_provider_error_mapping(): + error = RustUpstreamError(400, '{"message":"invalid model"}') + + mapped = ocr_main._map_rust_ocr_error( + error, + build_prepared_request(), + (RuntimeError, RustUpstreamError), + ) + + assert isinstance(mapped, BaseLLMException) + assert mapped.status_code == 400 + assert mapped.message == '{"message":"invalid model"}' + + def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), @@ -516,7 +524,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name: str) -> str | None: raise AssertionError(f"resolver should not be called for {name}") @@ -536,7 +544,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name): resolver_calls.append(name) @@ -559,7 +567,7 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -586,7 +594,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) def _resolver(name: str) -> str | None: return { @@ -610,7 +618,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -628,7 +636,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -649,7 +657,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=bridge) + rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( prepared_request=build_prepared_request( @@ -737,7 +745,7 @@ def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(ocr=RaisingBridge()) + rust_bridge._OCR.override(RaisingBridge()) with pytest.raises(CapturedException): litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -783,7 +791,7 @@ async def test_aocr_exception_type_uses_resolved_provider_context( monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) litellm.rust(True) - rust_bridge.set_rust_ocr(aocr=RaisingAsyncBridge()) + rust_bridge._AOCR.override(RaisingAsyncBridge()) with pytest.raises(CapturedException): await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") @@ -812,9 +820,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): """With the flag off, the bridge must not be consulted even if an impl exists.""" bridge = RecordingBridge() litellm.rust(False) - rust_bridge.set_rust_ocr(ocr=bridge) - - assert rust_bridge.rust_ocr_enabled() is False + rust_bridge._OCR.override(bridge) # The impl stays available for injection, but the disabled flag gates usage, # so ocr() never reaches the Rust path (asserted via the enabled-path test). assert bridge.calls == [] 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 index d67d0df4d0e..1b003e11993 100644 --- 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 @@ -476,16 +476,47 @@ def test_raise_public_plain_unauthorized_has_no_challenge(): @pytest.mark.parametrize( - "root_path, expected_prefix", + "root_path", [ - ("/", ""), # "/" means no prefix - ("", ""), # empty means no prefix - ("/api/v1", "/api/v1"), # a real root path is prepended verbatim + "/", # "/" means no prefix + "", # empty means no prefix ], ) -def test_oauth_protected_resource_path_honors_root_path(root_path, expected_prefix): +def test_oauth_protected_resource_path_no_prefix(root_path, monkeypatch): + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) path = oauth_protected_resource_path(root_path, _server(alias="my-srv")) - assert path == f"/.well-known/oauth-protected-resource{expected_prefix}/mcp/my-srv" + assert path == "/.well-known/oauth-protected-resource/mcp/my-srv" + + +def test_oauth_protected_resource_path_scalar_prefix_uses_rfc8414_insertion(monkeypatch): + # A scalar SERVER_ROOT_PATH deployment registers the well-known routes with + # the prefix inserted (via well_known_root_suffix at import time). The URL + # must match that insertion or a client fetching it 404s. + monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1") + path = oauth_protected_resource_path("/api/v1", _server(alias="my-srv")) + assert path == "/.well-known/oauth-protected-resource/api/v1/mcp/my-srv" + + +def test_oauth_protected_resource_path_per_request_prefix_goes_before_wellknown(monkeypatch): + # Per-request deployment: SERVER_ROOT_PATHS matched /tenant-a for this + # request but the scalar SERVER_ROOT_PATH is unset. Routes were registered + # without the well-known insertion, so the URL must place the prefix + # *before* .well-known — PerRequestRootPathMiddleware strips it and the + # router matches the un-inserted route. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + path = oauth_protected_resource_path("/tenant-a", _server(alias="my-srv")) + assert path == "/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv" + + +def test_oauth_protected_resource_path_dynamic_prefix_wins_over_scalar(monkeypatch): + # Both env vars configured: the middleware matched a SERVER_ROOT_PATHS + # prefix (/tenant-a) that differs from the scalar (/legacy). The URL must + # advertise /tenant-a — the prefix the client called — with no /legacy + # segment stacked onto it. Same review-fix invariant get_custom_url pins. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + path = oauth_protected_resource_path("/tenant-a", _server(alias="my-srv")) + assert path == "/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv" + assert "/legacy" not in path @pytest.mark.parametrize( @@ -510,7 +541,11 @@ def test_raise_user_oauth_challenge_points_at_per_server_prm(): ) -def test_raise_user_oauth_challenge_includes_server_root_path(): +def test_raise_user_oauth_challenge_includes_server_root_path(monkeypatch): + # The scalar deployment: routes are registered with the prefix inserted + # (via well_known_root_suffix at import time), so the challenge URL uses + # the RFC 8414 §3 insertion form. + monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1") with pytest.raises(HTTPException) as exc_info: raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/api/v1") assert ( @@ -519,6 +554,20 @@ def test_raise_user_oauth_challenge_includes_server_root_path(): ) +def test_raise_user_oauth_challenge_per_request_prefix_is_routable(monkeypatch): + # Per-request deployment (SERVER_ROOT_PATHS matched /tenant-a): the + # challenge URL must place /tenant-a before .well-known so the client's + # discovery fetch routes through the same middleware strip the original + # request went through. The scalar-inserted form would 404 here. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + with pytest.raises(HTTPException) as exc_info: + raise_user_oauth_challenge(_server(alias="my-srv"), root_path="/tenant-a") + assert ( + exc_info.value.headers["WWW-Authenticate"] + == 'Bearer resource_metadata="/tenant-a/.well-known/oauth-protected-resource/mcp/my-srv"' + ) + + def test_raise_token_exchange_challenge_is_rfc9728_invalid_token(): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, @@ -535,17 +584,30 @@ def test_raise_token_exchange_challenge_is_rfc9728_invalid_token(): assert "error_description=" in www -def test_raise_token_exchange_challenge_includes_server_root_path(): +def test_raise_token_exchange_challenge_includes_server_root_path(monkeypatch): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, ) + monkeypatch.setenv("SERVER_ROOT_PATH", "/api/v1") with pytest.raises(HTTPException) as exc_info: raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/api/v1") www = exc_info.value.headers["WWW-Authenticate"] assert 'resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/obo-srv"' in www +def test_raise_token_exchange_challenge_per_request_prefix_is_routable(monkeypatch): + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_token_exchange_challenge, + ) + + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + with pytest.raises(HTTPException) as exc_info: + raise_token_exchange_challenge(_server(alias="obo-srv"), root_path="/tenant-a") + www = exc_info.value.headers["WWW-Authenticate"] + assert 'resource_metadata="/tenant-a/.well-known/oauth-protected-resource/mcp/obo-srv"' in www + + def test_raise_token_exchange_challenge_static_form_is_unchanged_without_step_up(): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( raise_token_exchange_challenge, @@ -589,9 +651,7 @@ def test_id_jag_client_secret_maps_to_config(): # ID-JAG asserts the user's id_token; the access_token default maps to id_token. assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" assert isinstance(spec.config.client_auth, ClientSecretAuth) - assert spec.config.client_auth.client_secret.get_secret_value() == ( - "litellm-client-secret" - ) + assert spec.config.client_auth.client_secret.get_secret_value() == ("litellm-client-secret") def test_id_jag_private_key_maps_to_private_key_jwt_auth(): @@ -617,9 +677,7 @@ def test_id_jag_private_key_wins_over_client_secret(): def test_id_jag_honors_explicit_subject_token_type(): - spec = to_server_spec( - _id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2") - ) + spec = to_server_spec(_id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2")) assert spec is not None and isinstance(spec.config, IdJagConfig) assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index a9bb24bbef9..be4206a1faf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -10342,6 +10342,230 @@ async def test_upstream_resource_sent_on_dcr_bridge_relay_authorize(): assert query["client_id"] == ["caller-client"] +# --------------------------------------------------------------------------- +# Per-request root_path (SERVER_ROOT_PATHS / PerRequestRootPathMiddleware): +# one app fronting several client-visible URL path prefixes, each prefix's +# discovery documents emitting URLs under the prefix the client called +# (RFC 9728 §3 exact-match). A scalar PROXY_BASE_URL / SERVER_ROOT_PATH can +# encode at most one prefix per pod; these tests pin the N-prefix case. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _no_proxy_base_url(monkeypatch): + """Discovery must derive URLs from the request in these tests, so the + scalar env overrides are cleared.""" + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + + +@pytest.fixture +def _isolated_mcp_registry(): + """Fixture-owned registry state: snapshot the shared registry, hand the + test an empty one, restore afterwards so nothing leaks between cases.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + saved = dict(global_mcp_server_manager.registry) + global_mcp_server_manager.registry.clear() + try: + yield global_mcp_server_manager.registry + finally: + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry.update(saved) + + +def _prefixed_discovery_client(prefixes): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + ) + + app = FastAPI() + app.include_router(router) + app.add_middleware(PerRequestRootPathMiddleware, root_paths=prefixes) + return TestClient(app) + + +class TestPerRequestRootPathDiscovery: + def test_prefixed_wellknown_not_routable_without_middleware(self, _isolated_mcp_registry): + """Control: on a plain app (the only shape a scalar root_path can + express), a prefixed well-known request 404s before any discovery + builder runs — the routing gap this feature exists to close.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + + server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a") + _isolated_mcp_registry[server.server_id] = server + app = FastAPI() + app.include_router(router) + client = TestClient(app) + resp = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp/server_a") + assert resp.status_code == 404 + + def test_two_prefixes_one_app_each_resource_matches_the_called_url( + self, _no_proxy_base_url, _isolated_mcp_registry + ): + """The multi-origin case itself: two prefixes served by the same app, + each per-server document's ``resource`` equal to the URL its client + called — including the prefix.""" + for sid, name in (("srv_a", "server_a"), ("srv_b", "server_b")): + _isolated_mcp_registry[sid] = _create_oauth2_server(server_id=sid, name=name, server_name=name, alias=name) + client = _prefixed_discovery_client(["/tenant-a", "/tenant-b"]) + + resp_a = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp/server_a") + resp_b = client.get("/tenant-b/.well-known/oauth-protected-resource/mcp/server_b") + + assert resp_a.status_code == 200 + assert resp_a.json()["resource"] == "http://testserver/tenant-a/mcp/server_a" + assert resp_b.status_code == 200 + assert resp_b.json()["resource"] == "http://testserver/tenant-b/mcp/server_b" + + # Every URL the document advertises stays under the request's + # prefix, so it resolves on this same app. + for auth_server in resp_a.json()["authorization_servers"]: + assert auth_server.startswith("http://testserver/tenant-a/") + + def test_unprefixed_requests_unchanged_on_the_same_app(self, _no_proxy_base_url, _isolated_mcp_registry): + """Backward compat on the very same app: a root request emits the + document byte-identical to a deployment without the middleware.""" + server = _create_oauth2_server(server_id="srv_a", name="server_a", server_name="server_a", alias="server_a") + _isolated_mcp_registry[server.server_id] = server + client = _prefixed_discovery_client(["/tenant-a"]) + resp = client.get("/.well-known/oauth-protected-resource/mcp/server_a") + assert resp.status_code == 200 + assert resp.json()["resource"] == "http://testserver/mcp/server_a" + + def test_unlisted_prefix_404s(self, _no_proxy_base_url): + client = _prefixed_discovery_client(["/tenant-a"]) + assert client.get("/tenant-c/.well-known/oauth-protected-resource/mcp/server_a").status_code == 404 + + def test_aggregate_documents_and_as_endpoints_under_prefix(self, _no_proxy_base_url, _isolated_mcp_registry): + """Aggregate PRM/AS documents carry the prefix, and the advertised + authorize endpoint actually resolves under it — the 404 trap that + invalidated prefixing discovery URLs without per-request routing + (#35226 review round 1).""" + client = _prefixed_discovery_client(["/tenant-a"]) + + prm = client.get("/tenant-a/.well-known/oauth-protected-resource/mcp") + asm = client.get("/tenant-a/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/tenant-a/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/tenant-a/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/tenant-a/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize" + + # The prefixed authorize URL routes to the real handler (not 404): + # under per-request root_path the whole app is reachable per-prefix, + # so discovery may advertise prefixed AS endpoints safely. + assert client.get("/tenant-a/authorize").status_code != 404 + + def test_passthrough_challenge_metadata_url_carries_prefix(self, _no_proxy_base_url): + """The WWW-Authenticate resource_metadata URL a 401 advertises must + land under the request's prefix, or the client is bounced to a + document whose ``resource`` cannot match the URL it called.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + get_passthrough_resource_metadata_url, + ) + + def _scope(path, root_path=None): + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": [], + "scheme": "http", + "server": ("testserver", 80), + "client": ("1.2.3.4", 4444), + } + if root_path is not None: + scope["root_path"] = root_path + return scope + + prefixed = get_passthrough_resource_metadata_url( + scope=_scope("/tenant-a/mcp/github", root_path="/tenant-a"), + server_name="github", + ) + assert prefixed == "http://testserver/tenant-a/.well-known/oauth-protected-resource/mcp/github" + + # Regression guard: no root_path → today's URL, unchanged. + bare = get_passthrough_resource_metadata_url( + scope=_scope("/mcp/github"), + server_name="github", + ) + assert bare == "http://testserver/.well-known/oauth-protected-resource/mcp/github" + + def test_user_oauth_challenge_url_routes_and_resource_matches_client_url( + self, _no_proxy_base_url, _isolated_mcp_registry + ): + """The reviewer's expected end-state, pinned end-to-end: an MCP endpoint + raising ``raise_user_oauth_challenge`` under a per-request prefix must + emit a resource_metadata URL the client can actually fetch, and the + document it returns must carry the same prefix the client originally + called. If either half breaks the client's discovery is dead.""" + import re + + from fastapi import FastAPI, HTTPException, Request + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_user_oauth_challenge, + ) + from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + get_request_root_path, + ) + + server = _create_oauth2_server( + server_id="srv_a", name="server_a", server_name="server_a", alias="server_a" + ) + _isolated_mcp_registry[server.server_id] = server + + app = FastAPI() + app.include_router(router) + + @app.post("/mcp/{name}") + def _mcp(name: str, request: Request): + try: + raise_user_oauth_challenge(server, root_path=get_request_root_path()) + except HTTPException as exc: + return {"www_authenticate": exc.headers["WWW-Authenticate"]} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a", "/tenant-b"]) + client = TestClient(app) + + for prefix, mcp_url in ( + ("/tenant-a", "http://testserver/tenant-a/mcp/server_a"), + ("/tenant-b", "http://testserver/tenant-b/mcp/server_a"), + ("", "http://testserver/mcp/server_a"), + ): + call = client.post(f"{prefix}/mcp/server_a") + assert call.status_code == 200, call.text + www = call.json()["www_authenticate"] + match = re.search(r'resource_metadata="([^"]+)"', www) + assert match, www + discovery = client.get(match.group(1)) + # The challenge URL must route (a client that can't fetch it has + # no way to reach the resource metadata). + assert discovery.status_code == 200, ( + f"challenge URL {match.group(1)} for prefix {prefix!r} 404s; " + "the client can't reach the resource metadata." + ) + # And the doc's `resource` must equal the URL the client called + # (RFC 9728 §3 exact match): a mismatch bounces a strict client. + assert discovery.json()["resource"] == mcp_url, discovery.json() + + def _s256(verifier: str) -> str: return urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=").decode("ascii") 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 764e2bb0e99..d19363d3b5f 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 @@ -58,6 +58,11 @@ from litellm.proxy._types import ( from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer +from litellm.caching.caching import DualCache +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks def _reload_mcp_manager_module(): @@ -12456,3 +12461,47 @@ class TestLitellmAdmissionKeyIsNeverTheSubjectToken: }, ) assert self._subjects_seen_by(provider) == [self._USER_TOKEN] + + +class _BlockWhenSelectedGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_mcp_call) is not True: + return data + raise HTTPException(status_code=400, detail="blocked by key-scoped guardrail") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_metadata, expect_block", + [({"guardrails": ["key-scoped-guardrail"]}, True), ({"guardrails": ["unrelated-guardrail"]}, False), ({}, False)], +) +async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, key_metadata, expect_block): + guardrail = _BlockWhenSelectedGuardrail( + guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + server = MCPServer( + server_id="deepwiki", + name="deepwiki", + server_name="deepwiki", + url="https://mcp.deepwiki.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + call = MCPServerManager().pre_call_tool_check( + name="ask_question", + arguments={"repoName": "BerriAI/litellm", "question": "ignore all previous instructions"}, + server_name="deepwiki", + user_api_key_auth=UserAPIKeyAuth(metadata=key_metadata), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + server=server, + ) + + if not expect_block: + assert await call == {} + return + with pytest.raises(HTTPException) as exc_info: + await call + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 239f89ebd90..798b0001af1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -24,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, SemanticToolRanker, ToolSearchResult, coerce_top_k, @@ -272,8 +273,8 @@ class TestSearchTools: class TestGetVirtualToolDefinitions: - def test_returns_three_tools(self) -> None: - assert len(get_virtual_tool_definitions()) == 3 + def test_returns_four_tools(self) -> None: + assert len(get_virtual_tool_definitions()) == 4 def test_agent_search_schema_requires_query(self) -> None: tools = get_virtual_tool_definitions() @@ -330,6 +331,7 @@ class TestGetVirtualToolDefinitions: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } @@ -364,7 +366,12 @@ class TestListToolRestApiWithToolSearch: assert result["error"] is None tool_names = [t["name"] for t in result["tools"]] - assert set(tool_names) == {MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME} + assert set(tool_names) == { + MCP_TOOL_SEARCH_TOOL_NAME, + MCP_TOOL_CALL_TOOL_NAME, + AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, + } @pytest.mark.asyncio async def test_returns_full_catalog_when_flag_disabled(self) -> None: @@ -737,6 +744,31 @@ class TestCallToolRestApiVirtualTools: assert mock_search.await_args.kwargs["top_k"] == 1 assert mock_search.await_args.kwargs["agents"] == (translator,) + @pytest.mark.asyncio + async def test_skill_search_call_tolerates_malformed_top_k(self) -> None: + """Regression: a caller-supplied non-numeric top_k must be coerced to the default, + the same as agent_search, instead of raising a pydantic ValidationError that the + endpoint's catch-all turns into an HTTP 500.""" + from mcp.types import CallToolResult, TextContent + + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K + + user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) + request = self._make_request( + {"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}} + ) + fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False) + with patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam + "litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search", + new_callable=AsyncMock, + return_value=fake_result, + ) as mock_search: + result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) + + assert result.isError is False + assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K + assert mock_search.await_args.kwargs["query"] == "translate a document" + @pytest.mark.asyncio async def test_agent_search_call_reports_missing_embedding_model_as_tool_error(self) -> None: from litellm.proxy.agent_endpoints.agent_search import AgentSearchNotConfigured @@ -782,10 +814,15 @@ class TestCallToolRestApiVirtualTools: router = MagicMock() router.aembedding = AsyncMock(side_effect=fake_aembedding) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) with ( patch( # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does "litellm.proxy.proxy_server.llm_router", router ), + patch( # test-quality-ok: the proxy's key-limit hooks are a module global; the embedding call runs them like /embeddings does + "litellm.proxy.proxy_server.proxy_logging_obj", key_limits + ), patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", new_callable=AsyncMock, @@ -795,6 +832,8 @@ class TestCallToolRestApiVirtualTools: result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict + assert key_limits.pre_call_hook.await_args.kwargs["call_type"] == "aembedding" + assert key_limits.pre_call_hook.await_args.kwargs["data"]["model"] == "emb" assert result.isError is False assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name] @@ -810,7 +849,9 @@ class TestCallToolRestApiVirtualTools: assert "mcp_tool_search.embedding_model" in result.content[0].text @pytest.mark.asyncio - async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + async def test_mcp_tool_search_reports_invalid_settings_as_tool_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0}) user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) @@ -1196,6 +1237,7 @@ class TestHandleListToolsVirtual: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } 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 c007d22117f..bd692776c82 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 @@ -602,6 +602,140 @@ class TestTestConnection: route = _get_route("/mcp-rest/test/connection", "POST") assert _route_has_dependency(route, user_api_key_auth) + @staticmethod + def _capture_execute(monkeypatch) -> dict: + captured: dict = {} + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + captured["request"] = request + captured["mcp_auth_header"] = mcp_auth_header + captured["oauth2_headers"] = oauth2_headers + return {"status": "ok"} + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + return captured + + @staticmethod + def _oauth2_authorization_code_payload(**overrides) -> NewMCPServerRequest: + return NewMCPServerRequest( + server_name="github_mcp", + url="https://api.githubcopilot.com/mcp/", + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url="https://github.com/login/oauth/authorize", + token_url="https://github.com/login/oauth/access_token", + **overrides, + ) + + @pytest.mark.asyncio + async def test_forwards_staged_oauth2_bearer(self, monkeypatch): + """The just-authorized upstream token rides the request's Authorization header, exactly + as /test/tools/list receives it; dropping it makes every authorization_code server fail + the connection test that its tools preview passes.""" + from litellm.proxy._types import LitellmUserRoles + + captured = self._capture_execute(monkeypatch) + request = _build_request( + {"x-litellm-api-key": "sk-admin-session", "authorization": "Bearer upstream-oauth-token"}, + path="/mcp-rest/test/connection", + ) + + result = await rest_endpoints.test_connection( + request, + self._oauth2_authorization_code_payload(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["oauth2_headers"] == {"Authorization": "Bearer upstream-oauth-token"} + assert captured["mcp_auth_header"] is None + + @pytest.mark.asyncio + async def test_forwards_staged_auth_value(self, monkeypatch): + from litellm.proxy._types import LitellmUserRoles + + captured = self._capture_execute(monkeypatch) + request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection") + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com/mcp", + auth_type=MCPAuth.bearer_token, + credentials={"auth_value": "upstream-static-token"}, + ) + + result = await rest_endpoints.test_connection( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["mcp_auth_header"] == "upstream-static-token" + assert captured["oauth2_headers"] is None + + @pytest.mark.asyncio + async def test_inherits_stored_credentials_of_saved_server(self, monkeypatch): + """The edit form resends a saved server without its masked credential; the stored one + must be used, as /test/tools/list already does.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + captured = self._capture_execute(monkeypatch) + saved = MCPServer( + server_id="saved-server-id", + name="example", + url="https://example.com/mcp", + transport="http", + auth_type=MCPAuth.bearer_token, + authentication_token="stored-upstream-token", + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: saved if server_id == "saved-server-id" else None, + ) + request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection") + payload = NewMCPServerRequest( + server_id="saved-server-id", + server_name="example", + url="https://example.com/mcp", + auth_type=MCPAuth.bearer_token, + ) + + result = await rest_endpoints.test_connection( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["mcp_auth_header"] == "stored-upstream-token" + assert captured["request"].credentials == {"auth_value": "stored-upstream-token"} + + @pytest.mark.asyncio + async def test_does_not_forward_authorization_that_satisfied_admission(self, monkeypatch): + """With no x-litellm-api-key, the Authorization value is the caller's LiteLLM key and + must never reach the upstream.""" + from litellm.proxy._types import LitellmUserRoles + + captured = self._capture_execute(monkeypatch) + request = _build_request({"authorization": "Bearer sk-litellm-admission-key"}, path="/mcp-rest/test/connection") + + result = await rest_endpoints.test_connection( + request, + self._oauth2_authorization_code_payload(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result == {"status": "ok"} + assert captured["oauth2_headers"] is None + class TestTestToolsList: pytestmark = pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py index daca244c0a1..c02ed1f37e5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -211,11 +211,24 @@ class TestAgentSearchIndex: assert isinstance(outcome, AgentSearchEmbeddingFailed) +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + class TestSearchAgents: @pytest.mark.asyncio async def test_no_embedding_model_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex(), user_api_key_dict=CALLER + "q", + AGENTS, + 5, + router=MagicMock(), + embedding_model=None, + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) assert "agent_search_embedding_model" in outcome.reason @@ -223,7 +236,14 @@ class TestSearchAgents: @pytest.mark.asyncio async def test_no_router_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex(), user_api_key_dict=CALLER + "q", + AGENTS, + 5, + router=None, + embedding_model="m", + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) @@ -244,6 +264,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchHits) assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"] @@ -266,6 +287,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) metadata = router.aembedding.await_args.kwargs["metadata"] assert metadata["user_api_key"] == "hashed-caller-key" @@ -302,6 +324,7 @@ def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock: ) ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", _pass_through_key_limits()) monkeypatch.setattr(litellm, "agent_search_embedding_model", "text-embedding-3-small") return router diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py new file mode 100644 index 00000000000..ae6b1471c93 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py @@ -0,0 +1,175 @@ +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from openai import APIConnectionError + +import litellm +from litellm.llms.litellm_proxy.skills.skill_search import skill_search_text +from litellm.proxy._types import LiteLLM_SkillsTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.skills_endpoints import router +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + +TRANSLATOR: Final = LiteLLM_SkillsTable( + skill_id="translate-file", + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", +) +SQL_ANALYST: Final = LiteLLM_SkillsTable( + skill_id="warehouse-sql-analyst", + display_title="Warehouse SQL Analyst", + description="Runs SQL against the inventory database", +) +TRIP_PLANNER: Final = LiteLLM_SkillsTable( + skill_id="trip-planner", + display_title="Trip Planner", + description="Books flights and hotels", +) +SKILLS: Final = (TRANSLATOR, SQL_ANALYST, TRIP_PLANNER) + +VECTORS: Final = MappingProxyType( + { + "language translation": (1.0, 0.0, 0.0), + skill_search_text(TRANSLATOR): (0.9, 0.1, 0.0), + skill_search_text(SQL_ANALYST): (0.0, 1.0, 0.0), + skill_search_text(TRIP_PLANNER): (0.3, 0.0, 1.0), + } +) + + +def _client(role: LitellmUserRoles) -> TestClient: + app = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role) + return TestClient(app) + + +@pytest.fixture +def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + list_for_search = AsyncMock(return_value=list(SKILLS)) + monkeypatch.setattr( + "litellm.llms.litellm_proxy.skills.handler.LiteLLMSkillsHandler.list_skills_for_search", list_for_search + ) + return list_for_search + + +@pytest.fixture +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + embedding_router = MagicMock() + embedding_router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], + ) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", embedding_router) + monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") + return embedding_router + + +class TestGetSkillsQuery: + def test_query_ranks_and_scores_and_truncates( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation", "top_k": 2}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + body = response.json()["data"] + assert [skill["id"] for skill in body] == ["translate-file", "trip-planner"] + assert body[0]["search_score"] > body[1]["search_score"] + assert embedding_router.aembedding.await_args.kwargs["metadata"]["user_api_key_user_id"] == "u" + + def test_restricted_key_only_ranks_the_skills_it_can_access( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + accessible_skills.return_value = [SQL_ANALYST] + response = _client(LitellmUserRoles.INTERNAL_USER).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + assert [skill["id"] for skill in response.json()["data"]] == ["warehouse-sql-analyst"] + + def test_no_accessible_skills_is_a_no_match_empty_result( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + accessible_skills.return_value = [] + response = _client(LitellmUserRoles.INTERNAL_USER).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 200 + assert response.json()["data"] == [] + embedding_router.aembedding.assert_not_awaited() + + def test_query_is_unsupported_for_the_anthropic_passthrough_provider( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", params={"query": "anything"}, headers={"Authorization": "Bearer k"} + ) + assert response.status_code == 400 + assert response.json()["detail"]["error"] == "skill_search_unsupported_provider" + accessible_skills.assert_not_awaited() + + def test_missing_embedding_model_is_a_400( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "skill_search_embedding_model", None) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 400 + assert response.json()["detail"]["error"] == "skill_search_not_configured" + + def test_embedding_provider_failure_is_a_503( + self, accessible_skills: AsyncMock, embedding_router: MagicMock + ) -> None: + embedding_router.aembedding = AsyncMock(side_effect=APIConnectionError(request=MagicMock())) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 503 + assert response.json()["detail"]["error"] == "skill_search_unavailable" + + def test_a_key_over_its_rate_limit_gets_a_429_without_embedding( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, key_limits: MagicMock + ) -> None: + key_limits.pre_call_hook = AsyncMock(side_effect=ProxyRateLimitError(detail="rpm exceeded")) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 429 + embedding_router.aembedding.assert_not_awaited() + + def test_top_k_is_validated(self, accessible_skills: AsyncMock, embedding_router: MagicMock) -> None: + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "anything", "top_k": 0}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 422 diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f513f397b64..a996de4d40c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -428,6 +428,106 @@ def test_get_model_from_request_openai_deployment_route_still_works(): ) +def test_get_model_from_request_bedrock_converse_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/converse", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_invoke_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_v2_converse_stream_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/v2/model/us.anthropic.claude-sonnet-4-6/converse-stream", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_model_id_with_slashes(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/aws/anthropic/model-name/invoke", + ) + == "aws/anthropic/model-name" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/agents/some-agent-route", + ) + is None + ) + + +def test_get_model_from_request_bedrock_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_count_tokens_uses_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/v1/messages/count_tokens", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/agents/some-agent-route", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index d58683fd1e5..36bfc4c5dd3 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -857,6 +857,24 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + foundry_key = "azure_ai/gpt-6-astra" + local_entry = litellm.get_model_cost_map(url="")[foundry_key] + registered_before = foundry_key in litellm.azure_ai_models + try: + litellm.add_known_models(model_cost_map={foundry_key: local_entry}) + assert foundry_key in get_known_models_from_wildcard("azure_ai/*") + finally: + if not registered_before: + litellm.azure_ai_models.discard(foundry_key) + litellm.add_known_models(model_cost_map={}) + + def test_get_complete_model_list_drops_no_default_models_sentinel(): from litellm.proxy.auth.model_checks import get_complete_model_list diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 48926cb7bc2..5b15d4a7d5e 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3217,6 +3217,65 @@ def test_internal_user_blocked_from_search_tool_writes(route): assert "Your role=internal_user" in str(exc_info.value) +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_non_admin_can_open_vector_store_details(user_role): + """Regression for LIT-7132: the dashboard lists a vector store via /vector_store/list + (an LLM API route) but opened it via /vector_store/info, which no non-admin allowlist + granted, so the route gate 401'd before the handler's per-store access check ran.""" + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="user@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {} + + granted = RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route="/vector_store/info", + request=request, + valid_token=valid_token, + request_data={}, + ) + assert granted is None + + +@pytest.mark.parametrize( + "route", + ["/vector_store/new", "/vector_store/update", "/vector_store/delete"], +) +def test_internal_user_blocked_from_vector_store_writes(route): + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="user@example.com", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + def test_proxy_admin_viewer_can_read_another_users_info(): """Admin Viewer has read parity with Proxy Admin, so the /user/info key-ownership gate must not apply to it — the Users page reads every row.""" diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py index 4ea655b8871..609dd13afc2 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -4,7 +4,7 @@ selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow. """ from contextlib import asynccontextmanager -from datetime import date, datetime, timezone +from datetime import date, datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock import pytest @@ -19,7 +19,6 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( upcoming_partitions, ) - DDL_TIMEOUT_MS = 30000 @@ -28,19 +27,23 @@ def _budget(ms: "int | None" = DDL_TIMEOUT_MS): return lambda: ms -def _wire_tx(db) -> list[str]: +def _wire_tx(db) -> "tuple[list[str], list[int | timedelta | None]]": """ Model the prisma seam the partition DDL uses. Every statement this manager issues, DDL and catalog query alike, runs inside db.tx() so it can carry SET LOCAL timeouts. Those SET LOCAL statements are - collected in the returned list rather than forwarded, so assertions on - db.execute_raw and db.query_raw still see only the real statements. + collected in the first returned list rather than forwarded, so assertions on + db.execute_raw and db.query_raw still see only the real statements. The + second list records the interactive-transaction timeout each tx was opened + with. """ session_settings: list[str] = [] + tx_timeouts: list[int | timedelta | None] = [] @asynccontextmanager - async def _tx(): + async def _tx(*, max_wait: "int | timedelta | None" = None, timeout: "int | timedelta | None" = None): + tx_timeouts.append(timeout) tx = MagicMock() async def _execute_raw(sql, *args): @@ -57,7 +60,7 @@ def _wire_tx(db) -> list[str]: yield tx db.tx = _tx - return session_settings + return session_settings, tx_timeouts def test_period_start_per_interval(): @@ -227,7 +230,7 @@ async def test_partition_ddl_carries_a_statement_and_lock_timeout(): } ] ) - session_settings = _wire_tx(client.db) + session_settings, _ = _wire_tx(client.db) await mgr.ensure_partitions(client, _budget(7000)) await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget(7000)) @@ -249,7 +252,7 @@ async def test_catalog_queries_carry_a_statement_timeout(): mgr = SpendLogsPartitionManager() client = MagicMock() client.db.query_raw = AsyncMock(return_value=[]) - session_settings = _wire_tx(client.db) + session_settings, _ = _wire_tx(client.db) await mgr.is_partitioned(client, _budget(4000)) assert session_settings == ["SET LOCAL statement_timeout = 4000"], ( @@ -263,6 +266,34 @@ async def test_catalog_queries_carry_a_statement_timeout(): ) +@pytest.mark.asyncio +async def test_partition_transactions_outlive_their_statement_bound(): + """ + prisma's interactive transaction has its own timeout, 5s by default, which + keeps ticking while a statement waits on the partition lock. A tx shorter + than the SET LOCAL bound it carries is closed mid-lock-wait and the engine + then answers the next call with a 422, so the partition is silently not + created. A tx equal to the bound is closed too: the statement can consume + its whole bound waiting on the lock, then still needs to commit. + """ + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + client.db.query_raw = AsyncMock(return_value=[]) + _, tx_timeouts = _wire_tx(client.db) + + await mgr.is_partitioned(client, _budget()) + await mgr.ensure_partitions(client, _budget()) + await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget()) + + assert len(tx_timeouts) > 0 + for tx_timeout in tx_timeouts: + assert isinstance(tx_timeout, timedelta) + assert tx_timeout > timedelta(milliseconds=DDL_TIMEOUT_MS), ( + f"tx timeout {tx_timeout} does not outlive its {DDL_TIMEOUT_MS}ms statement bound" + ) + + @pytest.mark.asyncio async def test_partition_loops_stop_when_the_budget_runs_out_mid_way(): """ @@ -308,17 +339,15 @@ async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_s def test_unsupported_interval_raises(): - with pytest.raises(ValueError, match='Unsupported partition interval: year'): + with pytest.raises(ValueError, match="Unsupported partition interval: year"): period_start(date(2026, 6, 1), "year") - with pytest.raises(ValueError, match='Unsupported partition interval: year'): + with pytest.raises(ValueError, match="Unsupported partition interval: year"): next_period_start(date(2026, 6, 1), "year") def test_parse_partition_upper_bound_unparseable_to_value_is_none(): """A TO(...) value that is not a valid timestamp must not raise; return None.""" - assert ( - parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None - ) + assert parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 11ef911de3e..0bca7c9492c 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -7,6 +7,7 @@ import re from collections.abc import Callable from contextlib import asynccontextmanager from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -2936,7 +2937,9 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey "call_type, expects_flush", [("aresponses", True), ("responses", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back( + call_type: str, expects_flush: bool +): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. @@ -2957,6 +2960,303 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c PrismaClient.spend_log_flush_requested.clear() +def _batch_cost_payload() -> dict: + return { + **_minimal_spend_payload(), + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + } + + +def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock: + prisma = _tool_usage_prisma() + prisma.jsonify_object = lambda data: dict(data) + prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) + prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) + prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over) + return prisma + + +async def _update_database_with( + db_writer: DBSpendUpdateWriter, + prisma: MagicMock, + payload: dict, + disable_spend_logs: bool = False, + response_cost: float = 0.25, +) -> bool: + with ( + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.disable_spend_logs", disable_spend_logs + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.prisma_client", prisma + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), + patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=payload, + ), + ): + charged = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"}, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=response_cost, + ) + await asyncio.sleep(0) + return charged + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0), True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="failure", spend=0.0), True), + (0, SimpleNamespace(call_type="aembedding", status="success", spend=0.25), True), + (0, None, True), + ], + ids=[ + "first_retrieve_owns_the_row", + "another_retrieve_already_charged", + "an_older_proxy_left_a_zero_row_while_the_batch_ran", + "failed_retrieve_holds_the_row", + "client_chosen_call_id_holds_the_row", + "row_gone_between_insert_and_lookup", + ], +) +async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote_its_row( + inserted: int, existing: object, charged: bool +): + """ + Every retrieve of one batch shares one spend row, so the insert that lands first is + the charge and every later retrieve must leave the counters alone (LIT-7048). A row + that recorded no charge must not be able to take the charge away: neither one a + client planted under the batch id, nor the $0 row a pre-upgrade proxy wrote every + time it polled the batch while it was still running. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + claimed_rows = prisma.db.litellm_spendlogs.create_many.await_args.kwargs + assert claimed_rows["skip_duplicates"] is True + assert [(row["request_id"], row["spend"]) for row in claimed_rows["data"]] == [("batch_abc_batch_cost", 0.25)] + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("taken_over", "charged"), + [(1, True), (0, False)], + ids=["this_retrieve_takes_it_over", "another_one_got_there_first"], +) +async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left_at_zero( + taken_over: int, charged: bool +): + """ + A proxy without this fix wrote the batch's row at $0 on every poll of a running batch, + and the row outlives the upgrade, so the charge has to land on the row itself. Charging + without writing it there would charge again on every later retrieve (LIT-7048). + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing, taken_over) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + taken = prisma.db.litellm_spendlogs.update_many.await_args.kwargs + assert taken["where"] == { + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.0, + } + assert taken["data"]["spend"] == 0.25 + assert "request_id" not in taken["data"] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_leaves_a_batch_whose_zero_row_it_could_not_take_over_to_the_next_retrieve(): + """ + A DB that refuses the takeover leaves the row reading $0, so charging here would charge + the batch again on every later retrieve. The retrieve that does take the row over is the + one that charges. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing) + prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False + + assert db_writer._batch_database_updates.await_count == 0 + + +@pytest.mark.asyncio +async def test_update_database_leaves_a_batch_that_cost_nothing_to_the_retrieve_that_wrote_its_row(): + """ + A batch every line of which failed costs $0, so its row reads $0 for the honest reason + and the retrieve that wrote it is still the one that accounted it. Taking that row over + on every later retrieve would count one batch as many requests. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), response_cost=0.0) is False + + prisma.db.litellm_spendlogs.update_many.assert_not_called() + assert db_writer._batch_database_updates.await_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False), + ], + ids=["first_retrieve_owns_the_row", "another_retrieve_already_charged"], +) +async def test_update_database_charges_a_batch_once_even_with_spend_logs_disabled( + inserted: int, existing: object, charged: bool +): + """ + disable_spend_logs drops the per-request logs, not the batch's charge, so the one row + that makes a batch chargeable exactly once is still written and still read back. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), True) is charged + + assert prisma.db.litellm_spendlogs.create_many.await_count == 1 + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disabled(): + """The batch carve-out above stays a carve-out: every other row still goes unwritten.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = {**_batch_cost_payload(), "call_type": "acompletion"} + + assert await _update_database_with(db_writer, prisma, payload, True) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == 1 + + +_BATCH_CLAIM_FIELDS = {"request_id", "call_type", "status", "spend", "startTime", "endTime"} + + +def _logged_batch_cost_payload() -> dict: + return { + **_batch_cost_payload(), + "api_key": "0e5b0e9e5f", + "model": "gpt-5.6-luna", + "user": "test-user", + "metadata": '{"batch_models": ["gpt-5.6-luna"]}', + "requester_ip_address": "127.0.0.1", + "proxy_server_request": '{"headers": {"user-agent": "litellm-batch-cost-check"}}', + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("disable_spend_logs", "logs_the_request"), + [(False, True), (True, False)], + ids=["spend_logs_on", "spend_logs_off"], +) +async def test_update_database_claims_a_batch_without_logging_the_request_that_polled_it( + disable_spend_logs: bool, logs_the_request: bool +): + """ + disable_spend_logs has to keep meaning that no request gets logged, and the batch's cost + row is the one row it cannot drop, so with logging off that row carries only what tells + the retrieves apart: no metadata, no requester IP, no key, model, or token counts. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = _logged_batch_cost_payload() + + assert await _update_database_with(db_writer, prisma, payload, disable_spend_logs) is True + + claimed = prisma.db.litellm_spendlogs.create_many.await_args.kwargs["data"][0] + assert set(claimed) == (set(payload) if logs_the_request else _BATCH_CLAIM_FIELDS) + assert claimed["spend"] == 0.25 + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled(): + """A refused claim is retried through the queue, so what it queues has to stay unlogged too.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _logged_batch_cost_payload(), True) is True + + assert [set(row) for row in prisma.spend_log_transactions] == [_BATCH_CLAIM_FIELDS] + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): + """An unreachable DB must not drop the batch's only spend row, nor its charge.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + assert [row["request_id"] for row in prisma.spend_log_transactions] == ["batch_abc_batch_cost"] + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [{**_batch_cost_payload(), "call_type": "acompletion"}, {**_batch_cost_payload(), "status": "failure"}], + ids=["not_a_batch_retrieve", "failed_batch_retrieve"], +) +async def test_update_database_queues_every_other_spend_row_for_the_next_flush(payload: dict): + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + + assert await _update_database_with(db_writer, prisma, payload) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [payload] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize( "injected_deployment, attributed", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 4444cd693ff..d57a91d45bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -6,16 +6,19 @@ and allows requests with only registered servers. Covers both /chat/completions and /responses API paths (same pre_call_hook logic, different call_type). """ +from typing import Literal from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException +import litellm from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( MCPSecurityGuardrail, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams @pytest.fixture @@ -182,3 +185,23 @@ class TestMCPSecurityGuardrailPreCall: call_type="acompletion", ) assert result == data + + +class TestInitializeGuardrail: + @pytest.mark.parametrize( + "configured,expected", + [("block", "block"), ("alert", "alert"), (None, "alert"), ("warn", "alert"), ("end_session", "alert")], + ) + def test_on_violation_from_litellm_params( + self, + configured: Literal["block", "alert", "warn", "end_session"] | None, + expected: Literal["block", "alert"], + ): + litellm_params = LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation=configured) + guardrail = Guardrail(guardrail_name="mcp-security-block", litellm_params=litellm_params) + + result = initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail) + + assert isinstance(result, MCPSecurityGuardrail) + assert result.on_violation == expected + assert result in litellm.callbacks diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index fe2cd819717..530f8ffd854 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1361,6 +1361,22 @@ async def test_patch_guardrail_endpoint( assert "Failed to update" in str(mock_logger.warning.call_args) +@pytest.mark.asyncio +async def test_patch_guardrail_rejects_mcp_only_on_violation_with_422(mocker, mock_guardrail_registry): + mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) # test-quality-ok: endpoint has no DI seam + mocker.patch( # test-quality-ok: endpoint has no DI seam + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry + ) + request = PatchGuardrailRequest(litellm_params=BaseLitellmParams(on_violation="block")) + + with pytest.raises(HTTPException) as exc_info: + await patch_guardrail("test-guardrail-id", request, user_api_key_dict=MOCK_ADMIN_USER) + + assert exc_info.value.status_code == 422 + assert "only supported by guardrail='mcp_security'" in str(exc_info.value.detail) + mock_guardrail_registry.update_guardrail_in_db.assert_not_called() + + @pytest.mark.parametrize( "scenario,expected_result,expected_exception", [ diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index ab4e15ff423..e650f796f29 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -31,6 +31,7 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): "mode": "during_call", "default_on": True, "file_sanitization_fail_open": False, + "block_on_file_modify": False, }, } ], @@ -43,9 +44,11 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): assert registered[0].default_on is True assert registered[0].event_hook == "during_call" assert registered[0].file_sanitization_fail_open is False + assert registered[0].block_on_file_modify is False config_model = registered[0].get_config_model() assert config_model is not None assert config_model().file_sanitization_fail_open is True + assert config_model().block_on_file_modify is True def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch): @@ -379,6 +382,121 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch): assert result is not None +@pytest.mark.asyncio +async def test_file_sanitization_modify_blocks_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True + ) + csv_data = b"name,email\nAlice,alice@example.com\n" + item = { + "type": "file", + "file": { + "data": base64.b64encode(csv_data).decode(), + "mime_type": "text/csv", + }, + } + upload_response = Response( + json={"jobId": "modify-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={ + "status": "done", + "content": "name,email\nAlice,[REDACTED]\n", + "metadata": {"action": "modify", "violations": ["Sensitive Data"]}, + }, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail._process_document_item(item, None) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "Document blocked by Prompt Security. Violations: Sensitive Data" + + +@pytest.mark.asyncio +async def test_standalone_image_sanitization_modify_blocks_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", event_hook="pre_call", default_on=True + ) + guardrail.poll_interval = 0 + image_url = "data:image/png;base64," + base64.b64encode(b"image-content").decode() + upload_response = Response( + json={"jobId": "modify-image-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={ + "status": "done", + "content": "Email: [REDACTED]", + "metadata": {"action": "modify", "violations": ["Sensitive Data"]}, + }, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)): + with pytest.raises(HTTPException) as exc_info: + await guardrail._process_standalone_images([image_url], None) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "Image blocked by Prompt Security. Violations: Sensitive Data" + + +@pytest.mark.asyncio +async def test_file_sanitization_modify_can_rewrite_when_blocking_disabled(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + block_on_file_modify=False, + ) + csv_data = b"name,email\nAlice,alice@example.com\n" + item = { + "type": "file", + "file": { + "data": base64.b64encode(csv_data).decode(), + "mime_type": "text/csv", + }, + } + upload_response = Response( + json={"jobId": "modify-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={ + "status": "done", + "content": "name,email\nAlice,[REDACTED]\n", + "metadata": {"action": "modify", "violations": ["Sensitive Data"]}, + }, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)): + result = await guardrail._process_document_item(item, None) + + assert base64.b64decode(result["file"]["data"]) == b"name,email\nAlice,[REDACTED]\n" + + @pytest.mark.asyncio @pytest.mark.parametrize( "timeout", diff --git a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py index f716a8533d8..dc7ebf6e2ec 100644 --- a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py +++ b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py @@ -1,3 +1,7 @@ +import os +import subprocess +import sys +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -65,3 +69,19 @@ async def test_execute_code_loop_dispatches_litellm_skill_tool(): mock_exec.assert_awaited_once() assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME assert result is final_response + + +def test_importing_proxy_hooks_does_not_mutate_litellm_callbacks(): + script: Final = ( + "import litellm; " + "litellm.callbacks = []; " + "import litellm.proxy.hooks; " + "assert len(litellm.callbacks) == 0, f'mutated: {litellm.callbacks}'" + ) + result: Final = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}" diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index f3fd9f41cc2..eff892f2d80 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,4 +1,4 @@ - +import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -70,9 +70,7 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { - "request_id": "test_request_id" - } + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -336,9 +334,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs[ - "budget_reservation" - ] + mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -433,36 +429,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } + assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} - ) - is None - ) - assert ( - _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": UserAPIKeyAuth( - budget_reservation=budget_reservation - ) - } + metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": dict( - UserAPIKeyAuth(budget_reservation=budget_reservation) - ) - } + metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata( - metadata={"user_api_key_budget_reservation": budget_reservation} - ) + _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) is budget_reservation ) @@ -470,9 +451,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=Exception("db unavailable") - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -508,9 +487,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=db_exception - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -554,12 +531,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call( - "Failed to release budget reservation after database update failed" - ) - mock_log_exception.assert_any_call( - "Failed to invalidate budget reservation counters after release failed" - ) + mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") + mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") increment_spend_counters.assert_not_awaited() @@ -778,6 +751,107 @@ async def test_track_cost_callback_defers_in_progress_background_interaction(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +def _batch_retrieve_kwargs(call_type: str, reservation: dict | None = None) -> dict: + metadata = { + "user_api_key": "hashed_key", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + **({"user_api_key_budget_reservation": reservation} if reservation is not None else {}), + } + return { + "call_type": call_type, + "model": "gpt-5.6-luna", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": {"response_cost": 0.0, "request_tags": None}, + "stream": False, + } + + +def _retrieved_batch(status: str, output_file_id: str | None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"), + [ + ("aretrieve_batch", "in_progress", None, True, False, False), + ("aretrieve_batch", "completed", None, True, False, False), + ("aretrieve_batch", "completed", "file-out", False, True, False), + ("aretrieve_batch", "completed", "file-out", True, True, True), + ("aretrieve_batch", "failed", None, True, True, True), + ("acreate_batch", "validating", None, True, True, True), + ], + ids=[ + "retrieve_before_final", + "retrieve_completed_without_output_yet", + "retrieve_after_another_retrieve_charged", + "retrieve_first_final", + "retrieve_failed_batch", + "create_before_final", + ], +) +async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate + call_type, status, output_file_id, row_claimed, spend_written, charged +): + """ + A poll before the batch is final used to pin its shared spend row at $0, and every + completed retrieve after the first charged the key again (LIT-7048). Only retrieves + are gated, since creating a batch is its own billable request, and a retrieve that + charges nothing hands its budget reservation back instead. + """ + logger = _ProxyDBLogger() + budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} + kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) + + with ( + patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment_spend_counters, + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ) as mock_update_cache, + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the release is imported inside the callback's helper, no seam + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", new_callable=AsyncMock + ) as mock_release_budget_reservation, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=row_claimed) + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=_retrieved_batch(status, output_file_id), + start_time=datetime.now(), + end_time=datetime.now(), + ) + await asyncio.sleep(0) + + mock_proxy_logging.failed_tracking_alert.assert_not_called() + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if spend_written else 0) + assert mock_increment_spend_counters.await_count == (1 if charged else 0) + assert mock_update_cache.await_count == (1 if charged else 0) + if charged: + mock_release_budget_reservation.assert_not_awaited() + else: + mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation) + + def _in_progress_interaction_kwargs(reservation: dict) -> dict: return { "call_type": "acreate_interaction", @@ -1101,10 +1175,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert ( - call_kwargs["standard_logging_object"]["trace_id"] - == "trace-id-from-logging-obj" - ) + assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1691,9 +1762,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } with patch( @@ -1772,15 +1841,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs - ) + update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] - == "mcp-user@example.com" - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" @pytest.mark.asyncio @@ -1875,9 +1939,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request( - call_type, expect_spend_log -): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1923,9 +1985,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( - 1 if expect_spend_log else 0 - ) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) class _FakeDeploymentLookup: diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 21f8d985f22..35e76c96c14 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -489,6 +489,7 @@ class TestAutoRouterBenchmarks: monkeypatch: pytest.MonkeyPatch, rows: Sequence[Mapping[str, object]], model_list: Sequence[object], + api_key: str | None = None, ) -> AutoRouterBenchmarksResponse: from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -503,6 +504,7 @@ class TestAutoRouterBenchmarks: user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01", + api_key=api_key, ) ROW = _SessionAggRow( @@ -652,8 +654,9 @@ class TestAutoRouterBenchmarks: user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01", + api_key="key-hash", ) - assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00") + assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash") assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 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 8766b1a1868..a873a367eab 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 @@ -17692,6 +17692,253 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"] +@pytest.mark.asyncio +async def test_update_key_soft_budget_updates_existing_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123") + mock_db = MagicMock() + mock_db.litellm_budgettable.update = AsyncMock() + mock_db.litellm_budgettable.create = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=25.0, + changed_by="user-1", + ) + + assert result == "budget-123" + mock_db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-123"}, + data={"soft_budget": 25.0, "updated_by": "user-1"}, + ) + + +@pytest.mark.asyncio +async def test_update_key_soft_budget_clears_existing_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123") + mock_db = MagicMock() + mock_db.litellm_budgettable.update = AsyncMock() + mock_db.litellm_budgettable.create = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=None, + changed_by="user-1", + ) + + assert result == "budget-123" + mock_db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-123"}, + data={"soft_budget": None, "updated_by": "user-1"}, + ) + + +@pytest.mark.asyncio +async def test_update_key_soft_budget_creates_budget_row_when_key_has_none(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock() + created_row.budget_id = "budget-new" + mock_db = MagicMock() + mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row) + mock_db.litellm_budgettable.update = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=10.5, + changed_by="user-1", + ) + + assert result == "budget-new" + mock_db.litellm_budgettable.create.assert_awaited_once_with( + data={"soft_budget": 10.5, "created_by": "user-1", "updated_by": "user-1"} + ) + + +@pytest.mark.asyncio +async def test_update_key_soft_budget_noop_when_clearing_without_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + mock_db = MagicMock() + mock_db.litellm_budgettable.create = AsyncMock() + mock_db.litellm_budgettable.update = AsyncMock() + + result = await _update_key_soft_budget( + db=mock_db, + existing_key_row=existing_key, + soft_budget=None, + changed_by="user-1", + ) + + assert result is None + mock_db.litellm_budgettable.create.assert_not_awaited() + mock_db.litellm_budgettable.update.assert_not_awaited() + + +def test_update_key_request_accepts_soft_budget(): + request = UpdateKeyRequest(key="sk-test", soft_budget=42.0) + assert request.soft_budget == 42.0 + assert "soft_budget" in request.model_fields_set + + +@pytest.mark.parametrize("valid_value", [None, 0.0, 25.0]) +def test_validate_soft_budget_value_accepts_valid_values(valid_value): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_soft_budget_value, + ) + + assert _validate_soft_budget_value(valid_value) is None + + +@pytest.mark.parametrize("invalid_value", [-5.0, float("nan"), float("inf")]) +def test_validate_soft_budget_value_rejects_invalid_values(invalid_value): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_soft_budget_value, + ) + + with pytest.raises(HTTPException) as exc_info: + _validate_soft_budget_value(invalid_value) + + assert exc_info.value.status_code == 400 + assert "soft_budget must be a non-negative finite number" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_soft_budget_update_adds_budget_id_for_new_budget_row(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _apply_soft_budget_update, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock() + created_row.budget_id = "budget-created-456" + mock_db = MagicMock() + mock_db.litellm_budgettable.create = AsyncMock(return_value=created_row) + mock_db.litellm_budgettable.update = AsyncMock() + + result = await _apply_soft_budget_update( + data=UpdateKeyRequest(key="sk-test", soft_budget=25.0), + non_default_values={"soft_budget": 25.0}, + db=mock_db, + existing_key_row=existing_key, + changed_by="user-1", + ) + + assert dict(result) == {"budget_id": "budget-created-456"} + + +@pytest.mark.asyncio +async def test_apply_soft_budget_update_keeps_existing_budget_id_out_of_token_update(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _apply_soft_budget_update, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id="budget-123") + mock_db = MagicMock() + mock_db.litellm_budgettable.update = AsyncMock() + mock_db.litellm_budgettable.create = AsyncMock() + + result = await _apply_soft_budget_update( + data=UpdateKeyRequest(key="sk-test", soft_budget=40.0), + non_default_values={"soft_budget": 40.0, "max_budget": 100.0}, + db=mock_db, + existing_key_row=existing_key, + changed_by="user-1", + ) + + assert dict(result) == {"max_budget": 100.0} + mock_db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-123"}, + data={"soft_budget": 40.0, "updated_by": "user-1"}, + ) + + +@pytest.mark.asyncio +async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transaction(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_row_with_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock(budget_id="budget-new") + updated_row = MagicMock() + updated_row.model_dump.return_value = {"token": "hashed", "budget_id": "budget-new"} + tx = MagicMock() + tx.litellm_budgettable.create = AsyncMock(return_value=created_row) + tx.litellm_verificationtoken.update = AsyncMock(return_value=updated_row) + tx_context = MagicMock() + tx_context.__aenter__ = AsyncMock(return_value=tx) + tx_context.__aexit__ = AsyncMock(return_value=None) + prisma_client = MagicMock() + prisma_client.tx.return_value = tx_context + prisma_client.jsonify_object = lambda data: dict(data) + + result = await _update_key_row_with_soft_budget( + prisma_client=prisma_client, + key="sk-test", + data=UpdateKeyRequest(key="sk-test", soft_budget=25.0), + non_default_values={"soft_budget": 25.0}, + existing_key_row=existing_key, + changed_by="user-1", + ) + + assert set(result) == {"token", "data"} + assert result["data"] == {"token": "hashed", "budget_id": "budget-new"} + tx.litellm_verificationtoken.update.assert_awaited_once() + update_call = tx.litellm_verificationtoken.update.await_args + assert update_call.kwargs["where"] == {"token": result["token"]} + assert update_call.kwargs["data"]["budget_id"] == "budget-new" + assert "soft_budget" not in update_call.kwargs["data"] + + +@pytest.mark.asyncio +async def test_update_key_row_with_soft_budget_propagates_transaction_error(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _update_key_row_with_soft_budget, + ) + + existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) + created_row = MagicMock(budget_id="budget-new") + tx = MagicMock() + tx.litellm_budgettable.create = AsyncMock(return_value=created_row) + tx.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed")) + tx_context = MagicMock() + tx_context.__aenter__ = AsyncMock(return_value=tx) + tx_context.__aexit__ = AsyncMock(return_value=None) + prisma_client = MagicMock() + prisma_client.tx.return_value = tx_context + prisma_client.jsonify_object = lambda data: dict(data) + + with pytest.raises(RuntimeError, match="update failed"): + await _update_key_row_with_soft_budget( + prisma_client=prisma_client, + key="sk-test", + data=UpdateKeyRequest(key="sk-test", soft_budget=25.0), + non_default_values={"soft_budget": 25.0}, + existing_key_row=existing_key, + changed_by="user-1", + ) + + tx_context.__aexit__.assert_awaited_once() + assert tx_context.__aexit__.await_args.args[0] is RuntimeError + + def test_generate_key_request_blank_team_id_is_personal(): """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925).""" from litellm.proxy._types import RegenerateKeyRequest diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 63c08529b4d..d90338f8480 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4447,13 +4447,13 @@ class TestStrategyRouterWriteValidation: self.litellm_proxymodeltable = MagicMock( create=AsyncMock(), update=AsyncMock(), - find_many=AsyncMock( - return_value=tuple( - LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows - ) - ), + find_many=AsyncMock(side_effect=self._find_many), ) + async def _find_many(self, where: object = None) -> tuple[LiteLLM_ProxyModelTable, ...]: + json.dumps(where) # prisma serializes the filter with json.dumps and rejects a mappingproxy + return tuple(LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows) + @property def db(self) -> "TestStrategyRouterWriteValidation._FakeTx": return self diff --git a/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py b/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py new file mode 100644 index 00000000000..61eea00bf7d --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_per_request_root_path_middleware.py @@ -0,0 +1,302 @@ +"""Tests for PerRequestRootPathMiddleware (``SERVER_ROOT_PATHS``). + +One deployment fronting several client-visible URL path prefixes: the matched +prefix becomes that request's ``root_path``, so Starlette route matching and +``request.base_url`` — and therefore every URL the proxy emits, the MCP OAuth +discovery documents among them — resolve under the prefix the client called. +""" + +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient + +from litellm.proxy.middleware.per_request_root_path_middleware import ( + PerRequestRootPathMiddleware, + get_request_root_path, + get_server_root_paths, + normalize_root_paths, +) + + +class TestNormalizeRootPaths: + def test_strips_whitespace_and_trailing_slash(self): + assert normalize_root_paths([" /tenant-a/ ", "/tenant-b"]) == ( + "/tenant-a", + "/tenant-b", + ) + + def test_drops_empty_entries(self): + assert normalize_root_paths(["", " ", "/tenant-a"]) == ("/tenant-a",) + + def test_drops_entries_without_leading_slash(self): + # A typo'd entry must not silently match nothing at request time. + assert normalize_root_paths(["tenant-a", "/tenant-b"]) == ("/tenant-b",) + + def test_drops_bare_root(self): + # "/" would turn every request into a root_path rewrite; a + # root-mounted deployment needs no entry at all. + assert normalize_root_paths(["/", "/tenant-a"]) == ("/tenant-a",) + + def test_dedupes(self): + assert normalize_root_paths(["/t", "/t/", " /t "]) == ("/t",) + + def test_longest_first_for_nested_prefixes(self): + # Longest-first ordering is what makes the most-specific nested + # prefix win at match time. + assert normalize_root_paths(["/t", "/t/deep"]) == ("/t/deep", "/t") + + +class TestGetServerRootPaths: + def test_unset_env_is_empty(self, monkeypatch): + monkeypatch.delenv("SERVER_ROOT_PATHS", raising=False) + assert get_server_root_paths() == () + + def test_empty_env_is_empty(self, monkeypatch): + monkeypatch.setenv("SERVER_ROOT_PATHS", "") + assert get_server_root_paths() == () + + def test_comma_separated_entries(self, monkeypatch): + monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a, /tenant-b/") + assert get_server_root_paths() == ("/tenant-a", "/tenant-b") + + +def _capture_scope_middleware(root_paths): + """Middleware wired to a downstream that records the scope it received.""" + captured = {} + + async def downstream(scope, receive, send): + captured.update(scope) + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + return PerRequestRootPathMiddleware(downstream, root_paths=root_paths), captured + + +async def _run(mw, scope): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message): + pass + + await mw(scope, receive, send) + + +class TestPerRequestRootPathMiddleware: + @pytest.mark.asyncio + async def test_matched_prefix_becomes_root_path_path_untouched(self): + # Starlette's router strips root_path from the (unmodified) path at + # match time, so the middleware must NOT rewrite scope["path"]. + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/tenant-a/mcp/x", "method": "GET", "headers": []}) + assert captured["root_path"] == "/tenant-a" + assert captured["path"] == "/tenant-a/mcp/x" + + @pytest.mark.asyncio + async def test_exact_prefix_matches(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/tenant-a", "method": "GET", "headers": []}) + assert captured["root_path"] == "/tenant-a" + + @pytest.mark.asyncio + async def test_segment_boundary_prevents_sibling_match(self): + # /tenant-ab must not match the /tenant-a prefix. + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/tenant-ab/mcp", "method": "GET", "headers": []}) + assert "root_path" not in captured + + @pytest.mark.asyncio + async def test_unmatched_path_untouched(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "http", "path": "/chat/completions", "method": "GET", "headers": []}) + assert "root_path" not in captured + assert captured["path"] == "/chat/completions" + + @pytest.mark.asyncio + async def test_longest_nested_prefix_wins(self): + mw, captured = _capture_scope_middleware(["/t", "/t/deep"]) + await _run(mw, {"type": "http", "path": "/t/deep/mcp", "method": "GET", "headers": []}) + assert captured["root_path"] == "/t/deep" + + @pytest.mark.asyncio + async def test_matched_prefix_overrides_scalar_root_path(self): + # FastAPI(root_path=SERVER_ROOT_PATH) stamps the scalar before the + # middleware stack runs; a matched dynamic prefix wins for that + # request (combining both mechanisms is warned about at startup). + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run( + mw, + {"type": "http", "path": "/tenant-a/mcp", "root_path": "/legacy", "method": "GET", "headers": []}, + ) + assert captured["root_path"] == "/tenant-a" + + @pytest.mark.asyncio + async def test_unmatched_request_keeps_scalar_root_path(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run( + mw, + {"type": "http", "path": "/legacy/mcp", "root_path": "/legacy", "method": "GET", "headers": []}, + ) + assert captured["root_path"] == "/legacy" + + @pytest.mark.asyncio + async def test_websocket_scope_matched(self): + mw, captured = _capture_scope_middleware(["/tenant-a"]) + await _run(mw, {"type": "websocket", "path": "/tenant-a/ws", "headers": []}) + assert captured["root_path"] == "/tenant-a" + + @pytest.mark.asyncio + async def test_lifespan_scope_passes_through(self): + called = {} + + async def downstream(scope, receive, send): + called["scope"] = scope + + mw = PerRequestRootPathMiddleware(downstream, root_paths=["/tenant-a"]) + await _run(mw, {"type": "lifespan"}) + assert called["scope"] == {"type": "lifespan"} + + +def _routed_client(prefixes): + app = FastAPI() + + @app.get("/where") + def where(request: Request): + return { + "base_url": str(request.base_url), + "root_path": request.scope.get("root_path", ""), + } + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=prefixes) + return TestClient(app) + + +class TestEndToEndRouting: + def test_two_prefixes_route_on_one_app(self): + # The property a scalar SERVER_ROOT_PATH cannot provide: two + # client-visible prefixes served by the same app, each request + # reconstructing its own base URL. + client = _routed_client(["/tenant-a", "/tenant-b"]) + + resp_a = client.get("/tenant-a/where") + resp_b = client.get("/tenant-b/where") + + assert resp_a.status_code == 200 + assert resp_a.json() == { + "base_url": "http://testserver/tenant-a/", + "root_path": "/tenant-a", + } + assert resp_b.status_code == 200 + assert resp_b.json() == { + "base_url": "http://testserver/tenant-b/", + "root_path": "/tenant-b", + } + + def test_unprefixed_route_still_served(self): + client = _routed_client(["/tenant-a"]) + resp = client.get("/where") + assert resp.status_code == 200 + assert resp.json()["base_url"] == "http://testserver/" + + def test_unlisted_prefix_404s(self): + client = _routed_client(["/tenant-a"]) + assert client.get("/tenant-c/where").status_code == 404 + + +class TestGetRequestRootPath: + """``get_request_root_path`` is the accessor that plumbs the middleware's + resolved prefix to code that doesn't have scope in hand — the 401 challenge + builders in ``mcp_server_manager`` / ``server`` and ``get_custom_url`` on the + SSO callback path. Reading the SERVER_ROOT_PATH scalar there would emit URLs + under a prefix the client didn't call, and stack a second prefix onto ones it + did (the two review points this fixture pins).""" + + def test_falls_back_to_server_root_path_env_outside_a_request(self, monkeypatch): + # Outside a request the middleware's ContextVar is unset. The scalar + # env still owns the answer, so pre-middleware call sites (module-load + # UI URL builders, background tasks) behave exactly as they did. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + assert get_request_root_path() == "/legacy" + + def test_returns_empty_string_when_no_env_and_no_request(self, monkeypatch): + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + assert get_request_root_path() == "" + + def test_returns_matched_prefix_inside_a_request(self, monkeypatch): + # With both env vars set, a request matching a SERVER_ROOT_PATHS prefix + # must see that prefix — not the SERVER_ROOT_PATH scalar — so the URL + # it emits stays under the prefix the router will resolve it against. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + seen: list[str] = [] + app = FastAPI() + + @app.get("/where") + def where(): + seen.append(get_request_root_path()) + return {} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"]) + client = TestClient(app) + + assert client.get("/tenant-a/where").status_code == 200 + assert seen == ["/tenant-a"] + + def test_unmatched_request_falls_through_to_scope_scalar(self, monkeypatch): + # A request the middleware saw but did not match keeps whatever + # scope["root_path"] the app was mounted under (the scalar). The + # ContextVar still reflects the effective per-request answer, so + # emitted URLs and the router agree even on the fallback path. + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + seen: list[str] = [] + + app = FastAPI(root_path="/legacy") + + @app.get("/where") + def where(): + seen.append(get_request_root_path()) + return {} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"]) + client = TestClient(app) + + assert client.get("/legacy/where").status_code == 200 + assert seen == ["/legacy"] + + def test_each_request_sees_its_own_prefix(self, monkeypatch): + # Two sequential requests through the same app must each see the + # prefix they arrived under, so one tenant's client is never sent + # the URL of another tenant's origin. + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + seen: list[tuple[str, str]] = [] + app = FastAPI() + + @app.get("/where") + def where(tag: str): + seen.append((tag, get_request_root_path())) + return {} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a", "/tenant-b"]) + client = TestClient(app) + assert client.get("/tenant-a/where?tag=a").status_code == 200 + assert client.get("/tenant-b/where?tag=b").status_code == 200 + assert seen == [("a", "/tenant-a"), ("b", "/tenant-b")] + + def test_context_var_reset_after_request(self, monkeypatch): + # A ContextVar left set after the request finishes would poison the + # module-load-time callers that read it lazily (they'd think they were + # inside a request under the last-seen prefix). + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + app = FastAPI() + + @app.get("/where") + def where(): + return {"prefix": get_request_root_path()} + + app.add_middleware(PerRequestRootPathMiddleware, root_paths=["/tenant-a"]) + client = TestClient(app) + assert client.get("/tenant-a/where").json() == {"prefix": "/tenant-a"} + # After the request finishes, the scalar-env fallback owns the answer + # again — nothing was left stashed from the last request's scope. + assert get_request_root_path() == "/legacy" diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a518c9289ae..7db85dc6943 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -821,6 +821,27 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): ) +@pytest.mark.asyncio +async def test_tuning_baseline_v2_is_created_alongside_the_legacy_row(): + from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT + + prisma_client = MagicMock() + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_config.create = AsyncMock() + deployment = { + "model_name": "a", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": {}}, + } + + result = await ProxyStartupEvent._load_heuristic_v1_tuning_baselines(prisma_client, [deployment]) + + assert result == {'yaml:["a",[]]': DEFAULT_TUNING_FINGERPRINT} + assert prisma_client.db.litellm_config.create.await_args.kwargs["data"] == { + "param_name": "auto_router_tuning_baseline_v2", + "param_value": json.dumps(dict(result)), + } + + @pytest.mark.asyncio async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch): prisma_client = MagicMock() diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 4a19ad3541c..0d82ed778f5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1,5 +1,6 @@ import re from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -327,6 +328,85 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_chatgpt_provider_fields(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers = response.json() + + chatgpt = next((p for p in providers if p["provider"] == "CHATGPT"), None) + assert chatgpt is not None, "ChatGPT provider entry not found" + + assert chatgpt["provider_display_name"] == "ChatGPT Subscription" + assert chatgpt["litellm_provider"] == LlmProviders.CHATGPT.value + assert chatgpt["default_model_placeholder"].startswith("chatgpt/") + assert chatgpt["credential_fields"] == [] + + +ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( + { + "a2a", + "a2a_agent", + "amazon_nova", + "apertis", + "aws_polly", + "black_forest_labs", + "charity_engine", + "chutes", + "darkbloom", + "gdc", + "helicone", + "inception", + "langflow", + "langgraph", + "libertai", + "litellm_agent", + "manus", + "meta", + "modelscope", + "mongodb", + "nano-gpt", + "neosantara", + "parasail", + "pinstripes", + "poe", + "publicai", + "ragflow", + "reducto", + "s3_vectors", + "sagemaker_nova", + "scaleway", + "stability", + "synthetic", + "tencent", + "tensormesh", + "text-completion-inception", + "valkey", + "xiaomi_mimo", + "zai", + } +) + + +def test_every_backend_provider_is_listed_in_add_model_or_frozen_as_unlisted(): + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + listed = {p["litellm_provider"] for p in response.json()} + + unlisted = {provider.value for provider in LlmProviders} - listed + assert unlisted == ADD_MODEL_UNLISTED_PROVIDERS, ( + "Add Model dropdown drift: give the new provider an entry in provider_create_fields.json " + "rather than adding it to ADD_MODEL_UNLISTED_PROVIDERS" + ) + + def test_google_ai_studio_provider_fields_expose_api_base(): """The Google AI Studio (gemini) credential form must let admins set a custom api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index f65f68812a2..c7de8c943ad 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -5,7 +5,7 @@ import pytest from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request +from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request from litellm.proxy.utils import ProxyLogging TOKEN_COUNTING_ROUTES: Final = ( @@ -46,3 +46,32 @@ async def test_non_exempt_llm_route_still_reserves_budget(): assert reservation is not None assert reservation["reserved_cost"] > 0 + + +BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6" +CONVERSE_BODY: Final = { + "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}], + "inferenceConfig": {"maxTokens": 5}, +} +INVOKE_BODY: Final = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 5, + "messages": [{"role": "user", "content": "Reply with one word: pong"}], +} + + +def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): + converse_cost: Final = estimate_request_max_cost( + request_body=CONVERSE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/converse", + llm_router=None, + input_token_counts={}, + ) + invoke_cost: Final = estimate_request_max_cost( + request_body=INVOKE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/invoke", + llm_router=None, + input_token_counts={}, + ) + assert converse_cost is not None and invoke_cost is not None + assert invoke_cost < converse_cost < 2 * invoke_cost diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index 3f775d82b7f..cc8fdeb0160 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -1,4 +1,4 @@ - +from typing import Final import pytest @@ -121,6 +121,147 @@ def _caching_usage(read: int, written: int, text: int = 10, out: int = 100) -> d } +@pytest.mark.parametrize( + "model,provider,prompt,reads,writes_5m,writes_1h,tier,region,location", + [ + pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 20000, 0, None, None, None, id="5m"), + pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 0, 20000, None, None, None, id="1h"), + pytest.param("claude-sonnet-4-5", "anthropic", 50000, 20000, 12000, 8000, None, None, None, id="mixed-ttl"), + pytest.param("claude-sonnet-4-5", "anthropic", 199999, 80000, 20000, 0, None, None, None, id="below-200k"), + pytest.param("claude-sonnet-4-5", "anthropic", 200000, 80000, 20000, 0, None, None, None, id="exactly-200k"), + pytest.param("claude-sonnet-4-5", "anthropic", 200001, 80000, 20000, 0, None, None, None, id="above-200k"), + pytest.param("claude-sonnet-4-5", "anthropic", 250000, 80000, 12000, 8000, None, None, None, id="ttl-and-200k"), + pytest.param( + "claude-sonnet-4-5", "anthropic", 250000, 80000, 0, 20000, "priority", None, None, id="absent-tier" + ), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "priority", None, None, id="priority"), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "flex", None, None, id="flex"), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "batch", None, None, id="batch-resolver-fallback"), + pytest.param("gpt-5.5", "openai", 300000, 80000, 0, 0, "flex", None, None, id="flex-and-272k"), + pytest.param("gpt-5.5", "openai", 100000, 80000, 0, 0, "priority", "eu", None, id="priority-and-eu"), + pytest.param("gemini-2.5-pro", "vertex_ai", 250000, 80000, 20000, 0, None, None, None, id="variant-only-write"), + pytest.param("gemini-3.5-flash", "vertex_ai", 100000, 80000, 0, 0, None, None, "us-east5", id="vertex-region"), + pytest.param("gpt-5.5", "openai", 300000, 0, 0, 0, "priority", "eu", None, id="no-cache"), + ], +) +def test_caching_savings_agree_with_biller_on_the_request_pricing_basis( + model: str, + provider: str, + prompt: int, + reads: int, + writes_5m: int, + writes_1h: int, + tier: str | None, + region: str | None, + location: str | None, +) -> None: + pricing: Final = litellm.get_model_info(model=model, custom_llm_provider=provider) + usage: Final = Usage( + prompt_tokens=prompt, + completion_tokens=100, + total_tokens=prompt + 100, + prompt_tokens_details={ + "cached_tokens": reads, + "cache_creation_tokens": writes_5m + writes_1h, + "text_tokens": prompt - reads - writes_5m - writes_1h, + "cache_creation_token_details": { + "ephemeral_5m_input_tokens": writes_5m, + "ephemeral_1h_input_tokens": writes_1h, + }, + }, + ) + uncached: Final = Usage( + prompt_tokens=prompt, + completion_tokens=100, + total_tokens=prompt + 100, + prompt_tokens_details={"text_tokens": prompt, "cached_tokens": 0, "cache_creation_tokens": 0}, + ) + costs: Final = tuple( + sum( + generic_cost_per_token( + model=model, + usage=arm, + custom_llm_provider=provider, + model_info=pricing, + service_tier=tier, + data_residency=region, + vertex_location=location, + ) + ) + for arm in (uncached, usage) + ) + expected: Final = costs[0] - costs[1] + for attributed in (False, True): + result: Final = compute_savings_spend( + model=model, + custom_llm_provider=provider, + compression_saved_tokens=4389, + gateway_injected_cache=attributed, + usage_object=usage.model_dump(), + cost_breakdown={"service_tier": tier, "data_residency": region, "vertex_location": location}, + billed_at="2026-09-07T12:00:00+00:00", + ) + assert result.prompt_caching == pytest.approx(expected) + assert result.gateway_injected_caching == pytest.approx(expected if attributed else 0.0) + assert result.compression == pytest.approx(4389 * (pricing["input_cost_per_token"] or 0.0)) + assert result.autorouter == 0.0 + if reads + writes_5m + writes_1h == 0: + assert expected == 0.0 + + +def test_negative_ttl_counts_do_not_become_cache_write_credits() -> None: + results: Final = tuple( + compute_savings_spend( + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=True, + usage_object={ + "prompt_tokens": 6000, + "completion_tokens": 100, + "prompt_tokens_details": { + "text_tokens": 1000, + "cache_creation_tokens": 5000, + "cache_creation_token_details": { + "ephemeral_5m_input_tokens": short_count, + "ephemeral_1h_input_tokens": 5000, + }, + }, + }, + ) + for short_count in (-5000, 0) + ) + assert results[0] == results[1] + assert results[0].prompt_caching < 0 + + +def test_unpublished_one_hour_price_uses_the_ordinary_write_price() -> None: + model: Final = "claude-4-opus-20250514" + pricing: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic") + assert pricing.get("cache_creation_input_token_cost_above_1hr") is None + assert pricing["cache_creation_input_token_cost"] > pricing["input_cost_per_token"] + results: Final = tuple( + compute_savings_spend( + model=model, + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=True, + usage_object={ + "prompt_tokens": 6000, + "completion_tokens": 100, + "prompt_tokens_details": { + "text_tokens": 1000, + "cache_creation_tokens": 5000, + "cache_creation_token_details": ttl, + }, + }, + ) + for ttl in (None, {"ephemeral_1h_input_tokens": 5000}) + ) + assert results[0] == results[1] + assert results[0].prompt_caching < 0 + + def test_prompt_caching_savings_nets_out_the_cache_write_premium(): """A cache-writing request is only credited the read discount minus the write premium.""" input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5") diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 96d9b0c5a26..6acd9d7258e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2,7 +2,7 @@ import asyncio import copy import datetime import json -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -1809,6 +1809,146 @@ class TestCommonRequestProcessingHelpers: response = await create_response(mock_generator(), "text/event-stream", custom_headers) assert response.headers["x-custom-header"] == "TestValue" + async def test_create_streaming_response_refresh_headers_after_first_chunk(self): + """LIT-6767: headers a caller can only resolve once the first chunk exists. + + A pre-first-chunk fallback replaces the deployment while the response + headers are still uncommitted, so ``refresh_headers`` is consulted after + the first chunk is buffered and its result wins. + """ + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + refresh_headers: Final = AsyncMock( + return_value={"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"} + ) + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment", "llm_provider-x-request-id": "req-FAILED"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert refresh_headers.await_count == 1 + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + assert response.headers["llm_provider-x-request-id"] == "req-FALLBACK" + # the buffering headers are still applied on top of the refreshed set + assert response.headers["x-accel-buffering"] == "no" + assert response.headers["cache-control"] == "no-cache" + + async def test_create_streaming_response_refreshes_only_after_the_first_chunk(self): + """LIT-6767: the refresh has to be consulted after the generator produced a chunk. + + A pre-first-chunk fallback only repoints the response while that first chunk is + being produced, so a refresh consulted any earlier still describes the attempt + that failed and the headers go out wrong. + """ + first_chunk_produced: Final = asyncio.Event() + + async def mock_generator(): + first_chunk_produced.set() + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + served = "fallback-deployment" if first_chunk_produced.is_set() else "failed-deployment" + return {"x-litellm-model-id": served} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + + async def test_create_streaming_response_empty_stream_uses_refreshed_headers(self): + """LIT-6767: a fallback that served nothing still gets to name itself. + + The empty-generator branch returns its own StreamingResponse, so it needs the + refreshed headers too or the client is told the failed deployment answered. + """ + + async def mock_generator(): + return + yield # make it an async generator + + async def refresh_headers(): + return {"x-litellm-model-id": "fallback-deployment"} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + assert response.headers["x-accel-buffering"] == "no" + + async def test_create_streaming_response_without_refresh_headers_is_unchanged(self): + """LIT-6767: the default keeps the caller-supplied headers verbatim.""" + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + ) + assert response.headers["x-litellm-model-id"] == "failed-deployment" + + async def test_create_streaming_response_refresh_headers_failure_keeps_stream(self): + """LIT-6767: the first chunk is already paid for, so a failing refresh + falls back to the caller's headers instead of erroring the stream.""" + + async def mock_generator(): + yield 'data: {"content": "data"}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + raise RuntimeError("boom") + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, StreamingResponse) + assert response.status_code == status.HTTP_200_OK + assert response.headers["x-litellm-model-id"] == "failed-deployment" + assert await self.consume_stream(response) == [ + 'data: {"content": "data"}\n\n', + "data: [DONE]\n\n", + ] + + async def test_create_response_first_chunk_error_uses_refreshed_headers(self): + """LIT-6767: the JSON error response built from a bad first chunk carries + the refreshed headers too, so it cannot describe a deployment that no + longer served the request.""" + + async def mock_generator(): + yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' + yield "data: [DONE]\n\n" + + async def refresh_headers(): + return {"x-litellm-model-id": "fallback-deployment"} + + response = await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-model-id": "failed-deployment"}, + refresh_headers=refresh_headers, + ) + assert isinstance(response, JSONResponse) + assert response.headers["x-litellm-model-id"] == "fallback-deployment" + async def test_create_streaming_response_disables_proxy_buffering(self): """Regression for #28384: every StreamingResponse create_response returns must carry the headers that stop nginx/ingress/Envoy from buffering the @@ -8014,3 +8154,108 @@ class TestDetachedStreamFailureHook: await logging_obj._on_detached_stream_failure(failure) assert [call["original_exception"] for call in recorder.calls] == [failure] + + +class TestStreamingResponseHeadersFollowFallback: + """LIT-6767: the streaming branch has to publish the deployment that served the stream.""" + + @staticmethod + def _fallback_adopting_stream(): + class _Stream: + def __init__(self) -> None: + self._hidden_params = { + "model_id": "failed-deployment", + "api_base": "http://127.0.0.1:20769/v1", + "additional_headers": {"llm_provider-stale-marker": "failed-deployment"}, + } + self.fallback_headers_adopted = False + + def adopt(self) -> None: + self._hidden_params = { + "model_id": "served-deployment", + "api_base": "https://api.openai.com", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + self.fallback_headers_adopted = True + + return _Stream() + + @pytest.mark.asyncio + async def test_streaming_headers_name_the_deployment_that_served(self, monkeypatch): + """A pre-first-chunk fallback repoints the stream while the headers are still + uncommitted, so the published headers must describe the fallback, not the attempt + the Router picked first.""" + stream = self._fallback_adopting_stream() + + def select_data_generator(**kwargs): + async def generator(): + stream.adopt() + yield 'data: {"choices": [{"delta": {"content": "OK"}}]}\n\n' + yield "data: [DONE]\n\n" + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-6767-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "oa-midfail", "stream": True, "litellm_logging_obj": logging_obj} + ) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-header": "kept"} + ) + + async def fake_route_request(**kwargs): + async def call(): + return stream + + return call() + + monkeypatch.setattr( + litellm.proxy.common_request_processing, "route_request", fake_route_request + ) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, StreamingResponse) + assert result.headers["x-litellm-model-id"] == "served-deployment" + assert result.headers["x-litellm-model-api-base"] == "https://api.openai.com" + assert result.headers["llm_provider-x-request-id"] == "req-SERVED" + assert "llm_provider-stale-marker" not in result.headers + assert result.headers["x-callback-header"] == "kept" + + +class TestPassthroughHeadersAcceptImmutableMappings: + """LIT-6767: the streaming branch now hands the passthrough helpers an immutable mapping.""" + + def test_merge_passthrough_streaming_headers_accepts_a_read_only_mapping(self): + merged = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + response_headers=httpx.Headers({"content-type": "text/event-stream", "transfer-encoding": "chunked"}), + custom_headers=MappingProxyType({"x-litellm-model-id": "served-deployment"}), + ) + + assert merged["x-litellm-model-id"] == "served-deployment" + assert merged["content-type"] == "text/event-stream" + # the excluded hop-by-hop header is still dropped + assert "transfer-encoding" not in merged diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9393ec0f8e6..b7bb58378d4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9224,6 +9224,82 @@ class TestLazyFeatureMiddleware: else: assert loads == [], f"{case}: feature must not load" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "env_root_path,scope_root_path,request_path,should_load,case", + [ + # Per-request root_path (PerRequestRootPathMiddleware under + # SERVER_ROOT_PATHS) with no scalar env: strip and match. + ("", "/tenant-a", "/tenant-a/dummy/x", True, "per-request root_path strip"), + # scope root_path is authoritative over the cached env scalar. + ("/api/v1", "/tenant-a", "/tenant-a/dummy/x", True, "scope wins over env scalar"), + # Boundary check still applies to the per-request value. + ("", "/tenant-a", "/tenant-ab/dummy/x", False, "boundary check on scope root_path"), + # Empty scope root_path falls back to the env scalar. + ("/api/v1", "", "/api/v1/dummy/x", True, "empty scope falls back to env"), + ], + ) + async def test_per_request_root_path_handling( + self, monkeypatch, env_root_path, scope_root_path, request_path, should_load, case + ): + """ + ``scope["root_path"]`` must be stripped before prefix matching when + set — the scalar SERVER_ROOT_PATH lands there via + ``FastAPI(root_path=...)``, and PerRequestRootPathMiddleware + (SERVER_ROOT_PATHS) resolves a per-request prefix there. Otherwise + lazily-registered features — the MCP OAuth discovery router among + them — stay unloaded under a client-visible prefix and 404. + """ + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + monkeypatch.setenv("SERVER_ROOT_PATH", env_root_path) + + loads = [] + + def fake_register(app, module): + loads.append(getattr(module, "__name__", "?")) + + feat = LazyFeature( + name=f"dummy_prr_{case}", + module_path="json", + path_prefixes=("/dummy",), + register_fn=fake_register, + ) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message): + pass + + await mw( + { + "type": "http", + "path": request_path, + "root_path": scope_root_path, + "method": "GET", + "headers": [], + }, + receive, + send, + ) + if should_load: + assert loads == ["json"], f"{case}: expected feature to load" + else: + assert loads == [], f"{case}: feature must not load" + @pytest.mark.asyncio async def test_concurrent_first_requests_only_register_once(self): """ diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcf18e773e8..9462f2c8eb0 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1922,6 +1922,34 @@ async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monk assert "test_proxy_utils" in captured["async_traceback"] +@pytest.mark.parametrize( + "key_metadata, team_metadata, expected_to_run", + [ + ({"guardrails": ["key-scoped-guardrail"]}, None, True), + ({}, {"guardrails": ["key-scoped-guardrail"]}, True), + ({"guardrails": ["some-other-guardrail"]}, None, False), + ({}, None, False), + ], +) +def test_convert_mcp_to_llm_format_carries_key_and_team_guardrails(key_metadata, team_metadata, expected_to_run): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + guardrail = CustomGuardrail(guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False) + kwargs = { + "name": "ask_question", + "arguments": {"question": "hello"}, + "server_name": "deepwiki", + "user_api_key_auth": UserAPIKeyAuth(metadata=key_metadata, team_metadata=team_metadata), + } + request_obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs) + + with patch( # test-quality-ok: the key-guardrail premium gate reads this proxy_server module global and has no injection seam + "litellm.proxy.proxy_server.premium_user", True + ): + synthetic = proxy_logging._convert_mcp_to_llm_format(request_obj, kwargs) + + assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is expected_to_run + + class _TracebackRecordingLogger(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py index 31ea1bdce74..5f23c5fc20e 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py @@ -314,3 +314,76 @@ def test_normalize_route_for_root_path_error_path_when_route_not_under_root( _clear_url_env(monkeypatch) monkeypatch.setenv("SERVER_ROOT_PATH", "/proxy") assert normalize_route_for_root_path("/other/v1/chat") is None + + +# --------------------------------------------------------------------------- +# get_custom_url under a per-request prefix (SERVER_ROOT_PATHS): +# when a request lives under a dynamic prefix, request.base_url already +# contains it. Appending the SERVER_ROOT_PATH scalar on top produced +# double-prefixed SSO callback / login URLs (a path that doesn't exist on +# the deployment). These pin that the emitted URL now stays under one +# prefix — the one the request actually arrived on. +# --------------------------------------------------------------------------- + + +def test_get_custom_url_uses_per_request_prefix_when_middleware_ran(monkeypatch): + """The middleware stashes the effective per-request prefix in a ContextVar. + ``get_custom_url`` reads that in preference to the SERVER_ROOT_PATH scalar, + and ``join_paths``'s tail-dedup collapses the append so a request whose + ``base_url`` already ends in ``/tenant-a`` does not become + ``/tenant-a/legacy/route``.""" + _clear_url_env(monkeypatch) + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + from litellm.proxy.middleware.per_request_root_path_middleware import ( + _request_root_path_var, + ) + + token = _request_root_path_var.set("/tenant-a") + try: + # request.base_url already carries the tenant prefix; the scalar + # SERVER_ROOT_PATH must not be re-appended on top. + result = get_custom_url( + request_base_url="https://request.example.com/tenant-a/", + route="/v1/chat", + ) + finally: + _request_root_path_var.reset(token) + + assert result == "https://request.example.com/tenant-a/v1/chat" + + +def test_get_custom_url_no_double_prefix_when_both_env_vars_configured(monkeypatch): + """Regression guard for the review point: with SERVER_ROOT_PATH also set + (scalar-legacy) and the request matched by SERVER_ROOT_PATHS (per-request), + the emitted URL is under one prefix — the per-request one — never both.""" + _clear_url_env(monkeypatch) + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a") + + from litellm.proxy.middleware.per_request_root_path_middleware import ( + _request_root_path_var, + ) + + token = _request_root_path_var.set("/tenant-a") + try: + # No "/legacy" ever appears — the fix pins the reviewer's expected + # behavior: only one prefix should apply per request. + result = get_custom_url("https://api.example.com/tenant-a", "/sso/callback") + finally: + _request_root_path_var.reset(token) + + assert result == "https://api.example.com/tenant-a/sso/callback" + assert "/legacy" not in result + + +def test_get_custom_url_scalar_only_still_stamps_root_path(monkeypatch): + """Pre-middleware deployments have not opted into SERVER_ROOT_PATHS at all; + the ContextVar stays unset and the SERVER_ROOT_PATH scalar owns the answer + — the behavior every existing scalar-only deployment relies on.""" + _clear_url_env(monkeypatch) + monkeypatch.setenv("SERVER_ROOT_PATH", "/legacy") + + result = get_custom_url("https://request.example.com", "/v1/chat") + + assert result == "https://request.example.com/legacy/v1/chat" diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 93049b21460..dc4f2900038 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -244,3 +244,65 @@ async def test_list_vector_stores_dashboard_session_resolves_real_teams( ), ): assert await _listed_ids(alice) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_team_ids", "expected_status"), + [ + (["team_a"], 200), + (["team_b"], 403), + ([], 403), + ], +) +async def test_get_vector_store_info_dashboard_session_resolves_real_teams( + user_team_ids: list[str], expected_status: int +): + """Regression for LIT-7132: /vector_store/info must grant a dashboard session the same team-owned stores + /vector_store/list shows it, instead of judging the session's reserved litellm-dashboard team id.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + get_vector_store_info, + ) + from litellm.types.vector_stores import VectorStoreInfoRequest + + alice = UserAPIKeyAuth( + team_id="litellm-dashboard", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj(team_id=team_id) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=MagicMock(model_dump=lambda: dict(_TEAM_A_OWNED)) + ) + + async def outcome() -> int: + try: + response = await get_vector_store_info( + data=VectorStoreInfoRequest(vector_store_id="vs_team_a"), user_api_key_dict=alice + ) + except HTTPException as exc: + return exc.status_code + assert response["vector_store"]["vector_store_id"] == "vs_team_a" + return 200 + + with ( + patch( # test-quality-ok: the endpoint reads the store row through the module-level prisma client, no injection seam + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch( # test-quality-ok: the endpoint consults the module-level registry before the DB, no injection seam + "litellm.vector_store_registry", None + ), + patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam + "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object + ), + patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam + "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids", + new=AsyncMock(return_value=user_team_ids), + ), + ): + assert await outcome() == expected_status diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 758d379f22c..a890d7ceed0 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -4,10 +4,13 @@ Tests for gateway repository layer. import json from datetime import datetime -from typing import Any, Dict, List, Optional +from types import SimpleNamespace +from typing import Any, Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from prisma import models as prisma_models +from prisma.builder import QueryBuilder from litellm.models.base import DomainModel from litellm.models.budget import LiteLLM_BudgetTable @@ -307,6 +310,23 @@ class TestModelRepository: client = MockPrismaClient() return ModelRepository(client) + @pytest.mark.asyncio + async def test_find_all_except_serializes_exclusion_for_prisma(self) -> None: + find_many: Final = AsyncMock(return_value=[]) + client: Final = SimpleNamespace( + db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many)) + ) + + await ModelRepository(client).find_all_except("current-model") + + find_many.assert_awaited_once() + query: Final = QueryBuilder( + method="find_many", + model=prisma_models.LiteLLM_ProxyModelTable, + arguments=find_many.call_args.kwargs, + ).build_query() + assert 'where: { model_id: { not: "current-model" } }' in " ".join(query.split()) + def test_table_is_wrapped_for_config_sync(self, repo): from litellm.proxy.common_utils.config_sync_pubsub import ( _PublishOnWriteActions, diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index cb6efa21036..9d9eefdceb3 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -440,6 +440,56 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.text_tokens == 20 assert result.completion_tokens_details.audio_tokens is None + def test_transform_realtime_usage_partitions_reasoning_out_of_text_tokens(self): + """Realtime nests reasoning_tokens inside text_tokens; the stored text share excludes them.""" + usage = { + "input_tokens": 237, + "output_tokens": 70, + "total_tokens": 307, + "input_token_details": {"text_tokens": 43, "audio_tokens": 0, "image_tokens": 194, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens == 70 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 18 + assert result.completion_tokens_details.reasoning_tokens == 52 + assert result.completion_tokens_details.audio_tokens == 0 + + def test_transform_realtime_usage_partitions_reasoning_beside_audio_output(self): + """Audio output stays as reported; only the text share sheds the nested reasoning tokens.""" + usage = { + "input_tokens": 100, + "output_tokens": 70, + "total_tokens": 170, + "input_token_details": {"text_tokens": 100, "audio_tokens": 0, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 39, "audio_tokens": 31, "reasoning_tokens": 23}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 16 + assert result.completion_tokens_details.audio_tokens == 31 + assert result.completion_tokens_details.reasoning_tokens == 23 + + def test_transform_response_api_usage_keeps_partitioned_text_tokens(self): + """A provider already reporting text_tokens beside reasoning_tokens is stored as sent.""" + usage = { + "input_tokens": 10, + "output_tokens": 20, + "total_tokens": 30, + "output_tokens_details": {"text_tokens": 12, "reasoning_tokens": 5}, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 12 + assert result.completion_tokens_details.reasoning_tokens == 5 + def test_transform_response_api_usage_carries_extra_provider_fields(self): """Non-standard usage fields (e.g. xAI tool details) must survive chat normalization.""" details = {"web_search_calls": 2, "x_search_calls": 0} diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 4b446368dbe..74d96bda336 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -4,7 +4,6 @@ import pytest from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled from litellm.rust_bridge import configuration, responses_websocket -from litellm.types.router import GenericLiteLLMParams class _FakeNativeConnection: @@ -48,22 +47,12 @@ def reset_responses_websocket(): configuration.reset_rust_configuration() -def test_rust_websocket_bridge_is_disabled_without_flag() -> None: - assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) - assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True)) - assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) - - -def test_explicit_false_overrides_process_enable() -> None: +def test_rust_websocket_bridge_uses_process_enablement() -> None: + configuration.rust(False) + assert not _rust_responses_websocket_enabled("openai") configuration.rust(True) - - assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) - - -def test_process_enable_applies_without_request_override() -> None: - configuration.rust(True) - - assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) + assert _rust_responses_websocket_enabled("openai") + assert not _rust_responses_websocket_enabled("anthropic") @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py index c3395ab641f..fa7a96adb20 100644 --- a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -72,6 +72,9 @@ class TestTuningFingerprint: config["classifier_llm_config"] = {"model": "judge"} assert tuning_fingerprint(config) != DEFAULT_TUNING_FINGERPRINT + def test_explicit_empty_tier_model_configs_follow_omission(self) -> None: + assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT + def test_tier_model_overrides_change_the_fingerprint(self) -> None: plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}}) with_override = tuning_fingerprint( @@ -80,11 +83,39 @@ class TestTuningFingerprint: assert plain != with_override def test_non_tuning_fields_do_not_change_the_fingerprint(self) -> None: - assert tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT + assert ( + tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT + ) def test_invalid_config_has_no_fingerprint(self) -> None: assert tuning_fingerprint({"tier_boundaries": "not-a-mapping"}) is None + def test_a_release_changing_a_shipped_default_is_not_an_operator_edit(self, monkeypatch) -> None: + """An omitted setting follows the shipped default and stays off the quota when that default moves; + only what the operator wrote is fingerprinted, so an explicit value equal to the old default still counts.""" + import litellm.router_strategy.complexity_router.config as config_module + + untouched = _router("a", {}) + tiers_only = _router("b", {"tiers": _TIERS}) + pinned = _router("c", {"dimension_weights": dict(config_module.DEFAULT_DIMENSION_WEIGHTS)}) + baselines = snapshot_tuning_baselines([untouched, tiers_only, pinned]) + assert tuning_fingerprint(pinned["litellm_params"]["complexity_router_config"]) != DEFAULT_TUNING_FINGERPRINT + + monkeypatch.setattr( + config_module, + "DEFAULT_DIMENSION_WEIGHTS", + {**config_module.DEFAULT_DIMENSION_WEIGHTS, "codePresence": 0.99}, + ) + monkeypatch.setattr( + config_module, + "DEFAULT_TIER_BOUNDARIES", + {**config_module.DEFAULT_TIER_BOUNDARIES, "simple_medium": 0.42}, + ) + + assert tuning_fingerprint({}) == DEFAULT_TUNING_FINGERPRINT + assert mutable_tuned_identities([untouched, tiers_only, pinned], baselines) == frozenset() + assert mutable_tuned_identities([untouched], snapshot_tuning_baselines([])) == frozenset() + class TestRouterIdentity: def test_db_rows_key_on_model_id_and_yaml_rows_on_name_and_tags(self) -> None: @@ -161,12 +192,18 @@ class TestQuota: new_c = _router("c", {"tiers": _TIERS}) assert tuning_quota_violation(candidate=edited_a, others=[legacy_b], baselines=baselines, limit=1) is None - assert tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) is None + assert ( + tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) + is None + ) assert tuning_quota_violation(candidate=legacy_a, others=[edited_b], baselines=baselines, limit=1) is None assert tuning_quota_violation(candidate=edited_b, others=[edited_a], baselines=baselines, limit=1) is not None assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=1) is not None assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=None) is None - assert tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) is None + assert ( + tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) + is None + ) def test_reverting_to_baseline_frees_the_quota(self) -> None: legacy_a = _router("a", {"tiers": _TIERS}) @@ -174,7 +211,10 @@ class TestQuota: baselines = snapshot_tuning_baselines([legacy_a, legacy_b]) edited_b = _router("b", {"tiers": _TIERS}) assert tuning_quota_violation(candidate=edited_b, others=[legacy_a], baselines=baselines, limit=1) is None - assert tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) is None + assert ( + tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) + is None + ) def test_violation_message_names_the_limit_and_remedy(self) -> None: message = tuning_limit_violation(held=2, limit=1) diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index f181370455d..ccd6766b13a 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -389,14 +389,24 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "max", ) - @pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra"]) - def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model): - """Microsoft Foundry serves the same model but its API accepts reasoning_effort none - (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" + @pytest.mark.parametrize( + "model,custom_llm_provider", + [ + ("azure/gpt-6-astra", "azure"), + ("azure/us/gpt-6-astra", "azure"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ], + ) + def test_an_azure_hosted_deployment_advertises_none_but_not_max( + self, local_model_cost_map, model, custom_llm_provider + ): + """Microsoft hosts the same model with a different level set than OpenAI does. Verified live + on both Azure routes: none returns 200 with zero reasoning tokens and unlocks temperature, + which OpenAI's API rejects, while max returns 400 unsupported_value naming none through + xhigh as the levels it does take.""" from litellm.utils import _get_model_info_helper - model_info = dict(_get_model_info_helper(model=model, custom_llm_provider="azure")) + model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( "none", @@ -404,5 +414,4 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "medium", "high", "xhigh", - "max", ) diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 0489f4ff017..b2fd2e6dcc0 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -10,8 +10,8 @@ from __future__ import annotations import pytest import litellm -from litellm.rust_bridge import chat_completions as bridge from litellm.rust_bridge import configuration +from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -67,10 +67,11 @@ def _hide_native_bridge(monkeypatch): @pytest.fixture(autouse=True) -def reset_bridge(): +def reset_bridge(monkeypatch): """Every test starts with no injected callables, and leaves none behind.""" bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "1") yield bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) configuration.reset_rust_configuration() @@ -112,7 +113,7 @@ def _accepts(**overrides) -> bool: "messages": MESSAGES, "optional_params": {"max_tokens": 16}, "custom_llm_provider": "anthropic", - "litellm_params": {"rust": True}, + "litellm_params": {}, "stream": None, } kwargs.update(overrides) @@ -126,23 +127,16 @@ class TestGate: bridge.set_rust_chat_completions(decline=gate) assert _accepts(litellm_params={}) is False assert _accepts(litellm_params=None) is False - assert _accepts(litellm_params={"rust": False}) is False assert gate.calls == [], "the gate must not be consulted before opt-in" def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) assert _accepts() is True assert gate.calls[0]["model"] == "claude-sonnet-4-5" assert gate.calls[0]["custom_llm_provider"] == "anthropic" - def test_explicit_false_overrides_process_enable(self): - bridge.set_rust_chat_completions(decline=_RecordingDecline()) - configuration.rust(True) - - assert _accepts(litellm_params={"rust": False}) is False - def test_process_enable_applies_without_request_override(self): bridge.set_rust_chat_completions(decline=_RecordingDecline()) configuration.rust(True) @@ -155,7 +149,7 @@ class TestGate: assert _accepts(litellm_params={}) is True def test_declines_streaming_and_providers_off_the_path(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) assert _accepts(stream=True) is False @@ -170,10 +164,10 @@ class TestGate: handed `optional_params` only, so accepting here would send the request to Anthropic with the abuse-detection attribution silently missing. """ - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) - assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}) is False + assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of" # Bedrock's Converse transform reads no `user_id`, and an Anthropic request @@ -182,13 +176,13 @@ class TestGate: _accepts( custom_llm_provider="bedrock", model="bedrock/us-east-1/anthropic.claude-v2", - litellm_params={"rust": True, "metadata": {"user_id": "u-123"}}, + litellm_params={"metadata": {"user_id": "u-123"}}, ) is True ) - assert _accepts(litellm_params={"rust": True, "metadata": {"trace_id": "t-1"}}) is True - assert _accepts(litellm_params={"rust": True, "metadata": {"user_id": None}}) is True - assert _accepts(litellm_params={"rust": True, "metadata": None}) is True + assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True + assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True + assert _accepts(litellm_params={"metadata": None}) is True def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch): """`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the @@ -196,7 +190,7 @@ class TestGate: evicting a caller-supplied one. The core can do neither, so an operator who armed `bedrock_request_metadata_fields` keeps the Python path. """ - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") gate = _RecordingDecline() bridge.set_rust_chat_completions(decline=gate) bedrock = { @@ -213,17 +207,17 @@ class TestGate: assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone" def test_declines_when_the_core_declines(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming")) assert _accepts() is False def test_declines_when_the_bridge_is_unavailable(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") _hide_native_bridge(monkeypatch) assert _accepts() is False def test_declines_when_the_gate_itself_raises(self, monkeypatch): - monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.setenv("LITELLM_RUST", "1") def exploding(**_kwargs): raise RuntimeError("boom") diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py index 15f69f95335..aff9d5acac1 100644 --- a/tests/test_litellm/rust_bridge/test_configuration.py +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -10,7 +10,6 @@ from typing import Final import pytest from litellm.rust_bridge import configuration -from litellm.rust_bridge import ocr as rust_ocr @pytest.fixture(autouse=True) @@ -19,42 +18,31 @@ def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest ) -> Generator[None]: configuration.reset_rust_configuration() monkeypatch.delenv("LITELLM_RUST", raising=False) - monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False) - rust_ocr.set_rust_ocr(ocr=None, aocr=None) yield configuration.reset_rust_configuration() - rust_ocr.set_rust_ocr(ocr=None, aocr=None) @pytest.mark.parametrize( - ("request_override", "process", "environment", "legacy_environment", "release_default", "expected"), + ("process", "environment", "release_default", "expected"), ( - (False, True, True, True, True, False), - (True, False, False, False, False, True), - (None, False, True, True, True, False), - (None, True, False, False, False, True), - (None, None, False, True, True, False), - (None, None, True, False, False, True), - (None, None, None, False, True, False), - (None, None, None, True, False, True), - (None, None, None, None, False, False), - (None, None, None, None, True, True), + (False, True, True, False), + (True, False, False, True), + (None, False, True, False), + (None, True, False, True), + (None, None, False, False), + (None, None, True, True), ), ) def test_resolution_precedence( - request_override: bool | None, process: bool | None, environment: bool | None, - legacy_environment: bool | None, release_default: bool, expected: bool, ) -> None: assert ( configuration.resolve_rust_enabled( - request_override=request_override, process_override=process, environment_override=environment, - legacy_environment_override=legacy_environment, release_default=release_default, ) is expected @@ -71,7 +59,6 @@ def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) configuration.rust(True) assert configuration.rust_enabled() is True - assert configuration.rust_enabled(request_override=False) is False def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None: @@ -83,18 +70,8 @@ def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPat @pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: monkeypatch.setenv("LITELLM_RUST", value) - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") assert configuration.rust_enabled() is False - assert configuration.rust_ocr_enabled() is False - - -@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) -def test_invalid_legacy_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: - monkeypatch.setenv("LITELLM_USE_RUST_OCR", value) - - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_enabled() is False def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None: @@ -104,40 +81,20 @@ def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytes assert executor.submit(configuration.rust_enabled).result() is True configuration.rust(False) assert executor.submit(configuration.rust_enabled).result() is False - assert executor.submit(configuration.rust_ocr_enabled).result() is False configuration.reset_rust_configuration() assert executor.submit(configuration.rust_enabled).result() is True - assert executor.submit(configuration.rust_ocr_enabled).result() is True def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_RUST", "sometimes") - assert configuration.rust_enabled(request_override=False) is False configuration.rust(True) assert configuration.rust_enabled() is True -def test_legacy_ocr_environment_is_deprecated_and_global(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_enabled() is True - with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): - assert configuration.rust_ocr_enabled() is True - - -def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LITELLM_RUST", "0") - monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - - assert configuration.rust_enabled() is False - - -@pytest.mark.parametrize("environment_name", ("LITELLM_RUST", "LITELLM_USE_RUST_OCR")) @pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) -def test_environment_controls_startup(environment_name: str, value: str, expected: str) -> None: - environment: Final = {**os.environ, environment_name: value} +def test_environment_controls_startup(value: str, expected: str) -> None: + environment: Final = {**os.environ, "LITELLM_RUST": value} result: Final = subprocess.run( ( sys.executable, diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py new file mode 100644 index 00000000000..e1c66c8d9ea --- /dev/null +++ b/tests/test_litellm/skills/test_skills_main.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import litellm.skills.main as skills_main +from litellm.types.utils import LlmProviders + + +def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( + monkeypatch, +) -> None: + """The REST /v1/skills form endpoint passes description/instructions as top-level + kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch + branch of create_skill() dropped both, so every LiteLLM-hosted skill was created + with description=None and instructions=None regardless of what the caller sent.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Converts files from one language into another" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( + "Take an uploaded document and produce it in the target language" + ) + + +def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: + """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Warehouse SQL Analyst", + extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Runs SQL against the inventory database" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" + + +def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) + + assert handler.create_skill_handler.call_args.kwargs["description"] is None + assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index bbeb6c38f78..112464bda22 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -44,7 +44,7 @@ class AsyncBridge: def test_enabled_sync_bridge_receives_audio() -> None: bridge = SyncBridge() - rust_bridge.configure_rust_transcription(True, transcription=bridge) + rust_bridge.configure_rust_transcription(transcription=bridge) result = rust_bridge.transcription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, @@ -61,7 +61,7 @@ def test_enabled_sync_bridge_receives_audio() -> None: @pytest.mark.asyncio async def test_enabled_async_bridge() -> None: - rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge()) + rust_bridge.configure_rust_transcription(atranscription=AsyncBridge()) result = await rust_bridge.atranscription( model="mistral.voxtral-mini-3b-2507", audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 1945e5ffac5..f8fa2231597 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4526,6 +4526,18 @@ def test_explicit_pricing_precedes_private_provider_response_model( assert selected == expected +def test_cost_per_token_mistral_voxtral_tts_bills_per_input_character(_local_model_cost_map): + prompt_usd, completion_usd = cost_per_token( + model="voxtral-mini-tts-2603", + custom_llm_provider="mistral", + call_type="speech", + prompt_characters=1000, + ) + + assert prompt_usd == pytest.approx(1000 * 1.6e-05) + assert completion_usd == 0.0 + + def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map): """gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens.""" from litellm.cost_calculator import batch_cost_calculator @@ -4538,3 +4550,98 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m assert prompt_cost == pytest.approx(1000 * 5e-6) assert completion_cost == pytest.approx(500 * 2.5e-5) + + +def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once( + _local_model_cost_map: None, +) -> None: + """Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once.""" + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 260, + "input_tokens": 237, + "output_tokens": 23, + "input_token_details": { + "text_tokens": 43, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + "cached_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, + }, + "output_token_details": {"text_tokens": 23, "audio_tokens": 0, "reasoning_tokens": 18}, + } + }, + }, + ] + combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + total_cost = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=combined_usage_object, + custom_llm_provider="azure", + litellm_model_name="azure/gpt-realtime-2.1-mini", + ) + + info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure") + expected = ( + 43 * info["input_cost_per_token"] + + 194 * info["input_cost_per_image_token"] + + 23 * info["output_cost_per_token"] + ) + assert total_cost == pytest.approx(expected) + assert total_cost == pytest.approx(0.0002362) + + +def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> None: + """The combined usage that lands in spend logs keeps reasoning out of text_tokens for every turn.""" + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}}, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 307, + "input_tokens": 237, + "output_tokens": 70, + "input_token_details": { + "text_tokens": 43, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + }, + "output_token_details": {"text_tokens": 70, "audio_tokens": 0, "reasoning_tokens": 52}, + } + }, + }, + { + "type": "response.done", + "response": { + "usage": { + "total_tokens": 363, + "input_tokens": 300, + "output_tokens": 63, + "input_token_details": { + "text_tokens": 106, + "audio_tokens": 0, + "image_tokens": 194, + "cached_tokens": 0, + }, + "output_token_details": {"text_tokens": 63, "audio_tokens": 0, "reasoning_tokens": 43}, + } + }, + }, + ] + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + + assert combined.completion_tokens == 133 + assert combined.completion_tokens_details is not None + assert combined.completion_tokens_details.reasoning_tokens == 95 + assert combined.completion_tokens_details.text_tokens == 38 + assert combined.completion_tokens_details.audio_tokens == 0 diff --git a/tests/test_litellm/test_cost_map_guard.py b/tests/test_litellm/test_cost_map_guard.py new file mode 100644 index 00000000000..1b4330ed62c --- /dev/null +++ b/tests/test_litellm/test_cost_map_guard.py @@ -0,0 +1,195 @@ +import importlib.util +import json +import subprocess +import sys +from pathlib import Path +from types import ModuleType +from typing import Final + +import pytest + +ROOT: Final = Path(__file__).resolve().parents[2] +CI_CD: Final = ROOT / "ci_cd" + + +def _load(name: str) -> ModuleType: + spec = importlib.util.spec_from_file_location(name, CI_CD / f"{name}.py") + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +schema_module: Final = _load("generate_model_prices_schema") +guard: Final = _load("cost_map_guard") + +MAP_FILES: Final = (guard.COST_MAP_PATH,) +BOT_REF: Final = "litellm_cost_map_sync_2026-09-04T12-00Z" + + +def _entry(price: float = 1e-06, **extra: object) -> dict[str, object]: + return { + "input_cost_per_token": price, + "output_cost_per_token": price * 2, + "litellm_provider": "openrouter", + "mode": "chat", + "max_tokens": 4096, + **extra, + } + + +BASE_MAP: Final = { + "sample_spec": {"input_cost_per_token": "USD per prompt token"}, + "fallback_generalizations": {"rules": [{"name": "r", "pattern": "^x"}]}, + "openrouter/a": _entry(supports_vision=True), + "openrouter/b": _entry(2e-06), +} + + +def _serialize(cost_map: dict[str, object]) -> str: + return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" + + +def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object: + text = _serialize(cost_map) + rendered = schema_module.render(schema_module.build_schema(cost_map)) + return guard.Snapshot( + cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema + ) + + +BASE: Final = _snapshot(BASE_MAP) + + +def _failures(head: object, changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True) -> tuple[str, ...]: + return guard.guard_failures(BASE, head, changed_files, bot) + + +def test_in_sync_files_pass_for_humans_and_bots() -> None: + assert _failures(BASE, bot=False) == () + assert _failures(BASE, bot=True) == () + + +def test_bot_may_add_and_reprice_models() -> None: + head = _snapshot({**BASE_MAP, "openrouter/a": _entry(9e-06, supports_vision=True), "openrouter/c": _entry()}) + assert _failures(head) == () + + +def test_broken_json_is_reported() -> None: + head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}") + assert _failures(head, bot=False) == ( + f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", + ) + + +def test_non_object_root_is_reported() -> None: + head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}") + assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",) + + +def test_backup_drift_is_reported() -> None: + head = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) + assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)] + + +def test_schema_out_of_sync_is_reported() -> None: + head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema) + assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)] + + +def test_schema_validation_errors_are_reported() -> None: + head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)}) + prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c." + assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [ + "input_cost_per_token", + "output_cost_per_token", + ] + + +def test_unclassified_entry_key_is_reported() -> None: + text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)}) + head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema) + (failure,) = _failures(head, bot=False) + assert "Unclassified keys" in failure and "weird_thing" in failure + + +def test_bot_may_only_touch_the_cost_map_files() -> None: + changed = (*guard.GUARDED_PATHS, "litellm/utils.py", ".github/workflows/cost-map-guard.yml") + assert _failures(BASE, changed_files=changed, bot=False) == () + assert _failures(BASE, changed_files=changed) == ( + "bot PRs may only change the cost map files, not litellm/utils.py", + "bot PRs may only change the cost map files, not .github/workflows/cost-map-guard.yml", + ) + + +def test_bot_may_not_remove_models() -> None: + head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}) + assert _failures(head, bot=False) == () + assert _failures(head) == ("bot PRs may not remove models: openrouter/b",) + + +def test_bot_may_not_remove_fields() -> None: + head = _snapshot({**BASE_MAP, "openrouter/a": _entry()}) + assert _failures(head, bot=False) == () + assert _failures(head) == ("bot PRs may not remove fields: openrouter/a.supports_vision",) + + +def test_bot_may_not_change_special_root_keys() -> None: + head = _snapshot({**BASE_MAP, "fallback_generalizations": {"rules": []}}) + assert _failures(head, bot=False) == () + assert _failures(head) == ("bot PRs may not change fallback_generalizations",) + + +def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: + text = _serialize(cost_map) + (repo / guard.COST_MAP_PATH).write_text(text) + (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) + (repo / guard.BACKUP_PATH).write_text(text) + (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map))) + subprocess.run(("git", "add", "-A"), cwd=repo, check=True) + subprocess.run( + ("git", "-c", "user.name=t", "-c", "user.email=t@example.com", "commit", "-q", "-m", message), + cwd=repo, + check=True, + ) + return subprocess.run( + ("git", "rev-parse", "HEAD"), cwd=repo, check=True, capture_output=True, text=True + ).stdout.strip() + + +def _run_guard(repo: Path, base: str, head: str, head_ref: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + (sys.executable, str(CI_CD / "cost_map_guard.py"), "--base", base, "--head", head, "--head-ref", head_ref), + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + + +@pytest.mark.parametrize( + ("head_ref", "expected_code", "expected_line"), + [ + (BOT_REF, 1, "- bot PRs may not remove models: openrouter/b"), + ("litellm_fix_pricing", 0, "cost map guard passed (human PR, file checks only)"), + ], +) +def test_main_reads_both_revisions_from_git( + tmp_path: Path, head_ref: str, expected_code: int, expected_line: str +) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base = _commit(tmp_path, BASE_MAP, "base") + head = _commit(tmp_path, {key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}, "head") + result = _run_guard(tmp_path, base, head, head_ref) + assert result.returncode == expected_code, result.stdout + result.stderr + assert expected_line in result.stdout.splitlines() + + +def test_main_rejects_a_bot_pr_that_edits_code(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base = _commit(tmp_path, BASE_MAP, "base") + (tmp_path / "litellm" / "utils.py").write_text("print('hi')\n") + head = _commit(tmp_path, {**BASE_MAP, "openrouter/c": _entry()}, "head") + assert _run_guard(tmp_path, base, head, BOT_REF).returncode == 1 + assert _run_guard(tmp_path, base, head, "litellm_fix_pricing").returncode == 0 diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 715ca8672b2..038df3656fe 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3351,6 +3351,52 @@ def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeyp assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) +def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-fake-mp3-bytes" + mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + response_format="wav", + speed=2, + instructions="sound cheerful", + ) + + assert mock_route.called + request_body: Final = json.loads(mock_route.calls.last.request.content) + assert request_body == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" + assert response.content == audio_bytes + + +def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-gateway-bytes" + gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + api_base="https://mistral.gateway.internal", + ) + + assert gateway_route.called + assert response.content == audio_bytes + + FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 6114d1d8aba..c2c22c25998 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -8,6 +8,9 @@ from pathlib import Path import jsonschema import pytest +from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name +from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + REPO_ROOT = Path(__file__).parents[2] GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py" PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -173,3 +176,44 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict): "sync the tier keys so service-tier requests against pinned snapshots are not " "billed at standard rates:\n" + "\n".join(drifted) ) + + +OPENAI_REASONING_FAMILY_MARKERS = ("codex", "deep-research", "chat-latest") + + +def is_openai_o_series(name: str) -> bool: + return len(name) > 1 and name[0] == "o" and name[1].isdigit() + + +def is_openai_reasoning_family(name: str) -> bool: + base = name.split("/")[-1].removeprefix("ft:") + if "search-api" in base: + return False + return ( + is_openai_o_series(base) + or is_gpt_reasoning_series_name(base) + or any(marker in base for marker in OPENAI_REASONING_FAMILY_MARKERS) + ) + + +def test_openai_reasoning_family_entries_carry_supports_reasoning(prices: dict): + unflagged = [ + name + for name, entry in prices.items() + if isinstance(entry, dict) + and entry.get("litellm_provider") == "openai" + and is_openai_reasoning_family(name) + and entry.get("supports_reasoning") is not True + ] + assert unflagged == [], ( + "OpenAI o-series, gpt-5+, codex, deep-research, and chat-latest models are reasoning " + "models, and the Responses API drops the `reasoning` param for any mapped OpenAI model " + "whose entry lacks supports_reasoning; flag these entries:\n" + "\n".join(unflagged) + ) + + +def test_chat_latest_declares_the_one_effort_openai_accepts(prices: dict): + """OpenAI rejects every reasoning.effort on chat-latest except medium, and a reasoning entry + with no declared levels resolves to None, which lets /model_group/info and the dashboard offer + levels the upstream will 400 on.""" + assert resolve_supported_reasoning_efforts(prices["chat-latest"], deployment_is_mapped=True) == ("medium",) diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index aa2260e89ea..b84cb8aa657 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -11,6 +11,12 @@ import pytest ROOT = Path(__file__).resolve().parents[2] SCRIPT = ROOT / "scripts" / "pre_commit_lint.sh" +WHOLE_TREE_RUFF = "run --no-sync ruff check --config ruff-tests.toml tests" +TEST_TREE_RAN = "ran: test-tree lint (ruff-tests.toml + test-quality budget)" +TEST_TREE_SKIPPED = ( + "skipped: test-tree lint (ruff-tests.toml + test-quality budget) " + "(no tests/ Python files or test-tree lint inputs in scope)" +) BARRIER_HELPER = """barrier_sync() { touch "$STUB_BARRIER_DIR/$1.started" @@ -30,6 +36,7 @@ BARRIER_HELPER = """barrier_sync() { MAKE_STUB = """#!/bin/sh . "$STUB_BIN/barrier.sh" +[ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/make.args" case "$*" in lint) [ "${STUB_FAIL:-}" = "make-lint" ] && exit 1 @@ -40,6 +47,9 @@ case "$*" in sleep 60 fi ;; + lint-test-quality) + [ "${STUB_FAIL:-}" = "test-quality" ] && exit 1 + ;; esac exit 0 """ @@ -69,6 +79,10 @@ case "$*" in *orjson*) [ -n "${STUB_BARRIER_DIR:-}" ] && barrier_sync genapi "python dashboard" ;; + "run --no-sync ruff check --config ruff-tests.toml"*) + [ -n "${STUB_ARGS_DIR:-}" ] && echo "$*" >> "$STUB_ARGS_DIR/ruff_tests.args" + [ "${STUB_FAIL:-}" = "tests-ruff" ] && exit 1 + ;; esac exit 0 """ @@ -153,6 +167,13 @@ def _set_base_ref(repo: Path) -> None: ) +def _stage_file(repo: Path, relative: str, body: str) -> None: + path = repo / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body) + subprocess.run(["git", "add", relative], cwd=repo, check=True) + + def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") @@ -405,6 +426,7 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No assert "ran: dashboard lint (prettier + eslint + lint budgets)" in proc.stdout assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout + assert TEST_TREE_SKIPPED in proc.stdout assert "check: PASS" in proc.stdout assert "check: FAIL" not in proc.stdout @@ -412,20 +434,146 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") - tests_dir = repo / "tests" / "test_litellm" - tests_dir.mkdir(parents=True) - (tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n") - subprocess.run(["git", "add", "tests"], cwd=repo, check=True) + _stage_file(repo, "scripts/tool.py", "def main() -> None: ...\n") proc = _run(repo, bin_dir, {}) assert proc.returncode == 0, proc.stdout + proc.stderr assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout - assert "tests/test_litellm/test_x.py" in proc.stdout + assert "scripts/tool.py" in proc.stdout assert "a no-op, not a lint verdict" in proc.stdout assert "check: PASS" in proc.stdout assert "linting Python" not in proc.stdout log = (repo / ".git" / "pre_commit_lint.log").read_text() assert "check: summary" in log assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log + assert TEST_TREE_SKIPPED in log + + +def _recorded(args_dir: Path, name: str) -> list[str]: + path = args_dir / name + return path.read_text().splitlines() if path.exists() else [] + + +def test_tests_only_change_runs_the_whole_test_tree_ruff_and_the_quality_gate(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _stage_file(repo, "tests/fixtures/data.json", "{}\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + assert TEST_TREE_RAN in proc.stdout + assert "no gating lint check matches" not in proc.stdout + assert "linting Python" not in proc.stdout + assert "check: PASS" in proc.stdout + + +@pytest.mark.parametrize( + "changed", + [ + "ruff-tests.toml", + "test-quality-budget.json", + "scripts/check_test_quality.py", + "scripts/test_quality_gate.py", + "tests/e2e/test_x.py", + ], +) +def test_test_tree_lint_inputs_trigger_the_test_tree_checks(tmp_path: Path, changed: str) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, changed, "x = 1\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert "lint-test-quality" in _recorded(args_dir, "make.args") + assert TEST_TREE_RAN in proc.stdout + + +def test_nothing_staged_tests_only_working_tree_change_runs_the_test_tree_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + _set_base_ref(repo) + args_dir = tmp_path / "args" + args_dir.mkdir() + (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + + +def test_a_failing_test_tree_ruff_fails_the_run_and_still_runs_the_quality_gate(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "tests-ruff"}) + assert proc.returncode == 1 + assert "Test-tree ruff failed" in proc.stdout + proc.stderr + assert "check: FAIL" in proc.stdout + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + + +def test_a_failing_quality_gate_fails_a_tests_only_run(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_FAIL": "test-quality"}) + assert proc.returncode == 1 + assert "Test-quality budget failed" in proc.stdout + proc.stderr + assert "check: FAIL" in proc.stdout + + +def test_tests_changed_alongside_litellm_files_defer_to_make_lint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "litellm/foo.py", "x = 2\n") + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir), "STUB_FAIL": "test-quality"}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "linting Python" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [] + assert _recorded(args_dir, "make.args") == ["lint"] + assert TEST_TREE_RAN in proc.stdout + + +def test_deleted_test_file_still_runs_the_test_tree_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + subprocess.run(["git", "rm", "-q", "tests/test_a.py"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert _recorded(args_dir, "ruff_tests.args") == [WHOLE_TREE_RUFF] + assert _recorded(args_dir, "make.args") == ["lint-test-quality"] + assert TEST_TREE_RAN in proc.stdout + + +def test_partial_staging_warns_when_test_files_are_left_unstaged(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _stage_file(repo, "tests/test_a.py", "def test_a() -> None: ...\n") + _commit_all(repo, "base") + args_dir = tmp_path / "args" + args_dir.mkdir() + _stage_file(repo, "notes.md", "hi\n") + (repo / "tests" / "test_a.py").write_text("def test_a() -> None:\n assert True\n") + proc = _run(repo, bin_dir, {"STUB_ARGS_DIR": str(args_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "SKIPPED test-tree lint (ruff-tests.toml + test-quality budget)" in proc.stdout + assert "tests/test_a.py" in proc.stdout + assert _recorded(args_dir, "ruff_tests.args") == [] + assert _recorded(args_dir, "make.args") == [] def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7c044310e14..eed34c79a06 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -39,6 +39,7 @@ from litellm.router import ( _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, ) +from litellm.router_strategy import simple_shuffle from litellm.types.router import DeploymentTypedDict @@ -2386,6 +2387,580 @@ async def test_acompletion_streaming_iterator_preserves_hidden_params(): assert result._hidden_params.get("_response_ms") == 500.0 +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_preserves_response_headers(): + """LIT-6767: the returned wrapper must carry the provider's raw response headers. + + Proxy callbacks read ``_response_headers`` off the object the router hands + back. The wrapper used to be built without it, so every streaming chat + completion reported zero raw provider headers while the non-streaming path + reported the full set. + """ + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + async def _empty(): + return + yield # make it an async generator + + provider_headers = { + "x-request-id": "req-provider-123", + "x-ratelimit-remaining-requests": "42", + # a provider must never be able to spoof an internal header + "x-litellm-model-id": "spoofed", + } + source = CustomStreamWrapper( + completion_stream=_empty(), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers=provider_headers, + ) + + result = await router._acompletion_streaming_iterator( + model_response=source, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + assert result._response_headers == provider_headers + additional_headers = result._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "req-provider-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + # internal-header protection survives: the provider value is namespaced, never promoted + assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed" + assert "x-litellm-model-id" not in additional_headers + + +def test_completion_streaming_iterator_preserves_response_headers(): + """LIT-6767, sync counterpart of the async header-preservation test.""" + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"} + source = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers=provider_headers, + ) + + result = router._completion_streaming_iterator( + model_response=source, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + assert result._response_headers == provider_headers + assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync" + + +def test_adopt_fallback_response_headers_replaces_rather_than_merges(): + """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers. + + Values from the failed attempt must not survive, so the wrapper replaces both + ``_response_headers`` and ``_hidden_params`` instead of merging them. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + wrapper._hidden_params = { + "model_id": "failed-deployment", + "only_on_failed_attempt": "stale", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + fallback = MagicMock() + fallback._response_headers = {"x-request-id": "req-FALLBACK"} + fallback._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert "only_on_failed_attempt" not in wrapper._hidden_params + assert wrapper._hidden_params is not fallback._hidden_params + # the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built + # from it would still be stamped with the deployment that failed + assert wrapper._base_hidden_params["model_id"] == "fallback-deployment" + # the nested header dict is copied too, so a later mutation on the fallback + # response cannot reach headers the proxy has already published + assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"] + fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED" + assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} + + +def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): + """LIT-6767: adoption still returns the fallback's params once the wrapper is gone.""" + import weakref + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + fallback: Final = MagicMock() + fallback._response_headers = {"x-request-id": "req-FALLBACK"} + fallback._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + ) + live_ref: Final = weakref.ref(wrapper) + prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback) + assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"]) + assert wrapper.fallback_headers_adopted is True + assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} + + dead_ref: Final = weakref.ref(wrapper) + del wrapper + assert dead_ref() is None + assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared + + +def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace(): + """LIT-6767: a fallback that carries no raw provider headers publishes none. + + Keeping the failed attempt's raw headers would hand the client and the callbacks a + provider ``x-request-id`` for a request that deployment never served, which is the + leak this fix exists to close. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}} + + fallback = MagicMock() + fallback._response_headers = None + fallback._hidden_params = {"model_id": "fallback-deployment"} + + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers is None + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert wrapper.fallback_headers_adopted is True + + +def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none(): + """A fallback response carrying no hidden params keeps the identity headers. + + Publishing no ``x-litellm-*`` header at all for a request the fallback served is + worse than keeping what is there, so only the raw provider headers are dropped. + """ + from unittest.mock import MagicMock + + from litellm.router import FallbackAwareStreamWrapper, Router + + wrapper = FallbackAwareStreamWrapper( + completion_stream=iter([]), + model="gpt-4", + custom_llm_provider="openai", + logging_obj=MagicMock(), + _response_headers={"x-request-id": "req-FAILED"}, + ) + hidden_params_before = wrapper._hidden_params + + fallback = object() + wrapper.adopt_fallback_response_headers( + fallback, Router._prepare_fallback_hidden_params(fallback) + ) + + assert wrapper._response_headers is None + assert wrapper._hidden_params is hidden_params_before + assert wrapper.fallback_headers_adopted is True + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_adopts_fallback_response_headers(): + """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must + describe the deployment that served the stream, with no value left over + from the attempt that failed.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "api_base": "https://failed.example", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + "only_on_failed_attempt": "stale", + } + + def __aiter__(self): + return self + + async def __anext__(self): + raise failed_error + + class FallbackStream: + def __init__(self): + self._response_headers = {"x-request-id": "req-FALLBACK"} + self._hidden_params = { + "model_id": "fallback-deployment", + "api_base": "https://fallback.example", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._chunks) + except StopIteration: + raise StopAsyncIteration from None + + fallback_stream = FallbackStream() + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=fallback_stream, + ): + result = await router._acompletion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + # the failed attempt is what the wrapper is built from + assert result._response_headers == {"x-request-id": "req-FAILED"} + # the very first chunk the fallback produces must already be published under + # the fallback's identity: the proxy commits response headers once that chunk + # is buffered, so adopting any later is adopting too late + await result.__anext__() + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + async for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + assert result._hidden_params["api_base"] == "https://fallback.example" + assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} + # stale values are removed, not merged over + assert "only_on_failed_attempt" not in result._hidden_params + # and the wrapper holds its own copy, so later fallback mutations cannot leak in + assert result._hidden_params is not fallback_stream._hidden_params + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): + """LIT-6767: a fallback that itself fails over before its first chunk. + + The selected fallback still describes its own failed attempt at selection time, so + the wrapper has to re-read it once a chunk exists or it publishes a deployment that + produced no output. + """ + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error: Final = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + def __aiter__(self): + return self + + async def __anext__(self): + raise failed_error + + class NestedFallbackStream: + """A fallback that repoints itself at a third deployment as it yields.""" + + def __init__(self): + self._response_headers = {"x-request-id": "req-MIDDLE"} + self._hidden_params = { + "model_id": "middle-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + chunk = next(self._chunks) + except StopIteration: + raise StopAsyncIteration from None + self._response_headers = {"x-request-id": "req-SERVED"} + self._hidden_params = { + "model_id": "served-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + return chunk + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=NestedFallbackStream(), + ): + result = await router._acompletion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + first_chunk: Final = await result.__anext__() + # the proxy commits response headers once this chunk is buffered + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"} + # and the chunk itself carries the same deployment + assert first_chunk._hidden_params["model_id"] == "served-deployment" + async for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + + +def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): + """LIT-6767, sync counterpart of the nested-fallback adoption test.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error: Final = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + } + + def __iter__(self): + return self + + def __next__(self): + raise failed_error + + class NestedFallbackStream: + """A fallback that repoints itself at a third deployment as it yields.""" + + def __init__(self): + self._response_headers = {"x-request-id": "req-MIDDLE"} + self._hidden_params = { + "model_id": "middle-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, + } + self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) + + def __iter__(self): + return self + + def __next__(self): + chunk = next(self._chunks) + self._response_headers = {"x-request-id": "req-SERVED"} + self._hidden_params = { + "model_id": "served-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, + } + return chunk + + with patch.object(router, "function_with_fallbacks", return_value=NestedFallbackStream()): + result = router._completion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + first_chunk: Final = next(result) + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + assert first_chunk._hidden_params["model_id"] == "served-deployment" + for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-SERVED"} + assert result._hidden_params["model_id"] == "served-deployment" + + +def test_completion_streaming_iterator_adopts_fallback_response_headers(): + """LIT-6767, sync counterpart of the fallback-adoption test.""" + from unittest.mock import MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + failed_error = MidStreamFallbackError( + message="upstream died before the first chunk", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=True, + ) + + class FailedStream: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + self._response_headers = {"x-request-id": "req-FAILED"} + self._hidden_params = { + "model_id": "failed-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, + "only_on_failed_attempt": "stale", + } + + def __iter__(self): + return self + + def __next__(self): + raise failed_error + + class FallbackStream: + def __init__(self): + self._response_headers = {"x-request-id": "req-FALLBACK"} + self._hidden_params = { + "model_id": "fallback-deployment", + "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, + } + + def __iter__(self): + return iter([]) + + with patch.object(router, "function_with_fallbacks", return_value=FallbackStream()): + result = router._completion_streaming_iterator( + model_response=FailedStream(), + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + assert result._response_headers == {"x-request-id": "req-FAILED"} + for _ in result: + pass + + assert result._response_headers == {"x-request-id": "req-FALLBACK"} + assert result._hidden_params["model_id"] == "fallback-deployment" + assert "only_on_failed_attempt" not in result._hidden_params + + def test_completion_streaming_iterator_fallback_on_429(): """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback. @@ -12805,6 +13380,82 @@ class TestTierParamsTheTargetAccepts: assert accepted == {"reasoning_effort": "max"} +@pytest.mark.asyncio +async def test_router_aspeech_without_voice_dispatches_ref_audio_cloning(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603"}, + } + ] + ) + + response = await router.aspeech(model="voxtral-tts", input="clone me", ref_audio="ZmFrZQ==") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body == {"model": "voxtral-mini-tts-2603", "input": "clone me", "ref_audio": "ZmFrZQ=="} + assert response.content == audio_bytes + + +@pytest.mark.asyncio +async def test_router_aspeech_without_voice_keeps_deployment_default_voice(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603", "voice": "en_paul_neutral"}, + } + ] + ) + + await router.aspeech(model="voxtral-tts", input="use my default") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body["voice_id"] == "en_paul_neutral" + + +@pytest.mark.asyncio +async def test_router_aspeech_request_voice_overrides_deployment_default(respx_mock, monkeypatch): + import base64 + + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + audio_bytes = b"RIFFfake-wav-bytes" + respx_mock.post("https://api.mistral.ai/v1/audio/speech").respond( + json={"audio_data": base64.b64encode(audio_bytes).decode()} + ) + router = Router( + model_list=[ + { + "model_name": "voxtral-tts", + "litellm_params": {"model": "mistral/voxtral-mini-tts-2603", "voice": "en_paul_neutral"}, + } + ] + ) + + await router.aspeech(model="voxtral-tts", input="override me", voice="gb_oliver_neutral") + + request_body = json.loads(respx_mock.calls.last.request.content) + assert request_body["voice_id"] == "gb_oliver_neutral" + + class TestRequestReasoningEffortOverride: def test_drop_effort_from_nested_carrier_preserves_other_nested_values(self): params: dict[str, object] = {"output_config": {"effort": "high", "format": "json"}} @@ -13116,6 +13767,7 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo ({"DefaultRetries": 0}, 502, litellm.BadGatewayError, 1), ({"DefaultRetries": 0, "ServiceUnavailableErrorRetries": 1}, 503, litellm.ServiceUnavailableError, 2), ({"ServiceUnavailableErrorRetries": 0}, 502, litellm.BadGatewayError, 3), + ({"BadRequestErrorRetries": 2}, 400, litellm.BadRequestError, 3), ], ) async def test_router_retry_policy_controls_upstream_attempt_count( @@ -13152,6 +13804,323 @@ async def test_router_retry_policy_controls_upstream_attempt_count( assert upstream.call_count == expected_upstream_calls +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retry_policy,upstream_error", + [ + ( + {"BadRequestErrorRetries": 2}, + { + "message": "This model's maximum context length is 16385 tokens", + "type": "invalid_request_error", + "code": "context_length_exceeded", + }, + ), + ( + {"ContentPolicyViolationErrorRetries": 2}, + { + "message": "Your request was rejected as a result of our safety system", + "type": "invalid_request_error", + "code": "content_policy_violation", + }, + ), + ], +) +async def test_router_retry_policy_400_retries_on_sibling_deployment( + monkeypatch: pytest.MonkeyPatch, retry_policy, upstream_error +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://rejecting.local/v1", + "weight": 1, + }, + "model_info": {"id": "rejecting"}, + }, + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://accepting.local/v1", + "weight": 0, + }, + "model_info": {"id": "accepting"}, + }, + ], + num_retries=2, + retry_policy=retry_policy, + disable_cooldowns=True, + ) + + with respx.mock(assert_all_called=True) as respx_mock: + rejecting = respx_mock.post("https://rejecting.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": upstream_error}) + ) + accepting = respx_mock.post("https://accepting.local/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lit-7036", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + }, + ) + ) + response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert rejecting.call_count == 1 + assert accepting.call_count == 1 + assert response.choices[0].message.content == "hi back" + assert response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] == 1 + + +_UPSTREAM_400 = {"message": "upstream refused this request", "type": "invalid_request_error", "code": "bad_request"} + + +def _retry_skip_deployment(deployment_id, host, litellm_params=None, model_info=None): + return { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": f"https://{host}.local/v1", + **(litellm_params or {}), + }, + "model_info": {"id": deployment_id, **(model_info or {})}, + } + + +@pytest.mark.parametrize( + "status_code,failed_deployment_id,already_skipped,expected", + [ + (400, "rejecting", None, ("rejecting",)), + (403, "rejecting", None, ("rejecting",)), + (400, "second", ("first",), ("first", "second")), + (400, "first", ("first",), ("first",)), + (429, "rejecting", None, ()), + (503, "rejecting", None, ()), + (408, "rejecting", None, ()), + (400, None, None, ()), + (None, "rejecting", None, ()), + ("400", "rejecting", None, ()), + (400, "second", 7, ("second",)), + (400, "second", "first", ("second",)), + (400, "second", ["first"], ("second",)), + (400, "second", ("first", 7), ("first", "second")), + ], +) +def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_id, already_skipped, expected): + exception = Exception("upstream refused this request") + exception.status_code = status_code + exception.failed_deployment_id = failed_deployment_id + + assert litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped) == expected + + +@pytest.mark.parametrize( + "value,expected", + [ + (("first", "second"), ("first", "second")), + ((), ()), + (("first", 7, None, "second"), ("first", "second")), + (None, ()), + (7, ()), + ("first", ()), + (["first"], ()), + ({"first": True}, ()), + (object(), ()), + ], +) +def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected): + from litellm.router import _as_retry_skipped_deployment_ids + + assert _as_retry_skipped_deployment_ids(value) == expected + + +@pytest.mark.parametrize( + "deployment_ids,skipped,expected", + [ + (["rejecting", "sibling"], ("rejecting",), ["sibling"]), + (["rejecting"], ("rejecting",), ["rejecting"]), + (["rejecting", "sibling"], ("rejecting", "sibling"), ["rejecting", "sibling"]), + (["rejecting", "sibling"], (), ["rejecting", "sibling"]), + (["rejecting", "sibling"], None, ["rejecting", "sibling"]), + (["rejecting", "sibling"], ("absent",), ["rejecting", "sibling"]), + (["rejecting", "sibling"], 7, ["rejecting", "sibling"]), + (["rejecting", "sibling"], "rejecting", ["rejecting", "sibling"]), + (["rejecting", "sibling"], ["rejecting"], ["rejecting", "sibling"]), + (["rejecting", "sibling"], {"rejecting": True}, ["rejecting", "sibling"]), + (["rejecting", "sibling"], ("rejecting", 7), ["sibling"]), + ], +) +@pytest.mark.asyncio +async def test_router_healthy_deployments_keep_the_last_candidate_a_retry_skipped(deployment_ids, skipped, expected): + router = litellm.Router( + model_list=[_retry_skip_deployment(deployment_id, deployment_id) for deployment_id in deployment_ids], + disable_cooldowns=True, + ) + request_kwargs = {"_retry_skipped_deployment_ids": skipped} + + healthy_deployments = await router.async_get_healthy_deployments(model="gpt-5.6", request_kwargs=request_kwargs) + + assert sorted(deployment["model_info"]["id"] for deployment in healthy_deployments) == sorted(expected) + assert "_retry_skipped_deployment_ids" not in request_kwargs + + +@pytest.mark.parametrize("client_supplied", [7, "rejecting", ["rejecting"], {"rejecting": True}, object()]) +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_when_a_client_forges_the_skip_list( + monkeypatch: pytest.MonkeyPatch, client_supplied +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[_retry_skip_deployment("rejecting", "rejecting"), _retry_skip_deployment("sibling", "sibling")], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + respx_mock.post("https://rejecting.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + respx_mock.post("https://sibling.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + _retry_skipped_deployment_ids=client_supplied, + ) + + assert "upstream refused this request" in str(raised.value) + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_on_order_fallback_hop(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + _retry_skip_deployment("order1", "order1", litellm_params={"order": 1}), + _retry_skip_deployment("order2", "order2", litellm_params={"order": 2}), + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + order1 = respx_mock.post("https://order1.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + order2 = respx_mock.post("https://order2.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert "upstream refused this request" in str(raised.value) + assert "No deployments available" not in str(raised.value) + assert order1.call_count >= 1 + assert order2.call_count >= 1 + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_keeps_upstream_error_when_tags_narrow_the_group( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + _retry_skip_deployment( + "tagged", "tagged", litellm_params={"tags": ["free"]}, model_info={"enable_tag_filtering": True} + ), + _retry_skip_deployment("untagged", "untagged", model_info={"enable_tag_filtering": True}), + ], + num_retries=2, + retry_policy={"BadRequestErrorRetries": 2}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + tagged = respx_mock.post("https://tagged.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + untagged = respx_mock.post("https://untagged.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["free"]}, + ) + + assert "upstream refused this request" in str(raised.value) + assert "No deployments available" not in str(raised.value) + assert tagged.call_count == 3 + assert untagged.call_count == 0 + + +@pytest.mark.asyncio +async def test_router_retry_policy_400_never_returns_to_a_deployment_that_already_refused( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(simple_shuffle.random, "choice", lambda deployments: deployments[0]) + router = litellm.Router( + model_list=[ + _retry_skip_deployment("first-refuser", "first-refuser", litellm_params={"weight": 1}), + _retry_skip_deployment("second-refuser", "second-refuser", litellm_params={"weight": 0}), + _retry_skip_deployment("accepting", "accepting", litellm_params={"weight": 0}), + ], + num_retries=3, + retry_policy={"BadRequestErrorRetries": 3}, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + first = respx_mock.post("https://first-refuser.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + second = respx_mock.post("https://second-refuser.local/v1/chat/completions").mock( + return_value=httpx.Response(400, json={"error": _UPSTREAM_400}) + ) + accepting = respx_mock.post("https://accepting.local/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lit-7036", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hi back"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + }, + ) + ) + response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + + assert first.call_count == 1 + assert second.call_count == 1 + assert accepting.call_count == 1 + assert response.choices[0].message.content == "hi back" + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/test_litellm/test_router_order_fallback.py index fde870e5abe..93895bbde08 100644 --- a/tests/test_litellm/test_router_order_fallback.py +++ b/tests/test_litellm/test_router_order_fallback.py @@ -18,10 +18,10 @@ from litellm import Router from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.types.router import RouterRateLimitError -from litellm.utils import _get_deployment_order, _get_order_filtered_deployments +from litellm.utils import _get_deployment_order, get_order_filtered_deployments # --------------------------------------------------------------------------- -# Unit tests for _get_order_filtered_deployments +# Unit tests for get_order_filtered_deployments # --------------------------------------------------------------------------- @@ -42,7 +42,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), self._make_deployment(1, "c"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 assert all(d["model_info"]["id"] in ("a", "c") for d in result) @@ -52,7 +52,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), self._make_deployment(3, "c"), ] - result = _get_order_filtered_deployments(deps, target_order=2) + result = get_order_filtered_deployments(deps, target_order=2) assert len(result) == 1 assert result[0]["model_info"]["id"] == "b" @@ -61,7 +61,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(1, "a"), self._make_deployment(2, "b"), ] - result = _get_order_filtered_deployments(deps, target_order=99) + result = get_order_filtered_deployments(deps, target_order=99) assert result == [] def test_target_order_no_match_does_not_reselect_lower_order(self): @@ -70,7 +70,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(2, "b"), ] remaining_after_pre_call = [deps[0]] - result = _get_order_filtered_deployments(remaining_after_pre_call, target_order=2) + result = get_order_filtered_deployments(remaining_after_pre_call, target_order=2) assert result == [] def test_no_order_set_returns_all(self): @@ -78,11 +78,11 @@ class TestGetOrderFilteredDeployments: self._make_deployment(None, "a"), self._make_deployment(None, "b"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 def test_empty_list(self): - result = _get_order_filtered_deployments([]) + result = get_order_filtered_deployments([]) assert result == [] def test_single_order_returns_all_with_that_order(self): @@ -90,7 +90,7 @@ class TestGetOrderFilteredDeployments: self._make_deployment(1, "a"), self._make_deployment(1, "b"), ] - result = _get_order_filtered_deployments(deps) + result = get_order_filtered_deployments(deps) assert len(result) == 2 diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 162312a8c67..9f05654f23f 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -15,11 +15,11 @@ import pytest import litellm from litellm import Router -from litellm.utils import _get_excluded_filtered_deployments +from litellm.utils import get_excluded_filtered_deployments # --------------------------------------------------------------------------- -# Unit tests for _get_excluded_filtered_deployments +# Unit tests for get_excluded_filtered_deployments # --------------------------------------------------------------------------- @@ -37,17 +37,17 @@ def _make_dep(dep_id: str, weight: Optional[int] = None) -> dict: class TestGetExcludedFilteredDeployments: def test_no_excluded_returns_all(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) assert len(result) == 2 def test_empty_excluded_returns_all(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) assert len(result) == 2 def test_drops_excluded(self): deps = [_make_dep("a"), _make_dep("b"), _make_dep("c")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) ids = sorted(d["model_info"]["id"] for d in result) assert ids == ["a", "c"] @@ -57,12 +57,12 @@ class TestGetExcludedFilteredDeployments: # error. Returning the original list here would re-include the # just-failed deployment and let weighted failover re-pick it. deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) assert result == [] def test_excluded_set_with_unknown_ids(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) assert len(result) == 2 def test_handles_missing_model_info(self): @@ -70,7 +70,7 @@ class TestGetExcludedFilteredDeployments: {"model_name": "x", "litellm_params": {"model": "gpt-4o"}}, # no model_info _make_dep("b"), ] - result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + result = get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) assert len(result) == 1 diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/test_litellm/test_test_quality_gate.py index 8cce6bc735a..6652211a828 100644 --- a/tests/test_litellm/test_test_quality_gate.py +++ b/tests/test_litellm/test_test_quality_gate.py @@ -1,16 +1,21 @@ """Tests for scripts/test_quality_gate.py. -The gate's whole value is that it blames a change only for what it adds, that a limit -can never rise, and that a limit cannot stay above a count the branch pushed below it. -All three live in pure functions, so they are tested directly: `evaluate` for the blame -rule, `ratcheted_budget` for the one-way ratchet, `unratcheted` for the ceiling a branch -left behind, and `parse_changed_lines` for the diff scan that turns a breach into -file:line. +The gate's whole value is that it blames a change only for what it adds and that a +limit can never rise. Both live in pure functions, so they are tested directly: +`evaluate` for the blame rule, `ratcheted_budget` for the one-way ratchet, and +`parse_changed_lines` for the diff scan that turns a breach into file:line. """ import importlib.util +import os +import signal +import subprocess import sys +import time +from collections.abc import Callable +from contextlib import suppress from pathlib import Path +from typing import NamedTuple _REPO_ROOT = Path(__file__).resolve().parents[2] _MODULE_PATH = _REPO_ROOT / "scripts" / "test_quality_gate.py" @@ -23,6 +28,16 @@ _spec.loader.exec_module(gate) _BUDGET = {"TQ001": {"limit": 10}, "TQ003": {"limit": 5}} +_SCAN_BASE = ( + "import importlib.util, pathlib, sys\n" + "spec = importlib.util.spec_from_file_location('test_quality_gate', sys.argv[1])\n" + "gate = importlib.util.module_from_spec(spec)\n" + "sys.modules[spec.name] = gate\n" + "spec.loader.exec_module(gate)\n" + "gate.base_counts('HEAD', repo_root=pathlib.Path(sys.argv[2]), checker=pathlib.Path(sys.argv[3]))\n" +) +_SCAN_BASE_WITH_SIGHUP_IGNORED = "import signal\nsignal.signal(signal.SIGHUP, signal.SIG_IGN)\n" + _SCAN_BASE + def test_a_rule_within_its_limit_is_not_a_breach(): assert gate.evaluate({"TQ001": 10}, {"TQ001": 10}, _BUDGET) == () @@ -73,29 +88,6 @@ def test_ratchet_lowers_a_rule_introduced_on_this_branch_like_any_other(): assert updated["TQ001"]["limit"] == 4 -def test_a_branch_that_cleared_violations_must_lower_the_ceiling(): - stale = gate.unratcheted({"TQ001": 6}, {"TQ001": 10}, _BUDGET) - assert [(b.rule, b.total, b.cap, b.added) for b in stale] == [("TQ001", 6, 10, -4)] - - -def test_headroom_already_in_the_base_is_not_blamed_on_this_branch(): - assert gate.unratcheted({"TQ001": 6}, {"TQ001": 6}, _BUDGET) == () - - -def test_a_branch_that_cleared_down_to_the_ceiling_exactly_is_clean(): - assert gate.unratcheted({"TQ001": 10}, {"TQ001": 12}, _BUDGET) == () - - -def test_a_branch_that_added_violations_is_not_a_ratchet_finding(): - assert gate.unratcheted({"TQ001": 14}, {"TQ001": 10}, _BUDGET) == () - - -def test_the_ratchet_finding_survives_the_update_that_answers_it(): - cleared = {"TQ001": 6} - updated = gate.ratcheted_budget(_BUDGET, cleared, {"TQ001": 10}) - assert gate.unratcheted(cleared, {"TQ001": 10}, updated) == () - - def test_parse_changed_lines_groups_hunks_under_their_own_file(): diff = ( "diff --git a/tests/a.py b/tests/a.py\n" @@ -146,3 +138,99 @@ def test_the_shipped_budget_covers_every_rule_the_checker_can_emit(): budget = json.loads((_REPO_ROOT / "test-quality-budget.json").read_text()) assert set(budget) == {"TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ008"} assert all(spec["limit"] >= 0 for spec in budget.values()) + + +def _git(cwd: Path, *args: str) -> str: + proc = subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + return proc.stdout.strip() + + +def _committed_repo(tmp_path: Path) -> Path: + repo = tmp_path / "repo" + (repo / "tests").mkdir(parents=True) + (repo / "tests" / "test_seed.py").write_text("def test_seed():\n assert True\n") + _git(repo, "init", "-q", "-b", "main") + _git(repo, "config", "user.email", "gate@example.com") + _git(repo, "config", "user.name", "gate") + _git(repo, "config", "commit.gpgsign", "false") + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-m", "seed") + return repo + + +def _wait_until(predicate: Callable[[], bool], timeout_seconds: float) -> bool: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return predicate() + + +def _reap(process: subprocess.Popen[bytes]) -> None: + with suppress(subprocess.TimeoutExpired): + process.wait(timeout=10) + if process.poll() is None: + process.kill() + process.wait(timeout=10) + + +def _registered_worktrees(repo: Path) -> int: + listing = _git(repo, "worktree", "list", "--porcelain") + return sum(line.startswith("worktree ") for line in listing.splitlines()) + + +class _StalledScan(NamedTuple): + process: subprocess.Popen[bytes] + repo: Path + release: Path + temp_dir: Path + + +def _base_scan_stalled_in_its_checker(tmp_path: Path, driver: str) -> _StalledScan: + repo = _committed_repo(tmp_path) + scanning = tmp_path / "scanning" + release = tmp_path / "release" + slow_checker = tmp_path / "slow_checker.py" + slow_checker.write_text( + "import pathlib, time\n" + f"pathlib.Path({str(scanning)!r}).touch()\n" + f"while not pathlib.Path({str(release)!r}).exists():\n" + " time.sleep(0.05)\n" + ) + temp_dir = tmp_path / "tmp" + temp_dir.mkdir() + scan = subprocess.Popen( + [sys.executable, "-c", driver, str(_MODULE_PATH), str(repo), str(slow_checker)], + env={**os.environ, "TMPDIR": str(temp_dir)}, + ) + if not _wait_until(scanning.exists, 30): + _reap(scan) + raise AssertionError("the base scan never reached the checker") + return _StalledScan(scan, repo, release, temp_dir) + + +def test_a_terminated_base_scan_still_removes_its_worktree(tmp_path: Path) -> None: + stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE) + try: + stalled.process.send_signal(signal.SIGTERM) + assert stalled.process.wait(timeout=30) == 128 + signal.SIGTERM + finally: + _reap(stalled.process) + assert _registered_worktrees(stalled.repo) == 1 + assert list(stalled.temp_dir.iterdir()) == [] + + +def test_a_base_scan_keeps_ignoring_the_hangup_its_parent_ignored(tmp_path: Path) -> None: + stalled = _base_scan_stalled_in_its_checker(tmp_path, _SCAN_BASE_WITH_SIGHUP_IGNORED) + try: + stalled.process.send_signal(signal.SIGHUP) + time.sleep(1) + assert stalled.process.poll() is None, "a hangup the parent ignored killed the scan" + stalled.release.touch() + assert stalled.process.wait(timeout=30) == 0 + finally: + _reap(stalled.process) + assert _registered_worktrees(stalled.repo) == 1 + assert list(stalled.temp_dir.iterdir()) == [] diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d686b032ee0..8a56a84ade7 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -198,6 +198,11 @@ def test_get_model_info_resolves_provider_prefixed_model_ids(local_model_cost_ma assert via_provider["mode"] == "responses" +def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local_model_cost_map): + info = litellm.get_model_info(model="ft:gpt-4o-2024-08-06:my-org::abc123", custom_llm_provider="openai") + assert info["key"] == "ft:gpt-4o-2024-08-06" + + def test_provider_prefixed_lookup_never_outranks_an_existing_row(local_model_cost_map): """The provider-prefixed candidate is tried last, after every candidate that already existed, so no model that resolves today can change answer. `perplexity/sonar` @@ -5239,52 +5244,6 @@ def test_client_side_timeout_marker_never_reaches_the_provider(): ) -def test_rust_flag_not_forwarded_as_provider_param(): - forwarded = get_non_default_completion_params({"rust": True, "temperature": 0.5}) - assert "rust" not in forwarded - - -def test_completion_does_not_leak_rust_flag_into_provider_request_body(): - mock_response = MagicMock() - mock_response.model_dump.return_value = { - "id": "chatcmpl-1", - "object": "chat.completion", - "created": 1234567890, - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "hi"}, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - } - - mock_raw_response = MagicMock() - mock_raw_response.headers = {} - mock_raw_response.parse.return_value = mock_response - - mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = mock_raw_response - - litellm.completion( - model="openai/gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - rust=True, - api_key="sk-test", - client=mock_client, - ) - - create_kwargs = mock_client.chat.completions.with_raw_response.create.call_args.kwargs - assert "rust" not in create_kwargs - assert "rust" not in (create_kwargs.get("extra_body") or {}) - - class _RecordingDeploymentFailureLogger(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/test_litellm/types/test_guardrails_case_normalization.py b/tests/test_litellm/types/test_guardrails_case_normalization.py index 3e7a573ea8e..26c1d395320 100644 --- a/tests/test_litellm/types/test_guardrails_case_normalization.py +++ b/tests/test_litellm/types/test_guardrails_case_normalization.py @@ -2,6 +2,8 @@ Test case normalization in LitellmParams for all guardrail types """ +from typing import Literal + import pytest from pydantic import ValidationError @@ -93,6 +95,34 @@ class TestLitellmParamsCaseNormalization: assert params.on_disallowed_action.islower() +class TestOnViolationAcceptedValues: + """on_violation is shared by /v1/realtime guardrails and the mcp_security guardrail""" + + @pytest.mark.parametrize("action", ["block", "alert"]) + def test_mcp_security_policy_template_on_violation_is_accepted(self, action: Literal["block", "alert"]): + params = LitellmParams( + guardrail="mcp_security", + mode="pre_call", + default_on=True, + on_violation=action, + ) + assert params.on_violation == action + + @pytest.mark.parametrize("action", ["warn", "end_session"]) + def test_realtime_on_violation_still_accepted(self, action: Literal["warn", "end_session"]): + params = LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action) + assert params.on_violation == action + + @pytest.mark.parametrize("action", ["block", "alert"]) + def test_mcp_only_on_violation_is_rejected_for_other_guardrails(self, action: Literal["block", "alert"]): + with pytest.raises(ValidationError, match="only supported by guardrail='mcp_security'"): + LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action) + + def test_unknown_on_violation_is_rejected(self): + with pytest.raises(ValidationError): + LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation="ignore") + + class TestSensitiveDataRoutingValidation: """on_sensitive_data='route' requires a target model to be set""" diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 554604ab200..5f44ba1773e 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -3,7 +3,7 @@ from typing import Final import pytest -from litellm.types.utils import HiddenParams, all_litellm_params +from litellm.types.utils import HiddenParams, all_litellm_params, text_tokens_without_nested_reasoning def test_rust_is_a_known_litellm_param(): @@ -768,3 +768,32 @@ def test_image_response_keeps_background(): response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png") assert response.background == "transparent" assert response.model_dump()["background"] == "transparent" + + +@pytest.mark.parametrize( + ("completion_tokens", "text_tokens", "reasoning_tokens", "other_modality_tokens", "expected_text_tokens"), + ( + pytest.param(50, 30, 20, 0, 30, id="details_sum_to_completion_is_a_no_op"), + pytest.param(34, 30, 24, 0, 10, id="strip_is_capped_at_the_over_sum"), + pytest.param(100, 100, 10, 70, 90, id="only_the_reasoning_share_is_stripped_when_text_over_reports_further"), + pytest.param(10, 5, 20, 0, 0, id="text_never_goes_negative_when_reasoning_exceeds_it"), + ), +) +def test_text_tokens_without_nested_reasoning_clamps( + completion_tokens: int, + text_tokens: int, + reasoning_tokens: int, + other_modality_tokens: int, + expected_text_tokens: int, +) -> None: + """The strip never exceeds the reasoning share, the reported text, or the over-sum past completion_tokens.""" + + assert ( + text_tokens_without_nested_reasoning( + completion_tokens=completion_tokens, + text_tokens=text_tokens, + reasoning_tokens=reasoning_tokens, + other_modality_tokens=other_modality_tokens, + ) + == expected_text_tokens + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index ce0cd75cd36..e7c6adc478c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -22,7 +22,7 @@ vi.mock("@/components/shared/advanced_date_picker", () => ({ import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; -import AutoRouterBenchmarksTab from "./AutoRouterBenchmarksTab"; +import AutoRouterBenchmarksTab, { AutoRouterUsageView } from "./AutoRouterBenchmarksTab"; import type { AutoRouterBenchmarkGroup, AutoRouterBenchmarksResponse, @@ -359,13 +359,37 @@ describe("AutoRouterBenchmarksTab", () => { mockHook({ data: response([group()]) }); const { dateValue, onDateChange } = renderTab(); - expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue); + expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, undefined); expect(screen.getByText("Jul 6 – Aug 5 (UTC)")).toBeInTheDocument(); fireEvent.click(screen.getByTestId("date-picker")); expect(onDateChange).toHaveBeenCalledWith({ from: new Date(2026, 7, 1), to: new Date(2026, 7, 5) }); }); + it("scopes the query to one key when the usage view is mounted for a key", () => { + mockHook({ data: response([group()]) }); + const dateValue = { from: new Date(2026, 6, 6), to: new Date(2026, 7, 5) }; + const activity = { + dateValue, + onDateChange: vi.fn(), + results: [], + loading: false, + isFetchingMore: false, + progress: { currentPage: 1, totalPages: 1 }, + cancelled: false, + cancel: vi.fn(), + }; + render( + + + , + ); + + expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, "key-hash-1"); + expect(screen.getByText("Total estimated savings")).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Shadow Evals" })).not.toBeInTheDocument(); + }); + it("shows usage by default and mounts shadow evals only when its sub-tab is selected", () => { mockHook({ data: response([group()]) }); renderTab(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index 09a0cf0242b..39e6b0fd390 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -273,12 +273,13 @@ const BenchmarksBody: React.FC = ({ isPending, error, data, interface AutoRouterBenchmarksTabProps { accessToken: string | null; - activity: DailyActivityRange; + activity: Pick; + apiKey?: string; } -const UsageView: React.FC = ({ accessToken, activity }) => { +export const AutoRouterUsageView: React.FC = ({ accessToken, activity, apiKey }) => { const { dateValue, onDateChange } = activity; - const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue); + const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey); const [selectedKey, setSelectedKey] = useState(ALL_ROUTERS); const { data: autoRouters } = useAutoRouters(); @@ -347,7 +348,7 @@ const AutoRouterBenchmarksTab: React.FC = ({ acces - + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts index e8a858ef3a5..1258c967ec9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.test.ts @@ -3,7 +3,9 @@ import { describe, expect, it, vi } from "vitest"; vi.mock("@/components/networking", () => ({ formatDate: vi.fn() })); vi.mock("@/lib/http/api", () => ({ $api: { useQuery: vi.fn() } })); -import { benchmarksWindow } from "./useAutoRouterBenchmarks"; +import { $api } from "@/lib/http/api"; + +import { benchmarksWindow, useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks"; const localDay = (offsetHours: number) => @@ -44,3 +46,30 @@ describe("benchmarksWindow", () => { expect(benchmarksWindow({ to: now }, now, pacific)).toEqual({}); }); }); + +describe("useAutoRouterBenchmarks", () => { + const range = { from: new Date("2026-07-06T19:00:00Z"), to: new Date("2026-08-05T19:00:00Z") }; + + it("forwards the key hash as the endpoint's api_key filter", () => { + useAutoRouterBenchmarks("sk-test", range, "key-hash-1"); + + const [, path, init] = vi.mocked($api.useQuery).mock.calls.at(-1) as unknown as [ + string, + string, + { params: { query: Record } }, + ]; + expect(path).toBe("/auto_router/benchmarks"); + expect(init.params.query.api_key).toBe("key-hash-1"); + }); + + it("leaves the read deployment-wide when no key is given", () => { + useAutoRouterBenchmarks("sk-test", range); + + const [, , init] = vi.mocked($api.useQuery).mock.calls.at(-1) as unknown as [ + string, + string, + { params: { query: Record } }, + ]; + expect(init.params.query.api_key).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts index 155c9d0c696..284342f0aec 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useAutoRouterBenchmarks.ts @@ -23,10 +23,10 @@ export const benchmarksWindow = ( }; }; -export const useAutoRouterBenchmarks = (accessToken: string | null, range: DateRange) => +export const useAutoRouterBenchmarks = (accessToken: string | null, range: DateRange, apiKey?: string) => $api.useQuery( "get", "/auto_router/benchmarks", - { params: { query: benchmarksWindow(range, new Date()) } }, + { params: { query: { ...benchmarksWindow(range, new Date()), api_key: apiKey } } }, { enabled: Boolean(accessToken && range.from && range.to), retry: false }, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx index 78280e32eed..51dfe9fda90 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.integration.test.tsx @@ -24,12 +24,18 @@ describe("useScopedDailyActivityRange wiring", () => { ); global.fetch = mockFetch; - renderHook(() => useScopedDailyActivityRange("sk-token", { userId: "u1", apiKey: "hash-abc" })); + const activity = { + dateValue: { from: new Date(2026, 7, 1), to: new Date(2026, 7, 10) }, + onDateChange: vi.fn(), + }; + renderHook(() => useScopedDailyActivityRange("sk-token", { userId: "u1", apiKey: "hash-abc" }, activity)); await waitFor(() => expect(mockFetch).toHaveBeenCalled()); const url = String(mockFetch.mock.calls[0][0]); expect(url).toContain("user_id=u1"); expect(url).toContain("api_key=hash-abc"); + expect(url).toContain("start_date=2026-08-01"); + expect(url).toContain("end_date=2026-08-10"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx index 6c9281060a4..00902aa9fdd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx @@ -25,11 +25,19 @@ vi.mock("@/components/networking", () => ({ })); import { userDailyActivityAggregatedCall } from "@/components/networking"; -import { useDailyActivityRange } from "./useDailyActivityRange"; +import { useActivityDateRange, useDailyActivityRange } from "./useDailyActivityRange"; const argsOfLastCall = () => mockUsePaginatedDailyActivity.mock.calls.at(-1)?.[0].args as unknown[]; describe("useDailyActivityRange", () => { + it("offers date-range state without starting a daily-activity query", () => { + const { result } = renderHook(() => useActivityDateRange()); + + expect(result.current.dateValue.from).toBeInstanceOf(Date); + expect(result.current.dateValue.to).toBeInstanceOf(Date); + expect(mockUsePaginatedDailyActivity).not.toHaveBeenCalled(); + }); + it("queries every user's activity for an admin", () => { renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin")); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts index 8ed57e36a94..3435b57dbc8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts @@ -37,14 +37,20 @@ export interface DailyActivityScope { apiKey?: string | null; } -export const useScopedDailyActivityRange = ( - accessToken: string | null, - scope: DailyActivityScope, -): DailyActivityRange => { +export type ActivityDateRange = Pick; + +export const useActivityDateRange = (): ActivityDateRange => { const initialFrom = useMemo(() => new Date(new Date().getTime() - THIRTY_DAYS_MS), []); const initialTo = useMemo(() => new Date(), []); const [dateValue, setDateValue] = useState({ from: initialFrom, to: initialTo }); + return { dateValue, onDateChange: setDateValue }; +}; +export const useScopedDailyActivityRange = ( + accessToken: string | null, + scope: DailyActivityScope, + { dateValue, onDateChange }: ActivityDateRange, +): DailyActivityRange => { const startTime = dateValue.from ?? null; const endTime = dateValue.to ?? null; const { userId, apiKey = null } = scope; @@ -63,7 +69,7 @@ export const useScopedDailyActivityRange = ( return { dateValue, - onDateChange: setDateValue, + onDateChange, results: data.results as DailyData[], loading, isFetchingMore, @@ -77,7 +83,7 @@ export const useDailyActivityRange = ( accessToken: string | null, userId: string | null, userRole: string, -): DailyActivityRange => - useScopedDailyActivityRange(accessToken, { - userId: spendScopeUserId(userRole, userId), - }); +): DailyActivityRange => { + const dateRange = useActivityDateRange(); + return useScopedDailyActivityRange(accessToken, { userId: spendScopeUserId(userRole, userId) }, dateRange); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts index f9cab181e71..66786fac1c3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editServerPayload.differential.test.ts @@ -12,7 +12,7 @@ import { } from "@/components/mcp_tools/types"; import { AUTH_TYPES_REQUIRING_CREDENTIALS } from "./createServerPayload"; import { TOOL_DISPLAY_NAME_PATTERN, normalizeEnvVars } from "./utils"; -import { buildEditServerPayload, type EditServerUiState } from "./editServerPayload"; +import { buildEditServerPayload, type EditServerFormValues, type EditServerUiState } from "./editServerPayload"; import { CASES, baseUi } from "./editServerPayload.differential.cases"; // GENERATED by scratchpad/emit_test.py. The body below is machine-extracted from @@ -317,6 +317,47 @@ describe("buildEditServerPayload matches the pre-extraction handleSave body", () }); }); +const EDIT_FORM_VALUES: EditServerFormValues = { + server_name: "srv", + alias: "srv_alias", + description: "a server", + transport: "http", + url: "https://example.com/mcp", + auth_type: "none", + mcp_access_groups: [], + extra_headers: [], + static_headers: [], + env_vars: [], + allow_all_keys: false, + available_on_public_internet: true, +}; + +describe("buildEditServerPayload wire contract", () => { + it("carries an edited alias and the server identifier onto the wire", () => { + const result = buildEditServerPayload({ ...EDIT_FORM_VALUES, alias: "renamed" }, baseUi); + + expect(result).toMatchObject({ kind: "ok", payload: { server_id: "srv_1", alias: "renamed" } }); + }); + + it.fails( + "sends description as an explicit null when the field is cleared (expected to fail until the forms revamp, tri-state PATCH tracker: today the cleared field reaches the wire as an empty string)", + () => { + const result = buildEditServerPayload({ ...EDIT_FORM_VALUES, description: "" }, baseUi); + + expect(result).toMatchObject({ kind: "ok", payload: { description: null } }); + }, + ); + + it.fails( + "sends only the server identifier and the edited alias (expected to fail until the forms revamp, tri-state PATCH tracker)", + () => { + const result = buildEditServerPayload({ ...EDIT_FORM_VALUES, alias: "renamed" }, baseUi); + + expect(result).toStrictEqual({ kind: "ok", payload: { server_id: "srv_1", alias: "renamed" } }); + }, + ); +}); + void ADMIN_CONFIG_CREDENTIAL_KEYS; void AUTH_TYPE; void AUTH_TYPES_REQUIRING_CREDENTIALS; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx index 8fb94ce477e..2571eb344f5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx @@ -1,8 +1,11 @@ import { cleanup, fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../../../tests/test-utils"; +import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; import { UserEditView } from "./user_edit_view"; +import * as networking from "@/components/networking"; + +vi.mock("@/components/networking"); vi.mock("@/components/key_team_helpers/fetch_available_models_team_key", () => ({ getModelDisplayName: vi.fn((model: string) => model), @@ -59,6 +62,10 @@ describe("UserEditView", () => { beforeEach(() => { vi.clearAllMocks(); + testQueryClient.clear(); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]); }); afterEach(() => { @@ -579,6 +586,44 @@ describe("UserEditView", () => { expect(budgetInput.closest("form")).not.toHaveAttribute("novalidate"); }); + it("shows the tool matrix for servers the user reaches only through an access group or toolset", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: "srv-group", server_name: "Group Server", alias: "Group Server", mcp_access_groups: ["group-a"] }, + { server_id: "srv-toolset", server_name: "Toolset Server", alias: "Toolset Server" }, + ]); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue(["group-a"]); + vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([ + { + toolset_id: "toolset-a", + toolset_name: "Toolset A", + tools: [{ server_id: "srv-toolset", tool_name: "list_issues" }], + } as never, + ]); + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: [{ name: "list_issues", description: "List issues" }], + error: false, + }); + + renderWithProviders( + , + ); + + expect(await screen.findByText("Via access group: group-a")).toBeInTheDocument(); + expect(await screen.findByText("Via toolset: Toolset A")).toBeInTheDocument(); + expect(networking.listMCPTools).toHaveBeenCalledWith("test-token", "srv-group"); + expect(networking.listMCPTools).toHaveBeenCalledWith("test-token", "srv-toolset"); + }); + it("should send objects for the mcp keys seeded from objectPermission", async () => { const payload = await submittedPayload({ objectPermission: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx index 96771dc6dc4..b7a3486c78e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx @@ -336,6 +336,8 @@ export function UserEditView({ form.setValue("mcp_tool_permissions", toolPerms)} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx index 7137da201d8..3224496b13a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx @@ -45,7 +45,7 @@ describe("VectorStoreManagement loading state", () => { it("should resolve the loading state when accessToken is null instead of showing the skeleton forever", async () => { const user = userEvent.setup(); - render(); + render(); await openManageTab(user); expect(await screen.findByText("table-loaded")).toBeInTheDocument(); expect(mockVectorStoreListCall).not.toHaveBeenCalled(); @@ -59,7 +59,7 @@ describe("VectorStoreManagement loading state", () => { resolveFetch = resolve; }), ); - render(); + render(); await openManageTab(user); expect(screen.getByText("table-loading")).toBeInTheDocument(); @@ -69,6 +69,43 @@ describe("VectorStoreManagement loading state", () => { }); }); +describe("VectorStoreManagement create flow visibility", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockVectorStoreListCall.mockResolvedValue({ data: [] }); + mockCredentialListCall.mockResolvedValue({ credentials: [] }); + }); + + it.each([ + { label: "Internal User", userRole: "Internal User", isViewOnly: false }, + { label: "Internal Viewer", userRole: "Internal Viewer", isViewOnly: true }, + { label: "proxy_admin_viewer session (userRole Admin, isViewOnly)", userRole: "Admin", isViewOnly: true }, + { label: "Org Admin", userRole: "Org Admin", isViewOnly: false }, + ])( + "should hide the Create Vector Store tab and button and skip /credentials for $label", + async ({ userRole, isViewOnly }) => { + render( + , + ); + await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test")); + expect(screen.queryByRole("tab", { name: "Create Vector Store" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Manage Vector Stores" })).toHaveAttribute("aria-selected", "true"); + expect(await screen.findByText("table-loaded")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "+ Add Vector Store" })).not.toBeInTheDocument(); + expect(mockCredentialListCall).not.toHaveBeenCalled(); + }, + ); + + it("should keep the Create Vector Store tab and button and fetch /credentials for a proxy admin", async () => { + const user = userEvent.setup(); + render(); + await waitFor(() => expect(mockCredentialListCall).toHaveBeenCalledWith("sk-test")); + expect(screen.getByRole("tab", { name: "Create Vector Store" })).toHaveAttribute("aria-selected", "true"); + await openManageTab(user); + expect(screen.getByRole("button", { name: "+ Add Vector Store" })).toBeInTheDocument(); + }); +}); + describe("VectorStoreManagement Indexes tab", () => { beforeEach(() => { vi.clearAllMocks(); @@ -88,7 +125,7 @@ describe("VectorStoreManagement Indexes tab", () => { }, ], }); - render(); + render(); await user.click(screen.getByRole("tab", { name: "Indexes" })); expect(await screen.findByText("support-docs-index")).toBeInTheDocument(); expect(screen.getByText("support-docs-store")).toBeInTheDocument(); @@ -96,7 +133,7 @@ describe("VectorStoreManagement Indexes tab", () => { }); it("should not render the Indexes tab for an Admin Viewer", async () => { - render(); + render(); await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test")); expect(screen.getByRole("tab", { name: "Manage Vector Stores" })).toBeInTheDocument(); expect(screen.queryByRole("tab", { name: "Indexes" })).not.toBeInTheDocument(); @@ -125,7 +162,7 @@ describe("VectorStoreManagement Indexes tab", () => { }, ], }); - render(); + render(); await user.click(screen.getByRole("tab", { name: "Indexes" })); await user.click(await screen.findByRole("button", { name: "support-docs-store" })); expect(await screen.findByTestId("vector-store-info-view")).toHaveTextContent("vs-1"); @@ -135,7 +172,7 @@ describe("VectorStoreManagement Indexes tab", () => { it("should link to the feature docs and a GitHub issue for unsupported providers on the Indexes tab", async () => { const user = userEvent.setup(); mockIndexesListCall.mockResolvedValue({ object: "list", data: [] }); - render(); + render(); await user.click(screen.getByRole("tab", { name: "Indexes" })); expect(screen.getByRole("link", { name: "vector store index docs" })).toHaveAttribute( "href", @@ -149,7 +186,7 @@ describe("VectorStoreManagement Indexes tab", () => { }); it("should not call indexesListCall until the Indexes tab is clicked", async () => { - render(); + render(); await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test")); expect(screen.getByRole("tab", { name: "Indexes" })).toBeInTheDocument(); expect(mockIndexesListCall).not.toHaveBeenCalled(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index 1745c710c51..285a96520bd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -24,9 +24,10 @@ interface VectorStoreProps { accessToken: string | null; userID: string | null; userRole: string | null; + isViewOnly: boolean; } -const VectorStoreManagement: React.FC = ({ accessToken, userID, userRole }) => { +const VectorStoreManagement: React.FC = ({ accessToken, userID, userRole, isViewOnly }) => { const [vectorStores, setVectorStores] = useState([]); const [isLoadingVectorStores, setIsLoadingVectorStores] = useState(true); const [isCreateModalVisible, setIsCreateModalVisible] = useState(false); @@ -37,7 +38,9 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID const [selectedVectorStoreId, setSelectedVectorStoreId] = useState(null); const [editVectorStore, setEditVectorStore] = useState(false); const [isDeleting, setIsDeleting] = useState(false); - const { onTabChange, hasVisited } = useVisitedTabs("create"); + const canCreateVectorStores = isProxyAdminRole(userRole || "") && !isViewOnly; + const defaultTab = canCreateVectorStores ? "create" : "manage"; + const { onTabChange, hasVisited } = useVisitedTabs(defaultTab); const fetchVectorStores = async () => { if (!accessToken) { @@ -56,7 +59,7 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID }; const fetchCredentials = async () => { - if (!accessToken) return; + if (!accessToken || !canCreateVectorStores) return; try { const response = await credentialListCall(accessToken); setCredentials(response.credentials || []); @@ -153,11 +156,13 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID You can use vector stores to store and retrieve LLM embeddings.

- + - - Create Vector Store - + {canCreateVectorStores && ( + + Create Vector Store + + )} Manage Vector Stores @@ -171,14 +176,18 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID )} - - - + {canCreateVectorStores && ( + + + + )} - + {canCreateVectorStores && ( + + )}
; + const { accessToken, userRole, userId, isViewOnly } = useAuthorized(); + return ( + + ); } diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 5605a993ded..014854ac712 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -60,6 +60,20 @@ const optionByLabel = (label: string): HTMLElement | undefined => const isOptionDisabled = (option: HTMLElement): boolean => option.getAttribute("aria-disabled") === "true"; +const tierChips = (tier: string): HTMLElement => { + const placeholder = `Select model(s) for ${tier.toLowerCase()} queries`; + const chips = screen + .getAllByRole("toolbar") + .find((candidate) => within(candidate).queryByLabelText(placeholder) !== null); + if (!chips) throw new Error(`No tier row found for "${tier}"`); + return chips; +}; + +const expectTierModel = (tier: string, model: string): void => { + const chips = within(tierChips(tier)).getAllByLabelText(/.+/, { selector: '[data-slot="combobox-chip"]' }); + expect(chips.map((chip) => chip.getAttribute("aria-label"))).toEqual([model]); +}; + const selectTemplate = async (label: string): Promise => { await userEvent.click(optionByLabel(label)!); }; @@ -188,11 +202,10 @@ describe("AddAutoRouterTab", () => { const button = await screen.findByTestId("configure-automatically-button"); await userEvent.click(button); - expect( - screen.getByText( - /Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: claude-opus-5.*Reasoning: claude-opus-5/, - ), - ).toBeInTheDocument(); + expectTierModel("Simple", "gpt-5.6-luna"); + expectTierModel("Medium", "claude-sonnet-5"); + expectTierModel("Complex", "claude-opus-5"); + expectTierModel("Reasoning", "claude-opus-5"); expect(toast.success).not.toHaveBeenCalledWith(expect.stringContaining("Configured with")); }); @@ -209,9 +222,23 @@ describe("AddAutoRouterTab", () => { const button = await screen.findByTestId("configure-automatically-button"); await userEvent.click(button); - expect( - screen.getByText(/Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: gpt-5.6-sol.*Reasoning: gpt-5.6-sol/), - ).toBeInTheDocument(); + expectTierModel("Simple", "gpt-5.6-luna"); + expectTierModel("Medium", "claude-sonnet-5"); + expectTierModel("Complex", "gpt-5.6-sol"); + expectTierModel("Reasoning", "gpt-5.6-sol"); + }); + + it("opens Detailed Configuration on the tiers automatic setup just filled in", async () => { + const simpleModel = "gpt-5.6-luna"; + mockFetchAvailableModels.mockResolvedValue([...ALL_FAMILY_MODELS, { model_group: simpleModel, mode: "chat" }]); + renderWithProviders(); + + expect(screen.queryByText("Complexity Tier Configuration")).not.toBeInTheDocument(); + + await userEvent.click(await screen.findByTestId("configure-automatically-button")); + + expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument(); + expectTierModel("Simple", simpleModel); }); // Nothing is filled in, so there is nothing to submit. The button reports that itself instead of diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 60780f9cbe0..2d0d6bc2fd8 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -207,10 +207,6 @@ const AddAutoRouterTab: React.FC = ({ const [isSubmitting, setIsSubmitting] = useState(false); const [selectedPreset, setSelectedPreset] = useState(undefined); - // Closed by default: a caller opens it deliberately, either by clicking it or by choosing Custom - // (which expands it automatically, since there's nothing else to show them their config from). A - // preset re-collapses it after prefilling, offering the same "here's what got filled in, expand to - // change it" affordance. A caller can always toggle it manually at any point. const [detailsExpanded, setDetailsExpanded] = useState(false); const [isRoutingTestVisible, setIsRoutingTestVisible] = useState(false); @@ -335,7 +331,7 @@ const AddAutoRouterTab: React.FC = ({ if (automaticRouterConfig === null) return; setSelectedPreset(undefined); applyPrefill({ ...buildEmptyPrefill(), complexityRouterConfig: automaticRouterConfig }); - setDetailsExpanded(false); + setDetailsExpanded(true); toast.success("Automatic setup created", { description: tierConfigSummary(automaticRouterConfig) }); }; @@ -543,14 +539,15 @@ const AddAutoRouterTab: React.FC = ({ {!automaticSetupLoading && automaticRouterConfig && ( - +
+
+

Not sure where to start?

+

Let us pick models for each complexity tier.

+
+ +
)}
diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx index 69a761b4723..149d231fff6 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx @@ -19,6 +19,7 @@ describe("MCPToolPermissions", () => { vi.clearAllMocks(); testQueryClient.clear(); vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); }); it("should update tool permissions when user selects a tool", async () => { @@ -621,6 +622,45 @@ describe("MCPToolPermissions", () => { ); expect(await screen.findByText("Unable to load MCP servers")).toBeInTheDocument(); + expect(screen.queryByText(/has 0 servers/)).not.toBeInTheDocument(); + }); + + it("tells the admin when a loaded access group has no member servers", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([groupServer]); + vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]); + vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: groupTools, error: false }); + + renderWithProviders( + , + ); + + expect(await screen.findByText('Access group "ops_readonly" has 0 servers')).toBeInTheDocument(); + expect(screen.getByText("Group Server")).toBeInTheDocument(); + expect(screen.queryByText('Access group "production-group" has 0 servers')).not.toBeInTheDocument(); + }); + + it("does not call a group empty when its servers are only hidden from the caller's catalog", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue(["production-group"]); + + renderWithProviders( + , + ); + + expect(await screen.findByText('Access group "ops_readonly" has 0 servers')).toBeInTheDocument(); + expect(screen.queryByText('Access group "production-group" has 0 servers')).not.toBeInTheDocument(); }); it("warns when the selected toolsets cannot be resolved to servers", async () => { diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index c866e9cc011..e26f1a6f511 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -4,6 +4,7 @@ import { MCPTool } from "../mcp_tools/types"; import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPAccessGroups } from "../../app/(dashboard)/hooks/mcpServers/useMCPAccessGroups"; import { useMCPToolsets } from "../../app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import McpCrudPermissionPanel from "../mcp_tools/McpCrudPermissionPanel"; import { classifyToolOp } from "../../utils/mcpToolCrudClassification"; @@ -12,6 +13,7 @@ import { EffectiveMcpServer, McpGrantSource, applyToolPermissionWrite, + emptyMcpAccessGroups, mcpAllowedToolsFor, resolveEffectiveMcpServers, } from "./effectiveMcpServers"; @@ -55,7 +57,13 @@ const MCPToolPermissions: React.FC = ({ onChange, disabled = false, }) => { - const { data: allServers = [], isError: serversFailed, isLoading: serversLoading } = useMCPServers(); + const { + data: allServers = [], + isError: serversFailed, + isLoading: serversLoading, + isSuccess: serversLoaded, + } = useMCPServers(); + const { data: populatedAccessGroups = [], isSuccess: accessGroupsLoaded } = useMCPAccessGroups(); const { data: toolsets = [], isError: toolsetsFailed, isLoading: toolsetsLoading } = useMCPToolsets(); const [serverTools, setServerTools] = useState>({}); const [loadingTools, setLoadingTools] = useState>({}); @@ -181,6 +189,18 @@ const MCPToolPermissions: React.FC = ({
)} + {serversLoaded && + accessGroupsLoaded && + emptyMcpAccessGroups(allServers, populatedAccessGroups, selectedAccessGroups).map((group) => ( +
+

Access group "{group}" has 0 servers

+

+ No MCP server lists this group, so it grants nothing. A server defined in config.yaml joins a group + through its access_groups key; mcp_access_groups is ignored there +

+
+ ))} + {toolsetsFailed && selectedToolsets.length > 0 && (

Unable to load toolsets

diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts index b2ac5337245..487c6f9f55e 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts +++ b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect } from "vitest"; import { MCPServer, MCPToolset } from "../mcp_tools/types"; import { applyToolPermissionWrite, + emptyMcpAccessGroups, mcpAllowedToolsFor, mcpServersForIdentifier, mcpToolPermissionKeyFor, @@ -101,6 +102,29 @@ describe("mcpToolPermissionKeyFor", () => { }); }); +describe("emptyMcpAccessGroups", () => { + const grouped = server({ server_id: "srv-group", server_name: "Grouped", mcp_access_groups: ["prod"] }); + const objectGrouped = { + ...grouped, + server_id: "srv-obj", + mcp_access_groups: [{ name: "legacy" }], + } as unknown as MCPServer; + + it("names only the selected groups no loaded server belongs to", () => { + expect(emptyMcpAccessGroups([grouped, objectGrouped], [], ["prod", "legacy", "ops_readonly"])).toEqual([ + "ops_readonly", + ]); + }); + + it("names every selected group when no server is loaded and the registry is empty", () => { + expect(emptyMcpAccessGroups([], [], ["prod"])).toEqual(["prod"]); + }); + + it("trusts the group registry when the caller's catalog hides the member servers", () => { + expect(emptyMcpAccessGroups([], ["prod"], ["prod", "ops_readonly"])).toEqual(["ops_readonly"]); + }); +}); + describe("resolveEffectiveMcpServers", () => { const direct = server({ server_id: "srv-direct", server_name: "Direct" }); const grouped = server({ server_id: "srv-group", server_name: "Grouped", mcp_access_groups: ["prod"] }); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts index 3320dc135c9..b3ba24f3c59 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts +++ b/ui/litellm-dashboard/src/components/mcp_server_management/effectiveMcpServers.ts @@ -54,6 +54,19 @@ const accessGroupNamesOf = (server: MCPServer): readonly string[] => return [typeof parsed.data === "string" ? parsed.data : parsed.data.name]; }); +// The server catalog can be trimmed to the caller's grants, so the unfiltered group registry +// (GET /v1/mcp/access_groups) has to agree before a group is called empty. +export const emptyMcpAccessGroups = ( + allServers: readonly MCPServer[], + populatedAccessGroups: readonly string[], + selectedAccessGroups: readonly string[], +): readonly string[] => + selectedAccessGroups.filter( + (group) => + !populatedAccessGroups.includes(group) && + !allServers.some((server) => accessGroupNamesOf(server).includes(group)), + ); + // Which servers an identifier names, with the same precedence the backend's expand_permission_list // applies: a string that is a registry server id names exactly that server, and only a string that // is not falls back to server_name/alias, which can name several. Matching all three fields at once diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 768183907db..3db9418dfb9 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -1729,5 +1729,77 @@ describe("ModelInfoView", () => { expect(payload.litellm_params.cache_control_injection_points).toEqual([{ location: "message", index: "2" }]); }); }); + + const setInputCost = (value: string) => { + fireEvent.change(screen.getByPlaceholderText("Enter input cost"), { target: { value } }); + }; + + it("carries an edited input cost and the model identifier onto the wire", async () => { + const user = userEvent.setup(); + await enterEditMode(user); + setInputCost("5"); + const payload = await save(user); + + expect(mockModelPatchUpdateCall.mock.calls[0][2]).toBe("123"); + expect(payload.litellm_params.input_cost_per_token).toBe(5 / 1_000_000); + }); + + it.fails( + "sends only the edited input cost (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const user = userEvent.setup(); + await enterEditMode(user); + setInputCost("5"); + const payload = await save(user); + + expect(payload).toStrictEqual({ litellm_params: { input_cost_per_token: 5 / 1_000_000 } }); + }, + ); + + const savePayloadAfterCostEditOnResolvedModel = async () => { + const resolved = { + ...defaultModelData, + model_info: { + ...defaultModelData.model_info, + max_input_tokens: 128_000, + mode: "chat", + supports_vision: true, + supports_function_calling: true, + }, + }; + mockUseModelsInfo.mockReturnValue({ data: { data: [resolved] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [resolved] }); + const user = userEvent.setup(); + await enterEditMode(user); + setInputCost("5"); + return save(user); + }; + + it.fails( + "leaves max_input_tokens off the wire when only the input cost is edited (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const payload = await savePayloadAfterCostEditOnResolvedModel(); + + expect(payload.model_info).not.toHaveProperty("max_input_tokens"); + }, + ); + + it.fails( + "leaves mode off the wire when only the input cost is edited (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const payload = await savePayloadAfterCostEditOnResolvedModel(); + + expect(payload.model_info).not.toHaveProperty("mode"); + }, + ); + + it.fails( + "leaves every supports_ capability off the wire when only the input cost is edited (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const payload = await savePayloadAfterCostEditOnResolvedModel(); + + expect(Object.keys(payload.model_info).filter((key) => key.startsWith("supports_"))).toStrictEqual([]); + }, + ); }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 045dcfa3ceb..3471ef00eb0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -544,6 +544,37 @@ describe("CreateKey", () => { expect((await createdPayload()).metadata).toBe('{"team":"research"}'); }); + + it("carries the typed key alias and the chosen team onto the wire", async () => { + state.teams = [{ team_id: "team-1", team_alias: "Team One", models: [] }]; + await openModal({ teams: state.teams as unknown as Team[] }); + await nameTheKey("wire-alias"); + await userEvent.click(await screen.findByLabelText("Team")); + await userEvent.click(await screen.findByRole("option", { name: /Team One/ })); + await submit(); + + expect(await createdPayload()).toMatchObject({ key_alias: "wire-alias", team_id: "team-1" }); + }); + + it("sends team_id as an explicit null when no team is chosen", async () => { + await openModal(); + await nameTheKey(); + await submit(); + + expect(await createdPayload()).toHaveProperty("team_id", null); + }); + + it.fails( + "adds no keys for an Optional Settings section the user opened but never filled (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await submit(); + + expect(await createdPayload()).toStrictEqual(ALL_CLOSED_PAYLOAD); + }, + ); }); describe("key ownership", () => { 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 dfc737ddd45..4c68e302267 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -89,6 +89,16 @@ describe("provider_info_helpers", () => { expect(result.logo).toBe(providerLogoMap[Providers.BedrockMantle]); }); + it("should map the chatgpt slug and CHATGPT enum key to the ChatGPT Subscription name and OpenAI logo", () => { + const fromSlug = getProviderLogoAndName("chatgpt"); + expect(fromSlug.displayName).toBe("ChatGPT Subscription"); + expect(fromSlug.logo).toContain("openai_small"); + + const fromEnumKey = getProviderLogoAndName("CHATGPT"); + expect(fromEnumKey.displayName).toBe("ChatGPT Subscription"); + expect(fromEnumKey.logo).toContain("openai_small"); + }); + it("should handle provider values case-insensitively", () => { const result = getProviderLogoAndName("OPENAI"); expect(result.displayName).toBe(Providers.OpenAI); @@ -272,6 +282,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Cognition)).toBe("cognition/swe-1.7"); }); + it("should return a chatgpt/ placeholder for the CHATGPT dropdown key", () => { + expect(getPlaceholder("CHATGPT")).toBe("chatgpt/gpt-5.4"); + }); + it("should return default gpt-3.5-turbo placeholder for unknown provider", () => { expect(getPlaceholder("UnknownProvider" as any)).toBe("gpt-3.5-turbo"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index d01a6a34cbe..72ec5990557 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -84,6 +84,7 @@ export enum Providers { BASETEN = "Baseten", BYTEZ = "Bytez", Cerebras = "Cerebras", + CHATGPT = "ChatGPT Subscription", CLARIFAI = "Clarifai", CLOUDFLARE = "Cloudflare", CODESTRAL = "Codestral", @@ -198,6 +199,7 @@ export const provider_map: Record = { BedrockMantle: "bedrock_mantle", BYTEZ: "bytez", Cerebras: "cerebras", + CHATGPT: "chatgpt", CLARIFAI: "clarifai", CLOUDFLARE: "cloudflare", CODESTRAL: "codestral", @@ -314,6 +316,7 @@ export const providerLogoMap: Partial> = { [Providers.BedrockMantle]: bedrockLogo.src, [Providers.SageMaker]: bedrockLogo.src, [Providers.Cerebras]: cerebrasLogo.src, + [Providers.CHATGPT]: openaiSmallLogo.src, [Providers.CLOUDFLARE]: cloudflareLogo.src, [Providers.CODESTRAL]: mistralLogo.src, [Providers.Cohere]: cohereLogo.src, @@ -425,6 +428,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Azure]: "my-deployment", [Providers.Azure_AI_Studio]: "azure_ai/command-r-plus", [Providers.Bedrock]: "claude-3-opus", + [Providers.CHATGPT]: "chatgpt/gpt-5.4", [Providers.Cognition]: "cognition/swe-1.7", [Providers.Cursor]: "cursor/claude-4-sonnet", [Providers.DeepInfra]: "deepinfra/", diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 041ba4b548d..b286c8c1303 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -2052,6 +2052,19 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { expect(objectPermission.agent_access_groups).toStrictEqual([]); }); + it("sends an empty vector_stores array after the last vector store chip is removed", async () => { + const user = userEvent.setup({ delay: null }); + await openEditor(user); + + await user.click(within(screen.getByLabelText("vs-1")).getByRole("button")); + expect(screen.queryByLabelText("vs-1")).not.toBeInTheDocument(); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.vector_stores).toStrictEqual([]); + }); + it("resends every stored value once both sections are opened", async () => { const user = userEvent.setup({ delay: null }); await openEditor(user); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index efd69a79fc0..4947cec1441 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1039,7 +1039,7 @@ const TeamInfoView: React.FC = ({ delete values.agents_and_groups; // Handle vector stores permissions - if (values.vector_stores && values.vector_stores.length > 0) { + if (values.vector_stores) { updateData.object_permission.vector_stores = values.vector_stores; } diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx new file mode 100644 index 00000000000..8cfd9d1941e --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx @@ -0,0 +1,98 @@ +import { screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; +import KeyAutoRouterUsageTab from "./KeyAutoRouterUsageTab"; + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "test-token", userId: "admin-123", userRole: "Admin" }), +})); + +const jsonResponse = (body: unknown) => + new Response(JSON.stringify(body), { status: 200, headers: { "content-type": "application/json" } }); + +const cache = { + coverage_pct: 100, + hit_rate_pct: 50, + same_model: { turns: 2, hits: 1, hit_rate_pct: 50 }, + first_visit: { turns: 1, hits: 0, hit_rate_pct: 0 }, + return_to_tier: { turns: 1, hits: 1, hit_rate_pct: 100 }, + unordered_turns: 0, + return_misses_expired: 0, + return_misses_within_ttl: 0, + return_misses_unknown: 0, + ttl_5m_turns: 0, + ttl_1h_turns: 0, +}; + +const stats = { + sessions: 2, + turns: 4, + avg_turns_per_session: 2, + avg_session_seconds: 30, + avg_tokens_per_session: 100, + spend: 1.25, + saved_spend: 8.75, + baseline_spend: 10, + saved_pct: 87.5, + saved_per_session: 4.375, + cache, +}; + +const benchmarks = { + start_date: "2025-01-01", + end_date: "2025-01-31", + routers_in_scope: 2, + totals: stats, + groups: [ + { router_name: "router-one", router_type: "complexity", tier_turns: { SIMPLE: 4 }, ...stats }, + { + router_name: "router-two", + router_type: "complexity", + tier_turns: { SIMPLE: 1 }, + ...stats, + spend: 0.25, + saved_spend: 0.75, + baseline_spend: 1, + }, + ], +}; + +const noDeployments = { data: [], total_count: 0, current_page: 1, total_pages: 1, size: 1000 }; +const fetchMock = vi.fn(async (request: Request | string) => { + const url = typeof request === "string" ? request : request.url; + if (url.includes("/auto_router/benchmarks")) return jsonResponse(benchmarks); + return jsonResponse(noDeployments); +}); +const requestedUrls = () => + fetchMock.mock.calls.map(([request]) => (typeof request === "string" ? request : request.url)); + +describe("KeyAutoRouterUsageTab", () => { + beforeEach(() => { + vi.clearAllMocks(); + testQueryClient.clear(); + vi.stubGlobal("fetch", fetchMock); + }); + + it("renders this key's spend, baseline, savings and per-router filter", async () => { + const activity = { + dateValue: { from: new Date(2025, 0, 1), to: new Date(2025, 0, 31) }, + onDateChange: vi.fn(), + }; + renderWithProviders(); + + expect(await screen.findByText("$8.75")).toBeInTheDocument(); + expect(screen.getByText("Actual auto-router spend")).toBeInTheDocument(); + expect(screen.getByText("$1.25")).toBeInTheDocument(); + expect(screen.getByText("Estimated spend at highest-tier model")).toBeInTheDocument(); + expect(screen.getByText("$10.00")).toBeInTheDocument(); + expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument(); + expect(screen.getAllByText("50.0%").length).toBeGreaterThan(0); + expect(screen.getByText("All auto-routers")).toBeInTheDocument(); + + const benchmarkUrl = new URL(requestedUrls().find((url) => url.includes("/auto_router/benchmarks")) ?? ""); + expect(benchmarkUrl.searchParams.get("api_key")).toBe("key-hash-1"); + expect(benchmarkUrl.searchParams.get("start_date")).toBe("2025-01-01"); + expect(benchmarkUrl.searchParams.get("end_date")).toBe("2025-01-31"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx new file mode 100644 index 00000000000..9e1cee1ec2e --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.tsx @@ -0,0 +1,18 @@ +"use client"; + +import React from "react"; + +import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab"; +import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; + +interface KeyAutoRouterUsageTabProps { + accessToken: string | null; + keyToken: string; + activity: ActivityDateRange; +} + +const KeyAutoRouterUsageTab: React.FC = ({ accessToken, keyToken, activity }) => ( + +); + +export default KeyAutoRouterUsageTab; diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index 1ab7b9be52c..2bbbc2bd48b 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -1,7 +1,11 @@ import React from "react"; +import { Control } from "react-hook-form"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { CircleHelp } from "lucide-react"; +import { FormField } from "@/components/shared/form/FormField"; +import NumericalInput from "../shared/numerical_input"; +import { KeyEditFormValues } from "./keyEditFormValues"; export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => ( <> @@ -48,3 +52,27 @@ export const KeyTypeSelect = ({ ); + +export const KeyBudgetNumberField = ({ + control, + name, + label, + placeholder, +}: { + control: Control; + name: "max_budget" | "soft_budget"; + label: string; + placeholder: string; +}) => ( + + {({ ref: _ref, ...field }) => ( + + )} + +); diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 705679fe2e6..8a42ecff50c 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -1,5 +1,10 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { toast } from "@/lib/toast"; + +vi.mock("@/lib/toast", () => ({ + toast: { error: vi.fn(), success: vi.fn() }, +})); // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ---- const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => { @@ -483,3 +488,104 @@ describe("KeyInfoView handleKeyUpdate empty strings", () => { }); }); }); + +describe("KeyInfoView handleKeyUpdate soft_budget", () => { + const premiumAdminAuth = { + accessToken: "access_abc", + userId: "user_1", + userRole: "Admin", + premiumUser: true, + token: "token_123", + userEmail: "test@example.com", + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }; + + const renderWithSoftBudget = (softBudget: number | null) => { + mockUseAuthorized.mockReturnValue(premiumAdminAuth); + + return render( + {}} + keyData={ + { ...baseKeyData, litellm_budget_table: softBudget === null ? null : { soft_budget: softBudget } } as any + } + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + }; + + it("should send a changed soft_budget as a number", async () => { + renderWithSoftBudget(null); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + soft_budget: "25", + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.soft_budget).toBe(25); + }); + + it("should omit an unchanged soft_budget so unrelated edits skip the budget gate", async () => { + renderWithSoftBudget(25); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + soft_budget: 25, + key_alias: "renamed", + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect("soft_budget" in sentPayload).toBe(false); + }); + + it("should forward a cleared soft_budget as an explicit null the JSON body keeps", async () => { + renderWithSoftBudget(25); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + soft_budget: "", + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.soft_budget).toBeNull(); + expect(JSON.stringify({ ...sentPayload })).toContain('"soft_budget":null'); + }); + + it("should reject an overflowing soft_budget instead of silently clearing it", async () => { + renderWithSoftBudget(25); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + soft_budget: "1e309", + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(toast.error).toHaveBeenCalled()); + expect(keyUpdateCallMock).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx index 6c7d65ee566..385c3967d02 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx @@ -46,6 +46,11 @@ const mockActivity = ( const scopedRange = () => vi.spyOn(useScopedDailyActivityRangeModule, "useScopedDailyActivityRange"); +const activity = { + dateValue: { from: new Date(2025, 0, 1), to: new Date(2025, 0, 31) }, + onDateChange: vi.fn(), +}; + const renderTab = (props: Partial> = {}) => render( > = keyToken="key-abc123" userId="user-123" userRole="Internal User" + activity={activity} {...props} />, ); @@ -114,7 +120,7 @@ describe("KeySavingsTab", () => { renderTab({ userId: "user-456", userRole: "Internal User" }); - expect(hook).toHaveBeenCalledWith("test-token", { userId: "user-456", apiKey: "key-abc123" }); + expect(hook).toHaveBeenCalledWith("test-token", { userId: "user-456", apiKey: "key-abc123" }, activity); expect(screen.getByTestId("key-savings-scope-note")).toHaveTextContent("Showing your own requests"); }); @@ -123,7 +129,7 @@ describe("KeySavingsTab", () => { renderTab({ userId: "admin-123", userRole: "Admin" }); - expect(hook).toHaveBeenCalledWith("test-token", { userId: null, apiKey: "key-abc123" }); + expect(hook).toHaveBeenCalledWith("test-token", { userId: null, apiKey: "key-abc123" }, activity); expect(screen.queryByTestId("key-savings-scope-note")).not.toBeInTheDocument(); }); @@ -132,7 +138,7 @@ describe("KeySavingsTab", () => { renderTab({ userId: "org-admin-1", userRole: "Org Admin" }); - expect(hook).toHaveBeenCalledWith("test-token", { userId: "org-admin-1", apiKey: "key-abc123" }); + expect(hook).toHaveBeenCalledWith("test-token", { userId: "org-admin-1", apiKey: "key-abc123" }, activity); expect(screen.getByTestId("key-savings-scope-note")).toHaveTextContent("Showing your own requests"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx index 967c3c7eff0..c33529eb042 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx @@ -22,7 +22,10 @@ import { usd, withStartAnchor, } from "@/app/(dashboard)/cost-optimization/_components/costOptimizationUtils"; -import { useScopedDailyActivityRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; +import { + useScopedDailyActivityRange, + type ActivityDateRange, +} from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; interface KeySavingsTabProps { accessToken: string | null; @@ -30,19 +33,19 @@ interface KeySavingsTabProps { keyToken: string; userId: string | null; userRole: string; + activity: ActivityDateRange; } -const KeySavingsTab: React.FC = ({ accessToken, keyToken, userId, userRole }) => { +const KeySavingsTab: React.FC = ({ accessToken, keyToken, userId, userRole, activity }) => { // Proxy admins read the whole key. For anyone else the endpoint applies the caller's own user_id // alongside the key filter, so the figures cover only that viewer's requests on this key -- said // plainly in the scope note below rather than left to be misread as the key's total. const readsWholeKey = hasProxyWideSpendView(userRole); - const activity = useScopedDailyActivityRange(accessToken, { - userId: spendScopeUserId(userRole, userId), - apiKey: keyToken, - }); - - const { dateValue, onDateChange, results, loading, isFetchingMore } = activity; + const { dateValue, onDateChange, results, loading, isFetchingMore } = useScopedDailyActivityRange( + accessToken, + { userId: spendScopeUserId(userRole, userId), apiKey: keyToken }, + activity, + ); const startTime = dateValue.from ?? null; const endTime = dateValue.to ?? null; diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index ad90610b732..f24fb6e5a86 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -22,6 +22,7 @@ export interface KeyEditFormValues { models?: string[]; allowed_routes?: string; max_budget?: number | string | null; + soft_budget?: number | string | null; budget_duration?: string | null; tpm_limit?: number | string | null; tpm_limit_type?: string | null; @@ -67,6 +68,8 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 ? keyData.allowed_routes.join(", ") : "", max_budget: keyData.max_budget, + soft_budget: + (keyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ?? null, budget_duration: canonicalBudgetDuration(keyData.budget_duration), tpm_limit: keyData.tpm_limit, tpm_limit_type: (keyData as { tpm_limit_type?: string | null }).tpm_limit_type ?? null, @@ -117,6 +120,7 @@ export const keyEditFormSchema = z.object({ models: z.custom(), allowed_routes: z.custom(), max_budget: z.custom(), + soft_budget: z.custom(), budget_duration: z.custom(), tpm_limit: z.custom(), tpm_limit_type: z.custom(), @@ -168,6 +172,7 @@ export const toSubmittedValues = ( models: values.models, allowed_routes: values.allowed_routes, max_budget: values.max_budget, + soft_budget: values.soft_budget, budget_duration: values.budget_duration, tpm_limit: values.tpm_limit, tpm_limit_type: values.tpm_limit_type, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 8b506329595..35e9268e1c2 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1885,6 +1885,7 @@ describe("KeyEditView", () => { key_alias: "asdasdas", models: [], max_budget: 0, + soft_budget: null, budget_duration: "30d", tpm_limit: 10, tpm_limit_type: null, @@ -2300,5 +2301,57 @@ describe("KeyEditView", () => { }); expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("tag_rpm_limit", { "test-tag": 7 }); }); + + const setRpmLimit = (value: string) => { + fireEvent.change(screen.getByLabelText("RPM Limit"), { target: { value } }); + }; + + it("carries an edited RPM limit and the key identifier onto the wire", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + setRpmLimit("25"); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledTimes(1); + }); + expect(onSubmitMock.mock.calls[0][0]).toMatchObject({ token: "test-token-123", rpm_limit: "25" }); + }); + + it.fails( + "sends max_budget as an explicit null when the field is cleared (expected to fail until the forms revamp, tri-state PATCH tracker: today the view hands KeyInfoView an empty string and handleKeyUpdate maps it to null)", + async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + await userEvent.clear(screen.getByLabelText("Max Budget (USD)")); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledTimes(1); + }); + expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("max_budget", null); + }, + ); + + it.fails( + "sends only the key identifier and the edited RPM limit (expected to fail until the forms revamp, tri-state PATCH tracker)", + async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmitMock); + await screen.findByRole("button", { name: /save changes/i }); + + setRpmLimit("25"); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledTimes(1); + }); + expect(onSubmitMock.mock.calls[0][0]).toStrictEqual({ token: "test-token-123", rpm_limit: "25" }); + }, + ); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 0f955fee895..b1bbbeed6f9 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -32,7 +32,7 @@ import { modelSentinelOptions, parseAllowedRoutes, } from "./keyEditFieldNormalizers"; -import { KeyTypeSelect, labelWithHint } from "./KeyEditViewControls"; +import { KeyBudgetNumberField, KeyTypeSelect, labelWithHint } from "./KeyEditViewControls"; import { AgentsAndGroups, KeyEditFormValues, @@ -417,17 +417,19 @@ export function KeyEditView({ )} - - {({ ref: _ref, ...field }) => ( - - )} - + + + {({ value, onChange, id }) => ( diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 7817d4e7cea..b403255b329 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -26,6 +26,30 @@ vi.mock("./key_edit_view", () => ({ }, })); +import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; + +const AnalyticsDateControl = ({ activity, keyToken }: { activity: ActivityDateRange; keyToken: string }) => ( +
+ {keyToken} + {activity.dateValue.from?.toISOString()} + {[1, 10].map((day) => ( + + ))} +
+); + +vi.mock("./KeyAutoRouterUsageTab", () => ({ + default: (props: React.ComponentProps) => , +})); +vi.mock("./KeySavingsTab", () => ({ + default: (props: React.ComponentProps) => , +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: vi.fn(), })); @@ -176,6 +200,38 @@ describe("KeyInfoView", () => { await userEvent.click(await screen.findByRole("button", { name: /more key actions/i })); }; + it("shows key-scoped auto-router usage as its own admin tab", async () => { + vi.mocked(useAuthorized).mockReturnValue({ ...baseUseAuthorizedMock, userRole: "Admin" }); + renderWithProviders( {}} keyId="test-key-id" teams={[]} />); + + await userEvent.click(screen.getByRole("tab", { name: "Auto-router usage" })); + + expect(screen.getByTestId("key-auto-router-usage")).toHaveTextContent("test-token-123"); + }); + + it("preserves dates in both directions across unmounted analytics panels", async () => { + vi.mocked(useAuthorized).mockReturnValue({ ...baseUseAuthorizedMock, userRole: "Admin" }); + renderWithProviders( {}} keyId="test-key-id" teams={[]} />); + + expect(screen.queryByLabelText("Selected dates")).not.toBeInTheDocument(); + await userEvent.click(screen.getByRole("tab", { name: "Savings" })); + await userEvent.click(screen.getByRole("button", { name: "Select August 1" })); + await userEvent.click(screen.getByRole("tab", { name: "Auto-router usage" })); + expect(screen.getByLabelText("Selected dates")).toHaveTextContent("2026-08-01T00:00:00.000Z"); + await userEvent.click(screen.getByRole("button", { name: "Select August 10" })); + await userEvent.click(screen.getByRole("tab", { name: "Settings" })); + expect(screen.queryByLabelText("Selected dates")).not.toBeInTheDocument(); + await userEvent.click(screen.getByRole("tab", { name: "Savings" })); + expect(screen.getByLabelText("Selected dates")).toHaveTextContent("2026-08-10T00:00:00.000Z"); + }); + + it("does not offer the admin-only auto-router usage tab to an internal user", () => { + vi.mocked(useAuthorized).mockReturnValue({ ...baseUseAuthorizedMock, userRole: "Internal User" }); + renderWithProviders( {}} keyId="test-key-id" teams={[]} />); + + expect(screen.queryByRole("tab", { name: "Auto-router usage" })).not.toBeInTheDocument(); + }); + describe("last updated", () => { const renderWithTimestamps = (overrides: Partial) => { vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); @@ -1015,6 +1071,16 @@ describe("KeyInfoView", () => { expect(keyUpdateCall).toHaveBeenCalledWith(expect.anything(), expect.objectContaining({ policies: [] })); }); + + it("puts the key identifier and an explicit null max_budget on the wire when the edit view hands over a cleared budget", async () => { + await enterEditMode({ ...MOCK_KEY_DATA, user_id: "proxy-admin-user" } as KeyResponse); + await editViewMocks.onSubmit!({ token: MOCK_KEY_DATA.token, max_budget: "" }); + + expect(keyUpdateCall).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ key: "test-token-123", max_budget: null }), + ); + }); }); describe("MCP tool permissions on save", () => { diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 0dd0dd6d6af..969decb6613 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -16,8 +16,15 @@ import { modelGroupHref, teamDetailHref } from "@/utils/entityLinks"; import { BadgeLink } from "@/components/shared/BadgeLink"; import { KeyInfoHeader } from "./KeyInfoHeader"; import KeySavingsTab from "./KeySavingsTab"; +import KeyAutoRouterUsageTab from "./KeyAutoRouterUsageTab"; +import { useActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; import { useEffect, useState } from "react"; -import { isProxyAdminRole, isUserTeamAdminForSingleTeam, rolesWithWriteAccess } from "../../utils/roles"; +import { + hasProxyWideSpendView, + isProxyAdminRole, + isUserTeamAdminForSingleTeam, + rolesWithWriteAccess, +} from "../../utils/roles"; import { mapDisplayToInternalNames, mapInternalToDisplayNames } from "../callback_info_helpers"; import AutoRotationView from "../common_components/AutoRotationView"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; @@ -81,6 +88,7 @@ export default function KeyInfoView({ backButtonText = "Back to Keys", }: KeyInfoViewProps) { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); + const activityDateRange = useActivityDateRange(); const queryClient = useQueryClient(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); const { teams: teamsData } = useTeams(); @@ -210,6 +218,23 @@ export default function KeyInfoView({ // Handle max budget empty string formValues.max_budget = mapEmptyStringToNull(formValues.max_budget); + // soft_budget is a budget change server-side (admin-gated); only send it when it changed + // so a non-admin edit of unrelated fields isn't blocked by that gate. + const previousSoftBudget = + (currentKeyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ?? + null; + const nextSoftBudget = + formValues.soft_budget === "" || formValues.soft_budget == null ? null : Number(formValues.soft_budget); + if (nextSoftBudget !== null && !Number.isFinite(nextSoftBudget)) { + toast.error("Soft Budget must be a finite number"); + return; + } + if (nextSoftBudget === previousSoftBudget) { + delete formValues.soft_budget; + } else { + formValues.soft_budget = nextSoftBudget; + } + // Handle object_permission updates if (formValues.vector_stores !== undefined) { formValues.object_permission = { @@ -618,6 +643,11 @@ export default function KeyInfoView({ Savings + {hasProxyWideSpendView(userRole) && ( + + Auto-router usage + + )} Settings @@ -761,9 +791,20 @@ export default function KeyInfoView({ keyToken={currentKeyData.token} userId={userID} userRole={userRole} + activity={activityDateRange} /> + {hasProxyWideSpendView(userRole) && ( + + + + )} + {/* Settings Panel */} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1534c19670..6fb08445aff 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8017,6 +8017,11 @@ export interface paths { * Update Key Fn * @description Update an existing API key's parameters. * + * The body is a merge patch: a field left out keeps its stored value, and on the key's own columns + * an explicit null clears it. The metadata-backed fields below are the exception, merging into the + * stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself + * replaces the stored metadata wholesale. + * * Parameters: * - key: Optional[str] - The key to update. Either key or key_alias must be provided. * - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases) @@ -8034,7 +8039,7 @@ export interface paths { * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - * - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. + * - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget. * - max_parallel_requests: Optional[int] - Rate limit for parallel requests * - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} * - tpm_limit: Optional[int] - Tokens per minute limit @@ -19867,6 +19872,12 @@ export interface paths { * curl "http://localhost:4000/v1/skills?beta=true&limit=10" -H "Authorization: Bearer your-key" -H "x-litellm-model: claude-account-1" * ``` * + * Pass `?custom_llm_provider=litellm_proxy&query=` to rank the LiteLLM-hosted skills you can + * access by semantic similarity instead of paging through the whole registry: + * ```bash + * curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" -H "Authorization: Bearer your-key" + * ``` + * * Returns: ListSkillsResponse with list of skills */ get: operations["list_skills_v1_skills_get"]; @@ -23748,9 +23759,9 @@ export interface components { on_sensitive_data?: ("block" | "route") | null; /** * On Violation - * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. + * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. */ - on_violation?: ("warn" | "end_session") | null; + on_violation?: ("warn" | "end_session" | "block" | "alert") | null; /** * Only Scan New Messages * @description When True, the guardrail only scans messages that have not already been scanned earlier in the same session (identified by litellm_session_id / session_id). Message content is hashed per session and cached; only the diff (new or edited messages) is sent to the guardrail provider on follow-up calls. Falls back to a full scan when the request has no session id or the cache is unavailable. Intended for blocking/detection guardrails; not applied when mask_request_content is set. @@ -29562,8 +29573,6 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; - /** Rust */ - rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ @@ -30792,9 +30801,9 @@ export interface components { on_sensitive_data?: ("block" | "route") | null; /** * On Violation - * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. + * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. */ - on_violation?: ("warn" | "end_session") | null; + on_violation?: ("warn" | "end_session" | "block" | "alert") | null; /** * Only Scan New Messages * @description When True, the guardrail only scans messages that have not already been scanned earlier in the same session (identified by litellm_session_id / session_id). Message content is hashed per session and cached; only the diff (new or edited messages) is sent to the guardrail provider on follow-up calls. Falls back to a full scan when the request has no session id or the cache is unavailable. Intended for blocking/detection guardrails; not applied when mask_request_content is set. @@ -36047,12 +36056,16 @@ export interface components { Skill: { /** Created At */ created_at: string; + /** Description */ + description?: string | null; /** Display Title */ display_title?: string | null; /** Id */ id: string; /** Latest Version */ latest_version?: string | null; + /** Search Score */ + search_score?: number | null; /** Source */ source: string; /** @@ -37748,6 +37761,8 @@ export interface components { rpm_limit?: number | null; /** Rpm Limit Type */ rpm_limit_type?: ("guaranteed_throughput" | "best_effort_throughput" | "dynamic") | null; + /** Soft Budget */ + soft_budget?: number | null; /** Spend */ spend?: number | null; /** Tag Rpm Limit */ @@ -39729,8 +39744,6 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; - /** Rust */ - rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ @@ -41392,6 +41405,8 @@ export interface operations { start_date?: string | null; /** @description YYYY-MM-DD UTC, inclusive (defaults to today) */ end_date?: string | null; + /** @description Filter to one virtual key token hash */ + api_key?: string | null; }; header?: never; path?: never; @@ -64553,6 +64568,10 @@ export interface operations { after_id?: string | null; before_id?: string | null; custom_llm_provider?: string | null; + /** @description Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model. */ + query?: string | null; + /** @description With query: the maximum number of ranked skills to return. */ + top_k?: number; }; header?: never; path?: never;